Make MoE router-weight placement explicit in experts - #79
Draft
zianglih wants to merge 3 commits into
Draft
Conversation
This reverts commit 6d43fca.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
@HumansAnd
MoERouterWeightPlacementexpert API withFC2_INPUTandFC2_OUTPUTmodes.MEGATRON_MOE_APPLY_PROBS_ON_OUTPUT=1selects FC2-output placement.GroupedMLP,TEGroupedMLP, andSequentialMLP, while preserving the legacy three-argument call for disabled/custom-expert paths.Numerical contract
The FC2-output mode no longer passes unit probabilities through the weighted activation path. It runs SwiGLU and FC2 unweighted, promotes the materialized FC2 output, expert bias, and original router probability to at least FP32, computes
(fc2_output + bias) * router_prob, and casts the result once to the expert output dtype.This removes the unit-probability allocation and activation-sized multiply/reduction, and it avoids narrowing the router probability or creating a probability-weighted BF16 FC2 input.
This is not a grouped-GEMM epilogue: TE has already materialized and rounded the FC2 output to the expert output dtype. Applying the router probability to the GEMM accumulator would require a future TE/kernel API extension and is intentionally outside this PR.
Scope
moe_apply_probs_on_inputandmoe_combine_in_fp32remain incompatible with FC2-output placement.Validation
git diff --checkpython3 -m py_compileon all changed Python files5 passed)GPU validation used
radixark/miles@sha256:d9e01378d8820afd88824c798ea628b3b6cb87a6c6911db5165ef9d98187db55without reinstalling packages.