[ganyi1996ppo]

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_dp or distributed_dp_serving flags are required.

Motivation

The existing native DP implementation assumes every DP rank is local:

• global DP rank and local GPU rank are treated as the same value;
• one CoreManager launches every DP rank;
• EngineCore communication uses local IPC sockets;
• distributed initialization can select devices using global rank values.

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:

Argument	Description
--data-parallel-size, -dp	Total number of DP ranks across all nodes
--data-parallel-size-local	Number of DP ranks hosted on this node
--data-parallel-rank	First global DP rank hosted on this node
--data-parallel-master-ip	Routable IP address of the rank-0 node
--data-parallel-master-port	DP rendezvous and control-port base
--data-parallel-base-port	Model-runner distributed initialization port

--data-parallel-rank-local is assigned internally to EngineCore workers and
normally should not be set by users.

If --data-parallel-size-local is omitted, it defaults to the global DP size,
preserving existing single-node behavior.

Multi-node mode is inferred when:

data_parallel_size_local < data_parallel_size
or data_parallel_rank > 0

The topology is validated as follows:

data_parallel_size >= 1
data_parallel_size_local >= 1
data_parallel_rank >= 0
data_parallel_rank + data_parallel_size_local <= data_parallel_size

Coordinator/worker model

In multi-node mode:

• the node whose first global DP rank is 0 is the sole coordinator;
• only the coordinator serves the OpenAI-compatible HTTP API;
• the coordinator owns routing sockets for every global EngineCore rank;
• non-coordinator nodes launch only their local EngineCore workers;
• requests are routed round-robin unless a request carries an explicit DP-rank
  routing hint.

Single-node deployments continue using the existing local IPC transport.

Global and local ranks

Each EngineCore receives two DP ranks:

• global DP rank: used for distributed groups and collectives;
• local DP rank: used for GPU and NUMA placement.

For example, the second node in a DP16 deployment may be configured with:

global DP size = 16
local DP size  = 8
rank offset    = 8

Its rank mapping becomes:

global rank 8  -> local rank 0
global rank 9  -> local rank 1
...
global rank 15 -> local rank 7

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:

global_engine_rank = global_dp_rank * tp_size + tp_rank
local_engine_rank = local_dp_rank * tp_size + tp_rank

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:

input_port  = data_parallel_master_port + 100 + 2*r
output_port = data_parallel_master_port + 100 + 2*r + 1

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 8

Because the local DP size is omitted, it defaults to 8 and all ranks are
launched locally.

Verification

Completed static verification:

• Ruff on all modified Python files
• Python bytecode compilation
• git diff --check
• no RCCL/MoE backend implementation changes included

A full multi-node GPU serving and accuracy run is still required before merge.

## Usage examples

### Two nodes with one GPU each

Assume:

```bash
MODEL=/path/to/model
COORDINATOR_IP=10.0.0.10

The model path must be available on both nodes.

Node 0: coordinator, global DP rank 0

export AITER_LOG_LEVEL=WARNING

python -m atom.entrypoints.openai_server \
  --model "$MODEL" \
  -tp 1 \
  -dp 2 \
  --data-parallel-size-local 1 \
  --data-parallel-rank 0 \
  --data-parallel-master-ip "$COORDINATOR_IP" \
  --data-parallel-master-port 29500 \
  --data-parallel-base-port 29550 \
  --host 0.0.0.0 \
  --server-port 8000

Node 1: worker, global DP rank 1

export AITER_LOG_LEVEL=WARNING

python -m atom.entrypoints.openai_server \
  --model "$MODEL" \
  -tp 1 \
  -dp 2 \
  --data-parallel-size-local 1 \
  --data-parallel-rank 1 \
  --data-parallel-master-ip "$COORDINATOR_IP" \
  --data-parallel-master-port 29500 \
  --data-parallel-base-port 29550

Send requests only to node 0:

curl "http://${COORDINATOR_IP}:8000/v1/completions" \
  -H "Content-Type: application/json" \
  -d '{
    "model": "'"$MODEL"'",
    "prompt": "Explain data-parallel inference in one sentence.",
    "max_tokens": 64
  }'

Two nodes with TP4 on each node

This configuration creates two global DP replicas, each using four GPUs:

global DP size = 2
local DP size per node = 1
TP size = 4
total GPUs = DP × TP = 8

Node 0

export AITER_LOG_LEVEL=WARNING

python -m atom.entrypoints.openai_server \
  --model "$MODEL" \
  -tp 4 \
  -dp 2 \
  --data-parallel-size-local 1 \
  --data-parallel-rank 0 \
  --data-parallel-master-ip "$COORDINATOR_IP" \
  --data-parallel-master-port 29500 \
  --data-parallel-base-port 29550 \
  --host 0.0.0.0 \
  --server-port 8000

Node 1

export AITER_LOG_LEVEL=WARNING

python -m atom.entrypoints.openai_server \
  --model "$MODEL" \
  -tp 4 \
  -dp 2 \
  --data-parallel-size-local 1 \
  --data-parallel-rank 1 \
  --data-parallel-master-ip "$COORDINATOR_IP" \
  --data-parallel-master-port 29500 \
  --data-parallel-base-port 29550

Four nodes with two DP ranks per node

For global DP8, TP1, with two GPUs per node:

Node	Local size	Rank offset	Global ranks
Node 0	2	0	0–1
Node 1	2	2	2–3
Node 2	2	4	4–5
Node 3	2	6	6–7

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1603 --add-label <label>

[Copilot]

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.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (2)

atom/model_engine/engine_core_mgr.py:41

• PR description “EngineCore transport” specifies 2 ports per engine (base + 2*r), but the implementation and docs now allocate 3 ports per engine (request/output/control). Please update the PR description to match the current 3-port plan so future readers don’t misconfigure ports.

# Engine sockets for a multi-node run cannot use IPC paths, and the two ends
# must derive identical TCP ports without negotiating. Both compute this plan
# from the shared master port.
DP_SOCKET_PORT_OFFSET = 100
DP_PORTS_PER_ENGINE = 3

docs/distributed_guide.md:514

• The guide claims “explicit flags win”, but ParallelConfig.__post_init__ applies ATOM_DP_* environment overrides after initialization, so a set env var will override any CLI-provided value. Please update the docs (or adjust precedence in code) so users don’t get surprising DP topology values.

Each has an `ATOM_DP_*` environment equivalent (`ATOM_DP_SIZE_LOCAL`,
`ATOM_DP_RANK`, …) for launch scripts; explicit flags win.

[Copilot]

Pull request overview

Copilot reviewed 18 out of 18 changed files in this pull request and generated 1 comment.

Suppressed comments (3)

atom/utils/init.py:103

• This docstring says CoreManager wraps engine spawns in set_device_control_env_var, but set_device_control_env_var() is no longer called anywhere in the codebase. The reference to CoreManager is now misleading; update the wording to describe the current (rollout-only) consumer without implying the DP serving path uses it.

    Inert for the DP path that calls it, and must stay that way. It sets a
    variable named by the literal placeholder string; the only reader is
    ``RLHFModelRunner.DP_DEVICE_MAP_ENV`` (atom/rollout/model_runner_ext.py),
    which consults it solely when ``data_parallel_size <= 1``. ``CoreManager``
    only wraps a spawn in this when ``data_parallel_size > 1``, so the two
    never meet. Do not delete the placeholder string as dead -- it is a live

atom/model_engine/engine_core_mgr.py:40

• PR description (“EngineCore transport” section) still describes a 2-ports-per-rank plan, but the implementation and docs now use 3 ports per engine (input/output/control) via INTERNODE_DP_PORTS_PER_ENGINE = 3. Please update the PR description to avoid confusing reviewers/operators about required firewall openings and port coordination.

# Engine sockets for a multi-node run cannot use IPC paths, and the two ends
# must derive identical TCP ports without negotiating. Both compute this plan
# from the shared master port.
INTERNODE_DP_SOCKET_PORT_OFFSET = 100
INTERNODE_DP_PORTS_PER_ENGINE = 3

docs/distributed_guide.md:514

• This sentence claims “explicit flags win”, but ParallelConfig.__post_init__ currently applies ATOM_DP_* overrides whenever the env var is set, even if the value was provided via CLI flags/programmatic config. Either adjust the precedence behavior in code, or update this doc to match reality (env overrides when set).

Each has an `ATOM_DP_*` environment equivalent (`ATOM_DP_SIZE_LOCAL`,
`ATOM_DP_RANK`, …) for launch scripts; explicit flags win.

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (3)

docs/distributed_guide.md:514

• This claims CLI flags override env vars, but ParallelConfig.__post_init__ overwrites DP fields from ATOM_DP_* whenever those env vars are set. As written, a set env var will override an explicit flag value.

Each has an `ATOM_DP_*` environment equivalent (`ATOM_DP_SIZE_LOCAL`,
`ATOM_DP_RANK`, …) for launch scripts; explicit flags win.

docs/distributed_guide.md:545

• The PR description’s port formula uses 2 ports per rank (+ 2*r / + 2*r + 1), but the implementation and this doc section use 3 ports per engine (requests/outputs/control, base + 3r + {0,1,2}). Please update the PR description/design section to match the 3-port scheme so operators don’t misconfigure firewalls.

Engine sockets are derived from the DP master port; both ends compute them
independently and must agree. With `base = master_port + 100`, DP rank `r` uses:

| Port | Purpose |
|---|---|
| `base + 3r` | requests (ROUTER/DEALER) |
| `base + 3r + 1` | outputs (PULL/PUSH) |
| `base + 3r + 2` | control — utility commands, abort, shutdown |

atom/utils/init.py:104

• This docstring still states that CoreManager wraps engine spawns in this context manager when data_parallel_size > 1, but CoreManager no longer imports or uses set_device_control_env_var after the multi-node DP changes. That makes the current explanation about why RLHF and DP paths “never meet” inaccurate.

    ``RLHFModelRunner.DP_DEVICE_MAP_ENV`` (atom/rollout/model_runner_ext.py),
    which consults it solely when ``data_parallel_size <= 1``. ``CoreManager``
    only wraps a spawn in this when ``data_parallel_size > 1``, so the two
    never meet. Do not delete the placeholder string as dead -- it is a live
    protocol constant for the rollout adapter.

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (3)

docs/distributed_guide.md:527

• The multi-node example commands omit --data-parallel-base-port. In multi-node runs the default data_parallel_base_port is chosen dynamically (via get_open_port()), so each node can pick a different port and distributed initialization can hang. The example should set an explicit base port on both nodes (matching the PR description and the CLI help in atom/model_engine/arg_utils.py).

python -m atom.entrypoints.openai_server --model <model> -tp 2 \
  --data-parallel-size 8 --data-parallel-size-local 4 --data-parallel-rank 0 \
  --data-parallel-master-ip 10.0.0.1 --data-parallel-master-port 29500

docs/configuration_guide.md:259

• The table now documents data_parallel_size_local as None → data_parallel_size, but the Computed property section immediately below still lists world_size / world_size_across_dp. Those fields/properties were removed from ParallelConfig (see atom/config.py), so this section is now misleading and should be updated to reflect the current API (e.g., is_multinode_dp).

| `data_parallel_size_local` | `int \| None` | `None` → `data_parallel_size` | DP ranks on **this** node. Defaults to the global size, i.e. single-node. Set it lower to give a node one slice of a multi-node run; also reaches MoRI as `gpu_per_node`. Overridden by `ATOM_DP_SIZE_LOCAL` |
| `data_parallel_rank` | `int` | `0` | First **global** DP rank owned by this node; overridden by `ATOM_DP_RANK` |
| `data_parallel_rank_local` | `Optional[int]` | `None` | Local rank within the data-parallel group (SPMD mode); overridden by `ATOM_DP_RANK_LOCAL` |

docs/distributed_guide.md:542

• The PR description’s port formula describes 2 ports per rank (+ 2*r / + 2*r + 1), but the implementation and this guide use 3 ports per engine (requests/outputs/control via INTERNODE_DP_PORTS_PER_ENGINE = 3). Please reconcile the PR description with the actual 3-port plan so operators don’t open the wrong firewall rules.

Engine sockets are derived from the DP master port; both ends compute them
independently and must agree. With `base = master_port + 100`, DP rank `r` uses:

| Port | Purpose |
|---|---|

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (4)

docs/distributed_guide.md:527

• Same as node 0: the worker-node command should set --data-parallel-base-port to the same value as the coordinator so model-runner distributed initialization uses a consistent rendezvous endpoint across nodes.

python -m atom.entrypoints.openai_server --model <model> -tp 2 \
  --data-parallel-size 8 --data-parallel-size-local 4 --data-parallel-rank 4 \
  --data-parallel-master-ip 10.0.0.1 --data-parallel-master-port 29500

atom/config.py:805

• ParallelConfig.data_parallel_rank is now used as the first global DP rank owned by this node in multi-node mode (see is_multinode_dp logic and CLI/docs). The field docstring still describes it as the in-group rank, which is misleading for readers and future callers.

    data_parallel_rank: int = 0
    """Rank of the data parallel group."""

docs/distributed_guide.md:522

• The multi-node example omits --data-parallel-base-port. In multi-node runs the model-runner distributed init uses ParallelConfig.data_parallel_base_port, and if it’s left unset each node will pick a different ephemeral port, causing NCCL rendezvous failures/hangs. Include an explicit base port (must match on all nodes).

This issue also appears on line 525 of the same file.

python -m atom.entrypoints.openai_server --model <model> -tp 2 \
  --data-parallel-size 8 --data-parallel-size-local 4 --data-parallel-rank 0 \
  --data-parallel-master-ip 10.0.0.1 --data-parallel-master-port 29500

atom/model_engine/engine_core_mgr.py:40

• The PR description’s "EngineCore transport" section describes only 2 ports per engine rank (base + 2*r / base + 2*r + 1). The implementation (and docs) now use 3 ports per engine (requests/outputs/control) with INTERNODE_DP_PORTS_PER_ENGINE = 3. Please update the PR description to match the actual socket plan so the design doc doesn’t diverge from code.

# Engine sockets for a multi-node run cannot use IPC paths, and the two ends
# must derive identical TCP ports without negotiating. Both compute this plan
# from the shared master port.
INTERNODE_DP_SOCKET_PORT_OFFSET = 100
INTERNODE_DP_PORTS_PER_ENGINE = 3

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (6)

docs/distributed_guide.md:527

• Multi-node launches require a stable data_parallel_base_port shared across nodes (it’s used to build distributed_init_method). This example omits --data-parallel-base-port, which can lead to different nodes choosing different ports and hanging during distributed init.

python -m atom.entrypoints.openai_server --model <model> -tp 2 \
  --data-parallel-size 8 --data-parallel-size-local 4 --data-parallel-rank 4 \
  --data-parallel-master-ip 10.0.0.1 --data-parallel-master-port 29500

docs/distributed_guide.md:514

• Docs claim that DP-related environment variables are only equivalents and that explicit flags win, but ParallelConfig.__post_init__ applies env var overrides whenever they are set. This is likely to confuse launches (e.g., a stale ATOM_DP_RANK/ATOM_DP_SIZE_LOCAL in the environment silently overriding CLI args). Please update the precedence wording to match the actual behavior.

Each has an `ATOM_DP_*` environment equivalent (`ATOM_DP_SIZE_LOCAL`,
`ATOM_DP_RANK`, …) for launch scripts; explicit flags win.

docs/distributed_guide.md:522

• Multi-node launches require a stable data_parallel_base_port shared across nodes (it’s used to build distributed_init_method). This example omits --data-parallel-base-port, which can lead to different nodes choosing different ports and hanging during distributed init.

This issue also appears on line 525 of the same file.

python -m atom.entrypoints.openai_server --model <model> -tp 2 \
  --data-parallel-size 8 --data-parallel-size-local 4 --data-parallel-rank 0 \
  --data-parallel-master-ip 10.0.0.1 --data-parallel-master-port 29500

docs/distributed_guide.md:545

• The PR description's EngineCore port formula describes 2 ports per engine (+ 2*r), but the implementation/docs here use 3 ports per engine (requests/outputs/control). Please align the PR description (or the implementation) so operators don’t provision the wrong firewall rules/port ranges.

Engine sockets are derived from the DP master port; both ends compute them
independently and must agree. With `base = master_port + 100`, DP rank `r` uses:

| Port | Purpose |
|---|---|
| `base + 3r` | requests (ROUTER/DEALER) |
| `base + 3r + 1` | outputs (PULL/PUSH) |
| `base + 3r + 2` | control — utility commands, abort, shutdown |

atom/model_engine/arg_utils.py:663

• In inferred multi-node DP mode (data_parallel_size_local < data_parallel_size or data_parallel_rank > 0), leaving data_parallel_base_port unset means each node will fall back to a different ephemeral port, and distributed init can hang. Consider validating this at CLI/config construction time and raising a clear error unless --data-parallel-base-port (or ATOM_DP_BASE_PORT) is provided.

        # 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)

atom/model_engine/engine_core_mgr.py:68

• New multi-node DP helpers/logic here (deterministic TCP port planning, global↔local rank assignment, coordinator vs worker behavior) are not covered by existing unit tests (no references in tests/ to build_internode_dp_socket_plan, iter_dp_rank_assignments, or is_multinode_dp). Given the operational risk of hangs/misrouting, please add CPU-only tests for these helpers and for routing state sizing with global_engine_count.

def build_internode_dp_socket_plan(
    *, engine_count: int, master_port: int
) -> list[InternodeDPSocketPlan]:
    """Derive deterministic per-engine TCP ports from the DP master port.

    Multi-node only -- ``CoreManager.__init__`` calls this when
    ``is_multinode_dp`` and passes ``None`` otherwise, and that ``None`` is what
    selects the IPC path (and, with it, the connect/bind polarity) further down.

    Three ports per engine: requests, outputs, and control. Control is separate
    so the request socket keeps a single writer thread (see _send_request).
    """

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 3 comments.