Skip to content

[Refactor] Drop model placement attributes and parameters nothing reads - #41812

Merged
ch-wan merged 2 commits into
mainfrom
cheng/refactor/drop-unread-model-placement
Sep 30, 2026
Merged

ch-wan merged 2 commits into
mainfrom
cheng/refactor/drop-unread-model-placement

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

Model modules copy parallel ranks and sizes onto self (self.tp_size = get_parallel().tp_size and similar). In 97 places no method of the class, its subclasses or its bases reads the attribute back, and no other code reads it off an instance. A few constructors also accept placement parameters that no caller sets.

Modifications

  • Remove the 97 unread assignments across 53 model files. The Dots Note omni wrapper's tp_size, copied from its language model and itself never read, goes with the language model's. Three TODOs that only annotated a removed assignment go with it.
  • Step3p5MLP accepted tp_size / tp_rank only to forward them to its linear layers, and neither caller passes them, so the layers always got their own defaults. OlmoeMoE, Grok1MoE and MixtralMoE accepted a tp_size they never use. Remove these parameters; every caller already passes by keyword without them.

Code outside the tree that read one of the removed attributes, for example model.tp_size, should read get_parallel() instead.

Accuracy Tests

Not applicable: nothing that is read changes.

  • All 53 modified model files import.
  • 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 #36772717146
Latest PR Test (Extra): 🚫 Run #36772716531
Latest PR Test (AMD ROCm 10): ❌ Run #36772717248

@ch-wan
ch-wan force-pushed the cheng/refactor/drop-unread-placement branch from 52fc839 to 02fdd93 Compare September 30, 2026 04:24
@ch-wan
ch-wan force-pushed the cheng/refactor/drop-unread-model-placement branch from 0a98d71 to 791b0e2 Compare September 30, 2026 04:24
@ch-wan
ch-wan force-pushed the cheng/refactor/drop-unread-placement branch from 02fdd93 to 87b938b Compare September 30, 2026 05:05
@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/drop-unread-placement branch from 87b938b to 80e6449 Compare September 30, 2026 20:26
Base automatically changed from cheng/refactor/drop-unread-placement to main September 30, 2026 20:26
Model modules copied parallel ranks and sizes onto `self`
(`self.tp_size = get_parallel().tp_size` and similar) in 97 places where
no method of the class, its subclasses or its bases reads the attribute
back, and no other code reads it off an instance. Remove those
assignments across 53 model files. The Dots Note omni wrapper's
`tp_size`, copied from its language model and itself never read, goes
with the language model's.
`Step3p5MLP` accepted `tp_size` / `tp_rank` only to forward them to its
linear layers, and neither caller passes them, so the layers always got
their own defaults. `OlmoeMoE`, `Grok1MoE` and `MixtralMoE` accepted a
`tp_size` they never use. Remove the parameters; every caller already
passes by keyword without them.
@ch-wan
ch-wan force-pushed the cheng/refactor/drop-unread-model-placement branch from 134ee9d to afcf208 Compare September 30, 2026 20:27
@ch-wan
ch-wan merged commit a09aedf into main Sep 30, 2026
9 of 18 checks passed
@ch-wan
ch-wan deleted the cheng/refactor/drop-unread-model-placement branch September 30, 2026 20:27
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