feat(sglang): engine fault tolerance for rollouts by xiuhu17 · Pull Request #3187 · NVIDIA-NeMo/RL

xiuhu17 · 2026-07-14T02:24:05Z

What

Adds RolloutHealthMonitor: a daemon thread that health-checks each SGLang 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

Ft design is adapted from miles

Stack

Stacked on #2997 (reserved-band port allocation) — review only the last commit (feat(sglang): engine fault tolerance for rollouts); earlier commits belong to #2997. Follow-ups: #3190 (megatron refit) → #3188 (mxfp8) → #3189 (comm-split fix).

Split out of #3091.

copy-pr-bot · 2026-07-14T02:24:08Z

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 @xiuhu17 , I rebased main for easily review and left some comments.

yuki-97 · 2026-07-27T06:56:33Z

+            self.engine_gpu_offsets,
+        )
+
+    def clear_updatable_num_new_engines(self):


where will this function be called?

It is now called inside the refit. https://github.com/NVIDIA-NeMo/RL/pull/3187/changes#diff-ccc7bd9f2ab92a1a2c34c1899dda3e50ebf3558d605ffebf359a29e5db25418aR126

yuki-97 · 2026-07-27T06:56:44Z

    def check_weights(self, action: str):
        return self._make_request("weights_checker", {"action": action})

+    def _simulate_crash(self):


where will this function be called?

It is for testing purpose and used in test: test_crashed_engine_is_detected_recovered_and_generation_survives

yuki-97 · 2026-07-27T07:01:33Z

            ray.get(init_handles)

+        # Serializes weight refits against engine recovery across processes.
+        self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote()


looks it only be used in fault tolerance case, so shall we also init with None and do real init under if sglang_cfg["sglang_cfg"].get("use_fault_tolerance"): like self._health_monitor?

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>

nvpohanh · 2026-08-04T00:18:30Z

@xiuhu17 could you fix the DCO issue? thanks

…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>

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>

A hung or dead SGLang engine currently wedges the whole run: the refit's `ray.get` on its actor never returns, and the router keeps sending `/generate` to it. Add `RolloutHealthMonitor`, a daemon thread that polls `/health_generate` on every node-0 engine and, on failure, kills and restarts the actor. The refit picks the survivors up through the engine registry this PR's base already exposes (`get_updatable_engines_and_lock`), so a recovered engine is reconnected on the next weight update. The monitor must not run while the engines are offloaded: `/health_generate` always executes a real one-token generation (the `SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION=false` bypass only covers `/health`), so probing a released engine reports a false failure and kills a live actor. It is therefore wired to the generation lifecycle — resumed by `prepare_for_generation` once the KV cache is back, paused by `finish_generation` and for the duration of a refit. `pause()` blocks on an in-flight probe rather than only setting a flag, so a probe cannot overlap `release_memory_occupation`. `_kill_engine` bounds its graceful-shutdown `ray.get` and kills the actor even when that shutdown fails, which is the exact case the monitor exists to handle. Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>

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

This was referenced Jul 14, 2026

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

Open

megatron backend for sglang rollout #2446

Draft

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

Closed

xiuhu17 force-pushed the zhw/sglang_fault_tolerance branch from c20a694 to 3f62683 Compare July 14, 2026 02:29

xiuhu17 mentioned this pull request Jul 14, 2026

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

Open

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

Kh4L mentioned this pull request Jul 23, 2026

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

Draft

6 tasks

yuki-97 force-pushed the zhw/sglang_fault_tolerance branch from 3f62683 to 6760b97 Compare July 27, 2026 05:12

yuki-97 reviewed Jul 27, 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 27, 2026

xiuhu17 force-pushed the zhw/sglang_fault_tolerance branch 3 times, most recently from 72b52c2 to fafd4d4 Compare August 1, 2026 22:32

xiuhu17 requested review from a team as code owners August 1, 2026 22:32

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

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

xiuhu17 and others added 8 commits August 3, 2026 18:21

xiuhu17 force-pushed the zhw/sglang_fault_tolerance branch from fafd4d4 to 348b555 Compare August 4, 2026 01:26

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

Navigation Menu

Uh oh!

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

feat(sglang): engine fault tolerance for rollouts#3187
xiuhu17 wants to merge 8 commits into
NVIDIA-NeMo:mainfrom
xiuhu17:zhw/sglang_fault_tolerance

xiuhu17 commented Jul 14, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jul 14, 2026

Uh oh!

yuki-97 left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yuki-97 Jul 27, 2026

Uh oh!

xiuhu17 Aug 3, 2026

Uh oh!

yuki-97 Jul 27, 2026

Uh oh!

xiuhu17 Aug 3, 2026

Uh oh!

Uh oh!

yuki-97 Jul 27, 2026

Uh oh!

Uh oh!

nvpohanh commented Aug 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

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

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!

yuki-97 Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

xiuhu17 Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

yuki-97 Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

xiuhu17 Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yuki-97 Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

nvpohanh commented Aug 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

xiuhu17 commented Jul 14, 2026 •

edited

Loading