feat(sglang): megatron backend weight refit for sglang rollouts by xiuhu17 · Pull Request #3190 · NVIDIA-NeMo/RL

xiuhu17 · 2026-07-14T02:30:43Z

What

Replaces 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
drops stream_weights_via_http / set_rollout_num_gpus_per_engine
fix: 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 (is_initialized() and _get_default_group().bound_device_id) sees a lazily-initialized default group and never selects the split path. Root cause analysis by @Kh4L (co-authored).

Quantization-related parameters (target_precision, sglang_quantization_cfg) are inert bf16 plumbing here; the mxfp8 implementation lands in #3188.

Megatron broadcast design is adapted from miles

Result

Stack

Stacked on #3187 (fault tolerance) on #2997 (ports) — review only the last commit (feat(sglang): megatron backend weight refit for sglang rollouts). Follow-ups: #3188 (mxfp8)

Supersedes the refit part of #2446; split out of #3091.

copy-pr-bot · 2026-07-14T02:30:47Z

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.

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>

yuki-97

thanks for supporting this @xiuhu17 , I took a first round review and will review again after your updates.

Two scope notes before the per-line comments:

Skipping b09309918 per the stack note — that commit duplicates #3187 (fault tolerance); everything below reviews 26e1d60b..HEAD.
fetch_updatable_engines_with_recover should move to #3187. nemo_rl/models/policy/utils.py:959 is fault-tolerance glue — it reads sglang_cfg.use_fault_tolerance and gates the call to recover_updatable_engines. Its natural home is #3187; the three refit-driver callers (policy/utils.py, dtensor_sglang_refit.py, megatron_sglang_refit.py) can call get_updatable_engines_and_lock directly here, and the recover gate lives with the FT feature it belongs to.

Other refit-side uses of FT-provided primitives (pause_generation, num_new_engines, rollout_engine_lock acquire/release, clear_updatable_num_new_engines) are legitimate consumers of the FT interface and stay in this PR.

yuki-97 · 2026-07-31T10:05:57Z

+    return policy_generation.get_updatable_engines_and_lock()
+
+
+def broadcast_hf_buckets_via_distributed_impl(


out of curious how's this function different from the one vllm/trtllm using and is it possible to reuse?

…atch The SGLang refit needs an engine-side finalize after the last weight bucket lands: quantized weights carry derived state (swizzled blockscales aliased over their own storage, `g1/g2_alphas` products) that a plain `load_weights` leaves stale, and re-running the finalize per bucket is not idempotent. No tagged sglang release exposes a hook for it. sgl-project/sglang's `sglang-miles` branch does, as `begin_weight_update` / `end_weight_update` routed through `weight_updater` + `model_runner`. Pin that branch instead of PyPI 0.5.12.post1, and drop the one on-disk patch it makes redundant: - `_patch_sglang_custom_all_reduce_v2_tms_cudagraph` backported the "skip custom all-reduce v2 graph capture under torch_memory_saver" fix. The branch has it natively via `_graph_mode_allowed`. Worse than redundant: `set_cuda_graph_capture` no longer exists there, so the patch's anchor lookup raised and aborted the whole patch chain. - `_patch_sglang_file_replacements`, its only caller now gone. Kept: `_patch_sglang_safe_unpickler` (the branch's allowlist still has no `nemo_rl.` prefix, and the policy venv cannot import sglang to reuse its `tensor_bucket`), the imbalance-check env override, and the two megatron `hook_mode` patches. The branch is based on v0.5.16, which costs a port: - `sglang-kernel` 0.4.2.post2 -> 0.4.5, per the branch's own pyproject. - `disable_piecewise_cuda_graph` was replaced by a per-phase enum. All three shipped sglang configs set it and it is passed straight into `ServerArgs(**kwargs)`, so on v0.5.16 every run would die with `TypeError: unexpected keyword argument` at sglang_worker.py:91, before `launch_server` is spawned. Translated to `cuda_graph_backend_prefill: breakable` -- the non-piecewise graph mode -- to preserve why the flag existed (piecewise prefill hit an illegal memory access on torch 2.10). - `cuda_graph_bs` / `cuda_graph_max_bs` split per phase. No config set them, so they were dormant rather than broken; renamed to the `_decode`/`_prefill` pairs. All 39 kwargs `sglang_worker` passes resolve against the pinned rev, and the vendored `FlattenedTensorBucket` / `FlattenedTensorMetadata` are functionally identical to the branch's, so the IPC payload is unchanged. `uv.lock` is deliberately not hand-edited: a git-sourced sglang changes the package entry, its transitive deps and its hashes, and a plausible-but-wrong lock is worse than an obviously stale one. Run `uv lock` in the project container. Whether `breakable` is the right translation (v0.5.16 may have fixed the crash, in which case the key can go) needs one smoke test. Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>

yuki-97

thanks for the updates! some comments on the new added tests, and you'll need a rebase since there's some conflict with main branch.

yuki-97 · 2026-08-04T13:39:57Z

+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl-sglang-1n4g


project: nemo-rl-sglang-1n4g (and name below) are shadowed at runtime by the .sh CLI overrides logger.wandb.project=nemo-rl / logger.wandb.name=$EXP_NAME — Hydra CLI overrides beat YAML.

Sibling 1n8g-fsdp2tp1-sglang.yaml uses project: nemo-rl to match its .sh.

Suggested change

project: nemo-rl-sglang-1n4g

project: nemo-rl

yuki-97 · 2026-08-04T15:50:09Z

@@ -1,36 +1,20 @@
 #!/bin/bash


thanks for adding the two 1n4g tests, will also need to add them to tests/test_suites/nightly_gb200.txt.

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>

`_refit_sglang_dispatch` and `SGLangColocatedWeightSynchronizer._refit_colocated` run in the driver process but imported `megatron_policy_worker` / `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 is synced without the `mcore` / `automodel` extras, which is why worker imports otherwise go through `PY_EXECUTABLES`. Every SGLang refit therefore failed on its first attempt with `ModuleNotFoundError: No module named 'megatron'`, reproduced on a two node run. Move the four drivers to `nemo_rl/weight_sync/`, next to the synchronizer that calls them, and repoint both dispatch sites. The function bodies and names are unchanged, and the dispatch keeps its lazy `as _backend` import, so each call site changes only the module it imports from. The drivers call `policy` and `policy_generation` facade methods plus `ray`, and already imported `nemo_rl.models.policy.utils` lazily, so they need no extra. The test parses the sources instead of importing them, so it also fails in an environment that provides a backend, such as the Megatron unit test shards. Signed-off-by: Serge Panev <spanev@nvidia.com>

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> (cherry picked from commit e27d3f1)

…nizer The SGLang refit had two entry points: `grpo._refit_sglang_dispatch`, which the live GRPO/PPO/distillation loops used, and `SGLangColocatedWeightSynchronizer`, which nothing constructed. The offload and prepare/restore phases were duplicated across both, plus a third copy in `refit_policy_generation` itself. Collapse them onto the synchronizer: - `sglang_weight_synchronizer` keeps only the GPU phase transitions and delegates the refit itself (recover, connect, pause, KV invalidation, bucket transfer, post-process, continue) to the backend driver modules. `SGLangColocatedWeightSynchronizer` offloads the policy around the transfer; the new `SGLangDisaggregatedWeightSynchronizer` does not — the trainer keeps its own GPUs, and `prepare_for_training` onloads unconditionally on both the Megatron and DTensor workers, so there is no offload state to restore. `prepare_for_generation` still runs on both paths: `SGLangGeneration` gates it on `needs_offload` internally. - `create_weight_synchronizer` now returns the disaggregated synchronizer for non-colocated SGLang instead of raising. SGLang establishes its own weight-update group on the first refit, so it needs neither cluster handle. - `refit_policy_generation` hands SGLang to the factory and returns; the remaining branches are vLLM IPC/NCCL and the Megatron reshard. This covers Megatron colocated + disaggregated and DTensor colocated; DTensor disaggregated still raises from `dtensor_sglang_refit`. Also drop `http_weight_synchronizer`. It was dead (never constructed) and broken: `policy.stream_weights_via_http` was removed in 7b0f0bd along with `stream_weights_via_http_impl`. `weight_update_actor` now streams over the production Ray CUDA-IPC path (`connect_colocate_topology` + `send_hf_buckets_via_ipc_actor_impl`) rather than the deleted HTTP impl, so the real weight-update tests exercise the transport they claim to. The two `excluded_unit_tests.sh` deselects for the `test_refit_policy_generation_sglang_*` tests were stale — those tests went away in 7b0f0bd. Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>

The sglang extra tracks the sglang-miles branch, which requires transformers==5.12.1, kernels>=0.14.1 and flashinfer-python 0.6.14. Raising those forced three follow-on changes: - flashinfer-cubin/flashinfer-jit-cache lose their pins: cubin has no 0.6.14 release and jit-cache is only published on the flashinfer index, so both are left to resolve transitively. - llguidance drops its upper bound and xgrammar becomes a floor. vllm needs llguidance>=1.3 and sglang-miles needs >=1.7.6; the two are mutually exclusive extras, so no single pin serves both. xgrammar keeps the 0.1.33 floor for GHSA-7rgv-gqhr-fxg3 while allowing sglang's 0.2.1. - the outlines override is dropped, since it only existed to pull in Python 3.13 support that the current resolution already has. The deepseek_v3 tokenizer-blocklist patch asserted transformers < 5.12 on the assumption that the upstream fix would land there. It has not: MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS and TOKENIZER_MAPPING_NAMES both still carry deepseek_v3 in 5.12.1, so the patch stays load-bearing and the assert moves to < 5.13. Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>

The existing sglang recipes start at 1n8g, which needs a full node. These two run the same Qwen2.5-Math-1.5B GRPO workload on four GPUs, keeping every sglang setting identical to the 2n8g reference (tp_size 1, mem_fraction_static 0.6, ipc weight transfer, breakable prefill graph) so the only variable between them is the training backend. Measured on 4xB200: megatron averages 20.6s/step over steps 2-30, inside the 25s ceiling the suite already asserts for the larger sglang recipes. Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>

yuki-97 · 2026-08-10T10:22:24Z

+    # trainer has to agree with it, so this one has to land before the process
+    # group below picks a transport.
+    if rollout_backend == "sglang":
+        os.environ["NCCL_CUMEM_ENABLE"] = "0"


for other rollout backends which need to set NCCL_CUMEM_ENABLE=1 for non-colocatred refit, you have the comment "Deliberately after init_process_group: it targets the refit communicator, not the training process group."
why sglang need to put it before init_process_group? isn't sglang also only target the refit communicator?

if sglang should also do it after, I think let's still put it at validate_and_set_config since we set config there.

yuki-97 · 2026-08-10T10:47:14Z

+            policy=policy,
+            generation=policy_generation,
+            generation_backend=backend,
+            colocated=colocated_inference,


I see. currently seems other rollout backends only go here in disaggregated case so just curious with this.

I'm ok with this as is, cc @terrykong @saumishr could you help take a look on whether the implementation match your new weight_synchronizer design?

yuki-97

Follow-up on post-review commits (8660879 build + eac58ed 1n4g recipes): 3 nits — copyright year on new .py, wandb project shadow on the mcore sibling of an already-flagged yaml, and mcore nightly manifest registration for the mcore 1n4g sh.

yuki-97 · 2026-08-10T11:10:23Z

@@ -0,0 +1,96 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.


nit: copyright header year lags. This is a new file added by this PR, and per .claude/skills/copyright/SKILL.md (always use the current year) it should say 2026 — the sibling new file nemo_rl/weight_sync/sglang_weight_synchronizer.py in this same PR already uses 2026.

Suggested change

# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.

yuki-97 · 2026-08-10T11:10:23Z

+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl-sglang-1n4g


Same shadow issue as the fsdp2 sibling: the launcher .sh overrides logger.wandb.project=nemo-rl on the CLI, so this YAML value never takes effect. The other sglang recipes (1n8g/2n8g fsdp2, 2n8g megatron-noncolocated) all use nemo-rl.

Suggested change

project: nemo-rl-sglang-1n4g

project: nemo-rl

yuki-97 · 2026-08-10T11:10:23Z

@@ -0,0 +1,42 @@
+#!/bin/bash


For the megatron variant, also register in tests/test_suites/nightly_mcore_gb200.txt — mirrors how grpo-qwen2.5-math-1.5b-instruct-2n8g-megatrontp1-sglang-noncolocated.sh lands in both nightly.txt and nightly_mcore.txt in this PR.

github-actions Bot added the community-request label Jul 14, 2026

This was referenced Jul 14, 2026

feat(sglang): engine fault tolerance for rollouts #3187

Open

feat(sglang): mxfp8 quantized rollout and refit #3188

Open

fix(sglang): avoid ncclCommSplit hang in the refit weight-update group #3189

Closed

xiuhu17 marked this pull request as ready for review July 14, 2026 05:05

xiuhu17 requested review from a team as code owners July 14, 2026 05:05

svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Jul 16, 2026

This was referenced Jul 21, 2026

feat(sglang): add token-exact SGLang model server for RL training NVIDIA-NeMo/Gym#1787

Draft

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

Draft

xiuhu17 force-pushed the zhw/megatron_sglang_refit branch 2 times, most recently from 345428b to 3e010c0 Compare July 30, 2026 18:52

yuki-97 reviewed Jul 31, 2026

View reviewed changes

svcnvidia-nemo-ci added waiting-on-customer Waiting on the original author to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels Jul 31, 2026

xiuhu17 force-pushed the zhw/megatron_sglang_refit branch 2 times, most recently from c3c4211 to f266043 Compare August 1, 2026 22:32

svcnvidia-nemo-ci removed the waiting-on-customer Waiting on the original author to respond label Aug 4, 2026

xiuhu17 force-pushed the zhw/megatron_sglang_refit branch from 773210d to 9194ded Compare August 4, 2026 01:25

svcnvidia-nemo-ci added the waiting-on-customer Waiting on the original author to respond label Aug 4, 2026

yuki-97 reviewed Aug 4, 2026

View reviewed changes

anwithk mentioned this pull request Aug 4, 2026

NeMo RL v0.8.0 Roadmap #3493

Open

svcnvidia-nemo-ci added waiting-on-customer Waiting on the original author to respond and removed waiting-on-customer Waiting on the original author to respond labels Aug 5, 2026

xiuhu17 and others added 6 commits August 9, 2026 22:30

xiuhu17 force-pushed the zhw/megatron_sglang_refit branch from 9194ded to eac58ed Compare August 10, 2026 05:31

yuki-97 reviewed Aug 10, 2026

View reviewed changes

Kh4L mentioned this pull request Aug 12, 2026

feat(sglang): megatron backend weight refit for sglang rollouts #3612

Open

		return policy_generation.get_updatable_engines_and_lock()


		def broadcast_hf_buckets_via_distributed_impl(

		@@ -0,0 +1,96 @@
		# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

	# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
	# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.

Uh oh!

Conversation

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

Uh oh!

What

Result

Stack

Uh oh!

copy-pr-bot Bot commented Jul 14, 2026

Uh oh!

yuki-97 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yuki-97 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yuki-97 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

xiuhu17 commented Jul 14, 2026 •

edited

Loading