Skip to content

Add a lineage_expert_permutation config option for per-step expert - #5520

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_991959508
Open

copybara-service[bot] wants to merge 1 commit into
mainfrom
test_991959508

Conversation

@copybara-service

Copy link
Copy Markdown
Contributor

Add a lineage_expert_permutation config option for per-step expert
load balancing in the Lineage MoE path.

When set to "lpt", the decoder and each MTP layer keep a per-layer routed
expert permutation as runtime state (moe.ExpertPermutationVar, an
nnx.BatchStat initialised to the identity), which each forward pass
reads and overwrites with the permutation for the next step. Params-only
checkpoint loads skip it. The default (None) creates no state and leaves
behaviour unchanged.

@codecov

codecov Bot commented Oct 2, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 7.69231% with 12 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/layers/multi_token_prediction.py 0.00% 6 Missing ⚠️
src/maxtext/layers/nnx_decoders.py 0.00% 5 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

@copybara-service
copybara-service Bot force-pushed the test_991959508 branch 3 times, most recently from b547b37 to dca7669 Compare October 3, 2026 02:06
load balancing in the Lineage MoE path.

When set to "lpt", the decoder and each MTP layer keep a per-layer routed
expert permutation as runtime state (`moe.ExpertPermutationVar`, an
`nnx.BatchStat` initialised to the identity), which each forward pass
reads and overwrites with the permutation for the next step. Params-only
checkpoint loads skip it. The default (None) creates no state and leaves
behaviour unchanged.

PiperOrigin-RevId: 991959508

This branch has not been deployed

No deployments
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