[minimax_m3] gate AR+RMSNorm fusion on ATOM_ENABLE_ALLREDUCE_RMSNORM_…#1344
Draft
zejunchen-zejun wants to merge 1 commit into
Draft
[minimax_m3] gate AR+RMSNorm fusion on ATOM_ENABLE_ALLREDUCE_RMSNORM_…#1344zejunchen-zejun wants to merge 1 commit into
zejunchen-zejun wants to merge 1 commit into
Conversation
…FUSION Previously M3's target-model fusion was hardcoded on (fused_allreduce_gemma_rms_norm fused all-reduce + residual-add + Gemma RMSNorm whenever TP>1), ignoring the ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION env that the draft / deepseek_v2 / glm4_moe already honor — so the env couldn't disable M3's fusion for A/B testing. Make both halves of the fusion move together under the flag (default on): - layernorm.py fused_allreduce_gemma_rms_norm: also require envs.ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION; when off it returns a plain residual-add + Gemma RMSNorm (expects an already-all-reduced input). - minimax_m3.py: the RowParallel o_proj (dense + sparse attn), MoE experts, shared_experts and dense MLP down_proj now use reduce_results=not ENABLE_ALLREDUCE_RMSNORM_FUSION, so when the env is off they do their own all-reduce and the norm runs unfused. Default (=1) is byte-equivalent to before. Set ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION=0 to disable M3 target fusion (linears all-reduce, norm unfused). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
zufayu
requested review from
JiaoliangYu and
ZhangLirong-amd
and removed request for
ZhangLirong-amd
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
1 participant
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.
…FUSION
Previously M3's target-model fusion was hardcoded on (fused_allreduce_gemma_rms_norm fused all-reduce + residual-add + Gemma RMSNorm whenever TP>1), ignoring the ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION env that the draft / deepseek_v2 / glm4_moe already honor — so the env couldn't disable M3's fusion for A/B testing.
Make both halves of the fusion move together under the flag (default on):
Default (=1) is byte-equivalent to before. Set ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION=0 to disable M3 target fusion (linears all-reduce, norm unfused).
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist