fix(recover): warm up NCCL communicators before the first post-recover step - #1548
fix(recover): warm up NCCL communicators before the first post-recover step#1548Le8r0nJames wants to merge 3 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a communicator warmup mechanism to eagerly initialize NCCL communicators and allocate transport buffers during recovery restarts, preventing out-of-memory errors at peak occupancy. This is implemented across MegatronEngine, TrainController, and the recovery utility. Feedback highlights a potential crash during DP group warmup if the data-parallel size does not evenly divide the tensor size, and suggests wrapping the DP and PP warmup blocks in individual try-except blocks to improve robustness.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| dp_group = mpu.get_data_parallel_group(with_context_parallel=True) | ||
| dp_size = dist.get_world_size(group=dp_group) | ||
| if dp_size > 1: | ||
| shard = torch.empty( | ||
| large.numel() // dp_size, dtype=large.dtype, device=device | ||
| ) | ||
| dist.reduce_scatter_tensor(shard, large, group=dp_group) | ||
| dist.all_gather_into_tensor(large, shard, group=dp_group) | ||
| pp_size = mpu.get_pipeline_model_parallel_world_size() | ||
| if pp_size > 1: | ||
| pp_rank = mpu.get_pipeline_model_parallel_rank() | ||
| nxt = mpu.get_pipeline_model_parallel_next_rank() | ||
| prv = mpu.get_pipeline_model_parallel_prev_rank() | ||
| is_first = mpu.is_pipeline_first_stage(ignore_virtual=True) | ||
| is_last = mpu.is_pipeline_last_stage(ignore_virtual=True) | ||
| for buf in (small, large): | ||
| send_buf = buf.clone() | ||
| recv_buf = torch.empty_like(buf) | ||
| # Unbatched pair comms, both directions; even ranks send | ||
| # first so the chain cannot deadlock. | ||
| for send_peer, recv_peer, do_send, do_recv in ( | ||
| (nxt, prv, not is_last, not is_first), | ||
| (prv, nxt, not is_first, not is_last), | ||
| ): | ||
| if pp_rank % 2 == 0: | ||
| if do_send: | ||
| dist.send(send_buf, send_peer) | ||
| if do_recv: | ||
| dist.recv(recv_buf, recv_peer) | ||
| else: | ||
| if do_recv: | ||
| dist.recv(recv_buf, recv_peer) | ||
| if do_send: | ||
| dist.send(send_buf, send_peer) | ||
| ops = [] | ||
| if not is_last: | ||
| ops.append(dist.P2POp(dist.isend, send_buf, nxt)) | ||
| if not is_first: | ||
| ops.append(dist.P2POp(dist.irecv, recv_buf, prv)) | ||
| if ops: | ||
| for work in dist.batch_isend_irecv(ops): | ||
| work.wait() |
There was a problem hiding this comment.
Issue: Potential crash on non-power-of-two DP sizes and lack of exception handling
-
Size Mismatch in
reduce_scatter_tensor/all_gather_into_tensor:
largeis initialized with16 * 1024 * 1024(16,777,216) elements. If the data-parallel size (dp_size) is not a divisor of this number (e.g.,dp_size = 3,5,6, etc.),large.numel() // dp_sizewill result in a size mismatch duringreduce_scatter_tensorandall_gather_into_tensor, raising aRuntimeErrorand crashing the warmup. Slicinglargeto a multiple ofdp_size(large[:shard_size * dp_size]) resolves this. -
Robustness:
The DP and PP warmup sections are currently outside anytry...exceptblock. If any collective or parallel state query fails, it will crash the entire warmup process and fail the recovery load. Wrapping both sections in individualtry...exceptblocks ensures that a failure in one section does not prevent the other sections from running, and gracefully falls back to lazy connection with a warning.
try:
dp_group = mpu.get_data_parallel_group(with_context_parallel=True)
dp_size = dist.get_world_size(group=dp_group)
if dp_size > 1:
shard_size = large.numel() // dp_size
shard = torch.empty(
shard_size, dtype=large.dtype, device=device
)
dist.reduce_scatter_tensor(shard, large[:shard_size * dp_size], group=dp_group)
dist.all_gather_into_tensor(large[:shard_size * dp_size], shard, group=dp_group)
except Exception:
self.logger.warning(
"communicator warmup skipped DP group", exc_info=True
)
try:
pp_size = mpu.get_pipeline_model_parallel_world_size()
if pp_size > 1:
pp_rank = mpu.get_pipeline_model_parallel_rank()
nxt = mpu.get_pipeline_model_parallel_next_rank()
prv = mpu.get_pipeline_model_parallel_prev_rank()
is_first = mpu.is_pipeline_first_stage(ignore_virtual=True)
is_last = mpu.is_pipeline_last_stage(ignore_virtual=True)
for buf in (small, large):
send_buf = buf.clone()
recv_buf = torch.empty_like(buf)
# Unbatched pair comms, both directions; even ranks send
# first so the chain cannot deadlock.
for send_peer, recv_peer, do_send, do_recv in (
(nxt, prv, not is_last, not is_first),
(prv, nxt, not is_first, not is_last),
):
if pp_rank % 2 == 0:
if do_send:
dist.send(send_buf, send_peer)
if do_recv:
dist.recv(recv_buf, recv_peer)
else:
if do_recv:
dist.recv(recv_buf, recv_peer)
if do_send:
dist.send(send_buf, send_peer)
ops = []
if not is_last:
ops.append(dist.P2POp(dist.isend, send_buf, nxt))
if not is_first:
ops.append(dist.P2POp(dist.irecv, recv_buf, prv))
if ops:
for work in dist.batch_isend_irecv(ops):
work.wait()
except Exception:
self.logger.warning(
"communicator warmup skipped PP group", exc_info=True
)There was a problem hiding this comment.
Fixed in 6b9030e: the DP warmup now truncates the buffer to shard_numel * dp_size so reduce_scatter/all_gather shapes stay consistent for any DP size, and the DP block degrades to lazy connection on failure like the per-group sweep. The pipeline P2P block intentionally stays outside a try/except: a one-sided exception there would leave the peer blocked, which is worse than surfacing the error - and the warmup as a whole is now opt-in.
| those allocations happen up front; the buffers persist across | ||
| offload/onload for the lifetime of the process. | ||
| """ | ||
| if os.environ.get("AREAL_SKIP_COMM_WARMUP", "").strip() == "1": |
There was a problem hiding this comment.
It's better to change the default value to 1 here, do not silently change the behavior, as there is still a risk of hanging in the warm up section.
There was a problem hiding this comment.
Done in 6b9030e - the warmup is now opt-in: AREAL_COMM_WARMUP=1 enables it and the default keeps the previous lazy-connection behavior. The docstring and commit message now state the hang rationale explicitly.
| continue | ||
| dist.all_reduce(small.clone(), group=group) | ||
| dist.all_reduce(large.clone(), group=group) | ||
| chunk = 4096 |
There was a problem hiding this comment.
What does this hardcoding represent here? Can it be applicable to tensors of different sizes? Please provide detailed explanation in the docstring or comments.
There was a problem hiding this comment.
Added a comment block in 6b9030e: the two probes do not model training tensor shapes - they sit on each side of NCCL's message-size protocol thresholds (LL/LL128 vs Simple; 2KB vs 32MB bf16 payloads), and NCCL allocates transport buffers per communicator and protocol class on first use. It targets the message classes observed in training rather than guaranteeing coverage of every tensor size, since protocol selection also depends on collective type, topology, and NCCL settings.
6b9030e to
f1a48af
Compare
| def get_device_stats(self) -> DeviceRuntimeInfo: | ||
| return DeviceRuntimeInfo.get_current() | ||
|
|
||
| def warmup_communicators(self) -> None: |
There was a problem hiding this comment.
If the issue only occurs under the conditions of mcore version 0.11 and ppsize=2, add a judgment in the function content to avoid undefined behavior as much as possible.
…r step PyTorch builds 2-rank communicators for unbatched pipeline send/recv lazily and NCCL connects transport buffers per protocol on first use. On a recover restart the backward-direction pipeline pairs therefore connect inside the first ppo_update, when the PP-last-stage ranks are already at peak occupancy, and the ~10MB transport calloc fails. Add an opt-in communicator warmup (AREAL_COMM_WARMUP=1), run right after the recover checkpoint load while memory is still light: sweep every registered NCCL process group with both message-size protocol classes, plus the distributed-optimizer collective shapes and both unbatched and batched pipeline P2P directions. The transport buffers persist across offload/onload for the process lifetime. The warmup stays disabled by default because it issues collectives on every group and would hang on an already-unhealthy rank; individual warmup failures degrade to the previous lazy-connection behavior.
The warmup reached for torch.cuda directly while the rest of the engine goes through current_platform, which pins it to CUDA. That is the wrong way round here: eager communicator setup exists in this tree because HCCL is the backend most prone to failing on a lazy first collective. The engine already resolves its own device during setup, so take it from there instead of asking the runtime again, and synchronise through the platform like the other seven call sites in this file. Skip the warmup altogether on CPU-only platforms, matching warmup_process_groups.
The pipeline probes were issued without a group, so they landed on the default group. Megatron routes unbatched pipeline p2p through the pipeline group, and only at pp_size == 2 does one direction move to WORLD so the two transfers can overlap. Past two stages the warmup therefore touched a communicator the train step never uses, leaving the one that allocates at peak memory cold - the case it was written for. The sweep is now split into transport helpers next to warmup_process_groups, which they complement: that one forces the communicator to exist with a tiny all-reduce, these size the buffers each collective shape needs. Knowledge of mpu stays in the engine, so the helpers are reusable and testable without one. all-to-all gets its own helper rather than running on every group. Its buffers live as long as the process, so warming a group that never dispatches tokens would hold that memory for nothing. Collective failures are no longer swallowed. A rank that logged a warning and moved on would leave its peers blocked on an operation that never arrives, which is worse than surfacing the error. The knob becomes a config field, since an env-only switch cannot be expressed in an experiment config.
f1a48af to
c291222
Compare
| dist.all_reduce(tensor, group=group) | ||
|
|
||
|
|
||
| def nccl_process_groups() -> list[dist.ProcessGroup]: |
There was a problem hiding this comment.
The XCCL recovery triggers a single-sided entry into cross-engine all_reduce. Before recovery, the RLTrainer calls actor.connect_engine.XCCL to create a custom NCCL group containing the training PP head and inference workers. This group is registered into the process-level _world.pg_map via _new_process_group_helper. Subsequent recovery occurs in rl_trainer.py, where nccl_process_groups() enumerates the entire _world.pg_map, and warmup_collective_transports() performs all_reduce on each group. Warmup RPC is sent only to training workers, excluding inference workers within the same XCCL group. This results in the training PP head with the weight-update group indefinitely waiting for inference members until a process-group timeout; the exception handling in RecoverHandler cannot take effect before the RPC returns.
This is why the current tests haven't uncovered issues: tests/test_megatron_comm_warmup.py:80-83 always provides an empty pg_map.
Suggested fix: Avoid traversing the process-level NCCL registry. Explicitly enumerate train-step-owned groups, or record reliable owner/purpose for groups and exclude cross-engine weight-update groups. Add real multi-process recovery tests for XCCL connected rollouts.
| logger.info(f"Saved recover checkpoint to {path} (with_optim={with_optim})") | ||
|
|
||
| @staticmethod | ||
| def _warmup_communicators( |
There was a problem hiding this comment.
Ray/Slurm's generic RPC will retry exceptions, including engine HTTP 500. RecoverHandler._warmup_communicators() ultimately catches all exceptions, logs a warning, and then continues to recover.
If a rank fails to complete a collective due to allocation or synchronize failure after some work is done, other ranks may have already successfully returned. After the failed rank is retried individually, it will re-enter the first collective, while the other ranks no longer participate, causing a new deadlock. Even if NCCL errors are returned on all ranks, the communicator may already be damaged and cannot safely "degrade" to lazy connect as in the first training.
Suggested fixes:
- The warmup RPC must be non-retriable.
- Stopping recovery and propagate the error if any rank fails.
- Do not claim to be able to fallback to lazy connection unless you can coordinate the destruction and reconstruction of the group on all ranks.
| for name, engine_ in normalized_engine.items(): | ||
| self._load_checkpoint(engine_, name=name) | ||
|
|
||
| self._warmup_communicators(normalized_engine) |
There was a problem hiding this comment.
Warmup occurs before SGLang's pause, KV-cache offload, and weight offload. Therefore, the allocation of 32 MiB probe tensors for warmup and the persistence of transport buffers will happen while SGLang is still resident. This is contrary to the goal of "while memory is light" and may retrigger the CUDA allocation failure that this PR attempts to address.
Suggested fix: The AWEX colocated path should pause/offload rollout first, then perform communicator warmup, and then load the actor checkpoint. The non-colocated path can maintain an independent timeline.
|
This pull request has been automatically marked as stale because it has not had recent activity within the last 14 days. Please add a comment or push new commits to keep it active. Thank you for your contribution! |
2 participants
Description
PyTorch builds 2-rank communicators for unbatched pipeline send/recv lazily, and NCCL connects transport buffers per protocol class on first use. After a recover restart those connects happen inside the first
ppo_update, when the device is already at peak occupancy, and the ~10MB transport calloc can fail on the PP-last-stage ranks (NCCL 'Failed to CUDA calloc' at operation=Connect) — observed on memory-tight multi-node MoE restarts.This PR exercises every communicator the train step needs right after the recover checkpoint load, while memory is still light:
all_reduceandall_to_all_singlein both message-size protocol classes;The recover handler invokes the hook via duck-typing; engines without
warmup_communicatorsare unaffected. Warmup failures degrade to lazy connection with a warning instead of failing the restart.AREAL_SKIP_COMM_WARMUP=1disables the warmup.Related Issue
N/A
Type of Change
Checklist
Contributing Guide
pre-commit run --all-files)main/create-prAdditional Context
Surfaced on restarts where the first post-recover step runs near peak memory (colocated/offload-heavy setups); the lazy-connect mechanism itself is generic to any PP + recover configuration. Multi-GPU behavior was validated on a 32-GPU MoE recover restart: previously failing first-step connects now complete during the post-load window.