diff --git a/atom/config.py b/atom/config.py index 75e0230cad..bb58e3a606 100644 --- a/atom/config.py +++ b/atom/config.py @@ -1069,6 +1069,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. @@ -2207,6 +2215,61 @@ 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 + 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 ( + ("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. @@ -2380,6 +2443,7 @@ def __post_init__(self): ) assert 1 <= self.tensor_parallel_size <= 8 + 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/distributed/dcp_utils.py b/atom/distributed/dcp_utils.py index bf1df4b6e1..0f22d68041 100644 --- a/atom/distributed/dcp_utils.py +++ b/atom/distributed/dcp_utils.py @@ -57,14 +57,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( diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index 092be91e19..db0826ed1e 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 @@ -212,6 +214,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", @@ -929,22 +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"), - } - 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 5a831cca42..8bd4f8cc0c 100644 --- a/atom/model_engine/async_proc.py +++ b/atom/model_engine/async_proc.py @@ -46,7 +46,9 @@ from atom.utils import ( envs, get_mp_context, + get_open_port, get_open_zmq_ipc_path, + get_tcp_uri, init_exit_handler, make_zmq_socket, resolve_obj_by_qualname, @@ -155,7 +157,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 +180,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) @@ -464,11 +470,22 @@ class AsyncIOProcManager: ) def __init__( - self, finalizer, proc_num: int, runner: str, *args, on_worker_failure=None + self, + finalizer, + proc_num: int, + runner: str, + *args, + local_proc_num: int | None = None, + on_worker_failure=None, ): self.parent_finalizer = finalizer self.on_worker_failure = on_worker_failure 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] = [] @@ -477,7 +494,10 @@ def __init__( 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() @@ -486,12 +506,14 @@ def __init__( 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) ] @@ -499,7 +521,7 @@ def __init__( 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) ] @@ -509,7 +531,7 @@ def __init__( # 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] @@ -569,6 +591,32 @@ def __init__( 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 get_tcp_uri(host, get_open_port()) def exit(self): if not self.still_running: @@ -756,6 +804,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 dae9e23bb0..04d314b14b 100644 --- a/atom/model_engine/engine_core.py +++ b/atom/model_engine/engine_core.py @@ -133,6 +133,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, on_worker_failure=self._send_engine_dead, ) self._post_model_load_hook() diff --git a/atom/model_engine/engine_core_mgr.py b/atom/model_engine/engine_core_mgr.py index 63d8def114..36a66b983b 100644 --- a/atom/model_engine/engine_core_mgr.py +++ b/atom/model_engine/engine_core_mgr.py @@ -382,8 +382,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( @@ -482,7 +482,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 4c7bc3bc35..1720ee71cc 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -62,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 ( @@ -1020,7 +1024,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( @@ -1046,19 +1051,17 @@ 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() ): # 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))) - ) - store = torch.distributed.TCPStore( - config.parallel_config.data_parallel_master_ip, - config.parallel_config.data_parallel_base_port, - is_master=False, + "HIP_VISIBLE_DEVICES", ",".join(map(str, range(visible))) ) + store = managed_store_client(config) torch.distributed.init_process_group( backend="nccl", store=store, rank=global_rank, world_size=world_size ) @@ -1948,7 +1951,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..eb78a2952f --- /dev/null +++ b/atom/model_engine/tp_multinode.py @@ -0,0 +1,250 @@ +# 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 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 + 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 + + 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 (or " + "ATOM_DP_MASTER_IP and ATOM_DP_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/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/atom/mono/runtime/deployment.py b/atom/mono/runtime/deployment.py index 1670b62a09..39a6a98a59 100644 --- a/atom/mono/runtime/deployment.py +++ b/atom/mono/runtime/deployment.py @@ -140,6 +140,13 @@ def common_refusals( "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/atom/utils/envs.py b/atom/utils/envs.py index 10f7d4349a..45b52a8e33 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -752,6 +752,13 @@ 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, 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" + ), # --- 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 2c7510d1a8..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 diff --git a/tests/test_deepseek_v41_mono_dispatch.py b/tests/test_deepseek_v41_mono_dispatch.py index b21b837af4..05dbed994c 100644 --- a/tests/test_deepseek_v41_mono_dispatch.py +++ b/tests/test_deepseek_v41_mono_dispatch.py @@ -118,6 +118,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"), ( diff --git a/tests/test_deepseek_v4_mono_deployment.py b/tests/test_deepseek_v4_mono_deployment.py new file mode 100644 index 0000000000..b6db441b2a --- /dev/null +++ b/tests/test_deepseek_v4_mono_deployment.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: MIT +# 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 + +# 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: + 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)" diff --git a/tests/test_otel_scheduler.py b/tests/test_otel_scheduler.py index 0ea6ad3edc..827d3d57b2 100644 --- a/tests/test_otel_scheduler.py +++ b/tests/test_otel_scheduler.py @@ -368,7 +368,7 @@ def call(name, *a, **k): failure_callbacks = [] - def runner_manager(*args, on_worker_failure): + def runner_manager(*args, local_proc_num, on_worker_failure): failure_callbacks.append(on_worker_failure) return SimpleNamespace(call_func=call) @@ -383,6 +383,7 @@ def runner_manager(*args, on_worker_failure): config = MockConfig( parallel_config=SimpleNamespace(control_address="unused", data_parallel_size=1), tp_world_size=1, + is_multinode_tp=False, prefill_context_parallel_size=1, runner_qualname="unused", enforce_eager=True, 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:")