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

[khazic]

Closes #2424.

What

Make SGLang accelerate the EAGLE-3 target in both deployment shapes:

1. Co-located (new): target_model_backend: sglang runs the frozen target through SGLang's ModelRunner inside the training process, replacing the HF eager forward (single-process runs).
2. Remote: serve_target --engine sglang serves the target on dedicated GPUs while the draft trains elsewhere.

Both sit on an engine-agnostic target-backend contract so another engine (e.g. vLLM) can drop in later without touching the trainer or the remote server.

Why

EAGLE-3 trains the draft against the target's auxiliary hidden states plus its distribution. The target forward dominates each training step, and SGLang is the fastest serving path for mainstream architectures. Running the frozen target through SGLang keeps training and inference numerically consistent while cutting target-side latency.

What changed

• target_runner.py (new): a narrow TargetRunner protocol plus RunnerEagle3TargetModel, which owns the supervision contract (shift / aux concatenation / aux-layer defaulting) for any runner. This is the single engine-agnostic seam.
• sglang_target.py / sglang_runner.py (new): the SGLang adapter. SGLangEagle3TargetModel adds only SGLang construction on top of RunnerEagle3TargetModel; SGLangTargetRunner owns the SGLang-internal forward (wraps LogitsProcessor to return all-position full-vocab logits plus the three EAGLE-3 aux hidden states) and is imported lazily so the module stays importable without SGLang. ServerArgs.dtype receives SGLang's string form (sglang_dtype_str), and building inside an initialized process group requires world_size == tp_size with a clear error.
• train_eagle3.py: target_model_backend: sglang selects the co-located SGLang path. SGLang's weight + KV pool defaults to half the GPU so the draft trains in the remainder; recipe_args.sglang_args forwards ServerArgs overrides. Single-process only (SGLang's parallel state must own every rank of the process group); multi-GPU runs use the remote backend. Example config: examples/speculative/eagle3/llama_eagle3_sglang.yaml.
• target.py: extract default_eagle3_aux_layer_ids / validate_eagle3_aux_layer_ids so every backend defaults and validates aux layers identically (behavior-preserving).
• serve_target.py: add --engine {hf,sglang}; lazy-import the HF AutoModel so the sglang engine runs in a minimal sglang-only environment.
• remote/client.py + remote/transport.py: gate the NCCL handshake on local availability (nccl_transport_available()). A client without sglang now goes straight to the wire fallback instead of asking the server to start a rendezvous it can never complete (which previously blocked the server for ~120s before falling back).
• scripts/smoke_sglang_target.py (new): GPU smoke for the server: validates the ported SGLang forward (shapes, finiteness) and its numerical agreement with the HF backend, with and without a pre-initialized process group (the co-located case).

Validation

The contract layer and the recipe wiring are fully unit-tested on CPU. The GPU path was validated end-to-end on a server: a standalone SGLang target server (Qwen3-4B) plus a sglang-free training client over HTTP + wire returns the correct supervision shapes (target_probs [B,S,vocab], aux [B,S,3*hidden]). The co-located path is validated on the server with scripts/smoke_sglang_target.py --compare-hf --init-dist.

Tests

• tests/unit_tests/speculative/test_eagle3_sglang.py: SGLang supervision is byte-for-byte identical to the co-located HF backend; the backend is engine-agnostic; aux-layer defaulting/validation; runner surface without SGLang; dtype string mapping.
• tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py: backend dispatch, CUDA / single-process guards, and sglang_args flow into the SGLang kwargs.
• tests/unit_tests/speculative/test_eagle3_remote_coverage.py: an sglang-free client does not contact the server for NCCL.
• Full speculative + eagle recipe suites green, ruff clean.

Follow-up

A separate PR will switch the remote data model to ship hidden states and recompute the target distribution trainer-side (lower bandwidth, less target-side compute, and less engine coupling).

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[HuiyingLi]

/ok to test b1f6e19

[HuiyingLi]

/ok to test 02c4a63

[khazic]

GPU validation of the SGLang EAGLE-3 target backend

Ran the backend on a single A800-80GB GPU at this branch's HEAD (cebe206a4).

Setup (dedicated env, sglang installed out of band as the PR documents):

• sglang 0.5.9, torch 2.9.1+cu128, transformers 4.57.1
• target: Qwen3-4B (bf16), single process, --init-dist (matches the co-located target_model_backend: sglang path the recipe uses)

1. scripts/smoke_sglang_target.py Stage 1: well-formed supervision (PASS)

aux_layer_ids=[1, 17, 32] hidden_size=2560 vocab_size=151936
aux (2, 16, 7680) torch.bfloat16
logits (2, 16, 151936) torch.float32
STAGE 1 OK: SGLang target produces well-formed supervision

Shapes ([B, S, 3*hidden] aux, [B, S, vocab] logits), dtype, full-logits encoding, and finiteness all check out.

2. SGLang vs HF supervision equivalence

The metric that matters for EAGLE-3 is the aux hidden states the draft trains on, and they line up almost exactly with the HF co-located backend.

Random tokens (the script default):

logits argmax match rate (loss positions): 0.9333
aux hidden-state mean cosine similarity:    0.9999

Real text, 4x48, equal length, no padding (188 supervised tokens):

logits argmax match rate:        0.9840
SGLang top-1 in HF top-5 rate:   1.0000
aux hidden-state mean cosine:    0.9999

Aux cosine is 0.9999 in both cases. The argmax dip on random tokens (0.9333) is bf16 argmax flips at near-tie positions on garbage input, not a representational difference: on real text the argmax match rises to 0.9840, every SGLang top-1 is within the HF top-5 (1.0000), and the cosine stays 0.9999. SGLang and HF agree.

Note for maintainers: transformers version pin

--compare-hf builds the HF target via NeMoAutoModelForCausalLM, which imports AutoModelForMultimodalLM (transformers 5.x), while sglang==0.5.9 pins transformers==4.57.1. The two can't co-exist in one env, so --compare-hf is not runnable as written alongside sglang. For the equivalence numbers above I loaded the HF side with plain transformers.AutoModelForCausalLM, still wrapped in this PR's HFEagle3TargetModel, so only the loader differs. Might be worth a line in the smoke docstring or the sglang doc that the HF comparison needs a transformers that satisfies both.

The recipe-level packing guard added here is covered by the unit tests in tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py (CPU suite passes locally).

[HuiyingLi]

/ok to test 493a904

[HuiyingLi]

Note for maintainers: transformers version pin

--compare-hf builds the HF target via NeMoAutoModelForCausalLM, which imports AutoModelForMultimodalLM (transformers 5.x), while sglang==0.5.9 pins transformers==4.57.1. The two can't co-exist in one env, so --compare-hf is not runnable as written alongside sglang. For the equivalence numbers above I loaded the HF side with plain transformers.AutoModelForCausalLM, still wrapped in this PR's HFEagle3TargetModel, so only the loader differs. Might be worth a line in the smoke docstring or the sglang doc that the HF comparison needs a transformers that satisfies both.

The recipe-level packing guard added here is covered by the unit tests in tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py (CPU suite passes locally).

Is it possible to include the version information in a guide/tutorial in the future? Thank you so much.

[khazic]

Done. Added a version-compatibility note to the EAGLE guide's environment-setup step (docs/guides/speculative/eagle.mdx) spelling out the sglang==0.5.9 / transformers==4.57.1 constraint and how to run --compare-hf. Thanks for the suggestion!

[github-actions]

🌿 Preview your docs: https://nvidia-preview-khazic-feat-eagle3-sglang-target.docs.buildwithfern.com/nemo/automodel