feat(speculative): add remote target serving for EAGLE-3 training

examples/speculative/eagle3/llama_eagle3_remote.yaml

@@ -0,0 +1,64 @@
+recipe: TrainEagle3Recipe
+
+# Remote (train-inference disaggregated) EAGLE-3 training.
+#
+# The frozen target model runs as a standalone server (see
+# ``nemo_automodel.components.speculative.serve_target``); this training
+# process holds only the draft model and pulls precomputed supervision over
+# HTTP + NCCL. Start the server first, e.g.:
+#
+#   python -m nemo_automodel.components.speculative.serve_target \
+#       --target meta-llama/Llama-3.2-1B --host 0.0.0.0 --port 8001
+#
+# then launch this recipe pointing ``remote_urls`` at it. Everything else is
+# identical to the co-located MVP config (llama_eagle3_mvp.yaml).
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 30
+
+recipe_args:
+  target_model_name_or_path: meta-llama/Llama-3.2-1B
+
+  # --- remote backend ---
+  # ``colocated`` (default) loads the target on the training GPU; ``remote``
+  # talks to one or more serve_target instances.
+  target_model_backend: remote
+  # One URL per target server. With multiple servers the prefetch pipeline
+  # round-robins requests across them.
+  remote_urls: ["http://127.0.0.1:8001"]
+  # Prefetch depth overlaps target inference with draft training. Capped to the
+  # number of servers (one in-flight request per server). 0 disables prefetch
+  # (pure disaggregation). Requires the remote backend.
+  target_prefetch_depth: 1
+  # remote_timeout: 120
+  # remote_max_retries: 3
+
+  train_data_path: /path/to/train.jsonl
+  val_data_path: null
+  train_split: null
+  val_split: null
+  output_dir: ./outputs/eagle3_llama_remote
+  seq_length: 1024
+  micro_batch_size: 1
+  grad_accumulation_steps: 1
+  num_workers: 0
+  num_epochs: 1
+  ttt_steps: 4
+  draft_vocab_size: 8192
+  freeze_embeddings: true
+  trust_remote_code: false
+  shuffle_seed: 42
+  log_every_steps: 10
+  max_grad_norm: 1.0
+
+optimizer:
+  lr: 1.0e-4
+  betas: [0.9, 0.95]
+  weight_decay: 0.0
+
+checkpoint:
+  enabled: true
+  checkpoint_dir: ./outputs/eagle3_llama_remote/checkpoints
+  model_save_format: safetensors
+  save_consolidated: true

nemo_automodel/components/speculative/eagle/__init__.py

@@ -24,6 +24,7 @@
 Import them directly from those modules.
 """
 
+from nemo_automodel.components.speculative.eagle.backend import Eagle3TargetBackend
 from nemo_automodel.components.speculative.eagle.core import Eagle3TrainerModule
 from nemo_automodel.components.speculative.eagle.core_v12 import EagleTrainerModule
 from nemo_automodel.components.speculative.eagle.draft_llama import LlamaEagle3DraftModel
@@ -41,6 +42,7 @@
 __all__ = [
     "EagleTrainerModule",
     "Eagle3TrainerModule",
+    "Eagle3TargetBackend",
     "HFEagleTargetModel",
     "HFEagle3TargetModel",
     "LlamaEagleDraftModel",

nemo_automodel/components/speculative/eagle/backend.py

@@ -0,0 +1,102 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Backend abstraction for the EAGLE-3 target model.
+
+The frozen target model is a *supervision provider*: for every training batch
+it produces the auxiliary hidden states and the per-token target distribution
+the draft model is trained against. EAGLE-3 training never updates the target,
+so it does not have to share the training GPU. This interface lets the recipe
+consume the target uniformly whether it runs co-located in-process
+(``HFEagle3TargetModel``) or, in a later change, as a remote inference service
+on separate GPUs.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+
+if TYPE_CHECKING:
+    from nemo_automodel.components.speculative.eagle.target import Eagle3TargetBatch
+
+
+class Eagle3TargetBackend(ABC):
+    """Abstract contract every EAGLE-3 target-model backend implements.
+
+    Two supervision encodings are allowed, both consumed directly by
+    :meth:`Eagle3TrainerModule.forward`:
+
+    - **full logits** -- :attr:`Eagle3TargetBatch.logits` carries the target's
+      full-vocab logits and the draft-vocab projection happens trainer-side.
+      Cheap when co-located (the tensor never leaves the GPU); impractical to
+      ship over a wire because it is full-vocab sized.
+    - **precomputed** -- :attr:`Eagle3TargetBatch.target_probs` and
+      :attr:`Eagle3TargetBatch.position_mask` carry the already-projected
+      draft-vocab distribution, so a backend that computes them itself (e.g. a
+      remote server) only has to transfer draft-vocab-sized tensors.
+
+    A backend returns exactly one of the two encodings from
+    :meth:`generate_batch`; the recipe forwards whichever is present.
+    """
+
+    @abstractmethod
+    def generate_batch(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        loss_mask: torch.Tensor,
+    ) -> "Eagle3TargetBatch":
+        """Run the target and return the supervision for one training batch."""
+
+    @abstractmethod
+    def get_input_embeddings(self) -> nn.Module:
+        """Return the target input-embedding module (used to seed the draft)."""
+
+    def set_vocab_mapping(
+        self,
+        selected_token_ids: torch.Tensor,
+        selected_token_mask: torch.Tensor,
+    ) -> None:
+        """Provide the draft-vocab mapping needed to precompute supervision.
+
+        Co-located backends keep the mapping on the trainer module and derive
+        the distribution there, so the default is a no-op. A backend that
+        computes ``target_probs`` itself overrides this to receive the mapping.
+        """
+
+    @property
+    def supports_async(self) -> bool:
+        """Whether :meth:`generate_batch_async` is implemented (prefetch-capable)."""
+        return False
+
+    def generate_batch_async(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        loss_mask: torch.Tensor,
+    ):
+        """Submit an asynchronous :meth:`generate_batch` for prefetch pipelining.
+
+        Only backends that overlap target inference with draft training
+        implement this; the default signals a synchronous backend so callers
+        fall back to :meth:`generate_batch`.
+        """
+        raise NotImplementedError(f"{type(self).__name__} does not support async prefetch")
+
+    def close(self) -> None:
+        """Release backend resources (remote connections, server handles)."""

nemo_automodel/components/speculative/eagle/remote/__init__.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Train-inference disaggregation for EAGLE-3 target serving.
+
+Runs the frozen target model as a standalone inference server on separate
+GPU(s) while the draft model trains elsewhere. The training side talks to the
+server through :class:`RemoteEagle3TargetModel`, which implements the
+:class:`~nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend`
+contract, so the EAGLE-3 recipe consumes a remote target exactly like the
+co-located ``HFEagle3TargetModel``.
+
+- HTTP is the control plane (input_ids up, tensor metadata down).
+- NCCL is the data plane for the large supervision tensors (GPU-to-GPU,
+  NVLink intra-node / RDMA inter-node), with a compact binary wire format
+  fallback when NCCL is unavailable.
+"""
+
+from nemo_automodel.components.speculative.eagle.remote.client import RemoteEagle3TargetModel
+
+__all__ = ["RemoteEagle3TargetModel"]