From ce21227508538a03624e25ade35ce23387390eca Mon Sep 17 00:00:00 2001 From: Guanbao Yu Date: Fri, 16 Oct 2026 10:07:56 +0000 Subject: [PATCH 1/6] [TP] Run one tensor-parallel group across nodes -tp 8 on two 4-GPU nodes: the node that runs TP rank 0 runs the engine, the API server and ranks 0-3; the other runs only ranks 4-7. node0: -tp 8 --tensor-parallel-size-local 4 \ --data-parallel-master-ip --data-parallel-base-port node1: the same, plus --tensor-parallel-rank 4 The engine reaches the remote ranks through AsyncIOProcManager's existing channels: the broadcast queue gets remote readers (aiter's MessageQueue already serves them over ZMQ), and their reply channels listen on TCP. The manager publishes those endpoints and the runner's constructor arguments in the model-runner TCPStore CoreManager serves, waits for every remote reader to subscribe (a PUB socket drops what it sends before that), then releases them. The remote node spawns AsyncIOProc workers from that and stops them if one fails or the engine's store goes away. ModelRunner places each worker on its node's GPU (rank - tp_rank_start) and takes the minimum KV block count over the TP group. Barrier RPCs are refused: a multiprocessing barrier does not reach another node. DP, DP attention, PP, PCP, --fake-eplb and RapidServe are rejected with a cross-node TP group. --- atom/config.py | 61 ++++++++ atom/entrypoints/openai/api_server.py | 7 + atom/model_engine/arg_utils.py | 26 ++++ atom/model_engine/async_proc.py | 74 +++++++++- atom/model_engine/engine_core.py | 1 + atom/model_engine/engine_core_mgr.py | 8 +- atom/model_engine/model_runner.py | 12 +- atom/model_engine/tp_multinode.py | 202 ++++++++++++++++++++++++++ atom/utils/envs.py | 6 + docs/environment_variables.md | 1 + 10 files changed, 385 insertions(+), 13 deletions(-) create mode 100644 atom/model_engine/tp_multinode.py diff --git a/atom/config.py b/atom/config.py index 120a49687c..89905f8edf 100644 --- a/atom/config.py +++ b/atom/config.py @@ -1049,6 +1049,14 @@ class ParallelConfig: data_parallel_master_ip: str = "127.0.0.1" + tensor_parallel_size_local: int | None = None + """TP ranks (model-runner workers) this node runs. None means all of them, + the single-node case. Set it below tensor_parallel_size to spread one TP + group over several nodes; see atom/model_engine/tp_multinode.py.""" + tensor_parallel_rank: int = 0 + """First TP rank this node runs. The node that runs rank 0 also runs the + engine and the API server; the others run only their workers.""" + @property def is_multinode_dp(self) -> bool: """Whether this node owns only part of the global DP group. @@ -2142,6 +2150,57 @@ def tp_world_size(self) -> int: visible = torch.cuda.device_count() return visible if 0 < visible < tp else tp + @property + def tp_size_local(self) -> int: + """TP worker processes on this node: fewer than `tp_world_size` when + one TP group spans several nodes.""" + local = self.parallel_config.tensor_parallel_size_local + return self.tp_world_size if local is None else local + + @property + def tp_rank_start(self) -> int: + """First TP rank on this node; it runs `tp_size_local` ranks from here.""" + return self.parallel_config.tensor_parallel_rank + + @property + def is_multinode_tp(self) -> bool: + return self.tp_size_local < self.tp_world_size or self.tp_rank_start > 0 + + def _check_multinode_tp(self) -> None: + pc = self.parallel_config + local, start, tp = self.tp_size_local, self.tp_rank_start, self.tp_world_size + if local < 1 or tp % local or start % local or not 0 <= start < tp: + raise ValueError( + f"--tensor-parallel-size-local {local} and --tensor-parallel-rank " + f"{start} must cut -tp {tp} into whole, equal per-node slices" + ) + unsupported = [ + name + for name, used in ( + ("data parallelism", pc.data_parallel_size > 1), + ("--enable-dp-attention", self.enable_dp_attention), + ("pipeline parallelism", self.pipeline_parallel_size > 1), + ("prefill context parallelism", self.prefill_context_parallel_size > 1), + ("--fake-eplb", self.fake_eplb), + ("RapidServe", self.enable_rapidserve), + ) + if used + ] + if unsupported: + raise ValueError( + "tensor parallelism across nodes does not support " + + ", ".join(unsupported) + ) + if pc.data_parallel_base_port == 0 or pc.data_parallel_master_ip in ( + "127.0.0.1", + "localhost", + ): + raise ValueError( + "tensor parallelism across nodes needs --data-parallel-master-ip " + "set to the address of the node that runs TP rank 0, and the " + "same nonzero --data-parallel-base-port on every node" + ) + @property def capture_sizes(self) -> list[int]: """The declared CUDAGraph capture ladder, in batch sizes. @@ -2250,6 +2309,8 @@ def __post_init__(self): ) assert 1 <= self.tensor_parallel_size <= 8 + if self.is_multinode_tp: + self._check_multinode_tp() if self.decode_context_parallel_size > 1: assert self.tensor_parallel_size % self.decode_context_parallel_size == 0, ( f"tp_size ({self.tensor_parallel_size}) must be divisible by " diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 3b038e1ede..a3b3913df1 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -2915,6 +2915,13 @@ def main(): ) args = parser.parse_args() + if args.tensor_parallel_rank > 0: + # The node that runs TP rank 0 serves the API; this one only runs the + # workers of its own ranks. + from atom.model_engine.tp_multinode import serve_remote_tp_ranks + + raise SystemExit(serve_remote_tp_ranks(EngineArgs.from_cli_args(args))) + if args.request_log: _request_logger = logging.getLogger("atom.request_log") _request_logger.setLevel(logging.INFO) diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index a3f69fa263..c4044e887d 100644 --- a/atom/model_engine/arg_utils.py +++ b/atom/model_engine/arg_utils.py @@ -119,6 +119,8 @@ class EngineArgs: model: str = "Qwen/Qwen3-0.6B" trust_remote_code: bool = False tensor_parallel_size: int = 1 + tensor_parallel_size_local: int | None = None + tensor_parallel_rank: int = 0 decode_context_parallel_size: int = 1 pipeline_parallel_size: int = 1 prefill_context_parallel_size: int = 1 @@ -210,6 +212,28 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: default=1, help="Tensor parallel size.", ) + parser.add_argument( + "--tensor-parallel-size-local", + type=int, + default=None, + help=( + "Number of tensor-parallel ranks to run on THIS node. Defaults " + "to --tensor-parallel-size (single-node). Set it lower to " + "spread one TP group over several nodes; every node then needs " + "--data-parallel-master-ip of the node that runs TP rank 0 and " + "the same --data-parallel-base-port." + ), + ) + parser.add_argument( + "--tensor-parallel-rank", + type=int, + default=0, + help=( + "First tensor-parallel rank owned by this node. Node 0 uses 0 " + "and serves the API; the second node of a 2x4 run uses 4 and " + "runs only its model-runner workers." + ), + ) parser.add_argument( "--pipeline-parallel-size", "-pp", @@ -908,6 +932,8 @@ def _get_engine_kwargs(self) -> dict: "data_parallel_rank": kwargs.pop("data_parallel_rank"), "data_parallel_master_ip": kwargs.pop("data_parallel_master_ip"), "data_parallel_master_port": kwargs.pop("data_parallel_master_port"), + "tensor_parallel_size_local": kwargs.pop("tensor_parallel_size_local"), + "tensor_parallel_rank": kwargs.pop("tensor_parallel_rank"), } base_port = kwargs.pop("data_parallel_base_port") if base_port is not None: diff --git a/atom/model_engine/async_proc.py b/atom/model_engine/async_proc.py index 16efabf6e4..9004290d9b 100644 --- a/atom/model_engine/async_proc.py +++ b/atom/model_engine/async_proc.py @@ -46,6 +46,7 @@ from atom.utils import ( envs, get_mp_context, + get_open_port, get_open_zmq_ipc_path, init_exit_handler, make_zmq_socket, @@ -155,7 +156,7 @@ def __init__( dp_local_rank = cfg.parallel_config.data_parallel_rank_local if dp_local_rank is None: dp_local_rank = cfg.parallel_config.data_parallel_rank - gpu = dp_local_rank * cfg.tp_world_size + rank + gpu = dp_local_rank * cfg.tp_world_size + rank - cfg.tp_rank_start numa_bind_to_node(gpu, label) except Exception as e: # noqa: BLE001 - binding is an optimization # NUMA binding only affects locality, never correctness, so any @@ -178,6 +179,10 @@ def __init__( self.rpc_queue: queue.Queue | None = None self.rpc_broadcast_mq = MessageQueue.create_from_handle(input_shm_handle, rank) + if self.rpc_broadcast_mq._is_remote_reader: + from atom.model_engine.tp_multinode import await_release + + await_release(self.rpc_broadcast_mq) import atexit atexit.register(self._cleanup_shared_memory) @@ -463,9 +468,21 @@ class AsyncIOProcManager: {*_PROTOCOL_RPC_NAMES, *UTILITY_MANAGED_RUNNER_METHODS} ) - def __init__(self, finalizer, proc_num: int, runner: str, *args): + def __init__( + self, + finalizer, + proc_num: int, + runner: str, + *args, + local_proc_num: int | None = None, + ): self.parent_finalizer = finalizer self.proc_num = proc_num + # Ranks from local_proc_num on run on other nodes (tp_multinode.py) and + # reach this node over TCP. + self.local_proc_num = proc_num if local_proc_num is None else local_proc_num + n_remote = proc_num - self.local_proc_num + host = args[0].parallel_config.data_parallel_master_ip if n_remote else None io_addrs = [get_open_zmq_ipc_path(), get_open_zmq_ipc_path()] self.procs: list[multiprocessing.Process] = [] @@ -474,7 +491,10 @@ def __init__(self, finalizer, proc_num: int, runner: str, *args): self.label = f"AsyncIOProcManager({self.runner_label})" self.rpc_broadcast_mq = MessageQueue( - proc_num, proc_num, max_chunk_bytes=16 * 1024 * 1024 + proc_num, + self.local_proc_num, + max_chunk_bytes=16 * 1024 * 1024, + connect_ip=host, ) self._block_table_encoder = BlockTableDeltaEncoder() scheduler_output_handle = self.rpc_broadcast_mq.export_handle() @@ -483,12 +503,14 @@ def __init__(self, finalizer, proc_num: int, runner: str, *args): import atexit atexit.register(self._cleanup_shared_memory) - self.all_ranks_barrier = ctx.Barrier(proc_num) + # A multiprocessing barrier cannot reach another node; barrier RPCs are + # refused instead (see _collective_rpc). + self.all_ranks_barrier = None if n_remote else ctx.Barrier(proc_num) init_exit_handler(self) # KV output aggregation infrastructure self.kv_output_aggregator: KVOutputAggregator | None = None - self.kv_output_addrs = [get_open_zmq_ipc_path() for _ in range(proc_num)] + self.kv_output_addrs = [self._reply_addr(i, host) for i in range(proc_num)] self.kv_outputs_queues: list[queue.Queue] = [ queue.Queue() for _ in range(proc_num) ] @@ -496,7 +518,7 @@ def __init__(self, finalizer, proc_num: int, runner: str, *args): self.kv_output_threads: list[threading.Thread] = [] # Generic collective-RPC reply channels, one per rank - self.rpc_output_addrs = [get_open_zmq_ipc_path() for _ in range(proc_num)] + self.rpc_output_addrs = [self._reply_addr(i, host) for i in range(proc_num)] self.rpc_outputs_queues: list[queue.Queue] = [ queue.Queue() for _ in range(proc_num) ] @@ -506,7 +528,7 @@ def __init__(self, finalizer, proc_num: int, runner: str, *args): # rank -> exit code, written by the monitor before it tears down. self._dead_ranks: dict[int, int | None] = {} - for i in range(proc_num): + for i in range(self.local_proc_num): label = f"ModelRunner{i}/{proc_num}" # Only rank 0 gets the primary output address addrs = [None, io_addrs[1]] if i == 0 else [None, None] @@ -566,6 +588,32 @@ def __init__(self, finalizer, proc_num: int, runner: str, *args): self.rpc_output_threads.append(t) self.monitor_procs() + if n_remote: + from atom.model_engine.tp_multinode import ( + publish_bootstrap, + release_remote_readers, + ) + + publish_bootstrap( + args[0], + { + "proc_num": proc_num, + "handle": scheduler_output_handle, + "runner": runner, + "args": args, + "kv_output_addrs": self.kv_output_addrs, + "rpc_output_addrs": self.rpc_output_addrs, + }, + ) + release_remote_readers( + self.rpc_broadcast_mq, n_remote, lambda: self.still_running + ) + + def _reply_addr(self, rank: int, host: str | None) -> str: + """Where `rank` sends its replies: IPC from this node, TCP from another.""" + if rank < self.local_proc_num: + return get_open_zmq_ipc_path() + return f"tcp://{host}:{get_open_port()}" def exit(self): if not self.still_running: @@ -753,6 +801,18 @@ def _collective_rpc( for rank in range(self.proc_num) ] + if payload.barrier and getattr(self, "local_proc_num", self.proc_num) < ( + self.proc_num + ): + error = ( + f"cannot run barrier RPC {func_name!r}: the TP ranks on other " + f"nodes share no barrier with this one" + ) + return [ + RpcResult(payload.request_id, rank, error=error) + for rank in range(self.proc_num) + ] + logger.debug( f"{self.label}: collective_rpc {func_name} id={payload.request_id}" ) diff --git a/atom/model_engine/engine_core.py b/atom/model_engine/engine_core.py index 35a9efb931..d9aa468f26 100644 --- a/atom/model_engine/engine_core.py +++ b/atom/model_engine/engine_core.py @@ -132,6 +132,7 @@ def __init__(self, config: Config, input_address: str, output_address: str): config.tp_world_size * config.prefill_context_parallel_size, config.runner_qualname, config, + local_proc_num=config.tp_size_local if config.is_multinode_tp else None, ) self._post_model_load_hook() block_info = self.runner_mgr.call_func("get_num_blocks", wait_out=True) diff --git a/atom/model_engine/engine_core_mgr.py b/atom/model_engine/engine_core_mgr.py index 5f734f6921..6e405fc14b 100644 --- a/atom/model_engine/engine_core_mgr.py +++ b/atom/model_engine/engine_core_mgr.py @@ -323,8 +323,8 @@ def _start_distributed_store( pc = config.parallel_config if multinode and pc.data_parallel_base_port == 0: raise ValueError( - "Multi-node DP requires the same nonzero --data-parallel-base-port " - "(or ATOM_DP_BASE_PORT) on every node." + "Multi-node DP or TP requires the same nonzero " + "--data-parallel-base-port (or ATOM_DP_BASE_PORT) on every node." ) if not multinode or pc.data_parallel_rank == 0: store = TCPStore( @@ -423,7 +423,9 @@ def __init__(self, config: Config): local_dp_ranks = [] try: - self._start_distributed_store(config, multinode=multinode) + self._start_distributed_store( + config, multinode=multinode or config.is_multinode_tp + ) for engine_index in range(self.local_engine_count): assignment_index = engine_index // self.pp_size dp_rank, local_dp_rank = rank_assignments[assignment_index] diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 2f3c5f6758..78593f2d76 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -1000,7 +1000,8 @@ def _setup_device_and_distributed(self, rank: int, config: Config): # tp_world_size: how many GPUs this stage actually occupies. stage_span = config.tp_world_size * config.prefill_context_parallel_size engine_index = dp_rank_local * pp_size + pp_rank - local_device_rank = engine_index * stage_span + rank + # A TP group that spans nodes has its ranks from tp_rank_start on here. + local_device_rank = engine_index * stage_span + rank - config.tp_rank_start num_gpus = torch.cuda.device_count() if local_device_rank >= num_gpus: raise ValueError( @@ -1030,8 +1031,9 @@ def _setup_device_and_distributed(self, rank: int, config: Config): and not torch.distributed.is_initialized() ): # Preserve AITER's environment setup when preinitializing its group. + visible = config.tp_size_local if config.is_multinode_tp else world_size os.environ.setdefault( - "HIP_VISIBLE_DEVICES", ",".join(map(str, range(world_size))) + "HIP_VISIBLE_DEVICES", ",".join(map(str, range(visible))) ) store = torch.distributed.TCPStore( config.parallel_config.data_parallel_master_ip, @@ -1915,7 +1917,11 @@ def get_num_blocks(self) -> dict[str, object]: # single source for every entry count, so it must never disagree with # the number the pool is actually built at. num_kvcache_blocks = plan.paged_entries - if config.pipeline_parallel_size > 1 and torch.distributed.is_initialized(): + # Only rank 0's answer reaches the scheduler; ranks on another node can + # have less memory free. + if ( + config.pipeline_parallel_size > 1 or config.is_multinode_tp + ) and torch.distributed.is_initialized(): t = torch.tensor( [num_kvcache_blocks], dtype=torch.int64, device=self.device ) diff --git a/atom/model_engine/tp_multinode.py b/atom/model_engine/tp_multinode.py new file mode 100644 index 0000000000..c6e76e3f9e --- /dev/null +++ b/atom/model_engine/tp_multinode.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Tensor parallelism across nodes. + +The node that runs TP rank 0 runs the engine, the API server and its own ranks' +workers. Every other node runs only its workers (:func:`serve_remote_tp_ranks`). +The engine reaches them the way it reaches local ones, through +:class:`~atom.model_engine.async_proc.AsyncIOProcManager`'s broadcast queue and +reply channels, whose remote ends listen on TCP. The manager publishes those +endpoints, with the runner's constructor arguments, in the model-runner +TCPStore that CoreManager serves on ``data_parallel_master_ip`` and +``data_parallel_base_port``; the workers' process group meets there too. +""" + +import logging +import pickle +import signal +import time +from datetime import timedelta + +from atom.utils import envs, get_mp_context, shutdown_all_processes + +logger = logging.getLogger("atom") + +BOOTSTRAP_KEY = "atom/tp_multinode/bootstrap" +_READY = b"READY" +# What an XPUB socket in verbose mode receives when a subscriber joins. +_SUBSCRIBED = b"\x01" +# A short-timeout store client, so that losing the engine's store is noticed +# in seconds rather than after ATOM_MULTINODE_TP_TIMEOUT_S. +_WATCHDOG_TIMEOUT_S = 30 +_WATCHDOG_MISSES = 3 + + +def _store(ip: str, port: int, timeout_s: float): + from torch.distributed import TCPStore + + return TCPStore(ip, port, is_master=False, timeout=timedelta(seconds=timeout_s)) + + +def publish_bootstrap(config, info: dict) -> None: + """Engine side: hand the other nodes what their workers are built from.""" + pc = config.parallel_config + store = _store( + pc.data_parallel_master_ip, + pc.data_parallel_base_port, + envs.ATOM_MULTINODE_TP_TIMEOUT_S, + ) + store.set(BOOTSTRAP_KEY, pickle.dumps(info)) + + +def release_remote_readers(mq, n_remote: int, is_running) -> None: + """Engine side: wait for every remote reader to subscribe to the broadcast + queue, then let them read. A PUB socket drops whatever it sends before a + subscriber is connected.""" + sock = mq.remote_socket + timeout_s = envs.ATOM_MULTINODE_TP_TIMEOUT_S + deadline = time.monotonic() + timeout_s + subscribed, next_log = 0, 0.0 + while subscribed < n_remote: + now = time.monotonic() + if now >= next_log: + logger.info( + "Waiting for %d of %d TP ranks on other nodes to subscribe", + n_remote - subscribed, + n_remote, + ) + next_log = now + 60 + if not is_running(): + raise RuntimeError("shut down while waiting for TP ranks on other nodes") + if now > deadline: + raise TimeoutError( + f"{n_remote - subscribed} of {n_remote} TP ranks on other nodes " + f"did not subscribe within {timeout_s:.0f}s" + ) + if sock.poll(timeout=1000) and sock.recv()[:1] == _SUBSCRIBED: + subscribed += 1 + sock.send(_READY) + logger.info("All %d TP ranks on other nodes subscribed", n_remote) + + +def await_release(mq) -> None: + """Remote rank side of :func:`release_remote_readers`.""" + timeout_s = envs.ATOM_MULTINODE_TP_TIMEOUT_S + if not mq.remote_socket.poll(timeout=int(timeout_s * 1000)): + raise TimeoutError( + f"the engine did not release this rank's broadcast queue within " + f"{timeout_s:.0f}s" + ) + msg = mq.remote_socket.recv() + if msg != _READY: + raise RuntimeError( + f"expected {_READY!r} on the broadcast queue, got {msg[:32]!r}" + ) + + +def serve_remote_tp_ranks(engine_args) -> int: + """Run this node's TP ranks for the engine on another node, until they end. + + Returns 0 when every worker exited cleanly (the engine shut them down), + 1 otherwise. + """ + from atom.model_engine.async_proc import AsyncIOProc + + start = engine_args.tensor_parallel_rank + local = engine_args.tensor_parallel_size_local + ip = engine_args.data_parallel_master_ip + port = engine_args.data_parallel_base_port + if not local or start % local or not port or ip in ("127.0.0.1", "localhost"): + raise ValueError( + "a node that runs only TP workers needs --tensor-parallel-size-local, " + "a --tensor-parallel-rank that is a multiple of it, and the engine " + "node's --data-parallel-master-ip and --data-parallel-base-port" + ) + ranks = range(start, start + local) + logger.info( + "TP ranks %d-%d: waiting for the engine at %s:%d", start, ranks[-1], ip, port + ) + store = _store(ip, port, envs.ATOM_MULTINODE_TP_TIMEOUT_S) + payload = store.get(BOOTSTRAP_KEY) + # The engine's config holds classes that a trust_remote_code model defines + # in transformers' dynamic modules; loading the model's config registers + # them here, and the workers inherit the path through spawn. + from atom.config import get_hf_config + + get_hf_config(engine_args.model, trust_remote_code=engine_args.trust_remote_code) + info = pickle.loads(payload) + proc_num = info["proc_num"] + if ranks[-1] >= proc_num: + raise ValueError( + f"TP ranks {start}-{ranks[-1]} do not exist: the engine runs {proc_num}" + ) + config, *runner_args = info["args"] + # The engine's copy describes its own node. + config.parallel_config.tensor_parallel_size_local = local + config.parallel_config.tensor_parallel_rank = start + + ctx = get_mp_context() + procs = [] + for rank in ranks: + label = f"ModelRunner{rank}/{proc_num}" + proc = ctx.Process( + target=AsyncIOProc, + name=label, + args=( + label, + [None, None], + info["handle"], + info["runner"], + rank, + info["kv_output_addrs"][rank], + None, + config, + *runner_args, + ), + kwargs={"rpc_output_addr": info["rpc_output_addrs"][rank]}, + ) + proc.start() + procs.append(proc) + logger.info("TP ranks %d-%d started for the engine at %s", start, ranks[-1], ip) + return _supervise(procs, ip, port) + + +def _supervise(procs, ip: str, port: int) -> int: + """Wait for the workers. One that fails, a signal, or losing the engine's + store (it lives in the engine node's server process) stops all of them.""" + stopping = [] + for sig in (signal.SIGINT, signal.SIGTERM): + signal.signal(sig, lambda signum, frame: stopping.append(signum)) + + watchdog, misses = None, 0 + status = 1 + while not stopping: + exited = [p for p in procs if not p.is_alive()] + failed = [p for p in exited if p.exitcode != 0] + if failed: + logger.error( + "%s exited with %s; stopping this node's TP ranks", + failed[0].name, + failed[0].exitcode, + ) + break + if len(exited) == len(procs): + status = 0 + break + try: + if watchdog is None: + watchdog = _store(ip, port, _WATCHDOG_TIMEOUT_S) + watchdog.check([BOOTSTRAP_KEY]) + misses = 0 + except Exception as e: # noqa: BLE001 - any store failure counts as a miss + watchdog, misses = None, misses + 1 + logger.warning("Engine store at %s:%d unreachable: %s", ip, port, e) + if misses >= _WATCHDOG_MISSES: + logger.error("Lost the engine; stopping this node's TP ranks") + break + time.sleep(5) + if stopping: + logger.info("Signal %d: stopping this node's TP ranks", stopping[0]) + shutdown_all_processes(procs, allowed_seconds=5) + return status diff --git a/atom/utils/envs.py b/atom/utils/envs.py index f890344b2b..3bd797e3ac 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -645,6 +645,12 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: "ATOM_SHUTDOWN_TIMEOUT_S": lambda: _positive_float_env( "ATOM_SHUTDOWN_TIMEOUT_S", "5" ), + # How long the nodes of a TP group that spans nodes wait for each other at + # startup: a remote node for the engine's endpoints, the engine for every + # remote rank to subscribe to its broadcast queue. + "ATOM_MULTINODE_TP_TIMEOUT_S": lambda: _positive_float_env( + "ATOM_MULTINODE_TP_TIMEOUT_S", "1800" + ), # --- Profiling & Logging --- "ATOM_METRICS_UPDATE_INTERVAL_S": lambda: _positive_float_env( "ATOM_METRICS_UPDATE_INTERVAL_S", "1.0" diff --git a/docs/environment_variables.md b/docs/environment_variables.md index e394513031..368a92790c 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -435,6 +435,7 @@ the host, give each its own endpoints. | Variable | Type | Default | Description | |----------|------|---------|-------------| | **ATOM_SHUTDOWN_TIMEOUT_S** | float | 5 | Seconds that shutdown waits for child processes to exit on their own before terminating them (SIGTERM, then SIGKILL): the server process waits this long for its EngineCores, and each EngineCore for its ModelRunner workers. All the children of one process share the same deadline. Must be finite and positive. Raise it when workers need longer to finish on exit, for example to write a profiler's output (`rocprofv3` writing a thread trace can take minutes). Set it before starting the server; every process reads it. | +| **ATOM_MULTINODE_TP_TIMEOUT_S** | float | 1800 | When one TP group spans nodes (`--tensor-parallel-size-local`), how many seconds the nodes wait for each other at startup: a node that runs only workers waits this long for the engine's endpoints, and the engine waits this long for every remote rank to subscribe to its broadcast queue. Must be finite and positive. | ## Profiling & debugging From c59ffd5dfa33dea8e4a0f145ceddf94bede6d480 Mon Sep 17 00:00:00 2001 From: Guanbao Yu Date: Fri, 16 Oct 2026 12:34:49 +0000 Subject: [PATCH 2/6] [MLA][DCP] Run DCP decode persistent on gfx1250 with the PS1 kernel dcp_persistent_supported() said gfx950 only. On gfx1250 the PS1 FP8 decode (AITER_MLA_DECODE_PS1_FLYDSL=1) is persistent and returns the LSE that DCP's merge needs (aiter #5844 added DCP to it), so DCP decode can take it. Without it the decode falls to gfx1250's non-persistent ASM kernel, which takes only 64-token pages: -tp 8 -dcp 8 for Kimi-K3 failed its first decode capture with 'mla_decode_gfx1250_dispatch: only support page_size == 64'. --- atom/distributed/dcp_utils.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/atom/distributed/dcp_utils.py b/atom/distributed/dcp_utils.py index f3464d7e81..f523032f86 100644 --- a/atom/distributed/dcp_utils.py +++ b/atom/distributed/dcp_utils.py @@ -46,14 +46,21 @@ def get_dcp_rank() -> int: def dcp_persistent_supported() -> bool: """Whether DCP decode can run in *persistent* mode on this GPU. - Needs an lse-emitting ASM decode kernel; only gfx950 ships one (gfx942 - falls back to the triton stage2 reduce instead). Cache the result once - per ``__init__`` — calling this per-forward would graph-break on - ``get_gfx()``. + Needs an lse-emitting persistent decode kernel: gfx950's ASM one, and on + gfx1250 the PS1 FP8 one that AITER_MLA_DECODE_PS1_FLYDSL=1 selects (aiter + #5844 added DCP to it); without that, gfx1250's non-persistent decode takes + only 64-token pages. gfx942 falls back to the triton stage2 reduce instead. + Cache the result once per ``__init__`` — calling this per-forward would + graph-break on ``get_gfx()``. """ + import os + from aiter.jit.utils.chip_info import get_gfx - return get_gfx() == "gfx950" + gfx = get_gfx() + if gfx == "gfx1250": + return os.environ.get("AITER_MLA_DECODE_PS1_FLYDSL", "0") == "1" + return gfx == "gfx950" def mla_dcp_decode_is_persistent( From c3cc2367ece794555483c9e640813fab14113712 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Thu, 22 Oct 2026 12:57:43 +0000 Subject: [PATCH 3/6] fix(tp): wait for a lazily created NCCL id as long as the cross-node TP setup Every NCCL communicator created on first use takes its unique id from the CoreManager's TCPStore, and a rank's client of that store kept torch's default 300 s timeout. When the TP group spans nodes, the trays reach such a collective as far apart as their checkpoint loads differ, which is more than 300 s on a tray that cannot page-cache the checkpoint. Kimi-K3, -tp 8 over CT2 (2 TB RAM) + CT17 (1 TB RAM), no DCP: CT17's weight load drained for 772 s against CT2's 17 s. CT2's ranks 1-3 reached the lm_head all-gather of the warmup while rank 0 had not, waited 300 s for its ncclUniqueId ("store->get('0') got error: wait timeout after 300000ms") and the engine shut down. With DCP the post-load DCP all-gather happens to create that communicator while all eight ranks are still together, so the same trays started. A rank whose TP group spans nodes now creates its store client with ATOM_MULTINODE_TP_TIMEOUT_S (default 1800 s), the wait the rest of the cross-node TP setup already uses. Single-node ranks keep the default. Co-Authored-By: Claude Opus 5.5 --- atom/model_engine/model_runner.py | 17 +++++++++++++++++ tests/test_distributed_store.py | 24 ++++++++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 78593f2d76..5d63cb9116 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -9,6 +9,7 @@ import os import time from contextlib import contextmanager, nullcontext +from datetime import timedelta from functools import partial, wraps from typing import Any, ClassVar, NamedTuple @@ -159,6 +160,21 @@ } +def managed_store_client_kwargs(config) -> dict: + """Extra arguments for a rank's client of the CoreManager's TCPStore. + + Every NCCL communicator created on first use takes its unique id from this + store. When the TP group spans nodes, its trays reach such a collective + apart by as much as their checkpoint loads differ (13 minutes on a tray + that cannot page-cache the checkpoint), so torch's default 300 s wait aborts + the start. Such a rank waits as long as the rest of the cross-node TP setup + does instead. + """ + if config.is_multinode_tp: + return {"timeout": timedelta(seconds=envs.ATOM_MULTINODE_TP_TIMEOUT_S)} + return {} + + def max_schedulable_decode_bs( max_num_seqs: int, max_num_batched_tokens: int, full_q_len: int ) -> int: @@ -1039,6 +1055,7 @@ def _setup_device_and_distributed(self, rank: int, config: Config): config.parallel_config.data_parallel_master_ip, config.parallel_config.data_parallel_base_port, is_master=False, + **managed_store_client_kwargs(config), ) torch.distributed.init_process_group( backend="nccl", store=store, rank=global_rank, world_size=world_size diff --git a/tests/test_distributed_store.py b/tests/test_distributed_store.py index f88e0df287..94ceae2447 100644 --- a/tests/test_distributed_store.py +++ b/tests/test_distributed_store.py @@ -249,3 +249,27 @@ def test_standalone_model_runner_keeps_rank_zero_rendezvous(manager): _spawn_workers( _model_runner_worker, [(pc, (2, 1, 1, 1, 1), rank) for rank in range(2)] ) + + +@pytest.mark.parametrize("env, seconds", [(None, 1800), ("60", 60)]) +def test_multinode_tp_rank_waits_on_the_store_as_long_as_tp_setup( + monkeypatch, env, seconds +): + # A cross-node TP rank fetches lazily created NCCL ids from this store; the + # trays reach their first such collective as far apart as their loads differ. + pytest.importorskip("aiter") + from atom.model_engine.model_runner import managed_store_client_kwargs + + if env is None: + monkeypatch.delenv("ATOM_MULTINODE_TP_TIMEOUT_S", raising=False) + else: + monkeypatch.setenv("ATOM_MULTINODE_TP_TIMEOUT_S", env) + kwargs = managed_store_client_kwargs(SimpleNamespace(is_multinode_tp=True)) + assert kwargs == {"timeout": timedelta(seconds=seconds)} + + +def test_single_node_rank_keeps_the_default_store_timeout(): + pytest.importorskip("aiter") + from atom.model_engine.model_runner import managed_store_client_kwargs + + assert managed_store_client_kwargs(SimpleNamespace(is_multinode_tp=False)) == {} From 24049fed262e7a9440be483747ae79ff3efb5be8 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Thu, 22 Oct 2026 12:57:43 +0000 Subject: [PATCH 4/6] fix(mono): refuse mono decode when the TP group spans nodes Mono builds its peer buffers across the TP group with HIP IPC, which cannot map a GPU on another node, and nothing turned it off for a cross-node TP group: common_refusals checks TP size, EP, DP, PP and context parallelism, and _check_multinode_tp does not refuse mono. The cross-node runs validated so far all used DCP, which mono already refuses; a cross-node TP run without DCP (Kimi-K3 -tp 8 on two 4-GPU nodes) enables it. common_refusals now refuses a TP group that spans nodes, so mono turns itself off with a logged reason instead of failing at IPC handle exchange. A config without the attribute (plugin modes) reads as single-node, like compilation_config. Co-Authored-By: Claude Opus 5.5 --- atom/mono/runtime/deployment.py | 7 +++++++ tests/test_deepseek_v41_mono_dispatch.py | 1 + 2 files changed, 8 insertions(+) diff --git a/atom/mono/runtime/deployment.py b/atom/mono/runtime/deployment.py index f3e5ddf96a..3e0e7cebe2 100644 --- a/atom/mono/runtime/deployment.py +++ b/atom/mono/runtime/deployment.py @@ -84,6 +84,13 @@ def common_refusals( (atom_config.pipeline_parallel_size == 1, "PP > 1"), (atom_config.decode_context_parallel_size == 1, "decode CP"), (atom_config.prefill_context_parallel_size == 1, "prefill CP"), + # Mono's peer buffers are HIP IPC mappings, which cannot reach a GPU on + # another node. Read like `compilation_config` below: a config without + # the attribute (plugin modes) is a single-node one. + ( + not getattr(atom_config, "is_multinode_tp", False), + "TP group spans nodes (HIP IPC peer memory)", + ), (not (atom_config.enable_tbo or atom_config.enable_tbo_decode), "TBO"), (not (is_vllm() or is_sglang()), "plugin mode"), ( diff --git a/tests/test_deepseek_v41_mono_dispatch.py b/tests/test_deepseek_v41_mono_dispatch.py index b830dfdee4..d3fa7a8967 100644 --- a/tests/test_deepseek_v41_mono_dispatch.py +++ b/tests/test_deepseek_v41_mono_dispatch.py @@ -100,6 +100,7 @@ def test_per_request_verify_lengths_are_served(native, schedule, ragged): ({"pipeline_parallel_size": 2}, "PP"), ({"decode_context_parallel_size": 2}, "decode CP"), ({"prefill_context_parallel_size": 2}, "prefill CP"), + ({"is_multinode_tp": True}, "TP group spans nodes"), ({"enable_tbo": True}, "TBO"), ({"enable_tbo_decode": True}, "TBO"), ( From 2b10dd677da541d8adc904601e049359239b1e1a Mon Sep 17 00:00:00 2001 From: Guanbao Yu Date: Tue, 3 Nov 2026 09:27:52 +0000 Subject: [PATCH 5/6] fix(dsv4): refuse Mono decode when the TP group spans nodes V4 Mono builds its peer buffers with HIP IPC (EpochPeerBuffer), which cannot map a GPU on another node. It admits a deployment through its own deployment_refusal, not common_refusals, so the cross-node refusal there does not reach it, and a pure TP8 group over two 4-GPU nodes passes its TP8 check. Mono is on by default (#2586), so such a run would fail at IPC handle exchange. deployment_refusal now refuses a TP group that spans nodes, with the reason common_refusals gives, so V4 Mono keeps the native path instead. --- atom/models/deepseek_v4/mono/profile.py | 2 + tests/test_deepseek_v4_mono_deployment.py | 55 +++++++++++++++++++++++ 2 files changed, 57 insertions(+) create mode 100644 tests/test_deepseek_v4_mono_deployment.py diff --git a/atom/models/deepseek_v4/mono/profile.py b/atom/models/deepseek_v4/mono/profile.py index d9f788f814..9c66b44c73 100644 --- a/atom/models/deepseek_v4/mono/profile.py +++ b/atom/models/deepseek_v4/mono/profile.py @@ -170,6 +170,8 @@ def deployment_refusal( or atom_config.parallel_config.data_parallel_size != 1 ): return "only pure TP8 without TBO/context parallel is supported" + if getattr(atom_config, "is_multinode_tp", False): + return "TP group spans nodes (HIP IPC peer memory)" if atom_config.kv_cache_dtype != "fp8" or atom_config.index_cache_dtype != "fp4": return "FP8 KV / FP4 index is required" if getattr(atom_config, "enable_prefix_caching", False): diff --git a/tests/test_deepseek_v4_mono_deployment.py b/tests/test_deepseek_v4_mono_deployment.py new file mode 100644 index 0000000000..0abddb96cd --- /dev/null +++ b/tests/test_deepseek_v4_mono_deployment.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""Which deployments the DeepSeek-V4 Mono decode admits.""" + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("aiter") + +from atom.models.deepseek_v4.mono.profile import ORDINARY_PROFILE, deployment_refusal + + +class _Mode: + name = "FULL" + + def uses_tr_hybrid(self): + return False + + def decode_mode(self): + return self + + +_MODEL_ARGS = SimpleNamespace( + use_index_cache=False, index_topk_freq=1, index_topk_pattern=None +) + + +def _atom_config(**overrides): + fields = { + "tensor_parallel_size": 8, + "enable_expert_parallel": False, + "enable_dp_attention": False, + "enable_tbo": False, + "pipeline_parallel_size": 1, + "prefill_context_parallel_size": 1, + "decode_context_parallel_size": 1, + "parallel_config": SimpleNamespace(data_parallel_size=1), + "kv_cache_dtype": "fp8", + "index_cache_dtype": "fp4", + "compilation_config": SimpleNamespace(level=0, cudagraph_mode=_Mode()), + } + fields.update(overrides) + return SimpleNamespace(**fields) + + +def test_pure_tp8_on_one_node_is_admitted(): + config = _atom_config() + assert deployment_refusal(config, ORDINARY_PROFILE, model_args=_MODEL_ARGS) is None + + +def test_tp_group_across_nodes_is_refused(): + config = _atom_config(is_multinode_tp=True) + why = deployment_refusal(config, ORDINARY_PROFILE, model_args=_MODEL_ARGS) + assert why == "TP group spans nodes (HIP IPC peer memory)" From 30954117d03f4e34d370b939dfee4a65c0158fb5 Mon Sep 17 00:00:00 2001 From: Guanbao Yu Date: Tue, 3 Nov 2026 10:10:34 +0000 Subject: [PATCH 6/6] fix(tp): address review of the cross-node TP group - A node that runs only TP workers took the engine's address from its CLI flags, while the engine's node takes it through ParallelConfig, which applies ATOM_DP_MASTER_IP and ATOM_DP_BASE_PORT. Set through the environment, node0 started and node1 refused to. Both now resolve it through EngineArgs.parallel_config(). - The TP-slice check ran only once the slice was cross-node, so --tensor-parallel-size-local 16 with -tp 8 was silently ignored. It now runs whenever either flag is set; the cross-node refusals still apply only across nodes. - A remote rank's reply address was formatted by hand, without the brackets an IPv6 host needs. It uses get_tcp_uri. - Only the OpenAI API server turned a node with --tensor-parallel-rank > 0 into a worker-only one; the atomesh server and the examples built a partial engine there. EngineArgs.create_engine now does it for every entrypoint. The API server on such a node now loads its tokenizer first, as it does on a remote DP node. - ATOM_MULTINODE_TP_TIMEOUT_S reached the store client but no process group. The default group and aiter's TP groups are created without a timeout and take torch's defaults, 30 minutes for gloo and 10 for NCCL. A rank whose TP group spans nodes now raises both to it, never lowering them, the way #2817 raises the gloo default when DP spans nodes. - The store client moves to tp_multinode, which imports without aiter, so its tests run in CI; it builds the client itself, and the test asserts what TCPStore receives. The DeepSeek-V4 Mono test loads profile.py from its file for the same reason. - docs: ATOM_MULTINODE_TP_TIMEOUT_S gets a section of its own with its full scope. The section says that AITER_MLA_DECODE_PS1_FLYDSL must match on every node, and what happens when it does not. --- atom/config.py | 7 +- atom/entrypoints/openai/api_server.py | 7 - atom/model_engine/arg_utils.py | 50 +++++-- atom/model_engine/async_proc.py | 3 +- atom/model_engine/model_runner.py | 28 +--- atom/model_engine/tp_multinode.py | 58 +++++++- atom/utils/envs.py | 3 +- docs/environment_variables.md | 11 +- tests/test_deepseek_v4_mono_deployment.py | 19 ++- tests/test_distributed_store.py | 24 --- tests/test_tp_multinode.py | 174 ++++++++++++++++++++++ 11 files changed, 302 insertions(+), 82 deletions(-) create mode 100644 tests/test_tp_multinode.py diff --git a/atom/config.py b/atom/config.py index 8e15ac68e2..bb58e3a606 100644 --- a/atom/config.py +++ b/atom/config.py @@ -2233,12 +2233,16 @@ def is_multinode_tp(self) -> bool: def _check_multinode_tp(self) -> None: pc = self.parallel_config + if pc.tensor_parallel_size_local is None and pc.tensor_parallel_rank == 0: + return local, start, tp = self.tp_size_local, self.tp_rank_start, self.tp_world_size if local < 1 or tp % local or start % local or not 0 <= start < tp: raise ValueError( f"--tensor-parallel-size-local {local} and --tensor-parallel-rank " f"{start} must cut -tp {tp} into whole, equal per-node slices" ) + if not self.is_multinode_tp: + return unsupported = [ name for name, used in ( @@ -2439,8 +2443,7 @@ def __post_init__(self): ) assert 1 <= self.tensor_parallel_size <= 8 - if self.is_multinode_tp: - self._check_multinode_tp() + self._check_multinode_tp() if self.decode_context_parallel_size > 1: assert self.tensor_parallel_size % self.decode_context_parallel_size == 0, ( f"tp_size ({self.tensor_parallel_size}) must be divisible by " diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index e2dbadc622..fce9d4d7d4 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -3285,13 +3285,6 @@ def main(): ) args = parser.parse_args() - if args.tensor_parallel_rank > 0: - # The node that runs TP rank 0 serves the API; this one only runs the - # workers of its own ranks. - from atom.model_engine.tp_multinode import serve_remote_tp_ranks - - raise SystemExit(serve_remote_tp_ranks(EngineArgs.from_cli_args(args))) - if args.request_log: _request_logger = logging.getLogger("atom.request_log") _request_logger.setLevel(logging.INFO) diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index 1badc475a4..db0826ed1e 100644 --- a/atom/model_engine/arg_utils.py +++ b/atom/model_engine/arg_utils.py @@ -953,24 +953,46 @@ def _get_engine_kwargs(self) -> dict: # DP topology -> ParallelConfig. `data_parallel_size` stays in kwargs # too: LLMEngine still reads the loose kwarg on the legacy path. - parallel_config_kwargs = { - "data_parallel_size": kwargs["data_parallel_size"], - "data_parallel_size_local": kwargs.pop("data_parallel_size_local"), - "data_parallel_rank": kwargs.pop("data_parallel_rank"), - "data_parallel_master_ip": kwargs.pop("data_parallel_master_ip"), - "data_parallel_master_port": kwargs.pop("data_parallel_master_port"), - "tensor_parallel_size_local": kwargs.pop("tensor_parallel_size_local"), - "tensor_parallel_rank": kwargs.pop("tensor_parallel_rank"), - } - base_port = kwargs.pop("data_parallel_base_port") - if base_port is not None: - parallel_config_kwargs["data_parallel_base_port"] = base_port - kwargs["parallel_config"] = ParallelConfig(**parallel_config_kwargs) + for name in self._PARALLEL_CONFIG_FIELDS: + if name != "data_parallel_size": + kwargs.pop(name) + kwargs["parallel_config"] = self.parallel_config() logger.info(f"Engine kwargs: {kwargs}") return kwargs + _PARALLEL_CONFIG_FIELDS = ( + "data_parallel_size", + "data_parallel_size_local", + "data_parallel_rank", + "data_parallel_master_ip", + "data_parallel_master_port", + "data_parallel_base_port", + "tensor_parallel_size_local", + "tensor_parallel_rank", + ) + + def parallel_config(self) -> ParallelConfig: + """This node's ParallelConfig, with the ATOM_DP_* environment applied. + + A node that runs only TP workers reads the engine's address from it as + well, so both nodes resolve that address the same way. + """ + kwargs = {name: getattr(self, name) for name in self._PARALLEL_CONFIG_FIELDS} + if kwargs["data_parallel_base_port"] is None: + del kwargs["data_parallel_base_port"] + return ParallelConfig(**kwargs) + def create_engine(self, tokenizer=None) -> LLMEngine: - """Create and return an LLMEngine instance with the configured parameters.""" + """Create and return an LLMEngine instance with the configured parameters. + + A node whose TP slice starts past rank 0 (--tensor-parallel-rank) has no + engine to create: it runs its TP workers for the engine on another node + and exits the process with their status once they stop. + """ + if self.tensor_parallel_rank > 0: + from atom.model_engine.tp_multinode import serve_remote_tp_ranks + + raise SystemExit(serve_remote_tp_ranks(self)) return LLMEngine(self.model, tokenizer=tokenizer, **self._get_engine_kwargs()) diff --git a/atom/model_engine/async_proc.py b/atom/model_engine/async_proc.py index a2b73c5dcf..8bd4f8cc0c 100644 --- a/atom/model_engine/async_proc.py +++ b/atom/model_engine/async_proc.py @@ -48,6 +48,7 @@ get_mp_context, get_open_port, get_open_zmq_ipc_path, + get_tcp_uri, init_exit_handler, make_zmq_socket, resolve_obj_by_qualname, @@ -615,7 +616,7 @@ def _reply_addr(self, rank: int, host: str | None) -> str: """Where `rank` sends its replies: IPC from this node, TCP from another.""" if rank < self.local_proc_num: return get_open_zmq_ipc_path() - return f"tcp://{host}:{get_open_port()}" + return get_tcp_uri(host, get_open_port()) def exit(self): if not self.still_running: diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 44ae111597..1720ee71cc 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -10,7 +10,6 @@ import time from collections import OrderedDict from contextlib import contextmanager, nullcontext -from datetime import timedelta from functools import partial, wraps from typing import Any, ClassVar, NamedTuple @@ -63,6 +62,10 @@ new_block_table, ) from atom.model_engine.state_runtime import StateRuntime +from atom.model_engine.tp_multinode import ( + managed_store_client, + raise_timeout_if_tp_spans_nodes, +) from atom.model_loader.loader import load_model from atom.model_ops.attentions.pool_layout.entry_arena import plan_regions from atom.model_ops.attentions.pool_layout.sub_pool_spec import ( @@ -127,21 +130,6 @@ logger = logging.getLogger("atom") -def managed_store_client_kwargs(config) -> dict: - """Extra arguments for a rank's client of the CoreManager's TCPStore. - - Every NCCL communicator created on first use takes its unique id from this - store. When the TP group spans nodes, its trays reach such a collective - apart by as much as their checkpoint loads differ (13 minutes on a tray - that cannot page-cache the checkpoint), so torch's default 300 s wait aborts - the start. Such a rank waits as long as the rest of the cross-node TP setup - does instead. - """ - if config.is_multinode_tp: - return {"timeout": timedelta(seconds=envs.ATOM_MULTINODE_TP_TIMEOUT_S)} - return {} - - def max_schedulable_decode_bs( max_num_seqs: int, max_num_batched_tokens: int, full_q_len: int ) -> int: @@ -1063,6 +1051,7 @@ def _setup_device_and_distributed(self, rank: int, config: Config): dp_rank = config.parallel_config.data_parallel_rank global_rank = (dp_rank * pp_size + pp_rank) * stage_span + rank raise_timeout_if_dp_spans_nodes(config.parallel_config) + raise_timeout_if_tp_spans_nodes(config) if ( config.parallel_config._managed_distributed_store and not torch.distributed.is_initialized() @@ -1072,12 +1061,7 @@ def _setup_device_and_distributed(self, rank: int, config: Config): os.environ.setdefault( "HIP_VISIBLE_DEVICES", ",".join(map(str, range(visible))) ) - store = torch.distributed.TCPStore( - config.parallel_config.data_parallel_master_ip, - config.parallel_config.data_parallel_base_port, - is_master=False, - **managed_store_client_kwargs(config), - ) + store = managed_store_client(config) torch.distributed.init_process_group( backend="nccl", store=store, rank=global_rank, world_size=world_size ) diff --git a/atom/model_engine/tp_multinode.py b/atom/model_engine/tp_multinode.py index c6e76e3f9e..eb78a2952f 100644 --- a/atom/model_engine/tp_multinode.py +++ b/atom/model_engine/tp_multinode.py @@ -39,6 +39,52 @@ def _store(ip: str, port: int, timeout_s: float): return TCPStore(ip, port, is_master=False, timeout=timedelta(seconds=timeout_s)) +def managed_store_client(config): + """A model-runner rank's client of the TCPStore that CoreManager serves. + + Every NCCL communicator created on first use takes its unique id from this + store. When the TP group spans nodes, its trays reach such a collective + apart by as much as their checkpoint loads differ (13 minutes on a tray + that cannot page-cache the checkpoint), so torch's default 300 s wait aborts + the start. Such a rank waits as long as the rest of the cross-node TP setup + does instead. + """ + from torch.distributed import TCPStore + + pc = config.parallel_config + kwargs = {} + if config.is_multinode_tp: + kwargs["timeout"] = timedelta(seconds=envs.ATOM_MULTINODE_TP_TIMEOUT_S) + return TCPStore( + pc.data_parallel_master_ip, + pc.data_parallel_base_port, + is_master=False, + **kwargs, + ) + + +def raise_timeout_if_tp_spans_nodes(config) -> None: + """Apply ATOM_MULTINODE_TP_TIMEOUT_S to the process groups created without a + timeout when the TP group spans nodes. + + The default group and aiter's TP groups are. Across nodes their first + collectives wait for the slowest node's load, which torch's defaults (30 + minutes for gloo, 10 for NCCL) need not cover. Raises those defaults and + never lowers them. + """ + if not config.is_multinode_tp: + return + from torch.distributed import distributed_c10d + + from atom.utils.distributed.utils import raise_default_group_timeout + + seconds = envs.ATOM_MULTINODE_TP_TIMEOUT_S + raise_default_group_timeout(seconds) + nccl = distributed_c10d.default_pg_nccl_timeout + if nccl is not None: + distributed_c10d.default_pg_nccl_timeout = max(nccl, timedelta(seconds=seconds)) + + def publish_bootstrap(config, info: dict) -> None: """Engine side: hand the other nodes what their workers are built from.""" pc = config.parallel_config @@ -103,15 +149,17 @@ def serve_remote_tp_ranks(engine_args) -> int: """ from atom.model_engine.async_proc import AsyncIOProc - start = engine_args.tensor_parallel_rank - local = engine_args.tensor_parallel_size_local - ip = engine_args.data_parallel_master_ip - port = engine_args.data_parallel_base_port + pc = engine_args.parallel_config() + start = pc.tensor_parallel_rank + local = pc.tensor_parallel_size_local + ip = pc.data_parallel_master_ip + port = pc.data_parallel_base_port if not local or start % local or not port or ip in ("127.0.0.1", "localhost"): raise ValueError( "a node that runs only TP workers needs --tensor-parallel-size-local, " "a --tensor-parallel-rank that is a multiple of it, and the engine " - "node's --data-parallel-master-ip and --data-parallel-base-port" + "node's --data-parallel-master-ip and --data-parallel-base-port (or " + "ATOM_DP_MASTER_IP and ATOM_DP_BASE_PORT)" ) ranks = range(start, start + local) logger.info( diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 17fae6504c..45b52a8e33 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -754,7 +754,8 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: ), # How long the nodes of a TP group that spans nodes wait for each other at # startup: a remote node for the engine's endpoints, the engine for every - # remote rank to subscribe to its broadcast queue. + # remote rank to subscribe to its broadcast queue, a rank on the + # model-runner store and in its process groups' first collectives. "ATOM_MULTINODE_TP_TIMEOUT_S": lambda: _positive_float_env( "ATOM_MULTINODE_TP_TIMEOUT_S", "1800" ), diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 2f20eb0b51..dff7b28362 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -36,6 +36,16 @@ limitations](observability/README.md). | **ATOM_DP_PREFILL_BALANCE_SLACK** | float | 0.1 | The target cost is at most mean x (1 + slack) over the ranks that prefill. | | **ATOM_DP_PREFILL_BALANCE_MIN_TOKENS** | int | 4096 | Fewest tokens a shrunk rank still prefills in a step. | +## Tensor parallelism across nodes + +One TP group over several nodes (`--tensor-parallel-size-local`, `--tensor-parallel-rank`): the node that runs TP rank 0 runs the engine and the API server, the others only their workers. + +| Variable | Type | Default | Description | +|----------|------|---------|-------------| +| **ATOM_MULTINODE_TP_TIMEOUT_S** | float | 1800 | How many seconds the nodes wait for each other at startup: a node that runs only workers waits this long for the engine's endpoints, the engine waits this long for every remote rank to subscribe to its broadcast queue, and a rank waits this long on the model-runner store (for the NCCL id of a communicator created on first use). It is also the floor for the timeout of the process groups created without one, the default group and aiter's TP groups: it raises torch's defaults (30 minutes for gloo, 10 for NCCL) and never lowers them. Must be finite and positive. | + +Every node reads its own environment, so a variable that changes what a rank computes must have the same value on every node. In particular, on gfx1250 `AITER_MLA_DECODE_PS1_FLYDSL` decides whether DCP decode runs persistent. Nodes that disagree size the gathered heads and the work metadata differently, and the DCP all-gather and merge then disagree across the group: it hangs, returns wrong output, or aborts on half the ranks with `only support page_size == 64`. Nothing checks this at startup. + ## Prefill delayer (TP/DCP and DP attention) Coalesces waiting prefills while decode continues. DP attention enables it by @@ -478,7 +488,6 @@ the host, give each its own endpoints. | Variable | Type | Default | Description | |----------|------|---------|-------------| | **ATOM_SHUTDOWN_TIMEOUT_S** | float | 5 | Seconds that shutdown waits for child processes to exit on their own before terminating them (SIGTERM, then SIGKILL): the server process waits this long for its EngineCores, and each EngineCore for its ModelRunner workers. All the children of one process share the same deadline. Must be finite and positive. Raise it when workers need longer to finish on exit, for example to write a profiler's output (`rocprofv3` writing a thread trace can take minutes). Set it before starting the server; every process reads it. | -| **ATOM_MULTINODE_TP_TIMEOUT_S** | float | 1800 | When one TP group spans nodes (`--tensor-parallel-size-local`), how many seconds the nodes wait for each other at startup: a node that runs only workers waits this long for the engine's endpoints, and the engine waits this long for every remote rank to subscribe to its broadcast queue. Must be finite and positive. | ## Profiling & debugging diff --git a/tests/test_deepseek_v4_mono_deployment.py b/tests/test_deepseek_v4_mono_deployment.py index 0abddb96cd..b6db441b2a 100644 --- a/tests/test_deepseek_v4_mono_deployment.py +++ b/tests/test_deepseek_v4_mono_deployment.py @@ -2,13 +2,22 @@ # Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. """Which deployments the DeepSeek-V4 Mono decode admits.""" +import importlib.util +import sys +from pathlib import Path from types import SimpleNamespace -import pytest - -pytest.importorskip("aiter") - -from atom.models.deepseek_v4.mono.profile import ORDINARY_PROFILE, deployment_refusal +# Loaded from its file: the module needs only dataclasses, but importing it by +# name runs atom.models.deepseek_v4's __init__, which needs aiter. +_spec = importlib.util.spec_from_file_location( + "_deepseek_v4_mono_profile", + Path(__file__).parents[1] / "atom/models/deepseek_v4/mono/profile.py", +) +profile = importlib.util.module_from_spec(_spec) +sys.modules[_spec.name] = profile # @dataclass looks its module up by name +_spec.loader.exec_module(profile) +ORDINARY_PROFILE = profile.ORDINARY_PROFILE +deployment_refusal = profile.deployment_refusal class _Mode: diff --git a/tests/test_distributed_store.py b/tests/test_distributed_store.py index 8a2d411ac4..6910fd527e 100644 --- a/tests/test_distributed_store.py +++ b/tests/test_distributed_store.py @@ -251,30 +251,6 @@ def test_standalone_model_runner_keeps_rank_zero_rendezvous(manager): ) -@pytest.mark.parametrize("env, seconds", [(None, 1800), ("60", 60)]) -def test_multinode_tp_rank_waits_on_the_store_as_long_as_tp_setup( - monkeypatch, env, seconds -): - # A cross-node TP rank fetches lazily created NCCL ids from this store; the - # trays reach their first such collective as far apart as their loads differ. - pytest.importorskip("aiter") - from atom.model_engine.model_runner import managed_store_client_kwargs - - if env is None: - monkeypatch.delenv("ATOM_MULTINODE_TP_TIMEOUT_S", raising=False) - else: - monkeypatch.setenv("ATOM_MULTINODE_TP_TIMEOUT_S", env) - kwargs = managed_store_client_kwargs(SimpleNamespace(is_multinode_tp=True)) - assert kwargs == {"timeout": timedelta(seconds=seconds)} - - -def test_single_node_rank_keeps_the_default_store_timeout(): - pytest.importorskip("aiter") - from atom.model_engine.model_runner import managed_store_client_kwargs - - assert managed_store_client_kwargs(SimpleNamespace(is_multinode_tp=False)) == {} - - @pytest.mark.parametrize( "dp_size,dp_size_local,first_rank,worker_rank,raised", [ diff --git a/tests/test_tp_multinode.py b/tests/test_tp_multinode.py new file mode 100644 index 0000000000..e77f70df58 --- /dev/null +++ b/tests/test_tp_multinode.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""One tensor-parallel group across nodes (atom/model_engine/tp_multinode.py).""" + +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch.distributed as dist + +from atom.config import Config, ParallelConfig +from atom.model_engine import arg_utils, tp_multinode +from atom.model_engine.arg_utils import EngineArgs + + +@pytest.fixture(autouse=True) +def _clean_environment(monkeypatch): + for name in ( + "ATOM_DP_SIZE", + "ATOM_DP_SIZE_LOCAL", + "ATOM_DP_RANK", + "ATOM_DP_RANK_LOCAL", + "ATOM_DP_MASTER_IP", + "ATOM_DP_MASTER_PORT", + "ATOM_DP_BASE_PORT", + "ATOM_MULTINODE_TP_TIMEOUT_S", + ): + monkeypatch.delenv(name, raising=False) + + +def _worker_node_args(): + return EngineArgs( + tensor_parallel_size=8, tensor_parallel_size_local=4, tensor_parallel_rank=4 + ) + + +def _config(local, start, dp=1): + """A Config holding only what the TP-slice check reads.""" + config = Config.__new__(Config) + config.parallel_config = ParallelConfig( + data_parallel_size=dp, + data_parallel_master_ip="10.0.0.1", + data_parallel_base_port=29600, + tensor_parallel_size_local=local, + tensor_parallel_rank=start, + ) + config.tensor_parallel_size = 8 + config.fake_eplb = False + config.enable_dp_attention = False + config.pipeline_parallel_size = 1 + config.prefill_context_parallel_size = 1 + config.enable_rapidserve = False + return config + + +@pytest.mark.parametrize( + "local, start, dp, error", + [ + (None, 0, 1, None), + (8, 0, 2, None), # one node, spelled out: no cross-node refusal + (4, 0, 1, None), + (4, 4, 1, None), + (16, 0, 1, "must cut -tp 8"), + (4, 2, 1, "must cut -tp 8"), + (4, 0, 2, "does not support data parallelism"), + ], +) +def test_tp_slices_are_checked_once_either_flag_is_set(local, start, dp, error): + config = _config(local, start, dp) + if error is None: + config._check_multinode_tp() + else: + with pytest.raises(ValueError, match=error): + config._check_multinode_tp() + + +def test_a_worker_node_reads_the_engine_address_the_way_the_engine_does( + monkeypatch, +): + from aiter_stub import stubbed_aiter + + monkeypatch.setenv("ATOM_DP_MASTER_IP", "10.0.0.1") + monkeypatch.setenv("ATOM_DP_BASE_PORT", "29600") + reached = [] + + def store(ip, port, timeout_s): + reached.append((ip, port)) + raise ConnectionError("stop before the bootstrap") + + monkeypatch.setattr(tp_multinode, "_store", store) + with stubbed_aiter(), pytest.raises(ConnectionError): + tp_multinode.serve_remote_tp_ranks(_worker_node_args()) + assert reached == [("10.0.0.1", 29600)] + + +def test_create_engine_on_a_worker_node_serves_its_ranks_instead(monkeypatch): + served = [] + monkeypatch.setattr( + tp_multinode, "serve_remote_tp_ranks", lambda args: served.append(args) or 0 + ) + monkeypatch.setattr( + arg_utils, "LLMEngine", lambda *a, **k: pytest.fail("built an engine") + ) + args = _worker_node_args() + with pytest.raises(SystemExit) as stopped: + args.create_engine() + assert stopped.value.code == 0 + assert served == [args] + + +@pytest.mark.parametrize( + "spans_nodes, env, timeout", + [ + (True, None, timedelta(seconds=1800)), + (True, "60", timedelta(seconds=60)), + (False, "60", None), + ], +) +def test_a_rank_store_client_waits_as_long_as_tp_setup( + monkeypatch, spans_nodes, env, timeout +): + # A cross-node TP rank fetches lazily created NCCL ids from this store; the + # trays reach their first such collective as far apart as their loads differ. + if env is not None: + monkeypatch.setenv("ATOM_MULTINODE_TP_TIMEOUT_S", env) + clients = [] + monkeypatch.setattr(dist, "TCPStore", lambda *a, **k: clients.append((a, k))) + config = SimpleNamespace( + is_multinode_tp=spans_nodes, + parallel_config=SimpleNamespace( + data_parallel_master_ip="10.0.0.1", data_parallel_base_port=29600 + ), + ) + tp_multinode.managed_store_client(config) + expected = {"is_master": False} + if timeout is not None: + expected["timeout"] = timeout + assert clients == [(("10.0.0.1", 29600), expected)] + + +@pytest.mark.parametrize("spans_nodes", [True, False]) +def test_groups_without_a_timeout_wait_as_long_as_tp_setup(monkeypatch, spans_nodes): + from torch.distributed import distributed_c10d + + monkeypatch.setenv("ATOM_MULTINODE_TP_TIMEOUT_S", "3600") + monkeypatch.setattr(distributed_c10d, "default_pg_timeout", timedelta(minutes=30)) + monkeypatch.setattr( + distributed_c10d, "default_pg_nccl_timeout", timedelta(minutes=10) + ) + config = SimpleNamespace(is_multinode_tp=spans_nodes) + tp_multinode.raise_timeout_if_tp_spans_nodes(config) + # What the default group and aiter's groups, created without one, get. + gloo = distributed_c10d._get_default_timeout(dist.Backend.GLOO) + nccl = distributed_c10d._get_default_timeout(dist.Backend.NCCL) + if not spans_nodes: + assert (gloo, nccl) == (timedelta(minutes=30), timedelta(minutes=10)) + return + assert gloo == nccl == timedelta(hours=1) + # A longer default is kept. + monkeypatch.setenv("ATOM_MULTINODE_TP_TIMEOUT_S", "60") + tp_multinode.raise_timeout_if_tp_spans_nodes(config) + assert distributed_c10d._get_default_timeout(dist.Backend.NCCL) == nccl + + +def test_a_remote_rank_replies_to_an_ipv6_host(): + from aiter_stub import stubbed_aiter + + with stubbed_aiter(): + from atom.model_engine.async_proc import AsyncIOProcManager + + manager = SimpleNamespace(local_proc_num=4) + reply = AsyncIOProcManager._reply_addr + assert reply(manager, 4, "fd00::1").startswith("tcp://[fd00::1]:") + assert reply(manager, 4, "10.0.0.1").startswith("tcp://10.0.0.1:")