Skip to content

[Refactor] Let FusedMoE's weight-loading helpers read the layer's MoE-TP rank - #41814

Merged
ch-wan merged 1 commit into
mainfrom
cheng/refactor/moe-loaders-read-own-rank
Sep 30, 2026
Merged

ch-wan merged 1 commit into
mainfrom
cheng/refactor/moe-loaders-read-own-rank

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

The seven FusedMoE loading helpers (_load_w13, _load_w2, the scale and g_idx loaders, the GGUF and FP8 shared-expert paths) each take a tp_rank argument. Every call is inside FusedMoE, and every value traces back to tp_rank = self.moe_tp_rank in the two weight loaders.

Modifications

Read self.moe_tp_rank in the helpers, and drop the argument and the two locals.

A layer that shards differently stores that choice in self.moe_tp_rank itself: the Inkling shared experts and their loading helper keep the full-TP rank there. So the helpers see the same value as before.

Accuracy Tests

H200, compared with main at the top of this stack, which contains this PR; 4 greedy prompts × 64 tokens and GSM8K (200 questions):

model config greedy GSM8K main / this stack
Qwen3-30B-A3B --tp 2 --ep-size 2 identical 0.930 / 0.925
Qwen3-VL-30B-A3B-Instruct-AWQ (moe_wna16) --tp 2 identical 0.945 / 0.945
Qwen3-VL-30B-A3B-Instruct-AWQ (moe_wna16) --tp 2 --ep-size 2 fails to load on main / 0.935 (the load fix is #41807, below in this stack)
  • unit/models, unit/layers/moe, unit/layers/quantization and test_runtime_context: the same results as the parent commit.
  • 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 #36772813010
Latest PR Test (Extra): 🚫 Run #36772812682
Latest PR Test (AMD ROCm 10): ❌ Run #36772813054

@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/moe-loaders-read-own-rank branch from a89e8fe to 2917852 Compare September 30, 2026 04:24
@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/moe-loaders-read-own-rank branch from 2917852 to 73420b3 Compare September 30, 2026 05:05
@ch-wan
ch-wan force-pushed the cheng/refactor/runner-placement-snapshots branch from b996c34 to aecc187 Compare September 30, 2026 20:27
Base automatically changed from cheng/refactor/runner-placement-snapshots to main September 30, 2026 20:27
The seven loading helpers (`_load_w13`, `_load_w2`, the scale and g_idx
loaders, the GGUF and FP8-shared-expert paths) each took a `tp_rank`
argument. Every call is inside `FusedMoE`, and every value traces back to
`tp_rank = self.moe_tp_rank` in the two weight loaders. Read
`self.moe_tp_rank` in the helpers and drop the argument and the two
locals.

A layer that shards differently stores that in `self.moe_tp_rank`
itself -- the Inkling shared experts and their loading helper keep the
full-TP rank there -- so the helpers see the same value as before.
@ch-wan
ch-wan force-pushed the cheng/refactor/moe-loaders-read-own-rank branch from 73420b3 to 9035759 Compare September 30, 2026 20:27
@ch-wan
ch-wan merged commit c7fa37a into main Sep 30, 2026
4 of 12 checks passed
@ch-wan
ch-wan deleted the cheng/refactor/moe-loaders-read-own-rank 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

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant