feat(sglang): add end-to-end training and refit integration by Kh4L · Pull Request #3330 · NVIDIA-NeMo/RL

Kh4L · Jul 23, 2026

Summary

This integration branch makes SGLang a reproducible generation backend for
NeMo RL training and in-memory weight refit:

synchronous and asynchronous GRPO rollouts, including NeMo-Gym;
Megatron-to-SGLang non-colocated NCCL broadcast refit;
colocated CUDA-IPC refit for the supported policy backends;
fail-closed engine startup, refit deadlines, cleanup, and quarantine;
immutable-source, immutable-image, public two-node smoke and 16-node async
integration reproducers with machine-checked evidence.

There is no disk-refit fallback.

Dependency stack

This is an integration/reproducer PR. The component PRs remain authoritative
for their respective lanes:

feat(sglang): engine fault tolerance for rollouts #3187: SGLang engine fault tolerance;
feat(sglang): megatron backend weight refit for sglang rollouts #3190: Megatron-to-SGLang weight refit;
feat(sglang): mxfp8 quantized rollout and refit #3188: MXFP8 rollout/refit support;
fix(sglang): avoid ncclCommSplit hang in the refit weight-update group #3189: side-by-side communicator split fix;
feat(sglang): enable NeMo-Gym async-GRPO rollouts on the SGLang backend #3127: async-GRPO and NeMo-Gym enablement, reduced here to the residual work
not supplied by the refit stack;
feat(nemo-gym): opt-in truncation of non-contiguous rollout episodes #3063: configurable handling of non-contiguous Gym episodes;
feat(sglang): add token-exact SGLang model server for RL training Gym#1787: the native SGLang Responses adapter pinned by the Gym
submodule.

Refit dispatch fix

Running the two-node reproducer surfaced a defect that made every SGLang refit
fail on its first attempt:

NRL_SGLANG_REFIT_FAILURE transfer=broadcast refit=1 \
    phase=prepare_or_transfer error_type=ModuleNotFoundError

_refit_sglang_dispatch runs in the driver but imported a policy worker module
to reach refit_sglang_{colocated,distributed}. Those modules import
megatron.bridge and nemo_automodel at module scope, and the driver
environment can hold neither: sglang is declared mutually exclusive with
mcore, fsdp, automodel and vllm in [tool.uv] conflicts. Both transfer
modes were affected, and SGLangColocatedWeightSynchronizer._refit_colocated
carried an independent copy of the same import. The DTensor path had the
identical latent failure on nemo_automodel, which is why both backends are
touched.

Those four helpers are driver-side orchestration: they drive both actor groups
through the pause/transfer/quarantine protocol and call only policy and
policy_generation facades, so they were simply filed in the wrong module.
They move verbatim to nemo_rl/weight_sync/, beside the synchronizer that
calls them, and the worker modules re-export the names. A new test parses the
source rather than importing it, so it fails on the offending commit even in an
environment where every backend happens to be installed.

The fix is isolated in its own commit so it can be cherry-picked onto #3190.

Safety semantics

SGLang refit is a synchronization barrier; replay collection may be
pipelined, but generation does not overlap a SGLang weight update.
Initial engine launch, health checks, cache flush, router registration, and
driver readiness share one startup deadline.
Non-colocated broadcast refits have one end-to-end deadline. A failed or
uncertain partial update aborts the communicator and quarantines the engines
instead of resuming generation.
Colocated partial-update failures also quarantine engines. Trainer-side
gather and per-bucket IPC waits on that path are not yet deadline-bounded and
are documented as requiring job termination if they stall.

Validation

CPU-only:

Ruff check and format check across all changed Python files;
Python byte-compilation across all changed Python files;
git diff --check;
11/11 evidence-validator tests;
3/3 driver-import guard tests;
recipe inheritance and test-suite invariant checks;
shell syntax and dry-run checks for both reproducer drivers;
targeted startup/config/router/refit cleanup tests;
36/36 focused Gym SGLang-adapter tests.

GPU, two-node smoke (2 nodes, 4 GPUs per node, public Qwen2.5-Math-1.5B-Instruct
at the pinned revision), run against a driver environment built without the
mcore extra so the dispatch fix is exercised rather than worked around:

{
  "markers": {"world_size": 5, "engines": 4,
              "group_ready_count": 1, "refit_success_count": 3},
  "metrics": {"expected_max_step": 3, "max_recorded_step": 3},
  "status": "passed"
}

The five-rank communicator (trainer rank 0 plus four engine ranks) formed,
streamed weights, and generation resumed on the refitted policy three times
across three GRPO steps with no failure marker. Refit wall time was 5.9 s for
the first transfer, which builds the communicator, then 0.24 s.

This smoke was re-run after rebasing the branch onto current main, so the
evidence above describes the branch as it now stands. Both the pre-rebase and
post-rebase runs produced identical markers.

Caveats on that evidence, stated plainly:

it was produced from a source checkout mounted into an existing prebuilt
container, not from a digest-qualified image built from this commit, so
image_ref in the bundle is a placeholder and the dependency set is this
branch's lockfile over that container's base;
the 16-node asynchronous NeMo-Gym integration run has not been performed.

Reproduce

See docs/guides/sglang-refit.md.

The manual public evidence runs are:

A two-node, four-GPU-per-node synchronous smoke test. It requires the exact
world_size=5 engines=4 refit topology, at least two successful refits, no
fatal refit marker, and train/loss through step 3.
A 16-node, eight-GPU-per-node asynchronous NeMo-Gym integration test. It
requires the exact world_size=65 engines=32 topology, at least two
successful refits, no fatal refit marker, train/loss through step 3, and a
non-empty Gym rollout artifact.

Both drivers require an immutable model revision, a digest-qualified image, a
publicly fetchable pinned Gym commit, and a fresh shared initial
Hugging-Face-to-Megatron conversion cache.

Remaining publication gates

Publish the cryptographically signed Gym commit stack and pin its final
public SHA (feat(sglang): add token-exact SGLang model server for RL training Gym#1787; all commits signature-verified, pinned
by the Gym submodule).
Demonstrate the two-node refit path end to end on GPUs.
Rebase onto current main and re-run the smoke.
Prove a clean recursive clone from public HTTPS remotes end to end. A
git clone --recursive of the branch, run with no credentials and with
global and system git config disabled so no url.insteadOf rewrite can
apply, resolves all four submodules from their public remotes, including
the pinned Gym commit fetched from NVIDIA-NeMo/Gym. uv lock --check
then passes against the committed lockfile in that fresh clone.
Build and publish the exact digest-qualified image and re-run the smoke
against it.
Run and attach a sanitized 16-node async integration evidence summary.

copy-pr-bot · Jul 23, 2026

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.

github-actions · Jul 23, 2026

❌ Submodule Fast-Forward Check Failed

Check based on commit: 3daef99 (PR #3330 from integration/sglang-training-shareable)

❌ Submodules that need attention:

Gym: ❌ Commits have DIVERGED from a common ancestor
TARGET (main branch): https://github.com/NVIDIA-NeMo/Gym/commits/d67ad6611cfe21dbaeb301c59e59df32ce22ec50/
CURRENT (PR #3330 from integration/sglang-training-shareable): https://github.com/NVIDIA-NeMo/Gym/commits/65ddaa2ec46a823d9d6d9278a3ca53232ad87649/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

Add RolloutHealthMonitor: a daemon thread that health-checks each engine and restarts hung or dead actors during rollout, gated by policy.generation.sglang_cfg.use_fault_tolerance (off by default). - fault_tolerance.py: monitor lifecycle (start/stop/pause/resume), health_generate probing, engine kill/restart bookkeeping - SGLangGeneration: engine recovery (_recover, recover_updatable_engines, get_updatable_engines_and_lock), num_new_engines tracking, monitor wiring - ray_utils.Lock: cooperative Ray lock serializing weight refits against engine recovery - SGLangGenerationWorker._simulate_crash: test-only crash injection - config + exemplar YAML keys for the health checker Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>

Replace the HTTP weight-streaming refit with two first-class SGLang refit paths driven from grpo via _refit_sglang_dispatch: - colocated (weight_transfer_mode: ipc): Ray CUDA-IPC buckets via send_hf_buckets_via_ipc_actor_impl + SGLangColocatedWeightSynchronizer, generalized Gloo gather topology (connect_colocate_topology) supporting FSDP and Megatron layouts - disaggregate (weight_transfer_mode: broadcast): trainer-rank-0 NCCL weight-update group (side-by-side init_process_group, connect/ disconnect_rollout_engines_from_distributed) broadcasting AutoBridge- restored HF buckets; engines join via init_weights_update_group - MegatronSGLangHfWeightIterator: AutoBridge export walk bucketed by post-transformation size - engine-side worker endpoints (update_weights_from_distributed, pause/continue_generation, post_process_weights, weight versioning) - NCCL_CUMEM_ENABLE=0 alignment between trainer and sglang scheduler - drop stream_weights_via_http / set_rollout_num_gpus_per_engine Quantization-related parameters (target_precision, sglang_quantization_cfg) are inert bf16 plumbing here; the mxfp8 implementation lands in a follow-up PR. Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>

Enable scheme=mxfp8 under policy.generation.sglang_cfg.quantization: - mxfp8_setup.py: offline HF->MXFP8 checkpoint conversion with cache fingerprinting (ensure_mxfp8_checkpoint), run before SGLang boot - mxfp8_quantization_core.py: quantize_mxfp8 (flashinfer), dynamic skip-substring policy, scale-key helpers shared by offline conversion and online refit - MegatronSGLangHfWeightIterator: quantize eligible tensors during the AutoBridge walk, emitting (weight, weight_scale_inv) pairs sized into post-transformation buckets - wire target_precision/sglang_quantization_cfg through the megatron refit paths; SglangQuantizationConfig type + exemplar YAML block Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>

When the trainer's default process group is eager-initialized with a bound device id (device_id= passed to torch.distributed.init_process_group), torch's _new_process_group_helper silently sets Options.split_from to the default group's NCCL communicator, and the first collective on the side weight-update group then issues ncclCommSplit. The split is collective over the parent (trainer) communicator, but only trainer rank 0 is a member of the refit group, so rank 0 blocks forever in the split bootstrap all-gather (commGetSplitInfo) while the remaining trainer ranks hit the watchdog. Clear default_pg.bound_device_id for the duration of the _new_process_group_helper call (restored in finally) so the helper's split predicate sees a lazily-initialized default group and never selects the split path. Current main initializes the Megatron default PG without device_id, so this is latent hardening: any trainer base that binds the default group would otherwise trip it silently. Verified: torch 2.10/2.11 have the identical split predicate and a pure field-assignment setter; gloo interop + split-predicate regression tests pass locally. Co-authored-by: Serge Panev <3193578+Kh4L@users.noreply.github.com> Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>

Rebased onto the SGLang refit stack (NVIDIA-NeMo#3187 -> NVIDIA-NeMo#3190 -> NVIDIA-NeMo#3188 -> NVIDIA-NeMo#3189). The stack now owns the collective-skip and the real Megatron->SGLang refit, so the earlier NCCL-skip / NRL_SGLANG_SKIP_REFIT / set_rollout_num_gpus_per_engine scaffolding is dropped. This keeps only the async-GRPO replay-path enablement: - lift the async_grpo_train backend gate to include sglang - expose_http_server + _spinup_nemo_gym wiring for the sglang branch - SGLangGeneration.cfg parity property + rollouts.py sglang context_length - picklability (__getstate__/__setstate__) dropping the aiohttp client, async loop, and (new on the stack) health monitor for the collector actor Signed-off-by: Serge Panev <spanev@nvidia.com>

… presence An SGLang recipe inherits vllm_cfg from a vLLM base via deep-merge, so the prior '"vllm_cfg" in cfg' check shadowed the SGLang branch and read the vLLM max_model_len instead of the SGLang context_length. Dispatch on cfg['backend'], and fall back to max_total_sequence_length when the nullable context_length is unset. Found by an independent audit. Signed-off-by: Serge Panev <spanev@nvidia.com>

Signed-off-by: Serge Panev <spanev@nvidia.com>

_postprocess_nemo_gym_to_nemo_rl_result hard-asserts that each assistant turn's accumulated token IDs form a prefix of the next turn's prompt. In long multi-turn agentic rollouts, rare tokenization/re-render edge cases can break this contiguity; the assert then kills the rollout task, and under async GRPO the step stalls indefinitely waiting for a trajectory that never arrives. Add an opt-in NemoGymConfig knob, truncate_noncontiguous_episodes (default false = current assert, unchanged). When enabled, a non-contiguous turn truncates the episode at the last contiguous turn with a warning: the corrupted tail is dropped and the valid prefix stays trainable, analogous to overlong filtering. Wired through GRPO and distillation setup identically to invalid_tool_call_patterns. Signed-off-by: Serge Panev <spanev@nvidia.com> (cherry picked from commit c0075c4)

…odes Per config conventions, a NotRequired field must not carry a non-None default at the call site. Use .pop(key, None) (matching the sibling thinking_tags line); the NemoGym consumer already gates on a truthiness .get(), so absent -> off is preserved. Found by an independent audit. Signed-off-by: Serge Panev <spanev@nvidia.com>

Signed-off-by: Serge Panev <spanev@nvidia.com>

`_refit_sglang_dispatch` ran in the driver but imported `megatron_policy_worker` (or `dtensor_policy_worker_v2`) to reach `refit_sglang_{colocated,distributed}`. Those modules import `megatron.bridge` and `nemo_automodel` at module scope, and the driver environment has neither: `sglang` is declared mutually exclusive with `mcore`, `fsdp`, `automodel` and `vllm` in `[tool.uv] conflicts`, so the driver can never hold both the SGLang stack and a trainer backend. Every SGLang refit therefore died on its first attempt with NRL_SGLANG_REFIT_FAILURE transfer=broadcast refit=1 \ phase=prepare_or_transfer error_type=ModuleNotFoundError on a 2-node GB200 run. Both transfer modes were affected, and `SGLangColocatedWeightSynchronizer._refit_colocated` carried an independent copy of the same import. The four helpers are driver-side orchestration: they drive both actor groups through the pause/transfer/quarantine protocol and call only `policy` and `policy_generation` facades, so they were merely filed in the wrong module. Move them verbatim to `nemo_rl/weight_sync/`, beside the synchronizer that calls them, and repoint both dispatch sites. The worker modules re-export the names so existing callers keep working. Verified on 2 nodes x 4 GPUs with a driver environment built without `mcore`: `NRL_SGLANG_REFIT_GROUP_READY world_size=5 engines=4`, three successful refits, no failures, `train/loss` through step 3. The new test parses the source rather than importing it, so it fails on the offending commit even where every backend happens to be installed. Signed-off-by: Serge Panev <spanev@nvidia.com>

…lectable Four defects found while running the two-node reproducer end to end: - Both driver scripts were committed non-executable, so `uv run <driver>` failed with `Permission denied` and the `TEST_DRYRUN=1` invocation in `test_all_tests_can_find_config_if_dryrun` exited 126. - `tests/test_sglang_refit_checks.py` sat at the `tests/` root. Every CI shard passes an explicit `unit/...` path to `tests/run_unit.sh`, so `testpaths` never applied and the evidence validator's only coverage never ran. Move it under `tests/unit/`. - The provenance guard compared `NEMO_RL_COMMIT` against a `ROOT_COMMIT` that defaults to the same value, so an unset variable passed the check and the run recorded the container's commit instead of the source's. Fail loudly instead. - The preflight reported "Could not attach to an existing Ray cluster" for any exception raised by `ray.init`, including a version mismatch after a successful connection. Print the chained cause. Signed-off-by: Serge Panev <spanev@nvidia.com>

github-actions · Jul 25, 2026

✅ Submodule Fast-Forward Check Results

Check based on commit: da23bed (PR #3330 from integration/sglang-training-shareable)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

The async NeMo-Gym integration recipe hard-codes eight GPUs per node, and its preflight rejects a cluster that cannot supply sixteen nodes with at least that many. Hosts that expose four GPUs per node therefore cannot run it at all, regardless of how many nodes are free. Add a 32n4g variant describing the same run on that host shape: sixteen training nodes and sixteen generation nodes, for the same sixty-four training GPUs, sixty-four generation GPUs and thirty-two TP2 engines. The inherited parallelism is untouched, so only rank placement differs, and the evidence assertion stays exactly world_size=65 engines=32. Register it alongside the existing variant in disabled.txt, since it is equally unsuited to nightly CI, and document both shapes in the guide so a reader picks the one matching their cluster. Signed-off-by: Serge Panev <spanev@nvidia.com>

github-actions · Jul 25, 2026

✅ Submodule Fast-Forward Check Results

Check based on commit: d322b52 (PR #3330 from integration/sglang-training-shareable)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

Both async NeMo-Gym drivers pass env.nemo_gym.uv_venv_dir=/opt/gym_venvs as a plain Hydra override, but nothing defines that key: it is absent from the recipe, from the exemplar it inherits, and from that exemplar's parent, and the recipe's env block sets _override_, which replaces rather than merges. Hydra rejects an override of a key that does not exist, so the run aborts during configuration parsing: ConfigCompositionException: Could not override 'env.nemo_gym.uv_venv_dir'. Declare the key in both recipes, defaulting to a path under HF_HOME, so the drivers can redirect it to the image's prefetched virtual environments. This matches grpo-qwen3-1.7b-1n8g-megatron-super-swe1, which declares the key in its recipe and overrides it from its driver. Found by running the 32n4g recipe on a thirty-two node allocation, where it failed after three and a half minutes without reaching the training loop. Signed-off-by: Serge Panev <spanev@nvidia.com>

github-actions Bot added the Documentation Improvements or additions to documentation label Jul 23, 2026

xiuhu17 and others added 24 commits July 24, 2026 16:45

fix(sglang): harden async Gym config handling

58a6c41

Signed-off-by: Serge Panev <spanev@nvidia.com>

fix(nemo-gym): centralize rollout contiguity defaults

e2e84d4

Signed-off-by: Serge Panev <spanev@nvidia.com>

fix(sglang): preserve runtime ownership across serialization

186152b

Signed-off-by: Serge Panev <spanev@nvidia.com>

fix(sglang): require cache invalidation before refit

e3396ae

Signed-off-by: Serge Panev <spanev@nvidia.com>

feat(sglang): emit refit lifecycle markers

26342af

Signed-off-by: Serge Panev <spanev@nvidia.com>

test(sglang): cover async backend selection and refit loop

5850729

Signed-off-by: Serge Panev <spanev@nvidia.com>

fix(sglang): validate runtime configuration

8f9ff6f

Signed-off-by: Serge Panev <spanev@nvidia.com>

docs(sglang): add reproducible training and refit recipes

25e32bf

Signed-off-by: Serge Panev <spanev@nvidia.com>

fix(sglang): bound and recover distributed refits

c7c1611

Signed-off-by: Serge Panev <spanev@nvidia.com>

fix(sglang): harden async backend lifecycle

6c07857

Signed-off-by: Serge Panev <spanev@nvidia.com>

fix(sglang): fail closed across startup and refit

bd2461f

Signed-off-by: Serge Panev <spanev@nvidia.com>

docs(sglang): add strict public refit reproducers

dbec358

Signed-off-by: Serge Panev <spanev@nvidia.com>

fix(sglang): make public evidence revision-exact

dd5c30c

Signed-off-by: Serge Panev <spanev@nvidia.com>

chore(gym): pin NeMo-Gym submodule to the token-exact SGLang adapter

7833c77

Signed-off-by: Serge Panev <spanev@nvidia.com>

Kh4L force-pushed the integration/sglang-training-shareable branch from ea6383b to da23bed Compare July 25, 2026 00:02

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat(sglang): add end-to-end training and refit integration#3330

feat(sglang): add end-to-end training and refit integration#3330
Kh4L wants to merge 26 commits into
NVIDIA-NeMo:mainNVIDIA-NeMo/RL:mainfrom
Kh4L:integration/sglang-training-shareableKh4L/NemoRL:integration/sglang-training-shareableCopy head branch name to clipboard

Kh4L commented Jul 23, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jul 23, 2026

Uh oh!

github-actions Bot commented Jul 23, 2026

Uh oh!

github-actions Bot commented Jul 25, 2026

Uh oh!

github-actions Bot commented Jul 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Search code, repositories, users, issues, pull requests...

Uh oh!

Conversation

Kh4L commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Dependency stack

Refit dispatch fix

Safety semantics

Validation

Reproduce

Remaining publication gates

Uh oh!

copy-pr-bot Bot commented Jul 23, 2026

Uh oh!

github-actions Bot commented Jul 23, 2026

❌ Submodule Fast-Forward Check Failed

❌ Submodules that need attention:

Uh oh!

github-actions Bot commented Jul 25, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Jul 25, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Kh4L commented Jul 23, 2026 •

edited

Loading