diff --git a/atom/config.py b/atom/config.py index 451bb05a43..1eba025d8e 100644 --- a/atom/config.py +++ b/atom/config.py @@ -9,6 +9,7 @@ import re from contextlib import contextmanager from dataclasses import dataclass, field, fields +from datetime import timedelta from typing import TYPE_CHECKING, Any, ClassVar import torch @@ -833,6 +834,14 @@ class ParallelConfig: data_parallel_master_ip: str = "127.0.0.1" + dp_sync_timeout_s: float | None = None + """Bound on the DP group's rendezvous and collectives. None keeps the + backend default (~30 min for gloo), which is right on one node where the + peers are sibling processes. Across nodes the common failure -- the other + node never started, or has the wrong master IP -- presents as every barrier + waiting forever with nothing logged, so a shorter bound is what makes it + diagnosable.""" + @property def is_multinode_dp(self) -> bool: """Whether this node owns only part of the global DP group. @@ -857,9 +866,14 @@ def get_next_dp_init_port(self) -> int: can live in different processes. To avoid port conflicts, we pop a new port from the prepared port list each time we need to initialize a new process group related to data parallelism. + + Every rank has to walk the same sequence, or the second group is + formed on a different port per rank and never completes. Advancing by + data_parallel_rank did not: rank 0 stayed on the base port while the + others each moved by a different amount. """ answer = self.data_parallel_master_port - self.data_parallel_master_port += self.data_parallel_rank + self.data_parallel_master_port += 1 return answer @@ -877,9 +891,17 @@ def stateless_init_dp_group(self): self.data_parallel_rank, self.data_parallel_size, backend="gloo", + timeout=self.dp_sync_timeout, ) return dp_group + @property + def dp_sync_timeout(self) -> timedelta | None: + """``dp_sync_timeout_s`` as a timedelta; None means backend default.""" + if self.dp_sync_timeout_s is None: + return None + return timedelta(seconds=self.dp_sync_timeout_s) + @staticmethod def has_unfinished_dp(dp_group: ProcessGroup, has_unfinished: bool) -> bool: tensor = torch.tensor([has_unfinished], dtype=torch.int32, device="cpu") @@ -950,6 +972,16 @@ def __post_init__(self) -> None: f"({self.data_parallel_size}): this node's slice would run off " f"the end of the global DP group" ) + if self.data_parallel_size % self.data_parallel_size_local: + raise ValueError( + f"data_parallel_size ({self.data_parallel_size}) must be " + f"divisible by data_parallel_size_local " + f"({self.data_parallel_size_local}): nodes have to hold equal " + f"slices. EPLB's hierarchical placement asserts " + f"num_gpus % num_nodes == 0 (model_ops/eplb.py), so a ragged " + f"split starts fine and then fails inside a rebalance, " + f"thousands of steps later" + ) _DSPARK_DEFAULT_MAX_BLOCK = 16 @@ -1381,6 +1413,15 @@ class EPLBConfig: # "biased" -> fully replicate top-K hottest experts to all GPUs # (K = num_redundant // num_gpus, per-node in multi-node) placement_policy: str = "naive" + # Cross-node migration guard. A rearrange whose inter-node weight movement + # would stall longer than this is skipped. 0 = no cap, count and log only, + # which is the default until M0-B measures the link: a made-up threshold + # that silently stops EPLB from ever rebalancing is worse than no cap. + # Single-node runs move zero cross-node bytes, so the cap never fires there. + max_migration_stall_ms: float = 0.0 + # Assumed inter-node bandwidth the guard converts bytes into time with. + # Calibrate at M0-B; only meaningful when max_migration_stall_ms > 0. + migration_bandwidth_gbps: float = 50.0 def __post_init__(self): self.load_window_size = int(self.load_window_size) @@ -1413,6 +1454,14 @@ def __post_init__(self): "naive", "biased", }, "eplb.placement_policy must be one of {'naive','biased'}" + self.max_migration_stall_ms = float(self.max_migration_stall_ms) + assert ( + self.max_migration_stall_ms >= 0 + ), "eplb.max_migration_stall_ms must be >= 0 (0 disables the cap)" + self.migration_bandwidth_gbps = float(self.migration_bandwidth_gbps) + assert ( + self.migration_bandwidth_gbps > 0 + ), "eplb.migration_bandwidth_gbps must be > 0" @classmethod def from_dict(cls, cfg: dict | None) -> "EPLBConfig": @@ -1977,6 +2026,131 @@ def __post_init__(self): elif self.index_cache_dtype is None: self.index_cache_dtype = self.kv_cache_dtype + self._validate_multinode() + + def _validate_multinode(self) -> None: + """Reject what multi-node DP has not been shown to work with. + + Multi-node narrows the configuration surface rather than adding an + orthogonal dimension: most parallelism knobs reshape the rank space, + and none of the combinations below have run across a node boundary. + Deny by default and widen each as it is verified. The failure being + avoided is not a crash -- it is a run that starts, serves, and is + quietly wrong or hangs with no message. + """ + pc = self.parallel_config + if not pc.is_multinode_dp: + return + + nnodes = pc.data_parallel_size // pc.data_parallel_size_local + + def reject(what: str, why: str, fix: str) -> None: + raise ValueError( + f"Multi-node DP ({nnodes} nodes, " + f"data_parallel_size={pc.data_parallel_size} " + f"data_parallel_size_local={pc.data_parallel_size_local}) " + f"{what}\n Why: {why}\n Fix: {fix}" + ) + + if not self.enable_dp_attention: + reject( + "requires --enable-dp-attention.", + "WideEP grows the EP group through the DP dimension. TP across " + "nodes would replicate the MLA KV cache and put an all-reduce " + "on every layer's critical path over the slower link.", + "add --enable-dp-attention, or run independent single-node " + "instances behind a router.", + ) + if not self.enable_expert_parallel: + reject( + "requires --enable-expert-parallel.", + "spanning nodes only pays for itself by sharding experts " + "wider; without EP each node holds a full copy and the " + "inter-node link buys nothing.", + "add --enable-expert-parallel, or run independent single-node " + "instances.", + ) + if self.pipeline_parallel_size != 1: + reject( + f"cannot be combined with pipeline_parallel_size=" + f"{self.pipeline_parallel_size}.", + "the engine index space folds PP stages into DP ranks, and " + "that mapping has not been reconciled across nodes.", + "set --pipeline-parallel-size 1.", + ) + if self.prefill_context_parallel_size != 1: + reject( + f"cannot be combined with prefill_context_parallel_size=" + f"{self.prefill_context_parallel_size}.", + "ATOM_PCP_MOE_MERGE folds PCP into the EP flatten " + "(model_ops/moe.py), which stacks on the multi-node rank " + "space in a way nothing has exercised.", + "set --prefill-context-parallel-size 1.", + ) + if pc.decode_context_parallel_size != 1: + reject( + f"cannot be combined with decode_context_parallel_size=" + f"{pc.decode_context_parallel_size}.", + "same unexercised interaction with the EP flatten as PCP.", + "set --decode-context-parallel-size 1.", + ) + if self.moe_backend != "standard": + reject( + f"cannot be combined with --moe-backend {self.moe_backend}.", + "FlyDSLAll2AllManager raises NotImplementedError for internode " + "(aiter dist/device_communicators/all2all.py); the mega path " + "has no cross-node transport.", + "use --moe-backend standard.", + ) + if self.enable_rapidserve: + reject( + "cannot be combined with --enable-rapidserve.", + "rapidserve selects DisaggCoreManager (model_engine/" + "llm_engine.py), which spawns engines and assigns addresses " + "through its own path -- a second, unreconciled rank layout.", + "drop --enable-rapidserve, or disaggregate across instances " + "instead of within a GPU.", + ) + if self.plugin_config is not None: + reject( + "cannot be combined with the SGLang/vLLM plugin path.", + "the plugin brings its own multi-node world (SGLang's " + "--nnodes), and plugin/config.py assigns " + "data_parallel_size_local the global size, so gpu_per_node " + "comes out wrong on that path.", + "use the native path, or drive multi-node through the " + "plugin's own topology flags.", + ) + + if self.enable_tbo: + logger.warning( + "Multi-node DP with --enable-tbo: TBO builds its own MoRI op " + "and is not verified across nodes. Expect to validate it " + "yourself before trusting the numbers." + ) + + base_port_default = next( + f.default for f in fields(ParallelConfig) if f.name == "data_parallel_base_port" + ) + if pc.data_parallel_base_port == base_port_default: + reject( + "requires an explicit --data-parallel-base-port.", + "the default is get_open_port(), evaluated separately in each " + "process, so the nodes pick different ports and every rank " + "waits on an address no one is bound to -- a silent hang, not " + "a connection error.", + "pass the same --data-parallel-base-port on every node.", + ) + if pc.data_parallel_master_ip in ("127.0.0.1", "localhost", "::1"): + logger.warning( + "Multi-node DP with data_parallel_master_ip=%s: this only " + "resolves if every 'node' is this machine, i.e. the " + "single-box logical-node setup. Across real nodes the " + "non-coordinator ranks connect to their own loopback and hang " + "with no message. Pass --data-parallel-master-ip .", + pc.data_parallel_master_ip, + ) + def compute_hash(self) -> str: """ WARNING: Whenever a new field is added to this config, diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index f6ef9af709..dbacc08452 100644 --- a/atom/model_engine/arg_utils.py +++ b/atom/model_engine/arg_utils.py @@ -47,6 +47,7 @@ class EngineArgs: data_parallel_master_ip: str = "127.0.0.1" data_parallel_master_port: int = 29500 data_parallel_base_port: int | None = None + dp_sync_timeout_s: float | None = None enforce_eager: bool = False enable_prefix_caching: bool = True port: int = 8006 @@ -184,6 +185,18 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: "explicitly for multi-node launches." ), ) + parser.add_argument( + "--dp-sync-timeout", + type=float, + default=None, + dest="dp_sync_timeout_s", + help=( + "Seconds to wait on DP rendezvous and collectives before " + "failing. Default is the backend's (~30 min for gloo). Set it " + "for multi-node: a peer that never starts otherwise blocks " + "every barrier indefinitely with nothing logged." + ), + ) parser.add_argument( "--decode-context-parallel-size", "-dcp", @@ -688,6 +701,7 @@ def _get_engine_kwargs(self) -> dict: base_port = kwargs.pop("data_parallel_base_port") if base_port is not None: parallel_config_kwargs["data_parallel_base_port"] = base_port + parallel_config_kwargs["dp_sync_timeout_s"] = kwargs.pop("dp_sync_timeout_s") kwargs["parallel_config"] = ParallelConfig(**parallel_config_kwargs) logger.info(f"Engine kwargs: {kwargs}") diff --git a/atom/model_engine/engine_core_mgr.py b/atom/model_engine/engine_core_mgr.py index 67003b5e06..7e85773b9f 100644 --- a/atom/model_engine/engine_core_mgr.py +++ b/atom/model_engine/engine_core_mgr.py @@ -8,6 +8,7 @@ import os import pickle import queue +import time import weakref from dataclasses import dataclass from threading import Lock, Thread @@ -436,7 +437,12 @@ def __init__(self, config: Config): self.output_sockets.append(output_socket) self.shutdown_paths.append(get_open_zmq_inproc_path()) - self._wait_for_all_ready_signals() + self._wait_for_all_ready_signals( + timeout_s=( + pc.dp_sync_timeout_s if multinode else None + ), + engines_per_node=local_engine_count, + ) logger.info( f"{self.label}: All EngineCores are fully initialized and ready" ) @@ -477,8 +483,17 @@ def __init__(self, config: Config): self._output_handler_task = None self._asyncio_mode = config.asyncio_mode - def _wait_for_all_ready_signals(self): - """Wait for READY signals from all DP ranks in parallel (no timeout).""" + def _wait_for_all_ready_signals( + self, *, timeout_s: float | None = None, engines_per_node: int | None = None + ): + """Wait for READY from every engine. + + Unbounded is right on one node: the engines are this process's children + and their death shows up other ways. Across nodes the usual failure is + that the other node was never started, or was pointed at a different + master, and it presents here as a wait that never ends and never logs a + thing -- the hardest shape of failure to attribute. + """ poller = zmq.Poller() for dp_rank, output_socket in enumerate(self.output_sockets): poller.register(output_socket, zmq.POLLIN) @@ -486,9 +501,20 @@ def _wait_for_all_ready_signals(self): engine_count = len(self.output_sockets) ready_received = [False] * engine_count remaining = engine_count + deadline = None if timeout_s is None else time.monotonic() + timeout_s while remaining > 0: - socks = poller.poll() # Wait indefinitely + if deadline is None: + socks = poller.poll() # Wait indefinitely + else: + left_ms = int((deadline - time.monotonic()) * 1000) + socks = poller.poll(timeout=max(0, left_ms)) + if not socks: + raise TimeoutError( + self._ready_timeout_message( + ready_received, timeout_s, engines_per_node + ) + ) if not socks: continue @@ -516,6 +542,38 @@ def _wait_for_all_ready_signals(self): f"{self.label}: Expected READY signal from DP rank {dp_rank}, but got {request_type}" ) + def _ready_timeout_message( + self, + ready_received: list[bool], + timeout_s: float, + engines_per_node: int | None, + ) -> str: + """Name the engines that never reported, grouped by the node they sit on. + + Which node is missing is the whole diagnosis here, and the coordinator + is the only process that can see it -- the engines that did come up are + blocked on the same barrier and have nothing to say. + """ + missing = [i for i, ok in enumerate(ready_received) if not ok] + if engines_per_node: + by_node: dict[int, list[int]] = {} + for rank in missing: + by_node.setdefault(rank // engines_per_node, []).append(rank) + where = "; ".join( + f"node {node}: dp ranks {ranks}" for node, ranks in sorted(by_node.items()) + ) + else: + where = f"dp ranks {missing}" + return ( + f"{self.label}: {len(missing)} of {len(ready_received)} engines did " + f"not report READY within {timeout_s:g}s ({where}).\n" + f" Why: every engine has to join before the group forms, so one " + f"absent node blocks all of them.\n" + f" Fix: check that node's launcher started, that it was given the " + f"same --data-parallel-master-ip/--data-parallel-base-port as this " + f"one, and that its log does not end in a weight-loading stall." + ) + def _create_output_thread( self, dp_rank: int, output_socket: zmq.Socket, shutdown_path: str ) -> Thread: diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index acdf39cdbe..3350c9be43 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -936,6 +936,83 @@ def _setup_device_and_distributed(self, rank: int, config: Config): config.parallel_config.data_parallel_master_ip, config.parallel_config.data_parallel_base_port, ) + # Before init_dist_env, not after: MoRI reads its heap settings when the + # heap is created, inside the first collective. + from atom.model_engine.topology import ( + WideEPTopology, + format_startup_summary, + node_count, + ) + from atom.model_ops.fused_moe.mori_env import apply_mori_env + + nnodes = node_count(config) + mori_env = apply_mori_env(nnodes=nnodes) + + # weight_iterator splits the prefetch by node, because page cache is + # per node. It reads these, and without them falls back to (0, 1) -- + # every rank prefetching the whole checkpoint -- as soon as the world + # is wider than the visible devices, which is every multi-node run. + # An external launcher that already set them knows better than we do. + if nnodes > 1: + topo = WideEPTopology.from_config(config) + os.environ.setdefault("LOCAL_RANK", str(local_device_rank)) + os.environ.setdefault("LOCAL_WORLD_SIZE", str(topo.gpu_per_node)) + if nnodes > 1: + logger.info( + format_startup_summary( + WideEPTopology.from_config(config), + dp_rank=config.parallel_config.data_parallel_rank, + dp_rank_local=dp_rank_local, + device_rank=local_device_rank, + visible_devices=os.environ.get("HIP_VISIBLE_DEVICES") + or os.environ.get("CUDA_VISIBLE_DEVICES"), + mori_env=mori_env, + ) + ) + + # aiter's init_dist_env calls init_process_group with no timeout, so the + # NCCL world is the one barrier we cannot bound from the call site. + # init_distributed_environment skips creation when a world already + # exists (aiter dist/parallel_state.py), so build it here instead, with + # the same rank arithmetic it would have used. + # + # Opt-in: multi-node AND an explicit --dp-sync-timeout. Nothing about + # the single-node path changes, and a multi-node run that did not ask + # for a bound keeps today's behaviour rather than silently taking a new + # init path that no single-node test covers. + if ( + nnodes > 1 + and config.parallel_config.dp_sync_timeout is not None + and pp_size == 1 + and not torch.distributed.is_initialized() + ): + global_world = config.parallel_config.data_parallel_size * stage_span + global_rank = ( + config.parallel_config.data_parallel_rank * stage_span + rank + ) + if "HIP_VISIBLE_DEVICES" not in os.environ: + # aiter would have set it to range(world_size) here, which is + # wrong past one node -- a 16-rank world does not mean 16 local + # devices. Creating the world ourselves skips that, so say so. + logger.warning( + "HIP_VISIBLE_DEVICES is unset on a %d-node run; the " + "launcher should pin this node's devices explicitly.", + nnodes, + ) + logger.info( + "[wideep] pre-creating the torch world: rank=%d/%d timeout=%s", + global_rank, + global_world, + config.parallel_config.dp_sync_timeout, + ) + torch.distributed.init_process_group( + backend="nccl", + init_method=distributed_init_method, + world_size=global_world, + rank=global_rank, + timeout=config.parallel_config.dp_sync_timeout, + ) + # Both branches handle simulated TP: the PP path only to reject it, # since it would otherwise deadlock on a group sized for absent ranks. if config.pipeline_parallel_size > 1: @@ -1701,14 +1778,41 @@ def get_num_blocks(self) -> dict[str, object]: # result back into the plan before publishing anything: the plan is the # single source for every entry count, so it must never disagree with # the number the pool is actually built at. + # + # Multi-node needs the same reduction for a different reason. Under + # DP-attention each rank owns its own BlockManager, so unequal counts do + # not corrupt anything -- the tighter rank simply preempts and recomputes + # sooner while the rest idle waiting for it. Nodes rarely have byte-equal + # free memory, and ATOM logs no preemption, so the only trace of this is + # a throughput number that is lower than it should be for no visible + # reason. Levelling costs the difference in blocks; not levelling costs + # a diagnosis. + from atom.model_engine.topology import node_count + num_kvcache_blocks = plan.paged_entries - if config.pipeline_parallel_size > 1 and torch.distributed.is_initialized(): + level_blocks = config.pipeline_parallel_size > 1 or node_count(config) > 1 + if level_blocks and torch.distributed.is_initialized(): t = torch.tensor( [num_kvcache_blocks], dtype=torch.int64, device=self.device ) + local_blocks = num_kvcache_blocks torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.MIN) num_kvcache_blocks = int(t.item()) plan = plan.with_paged_entries(num_kvcache_blocks) + if local_blocks != num_kvcache_blocks: + logger.warning( + "[wideep] KV pool levelled down: this rank sized %d blocks, " + "the group's minimum is %d. Every rank now runs at the " + "tightest one; %d blocks of this GPU go unused.", + local_blocks, + num_kvcache_blocks, + local_blocks - num_kvcache_blocks, + ) + else: + logger.info( + "[wideep] KV pool blocks=%d, equal to the group minimum", + num_kvcache_blocks, + ) block_bytes = plan.entry_bytes[plan.paged_class] # The whole plan travels to the engine process; BlockManager, the diff --git a/atom/model_engine/topology.py b/atom/model_engine/topology.py new file mode 100644 index 0000000000..9d045eb7ee --- /dev/null +++ b/atom/model_engine/topology.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""WideEP topology view (M-TOPO). + +Derived, immutable view of the engine/EP layout. Pure data + arithmetic; no +processes, sockets, CUDA, or environment reads. + +This does NOT replace the DP topology fields on ``ParallelConfig`` -- those stay +authoritative and ``CoreManager`` still rewrites them into engine units. What +this adds is a named, testable expression of the quantities the MoE and EPLB +layers need (``nnodes``, ``gpu_per_node``, ``ep_size``) plus assertions that +turn the rewrite's implicit invariants into explicit ones. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +def parse_dist_init_addr(addr: str) -> tuple[str, int]: + """Parse ``HOST:PORT`` or ``[IPv6]:PORT``.""" + addr = addr.strip() + if not addr: + raise ValueError("dist_init_addr must not be empty") + if addr.startswith("["): + end = addr.index("]") + host = addr[1:end] + rest = addr[end + 1 :] + if not rest.startswith(":"): + raise ValueError(f"Invalid dist_init_addr: {addr!r}") + port = int(rest[1:]) + else: + host, _, port_str = addr.rpartition(":") + if not host or not port_str: + raise ValueError(f"Invalid dist_init_addr: {addr!r}") + port = int(port_str) + if not (0 < port < 65536): + raise ValueError(f"Invalid port in dist_init_addr: {port}") + return host, port + + +def format_startup_summary( + topo: WideEPTopology, + *, + dp_rank: int, + dp_rank_local: int, + device_rank: int, + visible_devices: str | None = None, + mori_env: dict[str, str] | None = None, +) -> str: + """One greppable line per worker: is this rank where it thinks it is? + + Multi-node misconfiguration is diagnosed by comparing ranks across nodes, + which means reading the same fields out of several logs. Emitting them + together, once, on a line with a fixed prefix is the difference between + one grep and a correlation exercise. Everything here is already known at + worker init; nothing is computed for the message. + """ + parts = [ + topo.describe(), + f"rank={dp_rank} local_rank={dp_rank_local}", + f"device=cuda:{device_rank}", + ] + if visible_devices: + parts.append(f"visible={visible_devices}") + if mori_env: + parts.append( + "mori: " + " ".join(f"{k}={v}" for k, v in sorted(mori_env.items())) + ) + return " | ".join(parts) + + +def node_count(config) -> int: + """How many nodes this Config runs on. + + The one place that answers the question for the simulated case, so callers + that only need a node count do not each re-derive the rule: ``--fake-eplb`` + runs a single box pretending to be a wider deployment, and the DP fields + alone cannot distinguish that from being one node of a real split. + """ + if getattr(config, "dp_logical_size", 0): + return 1 + return WideEPTopology.from_config(config).nnodes + + +@dataclass(frozen=True) +class WideEPTopology: + """Engine-unit view of the parallel layout. + + Deliberately stores no raw ``-tp`` / ``-dp`` values: after CoreManager folds + TP into DP those are unrecoverable, and a second copy of an authoritative + field is how the two drift apart. + """ + + nnodes: int + node_rank: int + dp_attention: bool + tp_size: int + """TP width per engine. 1 under DP-attention, which folds TP into DP.""" + global_dp_size: int + """Engines across all nodes.""" + local_engine_count: int + """Engines on this node.""" + dist_init_host: str | None + dist_init_base_port: int | None + + @classmethod + def create( + cls, + *, + dp_attention: bool, + raw_tp_size: int, + raw_dp_size: int, + nnodes: int = 1, + node_rank: int = 0, + dist_init_addr: str | None = None, + ) -> WideEPTopology: + """Build from pre-fold CLI values (``-tp`` / ``-dp``).""" + tp_size = 1 if dp_attention else raw_tp_size + global_dp_size = raw_tp_size * raw_dp_size if dp_attention else raw_dp_size + if nnodes < 1: + raise ValueError(f"nnodes must be >= 1, got {nnodes}") + return cls._build( + nnodes=nnodes, + node_rank=node_rank, + dp_attention=dp_attention, + tp_size=tp_size, + global_dp_size=global_dp_size, + local_engine_count=global_dp_size // nnodes, + dist_init_addr=dist_init_addr, + ) + + @classmethod + def from_parallel_config( + cls, + parallel_config, + *, + tensor_parallel_size: int, + dp_attention: bool, + dp_logical_size: int = 0, + ) -> WideEPTopology: + """Derive from ``ParallelConfig``, before or after CoreManager's fold. + + Both regimes work because the fold is a change of units -- it scales the + DP quantities by ``tp_size`` and divides TP by the same -- so the ratio + and the product this reads are invariant under it. + + ``dp_logical_size`` (``Config.dp_logical_size``, non-zero only under + ``--fake-eplb``) must be passed through: while simulating, the DP fields + describe a deployment wider than the box, so their ratio is not a node + count. Refuse rather than report a node count that does not exist. + """ + if dp_logical_size: + raise ValueError( + f"cannot derive a node topology while simulating a " + f"{dp_logical_size}-wide deployment (--fake-eplb): the DP fields " + f"describe the simulated width, not this box. Callers that " + f"support simulation must handle it explicitly." + ) + global_dp_size = parallel_config.data_parallel_size + local_engine_count = parallel_config.data_parallel_size_local + rank_offset = parallel_config.data_parallel_rank + # CoreManager gives each spawned worker its global DP rank. Recover + # the node's slice offset before deriving node_rank; the topology + # object describes the node, not the individual worker. + rank_local = getattr(parallel_config, "data_parallel_rank_local", None) + if rank_local is not None: + rank_offset -= rank_local + if local_engine_count is None: + local_engine_count = global_dp_size + if dp_attention: + # Pre-fold the DP fields still count replicas, not engines. + if tensor_parallel_size > 1: + global_dp_size *= tensor_parallel_size + local_engine_count *= tensor_parallel_size + rank_offset *= tensor_parallel_size + tp_size = 1 + else: + tp_size = tensor_parallel_size + + if local_engine_count < 1: + raise ValueError( + f"data_parallel_size_local must be >= 1, got {local_engine_count}" + ) + # EPLB's hierarchical placement asserts num_gpus % num_nodes == 0 + # (model_ops/eplb.py), so a ragged split has to fail here rather than + # deep inside a rebalance. + if global_dp_size % local_engine_count != 0: + raise ValueError( + f"data_parallel_size ({global_dp_size}) must be divisible by " + f"data_parallel_size_local ({local_engine_count}): nodes must " + f"hold equal slices for EPLB's hierarchical placement" + ) + nnodes = global_dp_size // local_engine_count + if rank_offset % local_engine_count != 0: + raise ValueError( + f"data_parallel_rank ({rank_offset}) must be a multiple of " + f"data_parallel_size_local ({local_engine_count}): it names the " + f"first global rank of this node's slice" + ) + + host = getattr(parallel_config, "data_parallel_master_ip", None) + port = getattr(parallel_config, "data_parallel_master_port", None) + return cls._build( + nnodes=nnodes, + node_rank=rank_offset // local_engine_count, + dp_attention=dp_attention, + tp_size=tp_size, + global_dp_size=global_dp_size, + local_engine_count=local_engine_count, + dist_init_addr=f"{host}:{port}" if nnodes > 1 and host else None, + ) + + @classmethod + def from_config(cls, config) -> WideEPTopology: + """Same as ``from_parallel_config``, reading the fields off a Config.""" + return cls.from_parallel_config( + config.parallel_config, + tensor_parallel_size=config.tensor_parallel_size, + dp_attention=config.enable_dp_attention, + dp_logical_size=getattr(config, "dp_logical_size", 0), + ) + + @classmethod + def _build( + cls, + *, + nnodes: int, + node_rank: int, + dp_attention: bool, + tp_size: int, + global_dp_size: int, + local_engine_count: int, + dist_init_addr: str | None, + ) -> WideEPTopology: + host: str | None = None + base_port: int | None = None + if nnodes > 1: + if dist_init_addr is None: + raise ValueError( + "nnodes>1 requires a rendezvous address " + "(dist_init_addr, or data_parallel_master_ip/_port)" + ) + host, base_port = parse_dist_init_addr(dist_init_addr) + topo = cls( + nnodes=nnodes, + node_rank=node_rank, + dp_attention=dp_attention, + tp_size=tp_size, + global_dp_size=global_dp_size, + local_engine_count=local_engine_count, + dist_init_host=host, + dist_init_base_port=base_port, + ) + topo._validate() + return topo + + def _validate(self) -> None: + if self.nnodes < 1: + raise ValueError(f"nnodes must be >= 1, got {self.nnodes}") + if not (0 <= self.node_rank < self.nnodes): + raise ValueError( + f"node_rank must satisfy 0 <= node_rank < nnodes " + f"({self.node_rank}, {self.nnodes})" + ) + if self.nnodes > 1 and not self.dp_attention: + raise ValueError("nnodes>1 requires dp_attention (TP does not span nodes)") + if self.local_engine_count < 1: + raise ValueError( + f"local_engine_count must be >= 1, got {self.local_engine_count}" + ) + if self.local_engine_count * self.nnodes != self.global_dp_size: + divisors = [ + n + for n in range(1, self.global_dp_size + 1) + if self.global_dp_size % n == 0 + ] + raise ValueError( + f"global_dp_size={self.global_dp_size} is not divisible by " + f"nnodes={self.nnodes}. Valid nnodes values: {divisors}" + ) + # Only DP-attention flattens every rank into one EP group; otherwise EP + # spans a TP group and there are global_dp_size independent groups, so + # the identity does not apply. + if self.dp_attention and self.ep_size != self.gpu_per_node * self.nnodes: + raise ValueError( + f"ep_size invariant violated: ep_size={self.ep_size} != " + f"gpu_per_node({self.gpu_per_node}) * nnodes({self.nnodes})" + ) + + @property + def ep_size(self) -> int: + """Ranks in one EP group (``FusedMoEParallelConfig.make``).""" + if self.dp_attention: + return self.global_dp_size * self.tp_size + return self.tp_size + + @property + def gpu_per_node(self) -> int: + """GPUs this node contributes; reaches MoRI as ``gpu_per_node``.""" + return self.local_engine_count * self.tp_size + + @property + def is_multinode(self) -> bool: + return self.nnodes > 1 + + def _require_rendezvous_base_port(self) -> int: + if self.dist_init_base_port is None: + raise ValueError( + "rendezvous ports require nnodes>1 and a rendezvous address" + ) + return self.dist_init_base_port + + @property + def rendezvous_port_world(self) -> int: + return self._require_rendezvous_base_port() + + @property + def rendezvous_port_dp_gloo(self) -> int: + return self._require_rendezvous_base_port() + 1 + + def rendezvous_port_reserved(self, i: int) -> int: + if not (0 <= i < 6): + raise ValueError(f"reserved port index must be in [0, 6), got {i}") + return self._require_rendezvous_base_port() + 2 + i + + def dp_rank(self, engine_index: int, *, pp_size: int = 1) -> int: + return self.node_rank * self.local_engine_count + engine_index // pp_size + + def dp_rank_local(self, engine_index: int, *, pp_size: int = 1) -> int: + return engine_index // pp_size + + def local_device_rank( + self, + engine_index: int, + tp_rank: int, + *, + pp_size: int = 1, + pcp_size: int = 1, + ) -> int: + """This node's physical device index; mirrors ModelRunner's formula.""" + stage_span = self.tp_size * pcp_size + pp_rank = engine_index % pp_size + engine_idx = ( + self.dp_rank_local(engine_index, pp_size=pp_size) * pp_size + pp_rank + ) + return engine_idx * stage_span + tp_rank + + def describe(self) -> str: + return ( + f"[wideep] nnodes={self.nnodes} node_rank={self.node_rank} | " + f"ep={self.ep_size} gpu_per_node={self.gpu_per_node} | " + f"dp: global={self.global_dp_size} local={self.local_engine_count}" + ) diff --git a/atom/model_ops/eplb.py b/atom/model_ops/eplb.py index 8a24247aed..3d78615251 100644 --- a/atom/model_ops/eplb.py +++ b/atom/model_ops/eplb.py @@ -888,6 +888,114 @@ def _select_source_rank_for_receiver( return _assign_sender_for_receiver(ranks_to_send, ranks_to_recv, recv_rank) +@dataclass(frozen=True) +class MigrationCost: + """What a planned rearrange will move, split at the node boundary.""" + + transfers: int + cross_node_transfers: int + bytes_per_expert: int + + @property + def cross_node_bytes(self) -> int: + return self.cross_node_transfers * self.bytes_per_expert + + @property + def total_bytes(self) -> int: + return self.transfers * self.bytes_per_expert + + def cross_node_stall_ms(self, bandwidth_gbps: float) -> float: + """Time to move the cross-node bytes at an assumed link bandwidth. + + Intra-node transfers are excluded: over XGMI they are roughly an order + of magnitude cheaper and are not what the guard is protecting against. + """ + if bandwidth_gbps <= 0: + return float("inf") + return self.cross_node_bytes / (bandwidth_gbps * 1e9) * 1e3 + + +def count_migration_transfers( + *, + old_p2l: torch.Tensor, + new_p2l: torch.Tensor, + num_local_physical_experts: int, + num_gpu_per_node: int, +) -> tuple[int, int]: + """``(transfers, cross_node_transfers)`` for a whole rearrange. + + Exact rather than estimated: it applies the same recv rule as + ``_plan_single_layer_migration`` (one P2P per rank/logical the rank gains + and did not already hold) and the same source selection, but derives them + once from the global maps instead of replaying per-rank planning + ``world_size`` times. ``test_migration_cost_matches_planner`` pins the two + together so the shortcut cannot drift. + + Every rank computes the same answer -- the maps are identical across ranks + -- so a decision taken on this needs no collective to stay consistent. + """ + assert old_p2l.shape == new_p2l.shape + num_layers, num_physical = old_p2l.shape + assert num_local_physical_experts > 0 + assert num_physical % num_local_physical_experts == 0 + world_size = num_physical // num_local_physical_experts + assert num_gpu_per_node > 0 and world_size % num_gpu_per_node == 0 + + old_rows = old_p2l.cpu().tolist() + new_rows = new_p2l.cpu().tolist() + transfers = 0 + cross_node = 0 + + for layer in range(num_layers): + old_l = old_rows[layer] + new_l = new_rows[layer] + + # Ascending rank order, matching the planner: _assign_sender_for_receiver + # indexes into these lists, so their order is part of the contract. + holders: dict[int, list[int]] = {} + seen: dict[int, set[int]] = {} + for gslot in range(num_physical): + logical = old_l[gslot] + if logical < 0: + continue + r = gslot // num_local_physical_experts + if logical not in holders: + holders[logical] = [] + seen[logical] = set() + if r not in seen[logical]: + holders[logical].append(r) + seen[logical].add(r) + + recv_by_logical: dict[int, list[int]] = {} + for r in range(world_size): + rbase = r * num_local_physical_experts + old_r = old_l[rbase : rbase + num_local_physical_experts] + new_r = new_l[rbase : rbase + num_local_physical_experts] + old_set = {x for x in old_r if x >= 0} + for logical in sorted({x for x in new_r if x >= 0} - old_set): + recv_by_logical.setdefault(logical, []).append(r) + + for logical, recv_ranks in recv_by_logical.items(): + send_ranks = holders.get(logical) + if not send_ranks: + # The planner asserts here; counting must not be the thing that + # raises, or the guard turns a bad plan into a crash. + continue + recv_node = [r // num_gpu_per_node for r in recv_ranks] + for r, node in zip(recv_ranks, recv_node): + src = _select_source_rank_for_receiver( + ranks_to_send=send_ranks, + ranks_to_recv=recv_ranks, + recv_rank=r, + num_gpu_per_node=num_gpu_per_node, + ) + transfers += 1 + if src // num_gpu_per_node != node: + cross_node += 1 + + return transfers, cross_node + + def _effective_p2p_chunk_size(*, requested: int, num_logical_experts: int) -> int: if num_logical_experts <= 1: return 1 @@ -1661,6 +1769,12 @@ def __init__( self._nnodes: int = 1 self._rebalance_layers_per_chunk: int = 64 self._p2p_batch_chunk_size: int = 32 + self._max_migration_stall_ms: float = 0.0 + self._migration_bandwidth_gbps: float = 50.0 + # Cumulative, so a guard that is quietly suppressing every rebalance is + # visible in the log rather than looking like EPLB having nothing to do. + self._skipped_rebalances: int = 0 + self._skipped_cross_node_bytes: int = 0 # True iff the DP group == the migration (EP) group. When True the DP # sync's `any_rank_has_prefill` already spans exactly the migration # collective, so the prefill gate reuses it for free (no extra collective). @@ -1799,26 +1913,52 @@ def _maybe_initialize_runtime(self, owner: Any) -> None: "manager-owned runtime metadata cannot safely rebalance" ) from exc - try: - from atom.config import get_current_atom_config + # Node count for hierarchical placement. Deliberately outside the + # best-effort block below: a wrong value silently spreads expert groups + # across nodes and makes every rebalance pay avoidable inter-node + # traffic, so it fails loud instead of defaulting to 1. + from atom.config import get_current_atom_config + from atom.model_engine.topology import WideEPTopology, node_count + + atom_config = get_current_atom_config() + self._nnodes = node_count(atom_config) + if not atom_config.dp_logical_size: + # Skipped while simulating, where the two legitimately disagree. + topo = WideEPTopology.from_config(atom_config) + if topo.ep_size != ep_size: + raise RuntimeError( + f"EPLB ep_size ({ep_size}) disagrees with the topology " + f"({topo.ep_size}); hierarchical placement would partition " + f"the wrong rank set. {topo.describe()}" + ) - cfg = get_current_atom_config().eplb_config + try: + cfg = atom_config.eplb_config self._rebalance_layers_per_chunk = int( getattr(cfg, "rebalance_layers_per_chunk", 64) ) self._p2p_batch_chunk_size = int(getattr(cfg, "p2p_batch_chunk_size", 32)) + self._max_migration_stall_ms = float( + getattr(cfg, "max_migration_stall_ms", 0.0) + ) + self._migration_bandwidth_gbps = float( + getattr(cfg, "migration_bandwidth_gbps", 50.0) + ) except Exception: # noqa: BLE001 - optional eplb_config, fall back to defaults self._rebalance_layers_per_chunk = 64 self._p2p_batch_chunk_size = 32 + self._max_migration_stall_ms = 0.0 + self._migration_bandwidth_gbps = 50.0 logger.info( "EPLB runtime initialized: layers=%d num_logical=%d " - "num_physical=%d ep_size=%d ep_rank=%d", + "num_physical=%d ep_size=%d ep_rank=%d nnodes=%d", len(layers), num_logical, num_physical, ep_size, ep_rank, + self._nnodes, ) self.monitor.initialize_for_metadata(self.live_metadata) # Initialize redundant physical slots the checkpoint loader left empty. @@ -2176,6 +2316,11 @@ def _execute_runtime_rebalance(self): max_num_replicas=self.live_metadata.max_num_replicas, ) + cost = self._migration_cost(new_p2l=p2l) + stall_ms = cost.cross_node_stall_ms(self._migration_bandwidth_gbps) + if self._skip_for_migration_cost(rc=_rc, cost=cost, stall_ms=stall_ms): + return + chunk_size = max(1, int(self._rebalance_layers_per_chunk)) layer_ids = sorted(self._moe_layers) num_chunks = (len(layer_ids) + chunk_size - 1) // chunk_size @@ -2229,6 +2374,83 @@ def _execute_runtime_rebalance(self): _chunk_ms, ) + def _migration_cost(self, *, new_p2l: torch.Tensor) -> MigrationCost: + """What moving to ``new_p2l`` would transfer, from the global maps.""" + assert self.live_metadata is not None + num_local = int(self.live_metadata.num_local_physical_experts) + ep_size = int(self.live_metadata.ep_size) + first = min(self._expert_weights_of_layer) + bytes_per_expert = sum( + w.numel() // w.shape[0] * w.element_size() + for w in self._expert_weights_of_layer[first] + ) + transfers, cross_node = count_migration_transfers( + old_p2l=self.live_metadata.physical_to_logical_map, + new_p2l=new_p2l, + num_local_physical_experts=num_local, + num_gpu_per_node=ep_size // self._nnodes, + ) + return MigrationCost( + transfers=transfers, + cross_node_transfers=cross_node, + bytes_per_expert=bytes_per_expert, + ) + + def _skip_for_migration_cost( + self, *, rc: int, cost: MigrationCost, stall_ms: float + ) -> bool: + """Whether the cross-node cost of this rearrange is worth paying. + + Every rank derives the same maps, so every rank reaches the same verdict + without a collective -- which matters, since half the ranks migrating + while the other half skip would leave the placement inconsistent. + """ + _mib = 1024.0 * 1024.0 + if self._max_migration_stall_ms <= 0: + logger.info( + "EPLB rebalance #%d migration: %d transfers (%d cross-node), " + "%.0f MiB (%.0f MiB cross-node, ~%.0fms at %.0f GB/s); no cap set", + rc, + cost.transfers, + cost.cross_node_transfers, + cost.total_bytes / _mib, + cost.cross_node_bytes / _mib, + stall_ms, + self._migration_bandwidth_gbps, + ) + return False + + if stall_ms <= self._max_migration_stall_ms: + logger.info( + "EPLB rebalance #%d migration: %d transfers (%d cross-node), " + "%.0f MiB (%.0f MiB cross-node, ~%.0fms) within %.0fms cap", + rc, + cost.transfers, + cost.cross_node_transfers, + cost.total_bytes / _mib, + cost.cross_node_bytes / _mib, + stall_ms, + self._max_migration_stall_ms, + ) + return False + + self._skipped_rebalances += 1 + self._skipped_cross_node_bytes += cost.cross_node_bytes + logger.warning( + "EPLB rebalance #%d SKIPPED: %d cross-node transfers would move " + "%.0f MiB (~%.0fms at %.0f GB/s) over the %.0fms cap. Placement is " + "unchanged. Cumulative: %d rebalance(s) skipped, %.0f MiB not moved", + rc, + cost.cross_node_transfers, + cost.cross_node_bytes / _mib, + stall_ms, + self._migration_bandwidth_gbps, + self._max_migration_stall_ms, + self._skipped_rebalances, + self._skipped_cross_node_bytes / _mib, + ) + return True + def _log_rebalance_metrics( self, *, rc: int, logical_load: torch.Tensor, logcnt: torch.Tensor ) -> None: diff --git a/atom/model_ops/fused_moe/mori_env.py b/atom/model_ops/fused_moe/mori_env.py new file mode 100644 index 0000000000..9225e71bcb --- /dev/null +++ b/atom/model_ops/fused_moe/mori_env.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""MoRI environment for multi-node EP (M-MOE). + +MoRI reads these once, when its shared-memory heap is created, which happens +inside the first collective. They therefore have to be in the environment +before ``init_dist_env`` runs -- setting them later is a no-op that looks like +it worked. + +Split into a pure planner and a thin applier so the decisions are testable +without a process to mutate. +""" + +from __future__ import annotations + +import logging +import os + +logger = logging.getLogger("atom") + +# InterNodeV1 stages through five buffers (ShmemBufsInterNodeV1) where the +# intra-node path uses one, so the default heap that serves a single node is +# not enough once the group spans a boundary. +_MULTINODE_SHMEM_HEAP = "8G" + +_LOOPBACK_IFNAMES = frozenset({"lo", "lo0", "localhost"}) + + +def plan_mori_env( + *, nnodes: int, current: dict[str, str] | None = None +) -> dict[str, str]: + """Variables to add for this topology, given what is already exported. + + Never overrides an existing value: an operator who set one did so against a + specific fabric, and this has less information than they do. Returns only + the additions, so the caller can log exactly what it changed. + """ + env = os.environ if current is None else current + if nnodes <= 1: + return {} + + additions: dict[str, str] = {} + if not env.get("MORI_SHMEM_HEAP_SIZE"): + additions["MORI_SHMEM_HEAP_SIZE"] = _MULTINODE_SHMEM_HEAP + if not env.get("MORI_EP_LAUNCH_CONFIG_MODE"): + additions["MORI_EP_LAUNCH_CONFIG_MODE"] = "AUTO" + return additions + + +def check_mori_env(*, nnodes: int, current: dict[str, str] | None = None) -> None: + """Reject settings that cannot work across nodes. + + Only the interface name is checked. It is the one that is both commonly + wrong -- a single-node run needs ``lo`` and that value tends to survive + into a multi-node launch script -- and silent when wrong: peers on other + hosts are simply never reached, so the run hangs in the first collective + with nothing logged. + """ + if nnodes <= 1: + return + env = os.environ if current is None else current + ifname = env.get("MORI_SOCKET_IFNAME", "") + if ifname.lower() in _LOOPBACK_IFNAMES: + raise ValueError( + f"MORI_SOCKET_IFNAME={ifname!r} cannot reach {nnodes} nodes.\n" + f" Why: loopback resolves to this host, so MoRI's bootstrap never " + f"sees the other nodes' peers and the first collective hangs with " + f"nothing in the log. Single-node runs do need this value, which " + f"is how it survives into a multi-node launcher.\n" + f" Fix: set MORI_SOCKET_IFNAME to the interface that carries " + f"inter-node traffic (`ip -o addr show` on the node)." + ) + + +def apply_mori_env(*, nnodes: int) -> dict[str, str]: + """Validate, then export. Returns what it set, for the startup summary.""" + check_mori_env(nnodes=nnodes) + additions = plan_mori_env(nnodes=nnodes) + for key, value in additions.items(): + os.environ[key] = value + if additions: + logger.info( + "[wideep] MoRI env for %d nodes: %s", + nnodes, + " ".join(f"{k}={v}" for k, v in sorted(additions.items())), + ) + return additions diff --git a/atom/model_ops/fused_moe/wideep_experts.py b/atom/model_ops/fused_moe/wideep_experts.py new file mode 100644 index 0000000000..a5ec473f79 --- /dev/null +++ b/atom/model_ops/fused_moe/wideep_experts.py @@ -0,0 +1,364 @@ +# SPDX-License-Identifier: Apache-2.0 +"""AITER TestWideEpMoe integration for the EP16 inter-node path. + +The AITER operator owns the complete routed-MoE step: FP8 dispatch through +MORI InterNodeV1LL, both MXFP4 expert GEMMs, and combine. It therefore plugs +into ATOM as a whole-pipeline ``fused_experts`` backend, beside MegaMoEV2, +rather than through the standard prepare/GEMM/finalize seam. +""" + +from __future__ import annotations + +import torch + +from atom.utils.custom_register import direct_register_custom_op + +_WIDEEP_CACHE: dict = {} + + +def build_wideep_weights(layer) -> None: + """Build TestWideEpMoe's expert-major A16W4 shuffled weight layout.""" + from aiter.ops.shuffle import shuffle_scale_a16w4, shuffle_weight_a16w4 + + w1 = layer.w13_weight.data + local_experts = int(w1.shape[0]) + if local_experts != layer.local_num_experts: + raise RuntimeError( + "WideEP local weight width disagrees with the dispatch layout: " + f"weights={local_experts}, dispatch={layer.local_num_experts}" + ) + layer._wideep_w1 = shuffle_weight_a16w4(w1, 16, True).contiguous() + + w1_scale = layer.w13_weight_scale.data + w1_scale = w1_scale.reshape(local_experts * w1_scale.shape[1], w1_scale.shape[2]) + layer._wideep_w1_scale = shuffle_scale_a16w4( + w1_scale, local_experts, True + ).contiguous() + + w2 = layer.w2_weight.data + layer._wideep_w2 = shuffle_weight_a16w4(w2, 16, False).contiguous() + w2_scale = layer.w2_weight_scale.data + w2_scale = w2_scale.reshape(local_experts * w2_scale.shape[1], w2_scale.shape[2]) + layer._wideep_w2_scale = shuffle_scale_a16w4( + w2_scale, local_experts, False + ).contiguous() + # TestWideEpMoe enters the generic fused_moe frontend, which uses this + # marker to recognize that both weights already have kernel-ready layout. + layer._wideep_w1.is_shuffled = True + layer._wideep_w2.is_shuffled = True + + +def get_or_build_wideep_moe( + *, + rank: int, + world_size: int, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + quant: str, + mtpr: int, + swiglu_limit: float, + w1: torch.Tensor, + w1_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, +): + """Return the process-wide operator for this shape and bind layer weights.""" + key = ( + rank, + world_size, + model_dim, + inter_dim, + experts, + topk, + quant, + mtpr, + swiglu_limit, + ) + op = _WIDEEP_CACHE.get(key) + if op is None: + try: + from aiter.ops.flydsl.test_wide_ep_moe import TestWideEpMoe + except ImportError as exc: + raise RuntimeError( + "The EP16 WideEP path requires AITER dev/test_wide_ep_moe " + "with aiter.ops.flydsl.test_wide_ep_moe.TestWideEpMoe" + ) from exc + + with torch.inference_mode(False), torch.no_grad(): + op = TestWideEpMoe( + rank=rank, + world_size=world_size, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + quant=quant, + w1=w1, + w1_scale=w1_scale, + w2=w2, + w2_scale=w2_scale, + max_tok_per_rank=mtpr, + gpu_per_node=8, + swiglu_limit=swiglu_limit, + activation="silu", + gate_mode="interleave", + ) + _WIDEEP_CACHE[key] = op + + # One operator is shared by shape-identical MoE layers. TestWideEpMoe + # reads these tensors at execution time, so bind the current layer on every + # call. The shuffled tensors keep stable storage across graph replay. + op.w1 = w1 + op.w1_scale = w1_scale + op.w2 = w2 + op.w2_scale = w2_scale + return op + + +def atom_wideep_forward_impl( + x_quant: torch.Tensor, + x_scale: torch.Tensor, + weights: torch.Tensor, + ids: torch.Tensor, + w1: torch.Tensor, + w1_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + rank: int, + world_size: int, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + mtpr: int, + swiglu_limit: float, +) -> torch.Tensor: + """Dynamo boundary with layer weights explicit in the graph.""" + op = get_or_build_wideep_moe( + rank=rank, + world_size=world_size, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + quant="a8w4", + mtpr=mtpr, + swiglu_limit=swiglu_limit, + w1=w1, + w1_scale=w1_scale, + w2=w2, + w2_scale=w2_scale, + ) + # The MORI combine result aliases its symmetric arena. Materialize a + # graph-owned output before another MoE layer reuses the shared transport. + return op.forward_prequant(x_quant, x_scale, weights, ids).clone() + + +def atom_wideep_forward_fake( + x_quant: torch.Tensor, + x_scale: torch.Tensor, + weights: torch.Tensor, + ids: torch.Tensor, + w1: torch.Tensor, + w1_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + rank: int, + world_size: int, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + mtpr: int, + swiglu_limit: float, +) -> torch.Tensor: + del ( + x_scale, + weights, + w1, + w1_scale, + w2, + w2_scale, + rank, + world_size, + inter_dim, + experts, + topk, + mtpr, + swiglu_limit, + ) + return x_quant.new_empty((ids.shape[0], model_dim), dtype=torch.bfloat16) + + +direct_register_custom_op( + op_name="atom_wideep_forward", + op_func=atom_wideep_forward_impl, + mutates_args=[], + fake_impl=atom_wideep_forward_fake, +) + + +def run_wideep_moe( + layer, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + mtpr: int, + swiglu_limit: float, + quant: str = "a8w4", +) -> torch.Tensor: + """Run the complete EP16 InterNodeV1LL routed-MoE pipeline.""" + from aiter.dist.parallel_state import get_ep_group + + # This property initializes MORI's symmetric heap as a required side effect. + all2all_manager = get_ep_group().device_communicator.all2all_manager + rank = int(all2all_manager.rank) + world_size = int(all2all_manager.world_size) + if world_size != 16: + raise RuntimeError( + f"WideEP was selected for EP16, but the live EP group has {world_size} ranks" + ) + if hidden_states.dtype != torch.bfloat16: + raise TypeError( + "WideEP requires BF16 hidden states before its internal FP8 quantization, " + f"got {hidden_states.dtype}" + ) + if int(hidden_states.shape[0]) > mtpr: + raise ValueError( + f"WideEP run_tokens={hidden_states.shape[0]} exceeds max_num_tokens={mtpr}" + ) + + op = get_or_build_wideep_moe( + rank=rank, + world_size=world_size, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + quant=quant, + mtpr=mtpr, + swiglu_limit=float(swiglu_limit), + w1=layer._wideep_w1, + w1_scale=layer._wideep_w1_scale, + w2=layer._wideep_w2, + w2_scale=layer._wideep_w2_scale, + ) + + hidden_states = hidden_states.contiguous() + weights = topk_weights.to(torch.float32).contiguous() + ids = topk_ids.to(torch.int32).contiguous() + with torch.inference_mode(False), torch.no_grad(): + x_quant, x_scale = op.quantize(hidden_states) + if torch.compiler.is_compiling(): + return torch.ops.aiter.atom_wideep_forward( + x_quant, + x_scale, + weights, + ids, + layer._wideep_w1, + layer._wideep_w1_scale, + layer._wideep_w2, + layer._wideep_w2_scale, + rank, + world_size, + model_dim, + inter_dim, + experts, + topk, + mtpr, + float(swiglu_limit), + ) + return op.forward_prequant(x_quant, x_scale, weights, ids) + + +class WideEpFusedExperts: + """Whole-pipeline ``fused_experts`` adapter for AITER TestWideEpMoe.""" + + def __init__( + self, + layer: torch.nn.Module, + *, + model_dim: int, + inter_dim: int, + experts: int, + mtpr: int, + quant: str = "a8w4", + ) -> None: + self._layer = layer + self._model_dim = model_dim + self._inter_dim = inter_dim + self._experts = experts + self._mtpr = mtpr + self._quant = quant + + def __call__( + self, + *, + hidden_states: torch.Tensor, + w1: torch.Tensor | None = None, + w2: torch.Tensor | None = None, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + global_num_experts: int = -1, + activation=None, + apply_router_weight_on_input: bool = False, + expert_map: torch.Tensor | None = None, + bias1: torch.Tensor | None = None, + bias2: torch.Tensor | None = None, + hidden_pad: int = 0, + intermediate_pad: int = 0, + **_ignored, + ) -> torch.Tensor: + from aiter import ActivationType + + if apply_router_weight_on_input: + raise NotImplementedError( + "WideEP does not support apply_router_weight_on_input=True" + ) + if activation is not None and activation != ActivationType.Silu: + raise NotImplementedError( + f"WideEP A8W4 requires SiLU activation, got {activation}" + ) + if bias1 is not None or bias2 is not None: + raise NotImplementedError("WideEP does not support expert bias") + if hidden_pad or intermediate_pad: + raise NotImplementedError( + "WideEP requires model and intermediate dimensions with no padding" + ) + if global_num_experts not in (-1, self._experts): + raise RuntimeError( + "WideEP routing width disagrees with its configured experts: " + f"routing={global_num_experts}, configured={self._experts}" + ) + + # w1/w2 are emptied staging parameters. The dedicated shuffled weights + # live on the layer; expert_map is already reflected in global topk ids. + del w1, w2, expert_map + local_weight_experts = int(self._layer._wideep_w1.shape[0]) + if local_weight_experts != self._layer.local_num_experts: + raise RuntimeError( + "WideEP weight/dispatch layout changed after preparation: " + f"weights={local_weight_experts}, " + f"expected={self._layer.local_num_experts}" + ) + + return run_wideep_moe( + self._layer, + hidden_states, + topk_weights, + topk_ids, + model_dim=self._model_dim, + inter_dim=self._inter_dim, + experts=self._experts, + topk=int(topk_ids.shape[1]), + mtpr=self._mtpr, + swiglu_limit=getattr(self._layer, "swiglu_limit", 0.0), + quant=self._quant, + ) diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 705df6ef60..a74839457d 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -297,6 +297,8 @@ def flatten_tp_across_dp(dp_rank: int): # parallel update tp_size, tp_rank, ep_size and ep_rank to reflect that. ep_size = tp_size ep_rank = tp_rank + local_ep_size = atom_config.parallel_config.data_parallel_size_local * tp_size_ + _assert_gpu_per_node(atom_config, ep_size=ep_size, gpu_per_node=local_ep_size) return FusedMoEParallelConfig( tp_size=1, tp_rank=0, @@ -309,8 +311,35 @@ def flatten_tp_across_dp(dp_rank: int): dp_logical_ratio=( dp_logical // dp_size if flatten_tp_across_dp_for_moe else 1 ), - local_ep_size=atom_config.parallel_config.data_parallel_size_local - * tp_size_, + local_ep_size=local_ep_size, + ) + + +def _assert_gpu_per_node(atom_config, *, ep_size: int, gpu_per_node: int) -> None: + """Check the value that reaches MoRI as ``gpu_per_node``. + + InterNodeV1 derives its RDMA and XGMI peer targets from it, so a wrong + value does not raise -- it addresses the wrong ranks. Before #1603 it was 1 + on every run and nothing noticed, because the intra-node kernel does not + consult it. That is the argument for checking rather than trusting. + + Only meaningful under DP-attention: that is the case where every rank joins + one EP group, so the group has to partition evenly into nodes. Without it + EP spans a TP group and there are dp_size independent groups, and the + identity does not hold. + """ + if not atom_config.enable_dp_attention: + return + from atom.model_engine.topology import node_count + + nnodes = node_count(atom_config) + if ep_size % gpu_per_node or ep_size // gpu_per_node != nnodes: + raise ValueError( + f"gpu_per_node={gpu_per_node} does not partition ep_size={ep_size} " + f"into {nnodes} node(s). MoRI takes this as the node width and " + f"computes peer targets from it, so a wrong value reaches the " + f"wrong ranks instead of failing. Check " + f"data_parallel_size_local against the GPUs this node actually has." ) @@ -1772,11 +1801,97 @@ def get_eplb_weight_views(self, layer: torch.nn.Module) -> list[torch.Tensor]: return views +def _is_wideep_topology(moe: FusedMoEConfig) -> bool: + parallel = getattr(moe, "moe_parallel_config", None) + return ( + parallel is not None + and getattr(parallel, "use_ep", False) + and getattr(parallel, "ep_size", 1) == 16 + and getattr(parallel, "local_ep_size", 1) == 8 + ) + + +def _validate_wideep_mxfp4_config( + quant_config: LayerQuantConfig, moe: FusedMoEConfig +) -> None: + """Reject EP16 combinations that TestWideEpMoe has not implemented.""" + problems = [] + parallel = moe.moe_parallel_config + atom_config = get_current_atom_config() + + if get_gfx() != "gfx950": + problems.append("GPU architecture must be gfx950") + if parallel.dp_logical_ratio != 1: + problems.append("logical DP simulation is unsupported") + if not parallel.use_all2all_kernels: + problems.append("the MORI all-to-all runtime is unavailable") + if not quant_config.is_dynamic: + problems.append("activation quantization must be dynamic") + if not envs.ATOM_MOE_GU_ITLV: + problems.append("ATOM_MOE_GU_ITLV=1 is required") + if getattr(atom_config, "eplb_enable", False): + problems.append("EPLB must be disabled") + if moe.expert_layout.mode is not SharedExpertMode.NONE: + problems.append("fused shared experts are unsupported") + if moe.has_bias: + problems.append("expert bias is unsupported") + if moe.in_dtype != torch.bfloat16: + problems.append(f"model dtype must be torch.bfloat16, got {moe.in_dtype}") + + if problems: + raise ValueError( + "The EP16 (2 nodes x 8 GPUs) MXFP4 configuration selects AITER " + "TestWideEpMoe, but its contract is not satisfied: " + "; ".join(problems) + ) + + +class WideEpMxfp4MoEMethod(Mxfp4MoEMethod): + """MXFP4 method backed by AITER's EP16 InterNodeV1LL pipeline.""" + + def __init__(self, quant_config: LayerQuantConfig, moe: FusedMoEConfig): + super().__init__(quant_config, moe) + _validate_wideep_mxfp4_config(quant_config, moe) + self.use_triton = False + self.use_triton_decode = False + + def _process_weight_layout_after_loading(self, layer) -> None: + from atom.model_ops.fused_moe.wideep_experts import build_wideep_weights + + build_wideep_weights(layer) + layer.w13_weight.data = torch.empty( + 0, dtype=layer.w13_weight.dtype, device=layer.w13_weight.device + ) + layer.w2_weight.data = torch.empty( + 0, dtype=layer.w2_weight.dtype, device=layer.w2_weight.device + ) + logger.info("Prepared WideEP weights for the EP16 InterNodeV1LL MoE path") + + def init_prepare_finalize(self, layer: torch.nn.Module): + from atom.model_ops.fused_moe.wideep_experts import WideEpFusedExperts + + if self.hidden_pad or self.intermediate_pad: + raise ValueError( + "WideEP does not support padded MoE dimensions: " + f"hidden_pad={self.hidden_pad}, intermediate_pad={self.intermediate_pad}" + ) + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + self.fused_experts = WideEpFusedExperts( + layer, + model_dim=self.hidden_size, + inter_dim=self.intermediate_size, + experts=self.moe.num_experts, + mtpr=self.moe.max_num_tokens, + quant="a8w4", + ) + + def _make_mxfp4_moe_method( quant_config: LayerQuantConfig, moe: FusedMoEConfig ) -> Mxfp4MoEMethod: if get_current_atom_config().moe_backend == "mega": return MegaMxfp4MoEMethod(quant_config, moe) + if _is_wideep_topology(moe): + return WideEpMxfp4MoEMethod(quant_config, moe) return Mxfp4MoEMethod(quant_config, moe) diff --git a/atom/utils/distributed/utils.py b/atom/utils/distributed/utils.py index a11e247b34..b1dca803c2 100644 --- a/atom/utils/distributed/utils.py +++ b/atom/utils/distributed/utils.py @@ -74,7 +74,12 @@ def init_gloo_process_group( def stateless_init_torch_distributed_process_group( - host: str, port: int, rank: int, world_size: int, backend: str + host: str, + port: int, + rank: int, + world_size: int, + backend: str, + timeout: timedelta | None = None, ) -> ProcessGroup: """ A replacement for `torch.distributed.init_process_group` that does not @@ -106,10 +111,16 @@ def stateless_init_torch_distributed_process_group( are the same as process 1 and 5, the main communication channel is always formed with process 1, 2, ..., 8, and the additional communication channel is formed with process 9 and 10. + + ``timeout`` bounds the rendezvous, the store, and the group handshake. + ``None`` keeps the backend default, so single-node behaviour is unchanged; + a shorter value is what turns "the other node never started" from an + indefinite hang into an error naming the group that did not form. """ init_method = get_tcp_uri(host, port) backend = Backend(backend) # it is basically string - timeout = _get_default_timeout(backend) + if timeout is None: + timeout = _get_default_timeout(backend) store, rank, world_size = next( rendezvous(init_method, rank, world_size, timeout=timeout) diff --git a/tests/aiter_stub.py b/tests/aiter_stub.py index b60b127482..1ef278a7b2 100644 --- a/tests/aiter_stub.py +++ b/tests/aiter_stub.py @@ -17,7 +17,12 @@ def stubbed_aiter(): afterwards, so it is removed as soon as the engine classes are bound. """ installed = [] - for name in ("aiter", "aiter.dist", "aiter.dist.shm_broadcast"): + for name in ( + "aiter", + "aiter.dist", + "aiter.dist.shm_broadcast", + "aiter.dist.parallel_state", + ): if name in sys.modules: continue module = types.ModuleType(name) diff --git a/tests/test_eplb_hierarchical.py b/tests/test_eplb_hierarchical.py new file mode 100644 index 0000000000..ec9f5dd7f4 --- /dev/null +++ b/tests/test_eplb_hierarchical.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Hierarchical (multi-node) EPLB placement. + +`rebalance_experts(enable_hierarchical=True, num_nodes>1)` has never executed in +production: `EPLBManager._nnodes` is hardcoded to 1, so every rebalance so far +took the flat branch. These tests cover the node-aware branch before WideEP +starts depending on it. + +CPU-only; no GPU or distributed setup. +""" + +import pytest +import torch +from aiter_stub import stubbed_aiter + +# Placement is pure index arithmetic. Stub AITER rather than skip without it, +# so a CPU box runs these for real instead of reporting green. +with stubbed_aiter(): + import atom.config # noqa: F401 (initialize before model_ops' __init__ chain) + from atom.model_ops.eplb import rebalance_experts + +# DSv4-Pro-shaped: 384 logical experts in 8 groups, EP16 over 2 nodes, +# num_redundant = ep_size (the TRT "num_slots = experts + EP size" tier). +DSV4 = { + "num_logical": 384, + "num_groups": 8, + "num_physical": 400, + "num_gpus": 16, + "num_nodes": 2, +} + + +def _weights(num_layers, num_logical, *, skew=None, seed=0): + g = torch.Generator().manual_seed(seed) + w = torch.rand((num_layers, num_logical), generator=g) + 0.1 + if skew is not None: + w[:, :skew] *= 100.0 + return w + + +def _run(cfg, *, num_layers=2, weight=None, hierarchical=True): + w = _weights(num_layers, cfg["num_logical"]) if weight is None else weight + return rebalance_experts( + w, + num_physical=cfg["num_physical"], + num_groups=cfg["num_groups"], + num_nodes=cfg["num_nodes"], + num_gpus=cfg["num_gpus"], + enable_hierarchical=hierarchical, + ) + + +def _assert_placement_valid(p2l, logcnt, cfg): + """Invariants any placement must satisfy, flat or hierarchical.""" + num_layers = p2l.shape[0] + num_logical = cfg["num_logical"] + assert p2l.shape == (num_layers, cfg["num_physical"]) + assert logcnt.shape == (num_layers, num_logical) + + for layer in range(num_layers): + ids = p2l[layer] + assert int(ids.min()) >= 0, "unfilled physical slot" + assert int(ids.max()) < num_logical, "physical slot maps outside logical space" + # Every logical expert is placed, or its tokens have nowhere to go. + counted = torch.bincount(ids.to(torch.int64), minlength=num_logical) + assert int(counted.min()) >= 1, "logical expert with no physical slot" + assert int(counted.sum()) == cfg["num_physical"] + # logcnt must agree with p2l: the dispatch map is built from logcnt + # while routing reads p2l, so a mismatch silently drops replicas. + assert torch.equal(counted.to(torch.int32), logcnt[layer]) + + +class TestHierarchicalPlacement: + def test_dsv4_ep16_two_nodes(self): + p2l, l2p, logcnt = _run(DSV4) + _assert_placement_valid(p2l, logcnt, DSV4) + assert l2p.shape[:2] == (2, DSV4["num_logical"]) + + def test_group_stays_on_one_node(self): + """The whole point of the hierarchical path. + + A group's experts are co-routed, so splitting one across nodes turns + intra-node XGMI traffic into RDMA traffic for every token hitting it. + """ + cfg = DSV4 + p2l, _, _ = _run(cfg) + group_size = cfg["num_logical"] // cfg["num_groups"] + phy_per_node = cfg["num_physical"] // cfg["num_nodes"] + + for layer in range(p2l.shape[0]): + for g in range(cfg["num_groups"]): + members = set(range(g * group_size, (g + 1) * group_size)) + nodes = { + int(i) // phy_per_node + for i, e in enumerate(p2l[layer].tolist()) + if e in members + } + assert len(nodes) == 1, f"group {g} split across nodes {nodes}" + + def test_each_node_owns_its_share_of_slots(self): + cfg = DSV4 + p2l, _, _ = _run(cfg) + phy_per_node = cfg["num_physical"] // cfg["num_nodes"] + for layer in range(p2l.shape[0]): + for n in range(cfg["num_nodes"]): + lo, hi = n * phy_per_node, (n + 1) * phy_per_node + assert int((p2l[layer][lo:hi] >= 0).sum()) == phy_per_node + + def test_skewed_load_keeps_groups_balanced(self): + """balanced_packing must still hand each node groups_per_node groups. + + The implementation asserts that; a heavily skewed layer is the case + most likely to break it. + """ + cfg = DSV4 + w = _weights(2, cfg["num_logical"], skew=48) # one whole group is hot + p2l, _, logcnt = _run(cfg, weight=w) + _assert_placement_valid(p2l, logcnt, cfg) + + def test_hot_experts_get_more_replicas(self): + cfg = DSV4 + w = _weights(1, cfg["num_logical"], skew=48) + _, _, logcnt = _run(cfg, num_layers=1, weight=w) + hot = logcnt[0, :48].float().mean() + cold = logcnt[0, 48:].float().mean() + assert hot > cold, f"hot={hot} not replicated more than cold={cold}" + + +class TestHierarchicalVsFlat: + def test_flat_and_hierarchical_both_valid(self): + """Same inputs, both branches; only the placement should differ.""" + w = _weights(2, DSV4["num_logical"]) + flat_p2l, _, flat_cnt = _run(DSV4, weight=w, hierarchical=False) + hier_p2l, _, hier_cnt = _run(DSV4, weight=w, hierarchical=True) + _assert_placement_valid(flat_p2l, flat_cnt, DSV4) + _assert_placement_valid(hier_p2l, hier_cnt, DSV4) + + def test_flat_branch_does_split_groups(self): + """Pins the discriminating power of test_group_stays_on_one_node. + + Without this, a hierarchical branch that silently degraded into the + flat one would still satisfy the co-location test if the flat placement + happened to keep groups together. Measured: flat splits 16/16 + group-layer pairs, hierarchical splits 0/16. + """ + cfg = DSV4 + w = _weights(2, cfg["num_logical"]) + p2l, _, _ = _run(cfg, weight=w, hierarchical=False) + group_size = cfg["num_logical"] // cfg["num_groups"] + phy_per_node = cfg["num_physical"] // cfg["num_nodes"] + split = sum( + len( + { + i // phy_per_node + for i, e in enumerate(p2l[layer].tolist()) + if g * group_size <= e < (g + 1) * group_size + } + ) + > 1 + for layer in range(p2l.shape[0]) + for g in range(cfg["num_groups"]) + ) + assert split > 0, "flat placement kept every group on one node by chance" + + def test_single_node_ignores_hierarchical_flag(self): + cfg = dict(DSV4, num_nodes=1, num_gpus=8, num_physical=392) + w = _weights(2, cfg["num_logical"]) + a = _run(cfg, weight=w, hierarchical=True) + b = _run(cfg, weight=w, hierarchical=False) + assert torch.equal(a[0], b[0]), "nnodes=1 must take the flat branch" + + +class TestHierarchicalConstraints: + @pytest.mark.parametrize( + "override, match", + [ + ({"num_groups": 5}, "num_logical must be divisible by num_groups"), + ({"num_nodes": 3}, "num_groups must be divisible by num_nodes"), + ({"num_gpus": 15}, "num_gpus must be divisible by num_nodes"), + ({"num_physical": 401}, "num_physical must be divisible by num_gpus"), + ], + ) + def test_rejects_incoherent_topology(self, override, match): + cfg = dict(DSV4, **override) + with pytest.raises(AssertionError, match=match): + _run(cfg) + + @pytest.mark.parametrize("num_nodes, num_gpus", [(2, 16), (4, 16), (8, 16)]) + def test_scales_to_more_nodes(self, num_nodes, num_gpus): + cfg = dict(DSV4, num_nodes=num_nodes, num_gpus=num_gpus) + p2l, _, logcnt = _run(cfg, num_layers=1) + _assert_placement_valid(p2l, logcnt, cfg) diff --git a/tests/test_eplb_migration_cost.py b/tests/test_eplb_migration_cost.py new file mode 100644 index 0000000000..827d24aefe --- /dev/null +++ b/tests/test_eplb_migration_cost.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Unit tests for the cross-node migration guard (M2.2). + +The guard decides whether a rearrange is worth its inter-node traffic. It reads +a count derived from the global maps rather than from planning; the first test +here is what keeps that shortcut honest. +""" + +import torch +from aiter_stub import stubbed_aiter + +# Everything under test is pure index arithmetic. Stub AITER rather than skip +# without it, so a CPU box runs these for real instead of reporting green. +with stubbed_aiter(): + import atom.config # noqa: F401 (initialize before model_ops' __init__ chain) + from atom.model_ops.eplb import ( + MigrationCost, + _plan_single_layer_migration, + count_migration_transfers, + ) + +# DSv4-Pro at the WideEP target: 384 logical experts, num_redundant = ep_size. +NUM_LOGICAL = 384 +NUM_PHYSICAL = 400 +WORLD_SIZE = 16 +NUM_LOCAL = NUM_PHYSICAL // WORLD_SIZE +GPU_PER_NODE = 8 # 2 nodes + + +def _make_p2l(num_layers: int, seed: int) -> torch.Tensor: + """A valid placement: every logical present, the spare slots as replicas.""" + g = torch.Generator().manual_seed(seed) + rows = [] + for _ in range(num_layers): + base = torch.randperm(NUM_LOGICAL, generator=g) + spare = torch.randint( + 0, NUM_LOGICAL, (NUM_PHYSICAL - NUM_LOGICAL,), generator=g + ) + rows.append(torch.cat([base, spare])) + return torch.stack(rows) + + +def _planner_transfers(old_p2l, new_p2l, *, gpu_per_node): + """Replay the real per-rank planner and total up its P2P receives.""" + total = cross = 0 + num_layers = old_p2l.shape[0] + for layer in range(num_layers): + for rank in range(WORLD_SIZE): + _, _, recv_actions, _ = _plan_single_layer_migration( + old_p2l_layer=old_p2l[layer], + new_p2l_layer=new_p2l[layer], + num_local_physical_experts=NUM_LOCAL, + num_gpu_per_node=gpu_per_node, + rank=rank, + world_size=WORLD_SIZE, + ) + for action in recv_actions: + total += 1 + if action.peer_rank // gpu_per_node != rank // gpu_per_node: + cross += 1 + return total, cross + + +class TestCountMatchesPlanner: + def test_migration_cost_matches_planner(self): + """The shortcut must equal what the planner actually issues. + + count_migration_transfers derives the answer once from the global maps; + the planner derives it per rank. They encode the same recv rule and the + same source selection, and this is the only thing stopping them from + drifting apart. + """ + old = _make_p2l(4, seed=0) + new = _make_p2l(4, seed=1) + assert count_migration_transfers( + old_p2l=old, + new_p2l=new, + num_local_physical_experts=NUM_LOCAL, + num_gpu_per_node=GPU_PER_NODE, + ) == _planner_transfers(old, new, gpu_per_node=GPU_PER_NODE) + + def test_matches_planner_on_one_node(self): + old = _make_p2l(2, seed=2) + new = _make_p2l(2, seed=3) + assert count_migration_transfers( + old_p2l=old, + new_p2l=new, + num_local_physical_experts=NUM_LOCAL, + num_gpu_per_node=WORLD_SIZE, + ) == _planner_transfers(old, new, gpu_per_node=WORLD_SIZE) + + +class TestCounts: + def test_no_change_moves_nothing(self): + p2l = _make_p2l(3, seed=4) + assert count_migration_transfers( + old_p2l=p2l, + new_p2l=p2l.clone(), + num_local_physical_experts=NUM_LOCAL, + num_gpu_per_node=GPU_PER_NODE, + ) == (0, 0) + + def test_single_node_never_crosses(self): + old = _make_p2l(3, seed=5) + new = _make_p2l(3, seed=6) + transfers, cross = count_migration_transfers( + old_p2l=old, + new_p2l=new, + num_local_physical_experts=NUM_LOCAL, + num_gpu_per_node=WORLD_SIZE, + ) + assert transfers > 0 + assert cross == 0 + + def test_same_node_holder_keeps_it_local(self): + # 4 ranks, 2 per node, 1 slot each. Rank 1 wants expert 0; rank 0 (same + # node) holds it, so the node-aware selector must not reach off-node. + old = torch.tensor([[0, 9, 0, 8]]) + new = torch.tensor([[0, 0, 0, 8]]) + transfers, cross = count_migration_transfers( + old_p2l=old, + new_p2l=new, + num_local_physical_experts=1, + num_gpu_per_node=2, + ) + assert (transfers, cross) == (1, 0) + + def test_only_off_node_holder_forces_a_crossing(self): + # Same shape, but expert 0 lives only on node 0 and node 1 needs it. + old = torch.tensor([[0, 9, 7, 8]]) + new = torch.tensor([[0, 9, 0, 8]]) + transfers, cross = count_migration_transfers( + old_p2l=old, + new_p2l=new, + num_local_physical_experts=1, + num_gpu_per_node=2, + ) + assert (transfers, cross) == (1, 1) + + def test_replicas_on_one_rank_cost_one_transfer(self): + # Rank 1 gains two slots of expert 0; the planner sends it once and + # copies locally for the second (the free-rider path). + old = torch.tensor([[0, 0, 9, 8, 7, 6]]) + new = torch.tensor([[0, 0, 0, 0, 7, 6]]) + transfers, cross = count_migration_transfers( + old_p2l=old, + new_p2l=new, + num_local_physical_experts=2, + num_gpu_per_node=1, + ) + assert (transfers, cross) == (1, 1) + + +class TestStallEstimate: + def test_only_cross_node_bytes_count(self): + cost = MigrationCost( + transfers=100, cross_node_transfers=10, bytes_per_expert=10**8 + ) + assert cost.total_bytes == 10**10 + assert cost.cross_node_bytes == 10**9 + # 1 GB at 50 GB/s = 20 ms. + assert cost.cross_node_stall_ms(50.0) == 20.0 + + def test_nothing_to_move_is_free(self): + cost = MigrationCost( + transfers=5, cross_node_transfers=0, bytes_per_expert=10**8 + ) + assert cost.cross_node_stall_ms(50.0) == 0.0 diff --git a/tests/test_mori_env.py b/tests/test_mori_env.py new file mode 100644 index 0000000000..581f92f1fa --- /dev/null +++ b/tests/test_mori_env.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""MoRI's multi-node environment (M-MOE). + +Pure planning and validation; nothing here touches a process. The applier is +covered through the plan it applies. +""" + +import pytest + +from atom.model_ops.fused_moe.mori_env import check_mori_env, plan_mori_env + + +class TestPlan: + def test_single_node_changes_nothing(self): + # The single-node defaults are what production has been running on. + assert plan_mori_env(nnodes=1, current={}) == {} + + def test_multinode_sets_heap_and_launch_mode(self): + assert plan_mori_env(nnodes=2, current={}) == { + "MORI_SHMEM_HEAP_SIZE": "8G", + "MORI_EP_LAUNCH_CONFIG_MODE": "AUTO", + } + + def test_operator_settings_win(self): + # Someone who exported these did it against a specific fabric; this + # has strictly less information than they do. + current = {"MORI_SHMEM_HEAP_SIZE": "16G", "MORI_EP_LAUNCH_CONFIG_MODE": "MANUAL"} + assert plan_mori_env(nnodes=4, current=current) == {} + + def test_partial_override_fills_only_the_gap(self): + current = {"MORI_SHMEM_HEAP_SIZE": "16G"} + assert plan_mori_env(nnodes=2, current=current) == { + "MORI_EP_LAUNCH_CONFIG_MODE": "AUTO" + } + + def test_empty_string_counts_as_unset(self): + assert "MORI_SHMEM_HEAP_SIZE" in plan_mori_env( + nnodes=2, current={"MORI_SHMEM_HEAP_SIZE": ""} + ) + + +class TestCheck: + @pytest.mark.parametrize("ifname", ["lo", "LO", "lo0", "localhost"]) + def test_loopback_across_nodes_is_rejected(self, ifname): + # A single-node run needs lo, and that value survives into the + # multi-node launcher. Wrong, it hangs the first collective silently. + with pytest.raises(ValueError, match="MORI_SOCKET_IFNAME"): + check_mori_env(nnodes=2, current={"MORI_SOCKET_IFNAME": ifname}) + + def test_loopback_on_one_node_is_correct(self): + check_mori_env(nnodes=1, current={"MORI_SOCKET_IFNAME": "lo"}) + + def test_real_interface_passes(self): + check_mori_env(nnodes=2, current={"MORI_SOCKET_IFNAME": "ens3"}) + + def test_unset_interface_passes(self): + # MoRI has its own discovery; only an actively wrong value is rejected. + check_mori_env(nnodes=2, current={}) + + def test_message_says_how_to_find_the_interface(self): + with pytest.raises(ValueError) as exc: + check_mori_env(nnodes=2, current={"MORI_SOCKET_IFNAME": "lo"}) + assert "Fix:" in str(exc.value) + assert "ip -o addr show" in str(exc.value) diff --git a/tests/test_wideep_config_validation.py b/tests/test_wideep_config_validation.py new file mode 100644 index 0000000000..77bd6de1fb --- /dev/null +++ b/tests/test_wideep_config_validation.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Multi-node DP's mutual-exclusion matrix (M-CFG). + +One case per row: the combination is rejected, and the message says what to do +about it. A validator that only reports "invalid" moves the work to whoever +reads the traceback, which for a hang-shaped failure is the expensive path. + +`_validate_multinode` reads attributes and nothing else, so it is exercised +against a namespace rather than a real Config -- building one loads an HF +config off disk, which has nothing to do with what is under test. +""" + +from types import SimpleNamespace + +import pytest +from import_guard import skip_if_dependency_missing + +try: + from atom.config import Config, ParallelConfig +except ImportError as _e: # transformers/aiter absent on a bare runner + skip_if_dependency_missing(_e, "requires atom.config import env") + +# Whatever get_open_port() returned when ParallelConfig was defined. Leaving the +# field at this value is what "the user never passed --data-parallel-base-port" +# looks like from inside __post_init__. +BASE_PORT_DEFAULT = next( + f.default for f in ParallelConfig.__dataclass_fields__.values() + if f.name == "data_parallel_base_port" +) + + +def _pc(**over): + """Node 1 of a two-node, eight-engines-per-node run.""" + kw = { + "data_parallel_size": 16, + "data_parallel_size_local": 8, + "data_parallel_rank": 8, + "data_parallel_base_port": 40000, + "data_parallel_master_ip": "10.0.0.1", + "decode_context_parallel_size": 1, + } + kw.update(over) + pc = SimpleNamespace(**kw) + pc.is_multinode_dp = ( + kw["data_parallel_size_local"] < kw["data_parallel_size"] + or kw["data_parallel_rank"] > 0 + ) + return pc + + +def _cfg(*, pc=None, **over): + """A config that passes every multi-node check unless a field is overridden.""" + kw = { + "parallel_config": pc if pc is not None else _pc(), + "enable_dp_attention": True, + "enable_expert_parallel": True, + "pipeline_parallel_size": 1, + "prefill_context_parallel_size": 1, + "moe_backend": "standard", + "enable_rapidserve": False, + "plugin_config": None, + "enable_tbo": False, + } + kw.update(over) + return SimpleNamespace(**kw) + + +def _reject(cfg) -> str: + with pytest.raises(ValueError) as exc: + Config._validate_multinode(cfg) + return str(exc.value) + + +class TestAllowed: + def test_the_target_configuration_passes(self): + Config._validate_multinode(_cfg()) + + def test_single_node_skips_every_check(self): + # Same violations, one node: none of them are multi-node's business. + pc = _pc(data_parallel_size=8, data_parallel_size_local=8, data_parallel_rank=0) + Config._validate_multinode( + _cfg( + pc=pc, + enable_dp_attention=False, + enable_expert_parallel=False, + moe_backend="mega", + enable_rapidserve=True, + ) + ) + + +class TestRejected: + """One row of the matrix each. Every message must carry a fix.""" + + @pytest.mark.parametrize( + "override, expected", + [ + ({"enable_dp_attention": False}, "--enable-dp-attention"), + ({"enable_expert_parallel": False}, "--enable-expert-parallel"), + ({"pipeline_parallel_size": 2}, "pipeline_parallel_size"), + ({"prefill_context_parallel_size": 2}, "prefill_context_parallel_size"), + ({"moe_backend": "mega"}, "moe-backend"), + ({"enable_rapidserve": True}, "rapidserve"), + ({"plugin_config": object()}, "plugin"), + ], + ) + def test_row_is_rejected_with_a_fix(self, override, expected): + message = _reject(_cfg(**override)) + assert expected in message + assert "Why:" in message and "Fix:" in message + + def test_decode_context_parallel_is_rejected(self): + message = _reject(_cfg(pc=_pc(decode_context_parallel_size=2))) + assert "decode_context_parallel_size" in message + assert "Fix:" in message + + def test_message_names_the_topology(self): + message = _reject(_cfg(enable_dp_attention=False)) + assert "2 nodes" in message + assert "data_parallel_size=16" in message + + +class TestRendezvous: + def test_unset_base_port_is_rejected(self): + # The default is per-process, so two nodes never agree on it and the + # symptom is a hang rather than a connection error. + message = _reject(_cfg(pc=_pc(data_parallel_base_port=BASE_PORT_DEFAULT))) + assert "data-parallel-base-port" in message + assert "silent hang" in message + + @pytest.mark.parametrize("ip", ["127.0.0.1", "localhost", "::1"]) + def test_loopback_warns_but_is_allowed(self, ip, caplog): + # Gate 1 runs both logical nodes on one box, where loopback is correct. + # Rejecting it would block our own single-machine verification path. + with caplog.at_level("WARNING"): + Config._validate_multinode(_cfg(pc=_pc(data_parallel_master_ip=ip))) + assert "loopback" in caplog.text or "this machine" in caplog.text + + def test_routable_ip_is_silent(self, caplog): + with caplog.at_level("WARNING"): + Config._validate_multinode(_cfg(pc=_pc(data_parallel_master_ip="10.0.0.1"))) + assert caplog.text == "" + + +class TestTboWarning: + def test_tbo_warns_and_continues(self, caplog): + with caplog.at_level("WARNING"): + Config._validate_multinode(_cfg(enable_tbo=True)) + assert "TBO" in caplog.text diff --git a/tests/test_wideep_mxfp4_method.py b/tests/test_wideep_mxfp4_method.py new file mode 100644 index 0000000000..7fb8089dae --- /dev/null +++ b/tests/test_wideep_mxfp4_method.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: MIT + +from types import SimpleNamespace + +import pytest +import torch +from import_guard import skip_if_dependency_missing + +try: + import atom.model_ops.moe as moe_mod +except ImportError as exc: + skip_if_dependency_missing(exc, "requires full atom import env") + + +def _parallel(**overrides): + values = { + "use_ep": True, + "ep_size": 16, + "local_ep_size": 8, + "dp_logical_ratio": 1, + "use_all2all_kernels": True, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _moe(**overrides): + values = { + "moe_parallel_config": _parallel(), + "a_quant_dtype": "fp8_e4m3", + "expert_layout": SimpleNamespace(mode=moe_mod.SharedExpertMode.NONE), + "has_bias": False, + "in_dtype": torch.bfloat16, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _quant(**overrides): + values = {"is_dynamic": True} + values.update(overrides) + return SimpleNamespace(**values) + + +def _patch_supported_runtime(monkeypatch): + monkeypatch.setattr(moe_mod, "get_gfx", lambda: "gfx950") + monkeypatch.setattr(moe_mod.envs, "ATOM_MOE_GU_ITLV", True) + monkeypatch.setattr( + moe_mod, + "get_current_atom_config", + lambda: SimpleNamespace(moe_backend="standard", eplb_enable=False), + ) + + +def test_ep16_two_node_selects_wideep(monkeypatch): + _patch_supported_runtime(monkeypatch) + standard_method = object() + wideep_method = object() + monkeypatch.setattr(moe_mod, "Mxfp4MoEMethod", lambda *_args: standard_method) + monkeypatch.setattr(moe_mod, "WideEpMxfp4MoEMethod", lambda *_args: wideep_method) + + # TestWideEpMoe accepts BF16 hidden states and performs its own FP8 + # per-1x32 quantization, so the model need not declare FP8 activations. + selected = moe_mod._make_mxfp4_moe_method(_quant(), _moe(a_quant_dtype=None)) + + assert selected is wideep_method + + +@pytest.mark.parametrize( + "parallel", + [ + _parallel(ep_size=8, local_ep_size=8), + _parallel(ep_size=16, local_ep_size=16), + _parallel(use_ep=False, ep_size=1, local_ep_size=1), + ], +) +def test_non_wideep_topology_keeps_standard_backend(monkeypatch, parallel): + _patch_supported_runtime(monkeypatch) + standard_method = object() + wideep_method = object() + monkeypatch.setattr(moe_mod, "Mxfp4MoEMethod", lambda *_args: standard_method) + monkeypatch.setattr(moe_mod, "WideEpMxfp4MoEMethod", lambda *_args: wideep_method) + + selected = moe_mod._make_mxfp4_moe_method( + _quant(), _moe(moe_parallel_config=parallel) + ) + + assert selected is standard_method + + +@pytest.mark.parametrize( + ("quant", "moe", "config", "env_interleave", "gfx", "message"), + [ + (_quant(is_dynamic=False), _moe(), {}, True, "gfx950", "dynamic"), + (_quant(), _moe(), {"eplb_enable": True}, True, "gfx950", "EPLB"), + ( + _quant(), + _moe( + expert_layout=SimpleNamespace( + mode=moe_mod.SharedExpertMode.LOCAL_REPLICA + ) + ), + {}, + True, + "gfx950", + "shared", + ), + (_quant(), _moe(has_bias=True), {}, True, "gfx950", "bias"), + (_quant(), _moe(), {}, False, "gfx950", "ATOM_MOE_GU_ITLV"), + (_quant(), _moe(), {}, True, "gfx942", "gfx950"), + ], +) +def test_unsupported_ep16_configuration_is_rejected( + monkeypatch, quant, moe, config, env_interleave, gfx, message +): + monkeypatch.setattr(moe_mod, "get_gfx", lambda: gfx) + monkeypatch.setattr(moe_mod.envs, "ATOM_MOE_GU_ITLV", env_interleave) + config_values = {"moe_backend": "standard", "eplb_enable": False} + config_values.update(config) + monkeypatch.setattr( + moe_mod, + "get_current_atom_config", + lambda: SimpleNamespace(**config_values), + ) + + with pytest.raises(ValueError, match=message): + moe_mod._validate_wideep_mxfp4_config(quant, moe) + + +def test_wideep_wrapper_quantizes_and_uses_prequant_entry(monkeypatch): + from atom.model_ops.fused_moe import wideep_experts + + calls = [] + + class FakeOp: + def quantize(self, hidden): + calls.append(("quantize", hidden)) + return "x_quant", "x_scale" + + def forward_prequant(self, x_quant, x_scale, weights, ids): + calls.append(("forward_prequant", x_quant, x_scale, weights, ids)) + return "output" + + monkeypatch.setattr( + wideep_experts, "get_or_build_wideep_moe", lambda **_kw: FakeOp() + ) + + manager = SimpleNamespace(rank=0, world_size=16) + group = SimpleNamespace( + device_communicator=SimpleNamespace(all2all_manager=manager) + ) + from aiter.dist import parallel_state + + monkeypatch.setattr(parallel_state, "get_ep_group", lambda: group) + hidden = torch.ones(2, 4, dtype=torch.bfloat16) + weights = torch.ones(2, 2, dtype=torch.bfloat16).t() + ids = torch.ones(2, 2, dtype=torch.int64).t() + layer = SimpleNamespace( + _wideep_w1="w1", + _wideep_w1_scale="w1_scale", + _wideep_w2="w2", + _wideep_w2_scale="w2_scale", + ) + + output = wideep_experts.run_wideep_moe( + layer, + hidden, + weights, + ids, + model_dim=4, + inter_dim=8, + experts=16, + topk=2, + mtpr=8, + swiglu_limit=10.0, + ) + + assert output == "output" + assert calls[0][0] == "quantize" + assert calls[0][1] is hidden + assert calls[1][0:3] == ("forward_prequant", "x_quant", "x_scale") + assert calls[1][3].dtype == torch.float32 + assert calls[1][3].is_contiguous() + assert calls[1][4].dtype == torch.int32 + assert calls[1][4].is_contiguous() + + +def test_wideep_cache_key_includes_limit_and_rebinds_weights(monkeypatch): + from atom.model_ops.fused_moe import wideep_experts + + built = [] + + class FakeTestWideEpMoe: + def __init__(self, **kwargs): + built.append(kwargs) + self.w1 = kwargs["w1"] + self.w1_scale = kwargs["w1_scale"] + self.w2 = kwargs["w2"] + self.w2_scale = kwargs["w2_scale"] + + fake_module = SimpleNamespace(TestWideEpMoe=FakeTestWideEpMoe) + monkeypatch.setitem( + __import__("sys").modules, + "aiter.ops.flydsl.test_wide_ep_moe", + fake_module, + ) + wideep_experts._WIDEEP_CACHE.clear() + common = { + "rank": 0, + "world_size": 16, + "model_dim": 16, + "inter_dim": 8, + "experts": 32, + "topk": 2, + "quant": "a8w4", + "mtpr": 256, + } + + first = wideep_experts.get_or_build_wideep_moe( + **common, + swiglu_limit=10.0, + w1="layer0_w1", + w1_scale="layer0_s1", + w2="layer0_w2", + w2_scale="layer0_s2", + ) + rebound = wideep_experts.get_or_build_wideep_moe( + **common, + swiglu_limit=10.0, + w1="layer1_w1", + w1_scale="layer1_s1", + w2="layer1_w2", + w2_scale="layer1_s2", + ) + different_limit = wideep_experts.get_or_build_wideep_moe( + **common, + swiglu_limit=0.0, + w1="layer2_w1", + w1_scale="layer2_s1", + w2="layer2_w2", + w2_scale="layer2_s2", + ) + + assert first is rebound + assert rebound.w1 == "layer1_w1" + assert rebound.w1_scale == "layer1_s1" + assert rebound.w2 == "layer1_w2" + assert rebound.w2_scale == "layer1_s2" + assert different_limit is not first + assert [call["swiglu_limit"] for call in built] == [10.0, 0.0] + + +def test_compiled_boundary_receives_and_rebinds_layer_weights(monkeypatch): + from atom.model_ops.fused_moe import wideep_experts + + captured = {} + + class FakeOp: + def forward_prequant(self, x_quant, x_scale, weights, ids): + captured["inputs"] = (x_quant, x_scale, weights, ids) + return torch.ones((ids.shape[0], 4), dtype=torch.bfloat16) + + def fake_get_or_build(**kwargs): + captured["builder"] = kwargs + return FakeOp() + + monkeypatch.setattr(wideep_experts, "get_or_build_wideep_moe", fake_get_or_build) + x_quant = torch.ones((2, 4), dtype=torch.float8_e4m3fn) + x_scale = torch.ones((2, 1), dtype=torch.uint8) + weights = torch.ones((2, 2), dtype=torch.float32) + ids = torch.ones((2, 2), dtype=torch.int32) + w1 = torch.ones((2, 2)) + w1_scale = torch.ones((2, 2)) + w2 = torch.ones((2, 2)) + w2_scale = torch.ones((2, 2)) + + output = wideep_experts.atom_wideep_forward_impl( + x_quant, + x_scale, + weights, + ids, + w1, + w1_scale, + w2, + w2_scale, + 0, + 16, + 4, + 8, + 32, + 2, + 256, + 10.0, + ) + + assert output.shape == (2, 4) + assert captured["builder"]["w1"] is w1 + assert captured["builder"]["w1_scale"] is w1_scale + assert captured["builder"]["w2"] is w2 + assert captured["builder"]["w2_scale"] is w2_scale + assert captured["builder"]["swiglu_limit"] == 10.0 diff --git a/tests/test_wideep_timeouts.py b/tests/test_wideep_timeouts.py new file mode 100644 index 0000000000..ffd091eb84 --- /dev/null +++ b/tests/test_wideep_timeouts.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Bounds on the barriers multi-node can hang at (M-DPSYNC). + +The failure being designed against is not an exception, it is silence: a node +that never started leaves every other rank waiting on a barrier that logs +nothing. What is tested here is the message that replaces the silence, and the +port sequence that has to agree across ranks for a group to form at all. +""" + +from types import SimpleNamespace + +import pytest +from import_guard import skip_if_dependency_missing + +try: + from atom.config import ParallelConfig + from atom.model_engine.engine_core_mgr import CoreManager +except ImportError as _e: # transformers/zmq absent on a bare runner + skip_if_dependency_missing(_e, "requires atom engine import env") + + +class TestReadyTimeoutMessage: + """The coordinator is the only process that can say which node is missing.""" + + def _message(self, ready, *, per_node=4, timeout_s=120.0): + mgr = SimpleNamespace(label="Engine Core Mgr") + return CoreManager._ready_timeout_message(mgr, ready, timeout_s, per_node) + + def test_groups_missing_ranks_by_node(self): + # 8 engines over 2 nodes; the whole second node never reported. + message = self._message([True] * 4 + [False] * 4) + assert "node 1: dp ranks [4, 5, 6, 7]" in message + assert "node 0" not in message + + def test_counts_against_the_total(self): + message = self._message([True] * 4 + [False] * 4) + assert "4 of 8 engines" in message + + def test_names_the_timeout(self): + message = self._message([False] * 8, timeout_s=90.0) + assert "90s" in message + + def test_partial_node_is_reported_as_such(self): + # One engine down on an otherwise healthy node is a different problem + # from a node that never launched, and the message has to distinguish. + message = self._message([True] * 5 + [False] + [True] * 2) + assert "node 1: dp ranks [5]" in message + + def test_falls_back_when_node_width_is_unknown(self): + mgr = SimpleNamespace(label="Engine Core Mgr") + message = CoreManager._ready_timeout_message( + mgr, [True, False, False], 30.0, None + ) + assert "dp ranks [1, 2]" in message + + def test_says_what_to_check(self): + message = self._message([True] * 4 + [False] * 4) + assert "Fix:" in message + assert "--data-parallel-master-ip" in message + + +class TestInitPortSequence: + """All ranks must walk the same port sequence or the group never forms.""" + + def _pc(self, rank): + return ParallelConfig( + data_parallel_size=8, + data_parallel_size_local=4, + data_parallel_rank=rank, + data_parallel_master_port=29500, + ) + + def test_every_rank_gets_the_same_first_port(self): + assert {self._pc(r).get_next_dp_init_port() for r in (0, 4)} == {29500} + + def test_every_rank_gets_the_same_second_port(self): + # The previous rule advanced by data_parallel_rank, so rank 0 stayed on + # the base port while every other rank moved a different distance -- + # the second group was formed on a different port per rank. + seconds = set() + for rank in (0, 4): + pc = self._pc(rank) + pc.get_next_dp_init_port() + seconds.add(pc.get_next_dp_init_port()) + assert seconds == {29501} + + def test_sequence_is_contiguous(self): + pc = self._pc(0) + assert [pc.get_next_dp_init_port() for _ in range(4)] == [ + 29500, + 29501, + 29502, + 29503, + ] + + +class TestDpSyncTimeout: + def test_none_means_backend_default(self): + pc = ParallelConfig(data_parallel_size=1) + assert pc.dp_sync_timeout is None + + def test_seconds_become_a_timedelta(self): + pc = ParallelConfig(data_parallel_size=1, dp_sync_timeout_s=90) + assert pc.dp_sync_timeout.total_seconds() == pytest.approx(90.0) diff --git a/tests/test_wideep_topology.py b/tests/test_wideep_topology.py new file mode 100644 index 0000000000..b186bf5079 --- /dev/null +++ b/tests/test_wideep_topology.py @@ -0,0 +1,421 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Unit tests for M-TOPO (WideEPTopology).""" + +from types import SimpleNamespace + +import pytest + +from atom.model_engine.topology import ( + WideEPTopology, + format_startup_summary, + node_count, + parse_dist_init_addr, +) + + +def _legacy_local_engine_count( + *, dp_attention: bool, raw_tp: int, raw_dp: int, pp_size: int = 1 +) -> int: + """Mirror CoreManager's engine-count fold for Gate 0 regression. + + engine_core_mgr.py: iter_dp_rank_assignments + the enable_dp_attention + branch. Excludes the --fake-eplb shrink, which the topology refuses to + describe (see TestSimulatedDeployment). + """ + if dp_attention: + assert pp_size == 1 + return raw_tp * raw_dp + return raw_dp * pp_size + + +class TestParseDistInitAddr: + def test_ipv4(self): + assert parse_dist_init_addr("10.0.0.5:29500") == ("10.0.0.5", 29500) + + def test_ipv6(self): + assert parse_dist_init_addr("[::1]:30000") == ("::1", 30000) + + def test_invalid(self): + with pytest.raises(ValueError, match="Invalid dist_init_addr"): + parse_dist_init_addr("no-port") + + +class TestGate0Regression: + """nnodes=1 values must match today's engine_core_mgr derivation.""" + + @pytest.mark.parametrize( + "raw_tp, raw_dp", + [(8, 1), (4, 2), (2, 4)], + ) + def test_dp_attention_matches_legacy(self, raw_tp, raw_dp): + topo = WideEPTopology.create( + dp_attention=True, + raw_tp_size=raw_tp, + raw_dp_size=raw_dp, + ) + expected_count = _legacy_local_engine_count( + dp_attention=True, raw_tp=raw_tp, raw_dp=raw_dp + ) + assert topo.tp_size == 1 + assert topo.global_dp_size == raw_tp * raw_dp + assert topo.local_engine_count == expected_count + assert topo.ep_size == topo.global_dp_size + assert topo.gpu_per_node == expected_count + assert not topo.is_multinode + assert topo.dist_init_host is None + assert topo.dist_init_base_port is None + + for i in range(expected_count): + assert topo.dp_rank(i) == i + assert topo.dp_rank_local(i) == i + assert topo.local_device_rank(i, tp_rank=0) == i + + @pytest.mark.parametrize("raw_dp", [1, 2, 4]) + def test_no_dp_attention_matches_legacy(self, raw_dp): + topo = WideEPTopology.create( + dp_attention=False, + raw_tp_size=8, + raw_dp_size=raw_dp, + ) + expected_count = _legacy_local_engine_count( + dp_attention=False, raw_tp=8, raw_dp=raw_dp + ) + assert topo.tp_size == 8 + assert topo.global_dp_size == raw_dp + assert topo.local_engine_count == expected_count + # Without the DP-attention flatten, FusedMoEParallelConfig.make leaves + # ep_size at tp_size: EP spans a TP group and there are raw_dp + # independent groups. Asserting the multinode identity here instead + # would only be checking the view against itself. + assert topo.ep_size == 8 + assert topo.gpu_per_node == raw_dp * 8 + + +class TestMultinodeLayout: + def test_gate1_logical_node(self): + topo0 = WideEPTopology.create( + nnodes=2, + node_rank=0, + dp_attention=True, + raw_tp_size=8, + raw_dp_size=1, + dist_init_addr="127.0.0.1:29500", + ) + topo1 = WideEPTopology.create( + nnodes=2, + node_rank=1, + dp_attention=True, + raw_tp_size=8, + raw_dp_size=1, + dist_init_addr="127.0.0.1:29500", + ) + assert topo0.ep_size == 8 + assert topo0.gpu_per_node == 4 + assert topo0.local_engine_count == 4 + assert topo0.rendezvous_port_world == 29500 + assert topo0.rendezvous_port_dp_gloo == 29501 + assert topo0.rendezvous_port_reserved(0) == 29502 + assert topo0.rendezvous_port_reserved(5) == 29507 + + assert topo0.dp_rank(0) == 0 + assert topo0.dp_rank(3) == 3 + assert topo1.dp_rank(0) == 4 + assert topo1.dp_rank(3) == 7 + assert topo0.dp_rank_local(2) == 2 + assert topo1.dp_rank_local(2) == 2 + + def test_ep16_two_nodes(self): + topo = WideEPTopology.create( + nnodes=2, + node_rank=0, + dp_attention=True, + raw_tp_size=8, + raw_dp_size=2, + dist_init_addr="192.168.1.10:30000", + ) + assert topo.ep_size == 16 + assert topo.gpu_per_node == 8 + assert topo.dist_init_host == "192.168.1.10" + + +class _FakeParallelConfig(SimpleNamespace): + """Duck-typed stand-in: atom.config pulls torch, which unit tests avoid.""" + + def __init__(self, **kw): + kw.setdefault("data_parallel_master_ip", "10.0.0.1") + kw.setdefault("data_parallel_master_port", 29500) + super().__init__(**kw) + + +class TestFromParallelConfig: + """Derivation from the DP fields CoreManager owns.""" + + def test_fold_is_a_change_of_units(self): + """Pre- and post-fold ParallelConfig must yield the same topology. + + CoreManager rewrites the DP fields into engine units (scaling them by + tp_size and setting tp_size to 1). Every quantity this view exposes is + a ratio or a product of those, so the rewrite must not move it. This + test is what makes that invariant executable rather than assumed. + """ + pre = _FakeParallelConfig( + data_parallel_size=2, data_parallel_size_local=1, data_parallel_rank=1 + ) + post = _FakeParallelConfig( + data_parallel_size=16, data_parallel_size_local=8, data_parallel_rank=8 + ) + a = WideEPTopology.from_parallel_config( + pre, tensor_parallel_size=8, dp_attention=True + ) + b = WideEPTopology.from_parallel_config( + post, tensor_parallel_size=1, dp_attention=True + ) + assert a == b + assert (a.nnodes, a.node_rank, a.ep_size, a.gpu_per_node) == (2, 1, 16, 8) + + def test_post_fold_worker_rank_is_normalized_to_node_offset(self): + pc = _FakeParallelConfig( + data_parallel_size=16, + data_parallel_size_local=8, + data_parallel_rank=13, + data_parallel_rank_local=5, + ) + topo = WideEPTopology.from_parallel_config( + pc, tensor_parallel_size=1, dp_attention=True + ) + assert (topo.nnodes, topo.node_rank) == (2, 1) + assert (topo.ep_size, topo.gpu_per_node) == (16, 8) + + def test_single_node_defaults_local_to_global(self): + pc = _FakeParallelConfig( + data_parallel_size=1, data_parallel_size_local=None, data_parallel_rank=0 + ) + topo = WideEPTopology.from_parallel_config( + pc, tensor_parallel_size=8, dp_attention=True + ) + assert not topo.is_multinode + assert (topo.nnodes, topo.ep_size, topo.gpu_per_node) == (1, 8, 8) + + def test_ragged_split_rejected(self): + """EPLB's hierarchical placement asserts num_gpus % num_nodes == 0.""" + pc = _FakeParallelConfig( + data_parallel_size=8, data_parallel_size_local=3, data_parallel_rank=0 + ) + with pytest.raises(ValueError, match="must be divisible by"): + WideEPTopology.from_parallel_config( + pc, tensor_parallel_size=1, dp_attention=True + ) + + def test_rank_offset_must_align_to_slice(self): + pc = _FakeParallelConfig( + data_parallel_size=8, data_parallel_size_local=4, data_parallel_rank=3 + ) + with pytest.raises(ValueError, match="must be a multiple of"): + WideEPTopology.from_parallel_config( + pc, tensor_parallel_size=1, dp_attention=True + ) + + +class TestSimulatedDeployment: + """--fake-eplb makes the DP fields describe a width, not a node split.""" + + def test_simulation_is_rejected_not_guessed(self): + # -tp 8 on a 4-GPU box: post-fold the fields read exactly like node 0 of + # a real 2-node deployment, so the ratio must not be taken as a node + # count. Callers that support simulation handle it explicitly. + pc = _FakeParallelConfig( + data_parallel_size=8, data_parallel_size_local=4, data_parallel_rank=0 + ) + with pytest.raises(ValueError, match="--fake-eplb"): + WideEPTopology.from_parallel_config( + pc, tensor_parallel_size=1, dp_attention=True, dp_logical_size=8 + ) + + def test_not_simulating_is_the_default(self): + pc = _FakeParallelConfig( + data_parallel_size=8, data_parallel_size_local=4, data_parallel_rank=0 + ) + topo = WideEPTopology.from_parallel_config( + pc, tensor_parallel_size=1, dp_attention=True, dp_logical_size=0 + ) + assert topo.nnodes == 2 + + +class TestStartupSummary: + def _topo(self): + return WideEPTopology.create( + nnodes=2, + node_rank=1, + dp_attention=True, + raw_tp_size=8, + raw_dp_size=2, + dist_init_addr="10.0.0.1:29500", + ) + + def test_carries_every_field_needed_to_compare_two_nodes(self): + line = format_startup_summary( + self._topo(), + dp_rank=9, + dp_rank_local=1, + device_rank=1, + visible_devices="0,1,2,3,4,5,6,7", + mori_env={"MORI_SHMEM_HEAP_SIZE": "8G"}, + ) + assert line.startswith("[wideep] ") + for field in ( + "nnodes=2", + "node_rank=1", + "ep=16", + "gpu_per_node=8", + "global=16", + "local=8", + "rank=9", + "local_rank=1", + "device=cuda:1", + "visible=0,1,2,3,4,5,6,7", + "MORI_SHMEM_HEAP_SIZE=8G", + ): + assert field in line, field + + def test_optional_parts_are_omitted_not_blank(self): + line = format_startup_summary( + self._topo(), dp_rank=9, dp_rank_local=1, device_rank=1 + ) + assert "visible=" not in line + assert "mori:" not in line + assert "device=cuda:1" in line + + def test_stays_on_one_line(self): + line = format_startup_summary( + self._topo(), + dp_rank=9, + dp_rank_local=1, + device_rank=1, + visible_devices="0,1", + mori_env={"A": "1", "B": "2"}, + ) + assert "\n" not in line + + +class TestNodeCount: + """The one place that answers "how many nodes" including the simulated case.""" + + def _config(self, pc, *, tp, dp_attention=True, dp_logical_size=0): + return SimpleNamespace( + parallel_config=pc, + tensor_parallel_size=tp, + enable_dp_attention=dp_attention, + dp_logical_size=dp_logical_size, + ) + + def test_reports_the_split(self): + pc = _FakeParallelConfig( + data_parallel_size=16, data_parallel_size_local=8, data_parallel_rank=8 + ) + assert node_count(self._config(pc, tp=1)) == 2 + + def test_simulation_reports_one_box(self): + # Same fields as node 0 of a real 2-node run; only dp_logical_size + # distinguishes them, which is why this rule lives in one place. + pc = _FakeParallelConfig( + data_parallel_size=8, data_parallel_size_local=4, data_parallel_rank=0 + ) + assert node_count(self._config(pc, tp=1, dp_logical_size=8)) == 1 + + def test_missing_field_is_not_simulating(self): + # Older Config objects and test doubles have no dp_logical_size. + pc = _FakeParallelConfig( + data_parallel_size=8, data_parallel_size_local=4, data_parallel_rank=0 + ) + cfg = SimpleNamespace( + parallel_config=pc, tensor_parallel_size=1, enable_dp_attention=True + ) + assert node_count(cfg) == 2 + + +class TestValidation: + def test_global_dp_not_divisible_suggests_nnodes(self): + with pytest.raises(ValueError, match="Valid nnodes values"): + WideEPTopology.create( + nnodes=3, + dp_attention=True, + raw_tp_size=8, + raw_dp_size=1, + dist_init_addr="10.0.0.1:29500", + ) + + def test_multinode_requires_dist_init_addr(self): + with pytest.raises(ValueError, match="dist_init_addr"): + WideEPTopology.create( + nnodes=2, + dp_attention=True, + raw_tp_size=8, + raw_dp_size=1, + ) + + def test_multinode_requires_dp_attention(self): + with pytest.raises(ValueError, match="dp_attention"): + WideEPTopology.create( + nnodes=2, + dp_attention=False, + raw_tp_size=8, + raw_dp_size=1, + dist_init_addr="10.0.0.1:29500", + ) + + def test_invalid_node_rank(self): + with pytest.raises(ValueError, match="node_rank"): + WideEPTopology.create( + nnodes=2, + node_rank=2, + dp_attention=True, + raw_tp_size=8, + raw_dp_size=1, + dist_init_addr="10.0.0.1:29500", + ) + + def test_rendezvous_ports_require_multinode(self): + topo = WideEPTopology.create(dp_attention=True, raw_tp_size=8, raw_dp_size=1) + with pytest.raises(ValueError, match="rendezvous ports"): + _ = topo.rendezvous_port_world + + def test_reserved_port_index_bounds(self): + topo = WideEPTopology.create( + nnodes=2, + dp_attention=True, + raw_tp_size=8, + raw_dp_size=1, + dist_init_addr="10.0.0.1:29500", + ) + with pytest.raises(ValueError, match="reserved port index"): + topo.rendezvous_port_reserved(6) + + +class TestLocalDeviceRank: + def test_with_pcp(self): + topo = WideEPTopology.create( + dp_attention=True, + raw_tp_size=4, + raw_dp_size=2, + ) + # tp_size=1, pcp=2 => stage_span=2 + assert topo.local_device_rank(1, tp_rank=0, pcp_size=2) == 2 + assert topo.local_device_rank(1, tp_rank=1, pcp_size=2) == 3 + + +class TestDescribe: + def test_one_line_summary(self): + topo = WideEPTopology.create( + nnodes=2, + node_rank=1, + dp_attention=True, + raw_tp_size=8, + raw_dp_size=2, + dist_init_addr="10.0.0.1:29500", + ) + assert topo.describe() == ( + "[wideep] nnodes=2 node_rank=1 | ep=16 gpu_per_node=8 | " + "dp: global=16 local=8" + )