[Kh4L]

Original work: #3190 by @xiuhu17. This is a re-cut opened at his request so the change can complete review. His commits retain their original authorship and Signed-off-by trailers.

Summary

• Routes SGLang weight refit through the weight-synchronizer interface.
• Uses Ray CUDA IPC for colocated Megatron/FSDP2 policies.
• Uses a trainer-rank-0 NCCL broadcast for non-colocated Megatron policies.
• Rejects unsupported non-colocated FSDP2/SGLang during setup; support is tracked in feat: support non-colocated (broadcast) SGLang refit for the FSDP2/DTensor policy #3745.
• Wraps bucket transfer in SGLang begin_weight_update / end_weight_update sessions so derived quantized layouts are finalized once after the last bucket.
• Removes the legacy HTTP transport. It was live on main; this PR supersedes it with the IPC/NCCL transports.

Runtime dependencies

• SGLang: 0.5.17.dev30+g3003d70f6 from pinned commit 3003d70f680d41c59d1b7acbf65cb47795dfd19e
• SGLang kernel: 0.4.5
• FlashInfer Python/cubin/CUDA-13 JIT cache: 0.6.14
• Transformers in the SGLang environment: 5.12.1

The pinned SGLang revision contains sglang#27948, replacing the local TMS CUDA-graph patch removed here. The matching FlashInfer cubin and JIT-cache wheels are explicit so engine startup does not fall back to compiling kernels.

Dependency versions are isolated per extra: only SGLang uses Transformers 5.12.1 and TileLang 0.1.11; Megatron/vLLM stay on Transformers 5.8.1, and base/non-SGLang extras stay on TileLang 0.1.9 where the package's x86_64 marker applies. This PR no longer depends on #3498.

Review fixes included

• Lazily initializes backend-neutral refit transport state and reuses the shared Megatron HF iterator/bucketer.
• Guarantees generation/policy phase restoration when pause, onload, or transfer fails.
• Declares the four SGLang policy methods on ColocatablePolicyInterface.
• Preserves SGLang's engine-side message error text.
• Cleans up stale Gloo subgroups and the Ray lock actor, bounds lock acquisition, and routes pause/continue through the common HTTP helper.
• Centralizes NCCL_CUMEM_ENABLE selection before the first NCCL process group.
• Removes stale test deselections and moves the ordering AST check to the CPU lane.

Validation

Current head: 9336ee7af808556dc884f65462a4d025f50a8f87

Local/current-head checks:

• uv lock --check: pass.
• Per-extra exports confirm SGLang-only Transformers/TileLang isolation.
• Changed recipe inheritance configs pass minimize-check.
• Relevant launchers pass Bash syntax and TEST_DRYRUN=1.
• pytest --collect-only succeeds with all 57 weight-synchronizer tests collected.
• Ruff/format hooks pass; Pyrefly reports no errors in the touched files.
• Exact-head CI requested with /ok to test 9336ee7af.

HSG hardware validation used immutable image digest sha256:3ef02100ebceb15790ad0e1e32728feb3d2d144f8e6da97daf954758acc74891, built by GitLab job 411470249 from 469bd63927650b00fca5ed4399926b11f46baa21.

• Qwen job 6518683: COMPLETED (0:0), 450/450 training steps and 450/450 successful refits, four SGLang engines, validator passed, median token_mult_prob_error=1.003821.
• Nemotron-H/Nano job 6518682: COMPLETED (0:0), 4/4 training steps and 4/4 successful refits, positive gradients on steps 1–2, minimum 4043 valid tokens, validator passed, median token_mult_prob_error=1.015443.
• Decode CUDA graphs remained enabled; only prefill capture was disabled. Full --disable-cuda-graph was not used.
• Runtime preflight payload in job 6526299 printed PR3612_RUNTIME_VERSIONS_OK for the immutable image on GB200/ARM64: the SGLang environment has SGLang 0.5.17.dev30+g3003d70f6, kernel 0.4.5, FlashInfer 0.6.14, Transformers 5.12.1, TileLang 0.1.11, and Torch 2.11.0+cu130; the Megatron environment has Transformers 5.8.1, Transformer Engine 2.15.0+42b84005, and Torch 2.11.0+cu130. Megatron's x86_64-only TileLang dependency is correctly absent on ARM64. The wrapper was cancelled after the success marker when container teardown remained stuck, releasing the node.

Qwen 50-step window means from the 450-step run:

Steps	Reward	Loss	token_mult_prob_error
1–50	0.2950	0.003241	1.003745
51–100	0.3200	0.004743	1.003776
101–150	0.3700	0.006685	1.004218
151–200	0.3225	0.009041	1.004265
201–250	0.3300	0.001671	1.004324
251–300	0.3825	0.004617	1.003998
301–350	0.3950	0.011951	1.004246
351–400	0.3700	0.003949	1.004182
401–450	0.4075	0.001219	1.003809

The only change after the hardware-tested commit is 9336ee7af, which replaces the invalid breakable prefill backend in shipped/test configs with disabled. That is the same prefill setting used by both successful HSG runs; runtime code and dependencies are unchanged.

Stack

This is the base of the linear follow-up stack:

• feat(sglang): megatron backend weight refit for sglang rollouts #3612 — Megatron/SGLang weight refit
• feat(sglang): engine fault tolerance for rollouts #3613 — engine fault tolerance
• feat(sglang): mxfp8 quantized rollout and refit #3614 — MXFP8 rollout/refit
• feat(sglang): nvfp4 quantized rollout and refit #3615 — NVFP4 rollout/refit

KV-cache preservation parity is tracked separately in #3747.

[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.

[Kh4L]

Self-review (/review-pr-team), per the contributor self-review policy

Ran the four reviewer roles the skill specifies — rl-expert, bug-finder, design-reviewer, test-agent — plus the mandatory devil's-advocate challenge, and the repo's own pre-commit suite.

39 raw findings → 21 CONFIRMED, 8 DOWNGRADED, 4 DISPUTED. The adversarial pass earned its place: it killed a flashinfer-cubin finding by checking PyPI (no 0.6.14 exists, so the removal from the sglang extra was forced, not a divergence), corrected two sets of wrong GPU-hour arithmetic, and upgraded one dependency finding after reading the lockfile.

Already fixed in b5356f0

1. grpo.py raised NameError on every SGLang setup. It still referenced worker_init_timing_metrics, which main replaced with the SetupTimingMetrics dataclass. It was the only surviving reference in the file, inside the elif backend == "sglang" branch. It survived the rebase precisely because that hunk merged cleanly while its two neighbours conflicted and got fixed. Routed to setup_timing_metrics.extras, the documented overflow bucket the vLLM sparse transports already use.
2. A failed refit left every engine paused forever. begin_weight_update() and invalidate_kv_cache() sat outside the try whose finally calls continue_generation(). Each acquired state now has its own guard, and end_weight_update() only runs for a session that actually opened.
3. The refit driver-import guard was blind to from ...workers.megatron_policy_worker import X — ast.ImportFrom puts the module on node.module and only the symbol on alias.name. Checked against all four import spellings: the old comprehension caught 2 of 4.
4. sglang_weight_synchronizer.py was missing from pyrefly.toml, so the largest new module in the PR escaped type-checking.
5. pre-commit hygiene the PR shipped red: taplo formatting and three unminimized recipe YAMLs.

Worth stating plainly: the only guaranteed crash was found by pre-commit, not by any of the four agents. That is an argument for running the real linter as part of the review rather than trusting the agent pass alone.

Environment caveats

• No GPU on the review machine, so nothing runtime-dependent was executed; every such finding is marked unverified rather than asserted. A GPU re-run would strengthen the test findings.
• CLAUDE_CODE_EXPERIMENTAL_AGENT_TEAMS=1 was not set, so the roles ran without inter-agent messaging. Devil's-advocate verified findings itself instead of interrogating the reporting agent across its usual two challenge rounds.

The five comments below are what survived and is not fixed — left for the reviewer's judgement rather than changed unilaterally.

Generated by Claude Code

[Kh4L]

Addendum to the self-review above — a rationale in this PR's own commit message is wrong, and it hides how big the change is.

refactor(sglang): route every SGLang refit through the weight synchronizer (now 7c78def after today's rebase) deletes nemo_rl/weight_sync/http_weight_synchronizer.py, justifying it as:

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.

Both halves are false against main:

• It is constructed — for every SGLang run. nemo_rl/weight_sync/factory.py on main reads if generation_backend == SGLANG_BACKEND: ... return HTTPWeightSynchronizer(...). That is the current shipped SGLang weight-sync path, not dead code.
• stream_weights_via_http is alive on main. It is called at nemo_rl/algorithms/grpo.py:2362, declared on PolicyInterface, implemented on lm_policy and dtensor_policy_worker_v2, and backed by stream_weights_via_http_impl in nemo_rl/models/policy/utils.py. That impl was added in feat: Add SGLang rollout backend and tests #1674 (2026-01-22) and the only commit that has ever touched the symbol is the one that added it.
• 7b0f0bd2 is not on main. It is a commit on this PR's own branch. So the reasoning is circular: this stack removed the HTTP path, and then removed the synchronizer because "it was removed in 7b0f0bd".

To be clear about what I am not saying: the deletion itself is correct and is the point of the PR. Replacing the HTTP transport with the IPC/broadcast synchronizers is exactly what this change is for, it is done coherently (no dangling imports — factory.py has no HTTP branch left), and I am not asking for it to be reverted.

What I am asking for is that the PR stop describing it as removing dead code. A maintainer who reads "it was dead (never constructed)" will not scrutinize what is in fact a swap of the live SGLang weight-transfer transport, and the PR body doesn't mention the swap at all. Anyone currently on SGLang gets a different transport after this merges.

Action: reword that paragraph of the commit body to say it replaces the transport rather than removing dead code, and add a line to the PR description noting the behavioural change for existing SGLang users.

Found while rebasing the stack onto current main — this is also the reason all four PRs went CONFLICTING today (upstream #3618 added tests to tests/unit/weight_sync/test_weight_synchronizer.py, the same file this change edits). All four are rebased and MERGEABLE again as of 6cc4448.

Generated by Claude Code

[yuki-97]

Re-cut of #3190, so I re-checked all 34 of my review threads there against this tree. This review is the #3190 carry-over only — comments on the new code will follow separately.

Reuse of the vLLM / TRT-LLM transfer functions — settled offline, no action needed. SGLang's refit differs enough that it can't easily reuse them as-is: it runs its own Ray CUDA-IPC NCCL group, needs an HF-weight conversion step and SGL-specific begin/apply/commit hooks, and has different pause / KV-cache / recovery semantics after a partial update. So the three threads asking for reuse (r3689520162, r3689584423, r3689651312) are not carried over here; please treat them as closed.

Carried over from #3190, details inline: 4 unresolved threads that still apply to this tree, and 2 asks that had no natural line to sit on, anchored to the nearest relevant file.

Everything else is either the reuse family above or genuinely addressed here.

[yuki-97]

Follow-up round on the re-cut. The fix commits stacked on top of the original seven (dffdc88b9 through b86aa8e4b) address 16 of the 20 open threads, which are now resolved.

Still open:

• quantization schema — declaring the TypedDict is good, but see my reply: let's keep .get("scheme", "bf16") rather than adding get_sglang_quantization_scheme.
• NCCL_CUMEM_ENABLE placement — see my reply on megatron/setup.py.
• Convergence evidence (#3612 (comment)) — the 2n8g non-colocated nightly is back at 450 steps / 150 min, but the description still says this tree has not been run on hardware. A reward + loss + token_mult_prob_error curve from a longer run would close it.

The rest is inline.

Two things this review does not cover: no CI has run on this tree, and every shipped SGLang recipe is PP=1, so the multi-rank bucket path is not exercised by anything in the suite.

[yuki-97]

/ok to test b86aa8e

[Kh4L]

/ok to test f4d36ef

[terrykong]

Team review of the SGLang Megatron refit path, run across six agents with an adversarial verification pass (12 candidate findings were dropped as unreachable, pre-existing, or already covered by a test). Reviewed at b86aa8e; swept against the rebased head f4d36ef — comments below say individually what that push already fixed.

On structure — this does follow the weight-synchronizer pattern, and it is the right direction. The PR does not add a parallel mechanism: it replaces HTTPWeightSynchronizer, which was already one. The two new classes subclass the same ABC, are built by create_weight_synchronizer, use the same timer key as all five siblings, mirror IPCWeightSynchronizer's is_stale semantics and buffer sizing, add their tests to the existing test_weight_synchronizer.py, and delete the SGLang isinstance branches that used to live inside refit_policy_generation. Two composition asks remain on the Megatron worker (constructor block, duplicated iterator) — both have small in-PR fixes spelled out inline, with the longer-term shape sketched separately.

Status after the f4d36ef push: the 9 pyrefly errors are fixed and the rebase cleared the staleness gate, but no head of this PR has ever produced test evidence — the only /ok to test died at lint. main has moved again (2 commits) and GitHub currently reports the branch CONFLICTING, so it likely needs one more rebase before a meaningful CI run.

Dependency changes in this PR, at a glance

Verified against uv.lock on this branch vs main (unchanged by the rebase — pyproject.toml is byte-identical between the two heads):

package	main	this PR
transformers base range	>=5.5.0,<5.9.0	>=5.5.0,<5.13.0
transformers sglang extra	==5.6.0	==5.12.1
transformers resolved for the mcore venv	5.8.1	5.12.1
transformers resolved for automodel/fsdp	5.5.0	5.5.0 (unchanged)
sglang	==0.5.12.post1 (PyPI release)	git rev 3003d70f, unreleased sglang-miles branch
sglang-kernel	0.4.2.post2	0.4.5
kernels	>=0.12.0,<0.13	>=0.14.1,<0.15
flashinfer-python	0.6.11.post1	0.6.14
flashinfer-cubin / flashinfer-jit-cache	0.6.11.post1	deleted
llguidance override	>=1.3.0,<1.4.0	>=1.3.0 (cap dropped)

⚠️ Risk we want a decision on before merge (cc @yuki-97): the transformers bump is not contained to sglang. The lock resolves 5.12.1 for the mcore extra on its own, so every Megatron training venv moves 5.8.1 → 5.12.1 — four minor versions — in a PR whose test plan only covers the sglang recipes. The new comment in policy/__init__.py claims "the driver and policy venvs stay on 5.5.x", which the lock contradicts. Resolution: this PR should wait for #3498 to land and rebase on it. That PR (Automodel r0.6.0) makes the same <5.9.0 → <5.13.0 move with an actual transformers-5.12 migration and its own test surface, so the bump arrives already exercised and a transformers regression bisects to one PR instead of two upgrades landing together.

Findings already raised in @yuki-97's open threads (weight_transfer_mode — since deleted outright at f4d36ef — the vLLM pause_generation stubs, and the automodel NCCL_CUMEM_ENABLE guard) are deliberately not duplicated here; our independent checks agreed with all three.

The remaining comments, in severity order: CI status, dependency blast radius, two try/finally placement issues, the ColocatablePolicyInterface regression, setup-time vs refit-time validation, an error-message regression against main, test coverage for the new pure helpers, two stale --deselect lines, the two composition asks on the Megatron worker, and a bundled block of low-severity nits. All were re-verified against f4d36ef; each one still present at that head says so.

Generated by Claude Code

[Kh4L]

/ok to test f322100

[Kh4L]

/ok to test 3b75899

[yuki-97]

Re-checked the four threads from my last pass at 3b75899 — all four look resolved. One follow-up on cc8a54e below.

[yuki-97]

/ok to test 09c4efc

[yuki-97]

CI:L1 passed in 3b75899, after that 09c4efc only changes the comments. so changes to CI:docs to merge.