chore: CI proxy for #2449 - #2822
Closed
HuiyingLi wants to merge 15 commits into
Closed
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>
_setup_online_target now reads cfg.get("distributed.cp_size") for the
context-parallelism gate (merged in from #2465). The test helper built the
recipe via __new__ without a cfg, so the four tests that dispatch through
_setup_online_target failed with AttributeError. Give the stub an empty
_RecipeCfg so the gate defaults to cp_size=1 (no CP).
Signed-off-by: khazic <khazzz1c@gmail.com>
SGLang 0.5.9 pins transformers==4.57.1 while NeMoAutoModelForCausalLM needs transformers 5.x (AutoModelForMultimodalLM), so the smoke --compare-hf check cannot run as written in one environment. Document this in the EAGLE guide's environment-setup step, per maintainer request on #2449. Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
Author
|
/ok to test b2b8e97 |
Contributor
|
🌿 Preview your docs: https://nvidia-preview-huiyingl-ci-run-pr-2449.docs.buildwithfern.com/nemo/automodel |
Contributor
Author
|
Closing CI-only proxy because source PR #2449 is merged. |
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
2 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.
CI-only proxy PR. Do not merge / do not review. Points at the exact head commit of #2449 (b2b8e97) so internal CI runs under the internal-contributor queue; results post back to #2449 via the shared SHA. Source of truth: #2449. Close once CI completes.