Skip to content

[Refactor] Keep only the TP and PP groups on the model runner - #41813

Merged
ch-wan merged 1 commit into
mainfrom
cheng/refactor/runner-placement-snapshots
Sep 30, 2026
Merged

ch-wan merged 1 commit into
mainfrom
cheng/refactor/runner-placement-snapshots

Conversation

@ch-wan

@ch-wan ch-wan commented Sep 30, 2026 •

Copy link
Copy Markdown
Collaborator

This PR is part of a stack (oldest at bottom):

Motivation

ModelRunner.init_torch_distributed copies fifteen placement values onto the runner. Only two of them have to outlive the scope the runner is built in:

  • speculative workers re-enter a draft's scope through the draft runner's TP group;
  • draft forwards read its PP group without entering the pipeline scope.

The other thirteen are second copies of values that get_parallel() already answers wherever they are read.

Modifications

  • Keep tp_group and pp_group. Drop the other thirteen: attention_tp_group, tp_rank, tp_size, dp_size, attn_dp_size, pp_rank, pp_size, attn_cp_rank, attn_cp_size, attn_dcp_rank, attn_dcp_size, moe_ep_size, dp_rank.
  • Each reader takes the value from where it runs:
    • Readers that run while the runner is built, during a forward or a graph capture, or only on the target read get_parallel(). In those scopes it answers what the runner copied.
    • Readers that can run on a draft runner outside its scope read the runner's own groups. The KV cache configurator is built while pools are allocated, and the PP proxy hidden size is also reached when graphs are recaptured; both use pp_group.world_size / rank_in_group. The DFLASH / DSpark rank-0 log gates inside the draft scope read the target runner's tp_group.rank_in_group.
    • The HiSparse coordinator is built with the pools and reads get_parallel().attn_tp_group. Outside the scope that is the target's attention-TP group, which is also the group an attention-owning draft runs on; a draft that does not own attention never narrows it.
  • The runner-owned tp_size / tp_rank arguments to build_load_config, dist_barrier_after_load, the tensor-dump hook, apply_torch_tp and the quantized-MoE check come from the context too. All of these run in the runner's constructor, except the barrier after an overlapped startup load, which the scheduler runs on the target's runners outside any draft scope.
  • Tests that stood in for the removed fields follow:
    • the MLX runner stub tests publish their attention-DP width;
    • the DFLASH sampler and phase-1 tests give their fake a TP group and publish a topology;
    • fake runners in the attention test kits and backend tests drop placement fields nothing reads;
    • the placement-freeze census asserts that it finds the TP and PP groups.

No value a reader sees changes. Code outside the tree that reads, for example, model_runner.tp_rank should read get_parallel().tp_rank.

Accuracy Tests

H200:

  • Qwen3-0.6B with dummy weights, --random-seed 42, --tp-size 2 and --tp-size 2 --dp-size 2 --enable-dp-attention: 4 greedy prompts × 32 tokens are identical to the parent commit.
  • Qwen3-8B with an EAGLE3 draft, TP2 and TP2 + DP attention (the attention-owning draft path): greedy text and spec_verify_ct are identical to main. Measured at the top of this stack, which contains this PR.
  • The MLX runner stub tests pass (21) with a stand-in mlx module, on main and on this PR.
  • test/registered/unit at this PR's head, compared with main: no new failures.

Speed Tests and Profiling

Not applicable.

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): 🚫 Run #36772756644
Latest PR Test (Extra): 🚫 Run #36772756255
Latest PR Test (AMD ROCm 10): 🚫 Run #36772756706

@ch-wan
ch-wan force-pushed the cheng/refactor/runner-placement-snapshots branch from eea0f19 to e5c3fef Compare September 30, 2026 04:24
@ch-wan
ch-wan force-pushed the cheng/refactor/drop-unread-model-placement branch from 791b0e2 to 134ee9d Compare September 30, 2026 05:05
@ch-wan
ch-wan force-pushed the cheng/refactor/runner-placement-snapshots branch from e5c3fef to b996c34 Compare September 30, 2026 05:05
@ch-wan
ch-wan force-pushed the cheng/refactor/drop-unread-model-placement branch from 134ee9d to afcf208 Compare September 30, 2026 20:27
Base automatically changed from cheng/refactor/drop-unread-model-placement to main September 30, 2026 20:27
`ModelRunner.init_torch_distributed` copied fifteen placement values
onto the runner. Only two of them have to outlive the scope the runner
is built in: speculative workers re-enter a draft's scope through its TP
group, and draft forwards read its PP group without the pipeline scope.
Keep those two and drop the other thirteen (`attention_tp_group`,
`tp_rank`, `tp_size`, `dp_size`, `attn_dp_size`, `pp_rank`, `pp_size`,
`attn_cp_rank`, `attn_cp_size`, `attn_dcp_rank`, `attn_dcp_size`,
`moe_ep_size`, `dp_rank`).

Each reader now takes the value from where it runs:

- readers that run while the runner is built, during a forward or a
  graph capture, or only on the target read `get_parallel()`, which in
  those scopes is what the runner saw when it copied the value;
- readers that can run on a draft runner outside its scope read the
  runner's own groups: the KV cache configurator (built while pools are
  allocated) and the PP proxy hidden size (also reached when graphs are
  recaptured) use `pp_group.world_size` / `rank_in_group`, and the
  DFLASH / DSpark rank-0 log gates inside the draft scope read the target
  runner's `tp_group.rank_in_group`;
- the HiSparse coordinator, built with the pools, reads
  `get_parallel().attn_tp_group`. Outside the scope that is the target's
  attention-TP group, which is the group an attention-owning draft runs
  on; a draft that does not own attention never narrows it.

The runner-owned `tp_size` / `tp_rank` arguments to `build_load_config`,
`dist_barrier_after_load`, the tensor-dump hook, `apply_torch_tp` and the
quantized-MoE check now come from the context too. All of them run in
the runner's constructor, except the barrier after an overlapped startup
load, which the scheduler runs on the target's runners outside any draft
scope.

No value a reader sees changes.

Tests that stood in for the retired fields follow: the MLX runner stub
tests publish their attention-DP width, the DFLASH sampler and phase-1
tests give the fake a TP group and publish a topology, fake runners in
the attention test kits and backend tests drop placement fields nothing
reads, and the placement-freeze census asserts it finds the TP and PP
groups.
@ch-wan
ch-wan force-pushed the cheng/refactor/runner-placement-snapshots branch from b996c34 to aecc187 Compare September 30, 2026 20:27
@ch-wan
ch-wan merged commit 499e86d into main Sep 30, 2026
8 of 19 checks passed
@ch-wan
ch-wan deleted the cheng/refactor/runner-placement-snapshots branch September 30, 2026 20:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant