fix(recover): warm up NCCL communicators before the first post-recover step

areal/api/cli_args.py

@@ -1296,6 +1296,10 @@ class TrainEngineConfig:
             "choices": ["disk", "xccl", "awex"],
         },
     )
+    warmup_communicators: bool = field(
+        default=False,
+        metadata={"help": "Pre-connect collective transports after a recover load."},
+    )
     fsdp: FSDPEngineConfig = field(default_factory=FSDPEngineConfig)
     archon: ArchonEngineConfig = field(default_factory=ArchonEngineConfig)
     megatron: MegatronEngineConfig = field(default_factory=MegatronEngineConfig)

areal/engine/core/distributed.py

@@ -1,11 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import os
+from collections.abc import Sequence
 from datetime import timedelta
 
 import torch
 import torch.distributed as dist
 
+# The two probe sizes do not model training tensor shapes. They sit on either
+# side of NCCL's message-size protocol thresholds (LL/LL128 vs Simple; 2KB vs
+# 32MB payloads in bf16), and NCCL allocates transport buffers per communicator
+# and protocol class on first use. Protocol selection also depends on
+# collective type, topology and NCCL settings, so this targets the message
+# classes seen in training rather than guaranteeing coverage of every size.
+_PROBE_NUMELS: tuple[int, ...] = (1024, 16 * 1024 * 1024)
+
 
 def patch_dist_group_timeout(timeout: timedelta):
     """
@@ -23,6 +32,39 @@ def patch_dist_group_timeout(timeout: timedelta):
         distributed_c10d.default_pg_nccl_timeout = timeout
 
 
+def _probe_device() -> torch.device:
+    """The device warmup probes are allocated on.
+
+    Prefer LOCAL_RANK (set by torchrun and most launchers); fall back to the
+    device the caller has already configured via ``set_device``. This keeps the
+    helper usable under custom launchers that don't export LOCAL_RANK. Note
+    that the LOCAL_RANK branch also *sets* the device, matching what engines do
+    during setup.
+    """
+    from areal.infra.platforms import current_platform
+
+    local_rank_env = os.environ.get("LOCAL_RANK")
+    if local_rank_env is not None:
+        local_rank = int(local_rank_env)
+        current_platform.set_device(local_rank)
+    else:
+        local_rank = current_platform.current_device()
+    return torch.device(current_platform.device_type, local_rank)
+
+
+def _warmup_device() -> torch.device | None:
+    """The probe device, or None when there is nothing to warm up."""
+    from areal.infra.platforms import current_platform
+
+    if not dist.is_initialized() or current_platform.device_type == "cpu":
+        return None
+    return _probe_device()
+
+
+def _unique(groups: tuple[dist.ProcessGroup | None, ...]) -> list[dist.ProcessGroup]:
+    return list(dict.fromkeys(g for g in groups if g is not None))
+
+
 def warmup_process_groups(*groups: dist.ProcessGroup | None) -> None:
     """Force eager initialization of the collective communicator for each group.
 
@@ -40,34 +82,175 @@ def warmup_process_groups(*groups: dist.ProcessGroup | None) -> None:
     or before ``dist.init_process_group``. Safe to call repeatedly;
     subsequent calls on already-initialized groups are cheap.
     """
-    # Deferred import to keep this module importable without a platform
-    # (e.g. during lightweight tooling or unit tests).
     from areal.infra.platforms import current_platform
 
     if not dist.is_initialized() or current_platform.device_type == "cpu":
         return
 
-    # Preserve argument order while dropping None and duplicates.
-    unique_groups = list(dict.fromkeys(g for g in groups if g is not None))
+    unique_groups = _unique(groups)
     if not unique_groups:
         return
 
-    # Prefer LOCAL_RANK (set by torchrun/most launchers); fall back to the
-    # device the caller has already configured via ``set_device``. This keeps
-    # the helper usable under custom launchers that don't export LOCAL_RANK.
-    local_rank_env = os.environ.get("LOCAL_RANK")
-    if local_rank_env is not None:
-        local_rank = int(local_rank_env)
-        current_platform.set_device(local_rank)
-    else:
-        local_rank = current_platform.current_device()
-
-    device = torch.device(current_platform.device_type, local_rank)
+    device = _probe_device()
     tensor = torch.zeros(1, device=device)
     for group in unique_groups:
         dist.all_reduce(tensor, group=group)
 
 
+def nccl_process_groups() -> list[dist.ProcessGroup]:
+    """Every collective-capable group this rank belongs to, in creation order.
+
+    Sweeping the registry rather than naming groups keeps subgroups minted by
+    other components (for example the MoE expert grad-bucket reduction) from
+    being missed; a missed group still connects lazily at peak memory.
+
+    Insertion order is the creation order, and non-member ranks get no entry,
+    so any two ranks agree on the relative order of the groups they share.
+    Only metadata is read here, so a raise cannot leave a peer waiting.
+    """
+    from torch.distributed.distributed_c10d import _world
+
+    if not dist.is_initialized():
+        return []
+
+    groups: list[dist.ProcessGroup] = []
+    for group in list(_world.pg_map):
+        if "nccl" not in str(dist.get_backend(group)).lower():
+            continue
+        if dist.get_world_size(group=group) <= 1:
+            continue
+        groups.append(group)
+    return groups
+
+
+def warmup_collective_transports(
+    *groups: dist.ProcessGroup | None,
+    numels: Sequence[int] = _PROBE_NUMELS,
+    dtype: torch.dtype = torch.bfloat16,
+) -> None:
+    """Pre-connect the all-reduce transport buffers of each group.
+
+    Complements :func:`warmup_process_groups`: that one forces the communicator
+    to be *created* with a single tiny all-reduce, which leaves the large Simple
+    protocol buffer to be allocated during the first train step, at peak memory.
+
+    Exceptions are deliberately not caught. Every statement below is a
+    collective, so a rank that swallowed a failure and moved on would leave its
+    peers blocked on an operation that never arrives.
+    """
+    device = _warmup_device()
+    if device is None:
+        return
+    for group in _unique(groups):
+        for numel in numels:
+            dist.all_reduce(torch.zeros(numel, dtype=dtype, device=device), group=group)
+
+
+def warmup_all_to_all_transports(
+    *groups: dist.ProcessGroup | None,
+    numels: Sequence[int] = _PROBE_NUMELS,
+    dtype: torch.dtype = torch.bfloat16,
+) -> None:
+    """Pre-connect the all-to-all transport buffers of each group.
+
+    Kept separate from :func:`warmup_collective_transports` because all-to-all
+    allocates its own buffers and they live for the lifetime of the process.
+    Warming a group that never dispatches tokens would hold that memory for
+    nothing, which is the opposite of the point.
+    """
+    device = _warmup_device()
+    if device is None:
+        return
+    for group in _unique(groups):
+        world_size = dist.get_world_size(group=group)
+        for numel in numels:
+            aligned = (numel // world_size) * world_size
+            if not aligned:
+                continue
+            src = torch.zeros(aligned, dtype=dtype, device=device)
+            dist.all_to_all_single(torch.empty_like(src), src, group=group)
+
+
+def warmup_sharded_transports(
+    group: dist.ProcessGroup | None,
+    *,
+    numel: int = _PROBE_NUMELS[-1],
+    dtype: torch.dtype = torch.bfloat16,
+) -> None:
+    """Pre-connect the distributed optimizer's reduce_scatter / all_gather."""
+    if group is None:
+        return
+    device = _warmup_device()
+    if device is None:
+        return
+    world_size = dist.get_world_size(group=group)
+    if world_size <= 1:
+        return
+    shard_numel = numel // world_size
+    if shard_numel == 0:
+        return
+
+    flat = torch.zeros(shard_numel * world_size, dtype=dtype, device=device)
+    shard = torch.empty(shard_numel, dtype=dtype, device=device)
+    dist.reduce_scatter_tensor(shard, flat, group=group)
+    dist.all_gather_into_tensor(flat, shard, group=group)
+
+
+def warmup_p2p_transports(
+    *groups: dist.ProcessGroup | None,
+    prev_rank: int,
+    next_rank: int,
+    has_prev: bool,
+    has_next: bool,
+    numels: Sequence[int] = _PROBE_NUMELS,
+    dtype: torch.dtype = torch.bfloat16,
+) -> None:
+    """Pre-connect the 2-rank pair communicators pipeline send/recv creates.
+
+    ``prev_rank`` and ``next_rank`` are *global* ranks, matching Megatron's own
+    convention: it resolves them with ``dist.get_global_rank(pp_group, ...)``
+    and passes them alongside an explicit ``group=``.
+
+    Both the unbatched and the batched form are exercised, because they can end
+    up on different communicators. Non-blocking posts are issued before any
+    wait, so no even/odd rotation is needed to stay deadlock-free.
+    """
+    device = _warmup_device()
+    if device is None:
+        return
+    unique = _unique(groups)
+    if not unique or not (has_prev or has_next):
+        return
+
+    for group in unique:
+        for numel in numels:
+            send_buf = torch.zeros(numel, dtype=dtype, device=device)
+            recv_prev = torch.empty(numel, dtype=dtype, device=device)
+            recv_next = torch.empty(numel, dtype=dtype, device=device)
+
+            reqs = []
+            if has_next:
+                reqs.append(dist.isend(send_buf, next_rank, group=group))
+            if has_prev:
+                reqs.append(dist.irecv(recv_prev, prev_rank, group=group))
+            if has_prev:
+                reqs.append(dist.isend(send_buf, prev_rank, group=group))
+            if has_next:
+                reqs.append(dist.irecv(recv_next, next_rank, group=group))
+            for req in reqs:
+                req.wait()
+
+            ops = []
+            if has_prev:
+                ops.append(dist.P2POp(dist.isend, send_buf, prev_rank, group))
+                ops.append(dist.P2POp(dist.irecv, recv_prev, prev_rank, group))
+            if has_next:
+                ops.append(dist.P2POp(dist.isend, send_buf, next_rank, group))
+                ops.append(dist.P2POp(dist.irecv, recv_next, next_rank, group))
+            for work in dist.batch_isend_irecv(ops):
+                work.wait()
+
+
 # Copy from pytorch and OpenRLHF to allow creating multiple main groups.
 # This is needed because torch.distributed.init_process_group() only creates
 # the default global group, and torch.distributed.new_group() only creates

areal/engine/megatron_engine.py

@@ -56,7 +56,12 @@
 )
 from areal.engine.core.distributed import (
     init_custom_process_group,
+    nccl_process_groups,
+    warmup_all_to_all_transports,
+    warmup_collective_transports,
+    warmup_p2p_transports,
     warmup_process_groups,
+    warmup_sharded_transports,
 )
 from areal.engine.core.model import (
     disable_dropout_in_model,
@@ -1721,6 +1726,53 @@ def _validate_fp8_consistency(self):
     def get_device_stats(self) -> DeviceRuntimeInfo:
         return DeviceRuntimeInfo.get_current()
 
+    def warmup_communicators(self) -> None:
+        """Pre-connect the train step's communicators while memory is light.
+
+        NCCL allocates transport buffers per (communicator, protocol) on first
+        use. Left alone, those allocations land inside the first ppo_update at
+        peak occupancy, where the per-peer calloc fails on the last pipeline
+        stage. The buffers live for the lifetime of the process, so one warmup
+        covers every later step.
+
+        Issuing collectives on every registered group means an already unhealthy
+        rank surfaces here rather than at the first train step, hence the opt-in.
+        """
+        if not self.config.warmup_communicators:
+            return
+        if not dist.is_initialized() or current_platform.device_type == "cpu":
+            return
+
+        warmup_collective_transports(*nccl_process_groups())
+        if mpu.get_expert_model_parallel_world_size() > 1:
+            warmup_all_to_all_transports(mpu.get_expert_model_parallel_group())
+        warmup_sharded_transports(
+            mpu.get_data_parallel_group(with_context_parallel=True)
+        )
+
+        pp_size = mpu.get_pipeline_model_parallel_world_size()
+        if pp_size > 1:
+            pp_group = mpu.get_pipeline_model_parallel_group()
+            # AReaL turns off batch_p2p_comm (see areal/models/mcore/registry.py),
+            # so Megatron takes the unbatched path. That path sends both
+            # directions over pp_group, except at pp_size == 2 where one
+            # direction moves to WORLD so the two transfers can overlap; the
+            # 'ucc' backend is excluded because WORLD is always nccl. See
+            # megatron/core/pipeline_parallel/p2p_communication.py::_p2p_ops.
+            p2p_groups: list[dist.ProcessGroup] = [pp_group]
+            if pp_size == 2 and str(dist.get_backend(pp_group)).lower() != "ucc":
+                p2p_groups.append(dist.group.WORLD)
+            warmup_p2p_transports(
+                *p2p_groups,
+                prev_rank=mpu.get_pipeline_model_parallel_prev_rank(),
+                next_rank=mpu.get_pipeline_model_parallel_next_rank(),
+                has_prev=not mpu.is_pipeline_first_stage(ignore_virtual=True),
+                has_next=not mpu.is_pipeline_last_stage(ignore_virtual=True),
+            )
+
+        current_platform.synchronize()
+        self.logger.info("Train-step communicator warmup complete")
+
     def start_memory_profile(self, max_entries: int = 100000) -> None:
         torch.cuda.memory._record_memory_history(max_entries=max_entries)
 

areal/infra/controller/train_controller.py

@@ -713,6 +713,10 @@ def init_awex_adapter(self, meta_server_addr: str | None = None):
             "init_awex_adapter", meta_server_addr=meta_server_addr
         )
 
+    def warmup_communicators(self):
+        """Eagerly build train-step NCCL communicators on all workers."""
+        self._custom_function_call("warmup_communicators")
+
     def step_lr_scheduler(self):
         """Step the learning rate scheduler.
 

areal/utils/recover.py

@@ -371,6 +371,8 @@ def load(
                 for name, engine_ in normalized_engine.items():
                     self._load_checkpoint(engine_, name=name)
 
+            self._warmup_communicators(normalized_engine)
+
             if inference_engine is not None:
                 assert weight_update_meta is not None
                 update_engine = normalized_engine[inference_engine_update_from]
@@ -440,6 +442,24 @@ def _save_checkpoint(
         engine.save(meta)
         logger.info(f"Saved recover checkpoint to {path} (with_optim={with_optim})")
 
+    @staticmethod
+    def _warmup_communicators(
+        normalized_engine: dict[str, TrainEngine | TrainController],
+    ) -> None:
+        for name, engine_ in normalized_engine.items():
+            warmup = getattr(engine_, "warmup_communicators", None)
+            if warmup is None:
+                continue
+            try:
+                warmup()
+            except Exception:
+                logger.warning(
+                    "Communicator warmup failed for engine %s; the first "
+                    "train step will connect lazily instead.",
+                    name,
+                    exc_info=True,
+                )
+
     def _load_checkpoint(
         self,
         engine: TrainEngine | TrainController,