feat(speculative): add SGLang target backend for EAGLE-3 training

docs/guides/speculative/eagle.mdx

@@ -89,6 +89,18 @@ For SGLang serving (Step 5), install it in the same environment:
 uv pip install "sglang>=0.5.9"
 ```
 
+<Note>
+**SGLang / transformers version compatibility.** SGLang `0.5.9` pins
+`transformers==4.57.1`. The SGLang target backend (`target_model_backend: sglang`)
+and `serve_sglang` run in that environment without issue. The `--compare-hf`
+check in `smoke_sglang_target.py` additionally builds the HuggingFace target
+through `NeMoAutoModelForCausalLM`, which imports `AutoModelForMultimodalLM`
+(available only in `transformers` 5.x), so that single-process comparison cannot
+run as written under `transformers==4.57.1`. To compare the two backends, use a
+`transformers` build that satisfies both, or load the HuggingFace side with plain
+`transformers.AutoModelForCausalLM` wrapped in `HFEagle3TargetModel`.
+</Note>
+
 ---
 
 ## Step 1 — Understand EAGLE Architecture

examples/speculative/eagle3/llama_eagle3_sglang.yaml

@@ -0,0 +1,59 @@
+recipe: TrainEagle3Recipe
+
+# Co-located EAGLE-3 training with the frozen target served through SGLang.
+#
+# Identical to the co-located MVP config (llama_eagle3_mvp.yaml) except the
+# target forward runs through SGLang's ModelRunner instead of the HuggingFace
+# eager forward, which is substantially faster for mainstream architectures.
+# SGLang carves its weight + KV pool out of the training GPU up front
+# (``sglang_args.mem_fraction_static``); the draft trains in the remainder.
+#
+# Single-process only: SGLang's parallel state must own every rank of the
+# process group. For multi-GPU runs serve the target separately
+# (``serve_target --engine sglang``) and use llama_eagle3_remote.yaml.
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 30
+
+recipe_args:
+  target_model_name_or_path: meta-llama/Llama-3.2-1B
+
+  # --- sglang backend ---
+  # ``colocated`` (default) runs the target through HF on this GPU; ``sglang``
+  # runs it through SGLang on this GPU; ``remote`` talks to serve_target.
+  target_model_backend: sglang
+  # Extra SGLang ServerArgs. ``mem_fraction_static`` is the fraction of GPU
+  # memory SGLang reserves for target weights + KV pool (default 0.5 here;
+  # raise it for big targets, lower it if draft training runs out of memory).
+  sglang_args:
+    mem_fraction_static: 0.5
+
+  train_data_path: /path/to/train.jsonl
+  val_data_path: null
+  train_split: null
+  val_split: null
+  output_dir: ./outputs/eagle3_llama_sglang
+  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_sglang/checkpoints
+  model_save_format: safetensors
+  save_consolidated: true

nemo_automodel/components/speculative/eagle/__init__.py

@@ -39,13 +39,16 @@
     resolve_eagle3_draft_spec,
 )
 from nemo_automodel.components.speculative.eagle.target import HFEagle3TargetModel
+from nemo_automodel.components.speculative.eagle.target_runner import RunnerEagle3TargetModel, TargetRunner
 from nemo_automodel.components.speculative.eagle.target_v12 import HFEagleTargetModel
 
 __all__ = [
     "EagleTrainerModule",
     "Eagle3TrainerModule",
     "PEagleTrainerModule",
     "Eagle3TargetBackend",
+    "RunnerEagle3TargetModel",
+    "TargetRunner",
     "HFEagleTargetModel",
     "HFEagle3TargetModel",
     "LlamaEagleDraftModel",

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

@@ -41,7 +41,7 @@
 
 from nemo_automodel.components.speculative.eagle.backend import Eagle3TargetBackend
 from nemo_automodel.components.speculative.eagle.remote import protocol, wire
-from nemo_automodel.components.speculative.eagle.remote.transport import NCCLTransport
+from nemo_automodel.components.speculative.eagle.remote.transport import NCCLTransport, nccl_transport_available
 from nemo_automodel.components.speculative.eagle.target import Eagle3TargetBatch
 
 logger = logging.getLogger(__name__)
@@ -107,6 +107,17 @@ def _init_nccl(self) -> bool:
             if self._nccl_attempted:
                 return False
             self._nccl_attempted = True
+
+            # Only ask the server to bring up NCCL if this process can actually
+            # join the group. A client without sglang (the common disaggregated
+            # case: sglang target server + sglang-free training client) cannot,
+            # and contacting the server anyway leaves it blocked on a rendezvous
+            # that never completes until it times out. Fall back to wire instead.
+            if not nccl_transport_available():
+                logger.info("NCCL unavailable in this process (sglang not importable); using wire format")
+                self._nccl = None
+                return False
+
             port = self._nccl_port()
             self._nccl = NCCLTransport(nccl_port=port, host=self._host(), is_server=False)
 

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

@@ -48,6 +48,19 @@
 
 _HAS_SGLANG_PG, _init_custom_process_group = safe_import_from("sglang.srt.utils.common", "init_custom_process_group")
 
+
+def nccl_transport_available() -> bool:
+    """Whether GPU-direct NCCL transfer is usable in this process.
+
+    NCCL transfer relies on sglang's ``init_custom_process_group``; without
+    sglang (e.g. a training client that intentionally keeps sglang out of its
+    env) NCCL cannot work and callers should use the wire fallback. Checking
+    this *before* asking the server to set up its NCCL side avoids leaving the
+    server blocked on a rendezvous the client can never complete.
+    """
+    return _HAS_SGLANG_PG
+
+
 # dtypes NCCL P2P does not support; transmitted as raw uint8 views.
 _NCCL_UNSUPPORTED_DTYPES = {torch.int16, torch.int8, torch.bool}
 _ELEMENT_SIZE = {torch.int16: 2, torch.int8: 1, torch.bool: 1}