multi-node dp support#1603
Open
ganyi1996ppo wants to merge 3 commits into
Open
Conversation
Signed-off-by: ganyi <ygan@amd.com>
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds native multi-node data-parallel (DP) serving support to ATOM by separating global DP ranks (for distributed groups/collectives) from local DP ranks (for GPU/NUMA placement), and introducing a coordinator/worker model where only global DP rank 0 serves the HTTP API while other nodes run EngineCore workers and connect back to rank 0.
Changes:
- Extends DP topology/config to support multi-node deployments (global vs local DP size/rank, master IP/ports) and validates topology in
ParallelConfig. - Updates EngineCore/CoreManager launch and routing to support multi-node socket addressing and rank mapping, including DP-attention rank flattening behavior.
- Adds CLI/kwargs backward-compat plumbing so callers can configure DP topology consistently via
ParallelConfig(with fallback for older loose kwargs).
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
atom/utils/envs.py |
Adds ATOM_DP_SIZE_LOCAL env var for local DP topology override. |
atom/model_engine/model_runner.py |
Passes local_rank into distributed init to ensure correct device selection under multi-node DP. |
atom/model_engine/llm_engine.py |
Maintains backward compatibility for older callers passing DP topology via loose kwargs instead of ParallelConfig. |
atom/model_engine/engine_core.py |
Updates DP init assertions to allow global/local rank decoupling. |
atom/model_engine/engine_core_mgr.py |
Implements multi-node DP socket planning, global/local rank assignment, and coordinator vs worker behavior. |
atom/model_engine/async_proc.py |
Fixes NUMA/GPU mapping to use local DP rank when available. |
atom/model_engine/arg_utils.py |
Adds CLI flags and builds a ParallelConfig encapsulating multi-node DP topology. |
atom/config.py |
Makes data_parallel_size_local optional with default-to-global behavior, adds multinode detection, hashing, env overrides, and topology validation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+135
to
+139
| master_ip = config.parallel_config.data_parallel_master_ip | ||
| addresses = { | ||
| "input_address": f"tcp://{master_ip}:{plan.input_port}", | ||
| "output_address": f"tcp://{master_ip}:{plan.output_port}", | ||
| } |
Comment on lines
+164
to
+168
| coordinator_addresses = [ | ||
| { | ||
| "input_address": f"tcp://0.0.0.0:{plan.input_port}", | ||
| "output_address": f"tcp://0.0.0.0:{plan.output_port}", | ||
| } |
Comment on lines
+208
to
+212
| if not self._distributed_coordinator: | ||
| for proc in self.engine_core_processes: | ||
| proc.join() | ||
| raise SystemExit(0) | ||
|
|
Comment on lines
+53
to
+59
| if "parallel_config" not in config_kwargs: | ||
| if "data_parallel_size" in kwargs: | ||
| config.parallel_config.data_parallel_size = kwargs["data_parallel_size"] | ||
| if "data_parallel_size_local" in kwargs: | ||
| config.parallel_config.data_parallel_size_local = kwargs[ | ||
| "data_parallel_size_local" | ||
| ] |
Comment on lines
+769
to
+773
| def is_multinode_dp(self) -> bool: | ||
| return ( | ||
| self.data_parallel_size_local < self.data_parallel_size | ||
| or self.data_parallel_rank > 0 | ||
| ) |
| "ATOM_DP_RANK": lambda: int(os.getenv("ATOM_DP_RANK", "0")), | ||
| "ATOM_DP_RANK_LOCAL": lambda: int(os.getenv("ATOM_DP_RANK_LOCAL", "0")), | ||
| "ATOM_DP_SIZE": lambda: int(os.getenv("ATOM_DP_SIZE", "1")), | ||
| "ATOM_DP_SIZE_LOCAL": lambda: int(os.getenv("ATOM_DP_SIZE_LOCAL", "1")), |
Comment on lines
+431
to
+435
| "data_parallel_master_port": data_parallel_master_port, | ||
| } | ||
| if data_parallel_base_port is not None: | ||
| parallel_config_kwargs["data_parallel_base_port"] = data_parallel_base_port | ||
| kwargs["parallel_config"] = ParallelConfig(**parallel_config_kwargs) |
Signed-off-by: ganyi <ygan@amd.com>
Comment on lines
+51
to
+74
| # Backward compatibility for callers that still pass DP topology as | ||
| # loose LLMEngine kwargs instead of a ParallelConfig. | ||
| if "parallel_config" not in config_kwargs: | ||
| if "data_parallel_size" in kwargs: | ||
| config.parallel_config.data_parallel_size = kwargs["data_parallel_size"] | ||
| if "data_parallel_size_local" in kwargs: | ||
| config.parallel_config.data_parallel_size_local = kwargs[ | ||
| "data_parallel_size_local" | ||
| ] | ||
| if "data_parallel_rank" in kwargs: | ||
| config.parallel_config.data_parallel_rank = kwargs["data_parallel_rank"] | ||
| if "data_parallel_rank_local" in kwargs: | ||
| config.parallel_config.data_parallel_rank_local = kwargs[ | ||
| "data_parallel_rank_local" | ||
| ] | ||
| if "data_parallel_master_ip" in kwargs: | ||
| config.parallel_config.data_parallel_master_ip = kwargs[ | ||
| "data_parallel_master_ip" | ||
| ] | ||
| if "data_parallel_master_port" in kwargs: | ||
| config.parallel_config.data_parallel_master_port = kwargs[ | ||
| "data_parallel_master_port" | ||
| ] | ||
| self.data_parallel_size = config.parallel_config.data_parallel_size |
Comment on lines
+53
to
+55
| def _iter_dp_rank_assignments(config: Config): | ||
| """Yield ``(global_dp_rank, local_dp_rank)`` for ranks owned by this node.""" | ||
| pc = config.parallel_config |
Comment on lines
122
to
124
| logger.info( | ||
| f"{self.label}: Creating EngineCore for DP rank {dp_rank}/{self.local_engine_count}" | ||
| ) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add native multi-node data-parallel serving to ATOM.
Global DP ranks can now be distributed across multiple nodes. Global DP rank 0
acts as the sole serving coordinator, while other nodes launch their local
EngineCore workers and connect back to rank 0.
Multi-node mode is inferred directly from the configured DP topology. No
additional
distributed_dpordistributed_dp_servingflags are required.Motivation
The existing native DP implementation assumes every DP rank is local:
CoreManagerlaunches every DP rank;These assumptions do not hold when DP ranks span multiple nodes. For example,
a node that owns global DP ranks 8–15 still needs to map those workers to local
GPUs 0–7.
Design
DP topology
The topology is configured with:
--data-parallel-size,-dp--data-parallel-size-local--data-parallel-rank--data-parallel-master-ip--data-parallel-master-port--data-parallel-base-port--data-parallel-rank-localis assigned internally to EngineCore workers andnormally should not be set by users.
If
--data-parallel-size-localis omitted, it defaults to the global DP size,preserving existing single-node behavior.
Multi-node mode is inferred when:
The topology is validated as follows:
Coordinator/worker model
In multi-node mode:
routing hint.
Single-node deployments continue using the existing local IPC transport.
Global and local ranks
Each EngineCore receives two DP ranks:
For example, the second node in a DP16 deployment may be configured with:
Its rank mapping becomes:
This prevents global ranks on later nodes from being interpreted as local GPU
indices.
DP-attention rank mapping
With DP attention enabled, DP and TP ranks are flattened into EngineCore ranks:
The existing PCP + DP-attention incompatibility is unchanged.
EngineCore transport
The rank-0 coordinator binds one request and one output port for every effective
global EngineCore rank.
For rank
r:The socket plan stores only these semantic port values. The coordinator derives
wildcard bind addresses, while workers derive connect addresses using
data_parallel_master_ip.Both single-node IPC and multi-node TCP coordinator sockets are initialized
through the same binding path before workers start.
Backward compatibility
Existing single-node commands continue to work:
python -m atom.entrypoints.openai_server \ --model "$MODEL" \ -dp 8Because the local DP size is omitted, it defaults to 8 and all ranks are
launched locally.
Verification
Completed static verification:
git diff --checkA full multi-node GPU serving and accuracy run is still required before merge.
The model path must be available on both nodes.
Node 0: coordinator, global DP rank 0
Node 1: worker, global DP rank 1
Send requests only to node 0:
Two nodes with TP4 on each node
This configuration creates two global DP replicas, each using four GPUs:
Node 0
Node 1
Four nodes with two DP ranks per node
For global DP8, TP1, with two GPUs per node: