Skip to content

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

Closed
xiuhu17 wants to merge 5 commits into
NVIDIA-NeMo:mainfrom
xiuhu17:zhw/fix_sglang_refit_comm_split
Closed

fix(sglang): avoid ncclCommSplit hang in the refit weight-update group#3189
xiuhu17 wants to merge 5 commits into
NVIDIA-NeMo:mainfrom
xiuhu17:zhw/fix_sglang_refit_comm_split

Conversation

@xiuhu17

@xiuhu17 xiuhu17 commented Jul 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

What

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, NCCL 2.28.9) while the remaining trainer ranks hit the watchdog.

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.

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.

Root cause analysis by @Kh4L (co-authored).

Verification

  • torch 2.10.0 and 2.11.0 have the identical split predicate; setBoundDeviceId is a pure field assignment accepting None in both (no propagation to live backends)
  • gloo interop smoke (this helper ↔ sglang init_custom_process_group replica): broadcast + destroy OK
  • split-predicate regression: with a device-bound default group, the raw helper path consults _get_split_source (positive control); the fixed path never does, and bound_device_id is restored

Stack

Stacked on #3188 (mxfp8) on #3190 (megatron refit) on #3187 (fault tolerance) on #2997 (ports) — review only the last commit (fix(sglang): avoid ncclCommSplit hang...).

@copy-pr-bot

copy-pr-bot Bot commented Jul 14, 2026

Copy link
Copy Markdown

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.

xiuhu17 and others added 5 commits July 13, 2026 19:29
@xiuhu17
fix: consolidate SGLang port allocation onto reserved-band helpers
08b3e91 
Replace the TOCTOU-prone ephemeral-port probing in the SGLang backend
(find_available_port / get_free_port random walks) with the shared
reserved-band helpers in virtual_cluster.py:

- new _get_free_consecutive_ports_local() scans a bounded band and
  threads a per-node cursor so blocks never overlap on a node
- engine server/NCCL/dist_init ports come from the generation band
  (policy.generation.port_range_low/high)
- router and Prometheus ports get dedicated hard-coded bands
- workers expose _get_current_free_port/_get_current_node_ip instead of
  _get_current_node_ip_and_free_port
- port_cursors/node_port_cursor use None sentinels (no mutable defaults)
- boundary tests for the new helper

Closes NVIDIA-NeMo#2874

Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>
@xiuhu17
feat(sglang): engine fault tolerance for rollouts
3f62683 
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>
@xiuhu17
feat(sglang): megatron backend weight refit for sglang rollouts
7b0f0bd 
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>
@xiuhu17
feat(sglang): mxfp8 quantized rollout and refit
d06ef21 
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>
@xiuhu17 @Kh4L
fix(sglang): avoid ncclCommSplit hang in the refit weight-update group
e27d3f1 
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>
@xiuhu17
xiuhu17 force-pushed the zhw/fix_sglang_refit_comm_split branch from c92770b to e27d3f1 Compare July 14, 2026 02:29
This was referenced Jul 14, 2026
Open
Open
Open
@xiuhu17
xiuhu17 marked this pull request as ready for review July 14, 2026 05:05
@xiuhu17
xiuhu17 requested review from a team as code owners July 14, 2026 05:05
yuki-97
yuki-97 reviewed Jul 14, 2026
View reviewed changes

@yuki-97 yuki-97 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hi @xiuhu17 @Kh4L, thanks for fixing! we also meet similar issue in #3073.

curious are we able to split this fix out instead of stack on previous PRs since it looks not based on the new features. so that we can get it in more quickly.

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-customer Waiting on the original author to respond label Jul 14, 2026
@xiuhu17

xiuhu17 commented Jul 14, 2026

Copy link
Copy Markdown
Contributor Author

Hi @yuki-97 the init process group is used in disaggregated mode(which is supported for megatron+Sglang backend). We could fix it first here but it might be a dead code.

@svcnvidia-nemo-ci svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed waiting-on-customer Waiting on the original author to respond labels Jul 14, 2026
@Kh4L Kh4L mentioned this pull request Jul 23, 2026
Draft
6 tasks
Kh4L added a commit to Kh4L/NemoRL that referenced this pull request Jul 25, 2026
@Kh4L
feat(sglang): enable async-GRPO NeMo-Gym rollouts on the SGLang backend
78e4982 
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>
Kh4L added a commit to Kh4L/NemoRL that referenced this pull request Jul 27, 2026
@Kh4L
feat(sglang): enable async-GRPO NeMo-Gym rollouts on the SGLang backend
8c568b2 
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>
@xiuhu17 xiuhu17 closed this Jul 28, 2026
Kh4L added a commit to Kh4L/NemoRL that referenced this pull request Jul 28, 2026
@Kh4L
feat(sglang): enable async-GRPO NeMo-Gym rollouts on the SGLang backend
1528cd8 
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>
@svcnvidia-nemo-ci svcnvidia-nemo-ci removed the waiting-on-maintainers Waiting on maintainers to respond label Jul 28, 2026
Kh4L added a commit to Kh4L/NemoRL that referenced this pull request Jul 28, 2026
@Kh4L
feat(sglang): enable async-GRPO NeMo-Gym rollouts on the SGLang backend
27c5771 
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>
@anwithk anwithk mentioned this pull request Aug 4, 2026
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@yuki-97 yuki-97 yuki-97 left review comments

Assignees

No one assigned

Labels

community-request

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@xiuhu17 @yuki-97 @svcnvidia-nemo-ci