Skip to content

[Refactor] Add make_pp_layers so models stop handing their PP position down - #41816

Merged
ch-wan merged 1 commit into
mainfrom
cheng/refactor/make-pp-layers
Sep 30, 2026
Merged

ch-wan merged 1 commit into
mainfrom
cheng/refactor/make-pp-layers

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

48 models build their decoder stack with make_layers(..., pp_rank=self.pp_group.rank_in_group, pp_size=self.pp_group.world_size, ...); one spells it with get_parallel().pp_group. Passing the pair both supplies the stage and switches make_layers to return (layers, start_layer, end_layer), so it is not redundant as such. But every caller passes the placement of the scope it is being built in.

Modifications

  • Add make_pp_layers(num_hidden_layers, layer_fn, prefix, ...), which reads the stage from get_parallel() and returns the tuple, and use it at all 48 sites.
  • Models are constructed inside the scope that states their placement (a draft inside its pipeline scope), where get_parallel() answers the same stage as self.pp_group.
  • test_make_pp_layers.py: each stage builds its own slice from the published placement, and a single-stage scope builds every layer.

Accuracy Tests

H200:

  • Qwen3-0.6B with dummy weights, --tp-size 2 and --pp-size 2: greedy output identical with and without this PR and the two before it.
  • upstage/solar-pro-preview-instruct, --tp-size 2 --pp-size 2: greedy output identical to the Solar TP2 run earlier in this stack; GSM8K (200 questions) 0.900.
  • 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 #36772881908
Latest PR Test (Extra): 🚫 Run #36772880959
Latest PR Test (AMD ROCm 10): ❌ Run #36772881796

@ch-wan
ch-wan force-pushed the cheng/refactor/models-drop-same-leaf-placement branch from 4b8856c to 6b0dfee Compare September 30, 2026 04:24
@ch-wan
ch-wan force-pushed the cheng/refactor/make-pp-layers branch from ebb248e to be4156e Compare September 30, 2026 04:25
@ch-wan
ch-wan force-pushed the cheng/refactor/models-drop-same-leaf-placement branch from 6b0dfee to a5f7cf4 Compare September 30, 2026 05:06
@ch-wan
ch-wan force-pushed the cheng/refactor/make-pp-layers branch from be4156e to e92827d Compare September 30, 2026 05:06
@ch-wan
ch-wan force-pushed the cheng/refactor/models-drop-same-leaf-placement branch from a5f7cf4 to eccd9b9 Compare September 30, 2026 20:28
Base automatically changed from cheng/refactor/models-drop-same-leaf-placement to main September 30, 2026 20:28
48 models built their decoder stack with
`make_layers(..., pp_rank=self.pp_group.rank_in_group,
pp_size=self.pp_group.world_size, ...)` (one spelled it with
`get_parallel().pp_group`). Passing the pair both supplies the stage and
switches `make_layers` to return `(layers, start_layer, end_layer)`, so it
was not a redundant argument as such -- but every caller passed the
placement of the scope it was being built in.

Add `make_pp_layers(num_hidden_layers, layer_fn, prefix, ...)`, which
reads the stage from `get_parallel()` and returns the tuple, and use it at
all 48 sites. Models are constructed inside the scope that states their
placement (a draft inside its pipeline scope), where `get_parallel()`
answers the same stage as `self.pp_group`.

A unit test checks that each stage builds its own slice from the
published placement, and that a single-stage scope builds every layer.
@ch-wan
ch-wan force-pushed the cheng/refactor/make-pp-layers branch from e92827d to cf8b583 Compare September 30, 2026 20:28
@ch-wan
ch-wan merged commit e7f6a99 into main Sep 30, 2026
2 checks passed
@ch-wan
ch-wan deleted the cheng/refactor/make-pp-layers branch September 30, 2026 20:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant