diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 117642a246..b7a277da35 100644 --- a/areal/api/cli_args.py +++ b/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) diff --git a/areal/engine/core/distributed.py b/areal/engine/core/distributed.py index 2b80a5cdf4..2a1b1fa14f 100644 --- a/areal/engine/core/distributed.py +++ b/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 diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index 7beb16ef69..0e218bfc02 100644 --- a/areal/engine/megatron_engine.py +++ b/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) diff --git a/areal/infra/controller/train_controller.py b/areal/infra/controller/train_controller.py index 45f1c4b340..75876ce519 100644 --- a/areal/infra/controller/train_controller.py +++ b/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. diff --git a/areal/utils/recover.py b/areal/utils/recover.py index 2ff898ea16..e22f4131e1 100644 --- a/areal/utils/recover.py +++ b/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, diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index 2ab6520d6f..4fc16c5fe7 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -367,6 +367,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | | `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | +| `warmup_communicators` | boolean | `False` | Pre-connect collective transports after a recover load. | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | @@ -443,6 +444,7 @@ Configuration for PPO critic model, a subclass of a TrainEngine. | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | | `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | +| `warmup_communicators` | boolean | `False` | Pre-connect collective transports after a recover load. | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | @@ -492,6 +494,7 @@ Core configuration for model training, including optimization and backend settin | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | | `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | +| `warmup_communicators` | boolean | `False` | Pre-connect collective transports after a recover load. | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | @@ -1010,6 +1013,7 @@ fields. | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | | `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | +| `warmup_communicators` | boolean | `False` | Pre-connect collective transports after a recover load. | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | diff --git a/docs/zh/cli_reference.md b/docs/zh/cli_reference.md index 8214bce1b1..82b86268c6 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -365,6 +365,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | | `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | +| `warmup_communicators` | boolean | `False` | Pre-connect collective transports after a recover load. | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | @@ -441,6 +442,7 @@ Configuration for PPO critic model, a subclass of a TrainEngine. | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | | `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | +| `warmup_communicators` | boolean | `False` | Pre-connect collective transports after a recover load. | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | @@ -490,6 +492,7 @@ Core configuration for model training, including optimization and backend settin | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | | `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | +| `warmup_communicators` | boolean | `False` | Pre-connect collective transports after a recover load. | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | @@ -1008,6 +1011,7 @@ fields. | `optimizer_dtype` | string | `"float32"` | Underlying parameter storage dtype, also the dtype of optimizer states (exp_avg, exp_avg_sq) since torch.optim.AdamW inherits dtype from model.parameters(). Default 'float32' maintains fp32 master weights matching DeepSpeed ZeRO-3 and Megatron precision-aware optimizer behavior. FSDP2's MixedPrecisionPolicy(param_dtype=`dtype`) will still cast forward/backward computation to `dtype` (e.g. bfloat16). Set to 'bfloat16' together with optimizer.type='adam_bf16' to reduce memory at the cost of needing Kahan summation for stability. Currently FSDP-only; Megatron uses use_precision_aware_optimizer instead and ignores this field. | | `optimizer` | [`OptimizerConfig`](section-optimizer) \| None | `None` | Optimizer configuration. None means no training. | | `weight_update_mode` | string | `"xccl"` | Weight update backend type. 'awex' requires a Megatron actor and an SGLang rollout, and targets colocated actor-rollout setups. **Choices:** `disk`, `xccl`, `awex` | +| `warmup_communicators` | boolean | `False` | Pre-connect collective transports after a recover load. | | `fsdp` | [`FSDPEngineConfig`](section-fsdp-engine) | **Required** | - | | `archon` | [`ArchonEngineConfig`](section-archon-engine) | **Required** | - | | `megatron` | [`MegatronEngineConfig`](section-megatron-engine) | **Required** | - | diff --git a/tests/test_comm_warmup_helpers.py b/tests/test_comm_warmup_helpers.py new file mode 100644 index 0000000000..aed0bb9365 --- /dev/null +++ b/tests/test_comm_warmup_helpers.py @@ -0,0 +1,216 @@ +"""Unit tests for the transport warmup helpers. + +Only metadata and tensor shapes are asserted here; the collectives are +stubbed, so these run without a distributed environment or a GPU. +""" + +from __future__ import annotations + +from contextlib import ExitStack, contextmanager +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.distributed as dist + +from areal.engine.core.distributed import ( + nccl_process_groups, + warmup_all_to_all_transports, + warmup_collective_transports, + warmup_p2p_transports, + warmup_sharded_transports, +) + + +class _FakeGroup: + def __init__(self, name: str) -> None: + self._name = name + + def __repr__(self) -> str: + return f"" + + +@pytest.fixture(autouse=True) +def no_local_rank(monkeypatch): + monkeypatch.delenv("LOCAL_RANK", raising=False) + + +@contextmanager +def _registry(*entries: tuple[_FakeGroup, str, int]): + backends = {group: backend for group, backend, _ in entries} + sizes = {group: size for group, _, size in entries} + world = SimpleNamespace(pg_map={group: (backends[group], None) for group in sizes}) + with ExitStack() as stack: + stack.enter_context(patch("torch.distributed.distributed_c10d._world", world)) + stack.enter_context(patch.object(dist, "is_initialized", return_value=True)) + stack.enter_context( + patch.object(dist, "get_backend", side_effect=lambda group: backends[group]) + ) + stack.enter_context( + patch.object(dist, "get_world_size", side_effect=lambda group: sizes[group]) + ) + yield + + +@contextmanager +def _cuda_platform(device_type="cuda", world_size=4): + with ExitStack() as stack: + platform = stack.enter_context(patch("areal.infra.platforms.current_platform")) + platform.device_type = device_type + platform.current_device.return_value = 0 + stack.enter_context(patch.object(dist, "is_initialized", return_value=True)) + stack.enter_context( + patch.object(dist, "get_world_size", return_value=world_size) + ) + yield stack + + +def test_enumeration_is_empty_before_init(): + with patch.object(dist, "is_initialized", return_value=False): + assert nccl_process_groups() == [] + + +def test_enumeration_skips_gloo_and_solo_groups(): + nccl = _FakeGroup("nccl") + gloo = _FakeGroup("gloo") + solo = _FakeGroup("solo") + with _registry((nccl, "nccl", 8), (gloo, "gloo", 8), (solo, "nccl", 1)): + assert nccl_process_groups() == [nccl] + + +def test_enumeration_preserves_creation_order(): + first = _FakeGroup("first") + second = _FakeGroup("second") + with _registry((first, "nccl", 4), (second, "nccl", 4)): + assert nccl_process_groups() == [first, second] + + +def test_collective_warmup_is_a_noop_on_cpu(): + with _cuda_platform(device_type="cpu") as stack: + all_reduce = stack.enter_context(patch.object(dist, "all_reduce")) + warmup_collective_transports(_FakeGroup("dp")) + all_reduce.assert_not_called() + + +def test_collective_warmup_dedupes_and_covers_every_probe_size(): + group = _FakeGroup("dp") + with _cuda_platform() as stack: + stack.enter_context(patch("torch.zeros", return_value=MagicMock())) + all_reduce = stack.enter_context(patch.object(dist, "all_reduce")) + warmup_collective_transports(group, None, group) + + assert all_reduce.call_count == 2 + assert {c.kwargs["group"] for c in all_reduce.call_args_list} == {group} + + +def test_all_to_all_warmup_truncates_to_a_multiple_of_the_world_size(): + group = _FakeGroup("ep") + sizes = [] + with _cuda_platform(world_size=3) as stack: + stack.enter_context( + patch( + "torch.zeros", + side_effect=lambda n, **kw: sizes.append(n) or MagicMock(), + ) + ) + stack.enter_context(patch("torch.empty_like", return_value=MagicMock())) + stack.enter_context(patch.object(dist, "all_to_all_single")) + warmup_all_to_all_transports(group) + + assert sizes + assert all(size % 3 == 0 for size in sizes) + + +def test_all_to_all_warmup_skips_probes_smaller_than_the_world_size(): + group = _FakeGroup("ep") + with _cuda_platform(world_size=4096) as stack: + stack.enter_context(patch("torch.zeros", return_value=MagicMock())) + stack.enter_context(patch("torch.empty_like", return_value=MagicMock())) + all_to_all = stack.enter_context(patch.object(dist, "all_to_all_single")) + warmup_all_to_all_transports(group, numels=(1024,)) + + all_to_all.assert_not_called() + + +def test_sharded_warmup_is_a_noop_without_a_group(): + with patch.object(dist, "reduce_scatter_tensor") as reduce_scatter: + warmup_sharded_transports(None) + reduce_scatter.assert_not_called() + + +def test_sharded_warmup_is_a_noop_for_a_single_rank_group(): + with _cuda_platform(world_size=1) as stack: + reduce_scatter = stack.enter_context( + patch.object(dist, "reduce_scatter_tensor") + ) + warmup_sharded_transports(_FakeGroup("dp")) + reduce_scatter.assert_not_called() + + +def test_sharded_warmup_shapes_the_shard_to_the_world_size(): + shapes = [] + with _cuda_platform(world_size=4) as stack: + stack.enter_context( + patch( + "torch.zeros", + side_effect=lambda n, **kw: shapes.append(n) or MagicMock(), + ) + ) + stack.enter_context( + patch( + "torch.empty", + side_effect=lambda n, **kw: shapes.append(n) or MagicMock(), + ) + ) + stack.enter_context(patch.object(dist, "reduce_scatter_tensor")) + stack.enter_context(patch.object(dist, "all_gather_into_tensor")) + warmup_sharded_transports(_FakeGroup("dp"), numel=1024) + + assert shapes == [1024, 256] + + +def test_p2p_warmup_is_a_noop_for_a_single_stage_pipeline(): + with _cuda_platform() as stack: + isend = stack.enter_context(patch.object(dist, "isend")) + warmup_p2p_transports( + _FakeGroup("pp"), + prev_rank=0, + next_rank=0, + has_prev=False, + has_next=False, + ) + isend.assert_not_called() + + +def test_p2p_warmup_passes_the_group_explicitly(): + group = _FakeGroup("pp") + with _cuda_platform() as stack: + stack.enter_context(patch("torch.zeros", return_value=MagicMock())) + stack.enter_context(patch("torch.empty", return_value=MagicMock())) + stack.enter_context(patch.object(dist, "batch_isend_irecv", return_value=[])) + stack.enter_context(patch.object(dist, "irecv", return_value=MagicMock())) + stack.enter_context(patch.object(dist, "P2POp", return_value=MagicMock())) + isend = stack.enter_context( + patch.object(dist, "isend", return_value=MagicMock()) + ) + warmup_p2p_transports( + group, + prev_rank=0, + next_rank=2, + has_prev=True, + has_next=True, + ) + + assert isend.call_count + assert {c.kwargs["group"] for c in isend.call_args_list} == {group} + + +def test_torch_device_is_built_from_the_platform_device_type(): + group = _FakeGroup("dp") + with _cuda_platform() as stack: + zeros = stack.enter_context(patch("torch.zeros", return_value=MagicMock())) + stack.enter_context(patch.object(dist, "all_reduce")) + warmup_collective_transports(group, numels=(8,)) + + assert zeros.call_args.kwargs["device"] == torch.device("cuda", 0) diff --git a/tests/test_megatron_comm_warmup.py b/tests/test_megatron_comm_warmup.py new file mode 100644 index 0000000000..19bd7a54c6 --- /dev/null +++ b/tests/test_megatron_comm_warmup.py @@ -0,0 +1,196 @@ +"""Unit tests for the communicator warmup's device and group selection. + +The collectives themselves need a distributed environment, but the platform +guards and the choice of target group can be exercised on CPU by driving the +unbound method with a stand-in engine. +""" + +from __future__ import annotations + +from contextlib import ExitStack +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.distributed as dist + +from areal.engine.megatron_engine import MegatronEngine + +_WORLD = object() + + +@pytest.fixture +def engine(): + return SimpleNamespace( + device=torch.device("cuda", 3), + cpu_group=object(), + logger=MagicMock(), + config=SimpleNamespace(warmup_communicators=True), + ) + + +@pytest.fixture(autouse=True) +def no_local_rank(monkeypatch): + monkeypatch.delenv("LOCAL_RANK", raising=False) + + +def _topology(pp_size, pp_group, ep_size=1): + mpu = MagicMock() + mpu.get_expert_model_parallel_world_size.return_value = ep_size + mpu.get_pipeline_model_parallel_world_size.return_value = pp_size + mpu.get_pipeline_model_parallel_group.return_value = pp_group + mpu.get_pipeline_model_parallel_rank.return_value = 1 + mpu.get_pipeline_model_parallel_prev_rank.return_value = 0 + mpu.get_pipeline_model_parallel_next_rank.return_value = 2 + mpu.is_pipeline_first_stage.return_value = False + mpu.is_pipeline_last_stage.return_value = False + return mpu + + +class _P2PRecorder: + def __init__(self): + self.groups = [] + + def send_like(self, tensor, peer, group=None, *args, **kwargs): + self.groups.append(group) + return MagicMock() + + def p2p_op(self, op, tensor, peer, group=None, *args, **kwargs): + self.groups.append(group) + return MagicMock() + + +def _platform_patches(stack, device_type): + for target in ( + "areal.engine.megatron_engine.current_platform", + "areal.infra.platforms.current_platform", + ): + platform = stack.enter_context(patch(target)) + platform.device_type = device_type + platform.current_device.return_value = 3 + + +def _run_warmup(engine, pp_size, pp_group, backend="nccl", ep_size=1): + recorder = _P2PRecorder() + patches = [ + patch( + "areal.engine.megatron_engine.mpu", _topology(pp_size, pp_group, ep_size) + ), + patch( + "torch.distributed.distributed_c10d._world", + SimpleNamespace(pg_map={}, default_pg=_WORLD), + ), + patch.object(dist, "is_initialized", return_value=True), + patch.object(dist, "get_backend", return_value=backend), + patch.object(dist, "get_world_size", return_value=pp_size), + patch.object(dist, "get_rank", return_value=1), + patch.object(dist, "barrier"), + patch.object(dist, "all_reduce"), + patch.object(dist, "all_to_all_single"), + patch.object(dist, "reduce_scatter_tensor"), + patch.object(dist, "all_gather_into_tensor"), + patch.object(dist, "batch_isend_irecv", return_value=[]), + patch.object(dist, "send", side_effect=recorder.send_like), + patch.object(dist, "recv", side_effect=recorder.send_like), + patch.object(dist, "isend", side_effect=recorder.send_like), + patch.object(dist, "irecv", side_effect=recorder.send_like), + patch.object(dist, "P2POp", side_effect=recorder.p2p_op), + patch("torch.zeros", return_value=MagicMock()), + patch("torch.empty", return_value=MagicMock()), + patch("torch.empty_like", return_value=MagicMock()), + ] + with ExitStack() as stack: + _platform_patches(stack, "cuda") + for p in patches: + stack.enter_context(p) + MegatronEngine.warmup_communicators(engine) + return recorder + + +def test_warmup_is_skipped_when_not_enabled(engine): + engine.config.warmup_communicators = False + with ExitStack() as stack: + _platform_patches(stack, "cuda") + stack.enter_context(patch.object(dist, "is_initialized", return_value=True)) + all_reduce = stack.enter_context(patch.object(dist, "all_reduce")) + MegatronEngine.warmup_communicators(engine) + + all_reduce.assert_not_called() + + +def test_warmup_is_skipped_on_cpu_platforms(engine): + with ExitStack() as stack: + _platform_patches(stack, "cpu") + stack.enter_context(patch.object(dist, "is_initialized", return_value=True)) + all_reduce = stack.enter_context(patch.object(dist, "all_reduce")) + isend = stack.enter_context(patch.object(dist, "isend")) + MegatronEngine.warmup_communicators(engine) + + all_reduce.assert_not_called() + isend.assert_not_called() + + +def test_pipeline_warmup_targets_the_pipeline_group(engine): + """Megatron routes unbatched pipeline p2p through the pipeline group. + + Warming the default group instead leaves the communicator that the train + step actually uses cold, which is the allocation this warmup exists to + front-load. + """ + pp_group = object() + + recorder = _run_warmup(engine, pp_size=4, pp_group=pp_group) + + assert recorder.groups + assert set(recorder.groups) == {pp_group} + + +def test_pipeline_warmup_covers_the_world_group_when_pp_size_is_two(engine): + """At pp_size == 2 Megatron sends one direction over the default group.""" + pp_group = object() + + recorder = _run_warmup(engine, pp_size=2, pp_group=pp_group) + + assert set(recorder.groups) == {pp_group, _WORLD} + + +def test_pipeline_warmup_stays_on_the_pipeline_group_for_ucc(engine): + pp_group = object() + + recorder = _run_warmup(engine, pp_size=2, pp_group=pp_group, backend="ucc") + + assert set(recorder.groups) == {pp_group} + + +def test_no_pipeline_warmup_without_pipeline_parallelism(engine): + recorder = _run_warmup(engine, pp_size=1, pp_group=object()) + + assert recorder.groups == [] + + +def test_all_to_all_is_warmed_only_when_experts_are_sharded(engine): + """all_to_all buffers persist, so a group that never dispatches pays for nothing.""" + with ExitStack() as stack: + _platform_patches(stack, "cuda") + stack.enter_context( + patch("areal.engine.megatron_engine.mpu", _topology(1, object(), ep_size=1)) + ) + stack.enter_context( + patch( + "torch.distributed.distributed_c10d._world", + SimpleNamespace(pg_map={}, default_pg=_WORLD), + ) + ) + stack.enter_context(patch.object(dist, "is_initialized", return_value=True)) + stack.enter_context(patch.object(dist, "get_world_size", return_value=4)) + stack.enter_context(patch.object(dist, "barrier")) + stack.enter_context(patch.object(dist, "reduce_scatter_tensor")) + stack.enter_context(patch.object(dist, "all_gather_into_tensor")) + stack.enter_context(patch("torch.zeros", return_value=MagicMock())) + stack.enter_context(patch("torch.empty", return_value=MagicMock())) + stack.enter_context(patch("torch.empty_like", return_value=MagicMock())) + all_to_all = stack.enter_context(patch.object(dist, "all_to_all_single")) + MegatronEngine.warmup_communicators(engine) + + all_to_all.assert_not_called()