feat(speculative): add vLLM target backend for EAGLE-3 training - #2798
Merged
HuiyingLi merged 21 commits intoJun 29, 2026
Conversation
Add a third Eagle3TargetBackend implementation that runs the frozen target
through SGLang, alongside the co-located HF and remote backends. SGLang is the
fastest serving path for mainstream architectures, so the remote target server
can hold the target on dedicated GPUs while the draft trains elsewhere.
The backend is split into two layers:
- SGLangEagle3TargetModel (sglang_target.py) owns the supervision contract and
assembles an Eagle3TargetBatch whose shift / aux-concatenation semantics are
byte-for-byte identical to HFEagle3TargetModel, so a SGLang run is
numerically equivalent to a co-located one. It depends only on a small runner
protocol, so it is unit-testable on CPU without SGLang.
- SGLangTargetRunner (sglang_runner.py) owns the SGLang-internal forward
(ModelRunner + CaptureHiddenMode.FULL + a logits-processor wrap that returns
all-position full-vocab logits plus the three concatenated aux hidden
states). It is lazily imported and validated on the GPU server; CPU tests
cover the surface that does not need SGLang.
serve_target gains an --engine {hf,sglang} flag (engines share a builder
signature via a dispatch map). SGLang is declared as an optional spec_sglang
extra pinned to 0.5.9 and kept out of the main training image. Shared
aux-layer-id default / validation helpers are extracted in target.py so all
backends default identically (behavior-preserving refactor).
Signed-off-by: khazic <khazzz1c@gmail.com>
Generalize the SGLang-specific runner protocol and target backend into an engine-agnostic seam so a second engine (vLLM) can plug in without touching the trainer, the remote server, or the supervision contract: - new target_runner.py: TargetRunner protocol + RunnerEagle3TargetModel, which owns the shift / aux-concatenation contract for any runner; - sglang_target.py: SGLangEagle3TargetModel now only adds SGLang construction on top of RunnerEagle3TargetModel (SGLangRunnerProtocol kept as a backwards-compatible alias of TargetRunner); - sglang_runner.py unchanged (still the GPU/SGLang forward). Behavior-preserving: the SGLang supervision is byte-for-byte identical, all existing CPU contract tests pass, plus a test locking that the backend is engine-agnostic. Signed-off-by: khazic <khazzz1c@gmail.com>
sglang>=0.5.9 made moe_ep_rank/moe_ep_size required positional args on ModelRunner.__init__; the target runner is single-process with no expert parallelism, so pass (0, 1). Signed-off-by: khazic <khazzz1c@gmail.com>
ModelRunner.init_torch_distributed already calls initialize_model_parallel in sglang>=0.5.9, so calling it ourselves trips 'tensor model parallel group is already initialized'. Bring up only the world process group here and let ModelRunner build the TP group. Signed-off-by: khazic <khazzz1c@gmail.com>
The sglang engine never loads the HF AutoModel, so importing NeMoAutoModelForCausalLM at module top forced the sglang target server to pull in Automodel's full model stack. Move it into _build_hf_target, matching the lazy sglang import in _build_sglang_target, so the sglang server runs in a minimal sglang-only environment. Signed-off-by: khazic <khazzz1c@gmail.com>
A client without sglang (the disaggregated case: sglang target server + sglang-free training client) cannot join the NCCL group, but _init_nccl still POSTed /init_nccl and let the server block on the rendezvous until its 120s timeout before both fell back to wire. Gate the request on a local nccl_transport_available() check so an sglang-free client goes straight to wire and never stalls the server. Also fix the serve_target test to patch the now lazily-imported NeMoAutoModelForCausalLM at its source. Signed-off-by: khazic <khazzz1c@gmail.com>
Drop the dead SGLangRunnerProtocol alias (the protocol and alias were both introduced on this branch, so there is no prior name to keep resolving), and cache the constant teacher-forcing SamplingParams on the runner instead of rebuilding it every extend. Signed-off-by: khazic <khazzz1c@gmail.com>
Add target_model_backend: sglang to the EAGLE-3 recipe: the frozen target runs through SGLang's ModelRunner on the training GPU (single-process only; SGLang's parallel state must own every rank, so multi-GPU runs keep using serve_target --engine sglang + the remote backend). SGLang's memory pool defaults to half the GPU here so the draft trains in the remainder, tunable via recipe_args.sglang_args. Also fix the ServerArgs dtype handling (SGLang compares dtype against string literals, so torch.dtype objects silently missed every branch; addresses the review comment), guard runner construction against a process-group/tp_size mismatch with a clear error, and add a GPU smoke script that validates the SGLang forward against the HF backend on the server, including the pre-initialized process group of the co-located path. Signed-off-by: khazic <khazzz1c@gmail.com>
Signed-off-by: khazic <khazzz1c@gmail.com>
sglang==0.5.9 hard-pins transformers==4.57.1, which conflicts with the project's transformers==5.8.1. Declaring it as the spec_sglang extra made uv's universal resolution unsatisfiable, failing every install-dependent CI job (lint, type-check, uv-lock, builds). The SGLang target backend is lazy-imported (safe_import), so the package needs no declared sglang dependency; install it in a separate dedicated SD venv/container as the serve_target docstring now documents. Signed-off-by: khazic <khazzz1c@gmail.com>
…argets The packed_sequence_size guard only blocked the remote backend, but the SGLang runner processes each row as one full causal sequence with no per-document masking, so packing + sglang silently leaked supervision across document boundaries. Gate packing on backend != 'colocated' and hoist the backend-name validation ahead of the guard so a misspelled backend still reports the clearer 'unknown backend' error. Also fix the copyright year in the new test (2026 -> 2025) and add coverage for the packing guard. Signed-off-by: khazic <khazzz1c@gmail.com>
Mirror the SGLang target backend with a vLLM engine adapter built on the shared engine-agnostic TargetRunner contract. VLLMTargetRunner uses vLLM's native extract_hidden_states speculative method to capture the three EAGLE-3 aux hidden states (plus the final pre-norm hidden) over a prefill, then rebuilds full-vocab logits in-process via the target's final RMSNorm + LM head. VLLMEagle3TargetModel is the thin RunnerEagle3TargetModel adapter, matching SGLangEagle3TargetModel. Recipe dispatch, remote serving, example config, smoke script and tests follow once the forward is GPU-validated. Signed-off-by: khazic <khazzz1c@gmail.com>
HFEagle3TargetModel hooks the output of decoder layer aux_layer_id (HF hidden_states[aux_layer_id + 1]), but vLLM's capture id k records the residual stream entering layer k (output of layer k-1). Shift the vLLM aux capture ids by +1 so both backends capture the same hidden states; GPU validation shows per-layer cosine 0.9999-1.0 vs HF after the shift. Signed-off-by: khazic <khazzz1c@gmail.com>
…ig, tests Add target_model_backend='vllm' dispatch (_setup_vllm_target) to the EAGLE-3 recipe, an --engine vllm builder to serve_target for the remote path, an example llama_eagle3_vllm.yaml, a GPU smoke script (scripts/smoke_vllm_target.py), and CPU unit tests for the recipe-side wiring. Update the unknown-backend error message (and the sglang test asserting it) to include 'vllm'. Signed-off-by: khazic <khazzz1c@gmail.com>
Add test_eagle3_vllm.py covering the parts that do not need vLLM: the engine-agnostic backend subclassing, vllm_dtype_str mapping, the model shim, and serve_target's --engine vllm routing / _build_vllm_target delegation. Mark _load_target_head_weights (reads model safetensors on the GPU server) as no-cover. Signed-off-by: khazic <khazzz1c@gmail.com>
…lision Cache the SamplingParams, the hidden-states connector, and the fp32/transposed LM-head + fp32 norm once instead of rebuilding them every forward_eagle3 call; do a single device-to-host copy of the batch token ids; factor the lazy weight load into _ensure_weights_loaded. Raise at config time if an aux layer id + 1 collides with the final-hidden capture id. No behavior change on the validated path. Signed-off-by: khazic <khazzz1c@gmail.com>
CI runs the PR against its merge with main, which carries the context- parallelism gate (NVIDIA-NeMo#2465): _setup_online_target reads cfg.get("distributed.cp_size"). The vLLM and (stacked) SGLang test stubs build the recipe via __new__ without a cfg, so every test dispatching through _setup_online_target fails with AttributeError on the merged tree. Give both stubs an empty _RecipeCfg so the gate defaults to cp_size=1 (no CP). Signed-off-by: khazic <khazzz1c@gmail.com>
The vLLM target runner's __init__ (HF-config parsing + lazy-default wiring) had no CPU coverage, dropping the PR's patch coverage to 69.56% (target 80%). Add CPU unit tests that stub the CUDA/HF-config touch points and exercise the constructor's dim parsing, lazy-init defaults, TMPDIR-rooted storage path, and explicit-arg handling, bringing the new vLLM runner module to 100%. Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| @@ -0,0 +1,356 @@ | |||
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |||
Contributor
There was a problem hiding this comment.
LGTM — clean, well-tested PR. The vLLM backend follows the SGLang pattern exactly, the +1 layer-id shift logic is correct and guarded against the collision case, and the CPU test suite covers all the non-GPU surfaces. No bugs, no missing capabilities flags, no dtype hazards.
HuiyingLi
approved these changes
Jun 29, 2026
HuiyingLi
enabled auto-merge (squash)
June 29, 2026 02:47
svcnvidia-nemo-ci
removed
the
waiting-on-maintainers
24 tasks
akoumpa
added a commit
that referenced
this pull request
Aug 30, 2026
* build: bump FA3/FA4 pin to pick up CuTe compile-key fixes The pinned flash-attention ref 002cce0a1 (2026-07-03) predates both CuTe compile-key fixes, so every forward/backward with a tensor max_seqlen -- i.e. every varlen/packed-sequence step -- rebuilt the compile key and recompiled the kernel. Reported against a Qwen3.5-VL packed-sequence finetune on B300. Bump to ce088ab9ce0f (main), which adds over the old pin: - #2507 [CuTe, Bwd] fix backward compile key churn (max_seqlen tensor) - #2762 [CuTe, Fwd] stabilize tensor max_seqlen compile key - #2745 [CuTe] fix forward dynamic-shape correctness - #2819 [CuTe] speed up scalar SM100 mask compilation #2798 raised the cute requirement to nvidia-cutlass-dsl>=4.6.2, so pin the FA4 CUTLASS DSL install to 4.6.2 to match. Kept exact rather than floating so the image stays reproducible. Note this ref also feeds the FA3 (Hopper) wheel, which is built by default, so the bump affects the default x86 image and not just INSTALL_FA4=true builds. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * build: pin full flash-attention commit Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * build: bump cutlass-dsl to 4.6.2, ffpa-attn to 0.2.3, quack-kernels to 0.6.4 flash_attn.cute (FA4) at the pinned FLASH_ATTN_REF builds against nvidia-cutlass-dsl 4.6.2. quack-kernels and ffpa-attn both import the same `cutlass` package, so all three have to agree on one version. - quack-kernels 0.6.1 -> 0.6.4 (pins cutlass-dsl 4.6.2; 0.6.1 pinned 4.6.0) - ffpa-attn 0.2.2 -> 0.2.3 (first release pinning cutlass-dsl 4.6.2 and quack-kernels 0.6.4; 0.2.2 hard-pinned 4.6.0/0.6.1) - ffpa extra's explicit cutlass-dsl pin 4.6.0 -> 4.6.2 Because ffpa-attn 0.2.3 already declares the versions FA4 needs, the resolution is conflict-free -- no override-dependencies entries required. Corrects the apache-tvm-ffi note as well: FA4 declares apache-tvm-ffi>=0.1.12 but never imports tvm-ffi, so the <=0.1.11 cap that keeps tilelang working is not mutually exclusive with FA4. Both lock files regenerated with uv 0.8.22 (the version CI pins); the diff is confined to the six bumped packages. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * build: enable FA4 (INSTALL_FA4=true) by default Flips the default in both stages that declare the ARG: wheel_builder (gates building the flash_attn/cute wheel) and automodel_final (gates the nvidia-cutlass-dsl[cu13] install and the flash_attn/cute symlink). Docker does not inherit ARGs across stages, so flipping only the first would build the wheel while skipping the DSL install and the symlink -- a half-installed FA4. Both FA4 wheels are installed --no-deps, so flash-attn-4's declared apache-tvm-ffi>=0.1.12 bound is never resolved and the <=0.1.11 cap that keeps the tilelang kernels working still holds. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * build: correct the apache-tvm-ffi comment The note claimed flash-attn-4 never imports tvm-ffi. That is wrong: flash_attn/cute/cache_utils.py has a bare top-level `import tvm_ffi`, reached eagerly through flash_attn.cute.__init__ -> interface.py, so importing flash_attn.cute hard-requires it. The cap itself is still right, but for a different reason. FA4 is installed --no-deps, so its declared >=0.1.12 floor is never resolved, and the tvm_ffi that is present comes via quack-kernels (>=0.1.6,<0.2) pinned here to 0.1.11. FA4 therefore runs against an ffi older than it declares; the two symbols it references, tvm_ffi.Function and tvm_ffi.__version__, both exist in 0.1.11. Note this is unexercised: nothing selects attn_implementation="flash_attention_4" until the FA4 backend lands, so flash_attn.cute is never imported by any test here. Worth knowing given docker/Dockerfile sets FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1, which turns on the AOT cache path in that same cache_utils.py. Verified against flash-attention ce088ab9, the pinned FLASH_ATTN_REF. Comment only; no dependency or lock change. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> --------- Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #2425.
What
Adds a vLLM target backend for EAGLE-3 training, the exact analog of the SGLang
backend from #2449. It plugs into the same engine-agnostic
TargetRunnercontract (
target_runner.py), so the trainer, the remote server, and thesupervision semantics are unchanged: only a new engine is added.
New / changed:
eagle/vllm_runner.py:VLLMTargetRunner, the vLLM-internal forward (theanalog of
sglang_runner.py).eagle/vllm_target.py:VLLMEagle3TargetModel, the thinRunnerEagle3TargetModeladapter (analog ofsglang_target.py).recipes/llm/train_eagle3.py:target_model_backend: vllmdispatch(
_setup_vllm_target).serve_target.py:--engine vllmfor the remote path.examples/speculative/eagle3/llama_eagle3_vllm.yaml: example config.scripts/smoke_vllm_target.py: GPU smoke (not in CI).test_eagle3_vllm_backend.py,test_eagle3_vllm.py).How
vLLM does not expose SGLang's "one
ModelRunner.forwardreturns(logits, aux)" surface, and its v1 engine runs the model in a separateprocess. So the runner uses vLLM's native
extract_hidden_statesspeculativemethod (the supported, version-stable path):
LLMwithspeculative_config={"method": "extract_hidden_states", ...}requesting thethree EAGLE-3 capture layers plus
num_hidden_layers(the final pre-normhidden), chunked prefill disabled.
generate(max_tokens=1)over the batch; vLLM writes the captured hiddenstates to disk via
ExampleHiddenStatesConnectorand returns the path.[B, S, 3*hidden], andrebuild full-vocab logits in-process from the final hidden via the target's
RMSNorm + LM head.
One subtlety worth calling out: vLLM's capture id
krecords the residual streamentering layer
k(HFhidden_states[k]), whereasHFEagle3TargetModelhooksthe output of layer
k(HFhidden_states[k+1]). The runner shifts the auxids by
+1so both backends capture the identical hidden states.vLLM is not a NeMo-AutoModel dependency; install it out of band in a
dedicated environment (
uv pip install vllm==0.23.0), same as sglang.Validation (single A800-80GB, Qwen3-4B, vLLM 0.23.0)
Through the real
VLLMEagle3TargetModel.generate_batchvsHFEagle3TargetModel.generate_batch:top-5 1.0000, aux hidden-state mean cosine 1.0000.
scripts/smoke_vllm_target.py --compare-hf(random tokens, 4x48):STAGE 1 OK, argmax match 0.9681, aux cosine 1.0000,STAGE 2 OK.Per-layer cosine of the captured aux vs HF is 0.9999-1.0 after the
+1shift.CPU unit suite: the eagle3 tests pass locally (the vLLM forward itself is
GPU-only and exercised by the smoke, mirroring how the sglang forward is tested).