feat(wideep): multi-node wide expert parallelism on top of #1603

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 <routable-ip>.",
+                pc.data_parallel_master_ip,
+            )
+
     def compute_hash(self) -> str:
         """
         WARNING: Whenever a new field is added to this config,

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}")

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,18 +483,38 @@ 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)
 
         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: