Include the expert gradients in the unfused fp16 clip norm - #8613
Open
vineethsaivs wants to merge 1 commit into
Open
vineethsaivs wants to merge 1 commit into
vineethsaivs wants to merge 1 commit into
Conversation
Fold them in the same way, and pass the engine's `has_moe_layers` through so it can. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
vineethsaivs
requested review from
loadams,
tjruwase and
tohtana
as code owners
September 20, 2026 19:42
This branch has not been deployed
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.
With MoE and the fp16 unfused optimizer, gradient clipping is decided from a norm built out of the shared parameters alone, so the clip coefficient is too large and every gradient, expert ones included, is under-clipped.
stepcallssplit_params_grads_into_shared_and_expert_paramsand discards the expert half;step_fused_lambspends itsexpert_norm_groupson the overflow check only. The split exists because expert parameters are replicated across the expert-parallel group, so their norm needs a reduction over that group thatget_weight_normcannot do.FP16_Optimizerfolds them in withget_norm_with_moe_layers, and this class already carries thenorm_typethat call wants but never uses it. Fold them in the same way, and pass the engine'shas_moe_layersthrough so it can.Test:
TORCHDYNAMO_DISABLE=1 DS_ACCELERATOR=cpu PYTHONPATH=.:tests python -m pytest tests/unit/runtime/half_precision/ -q. Two new cases fail before, both spellings of the step: four shared gradients of 3.0 and four expert gradients of 4.0 give a reported norm of 6.0 where 10.0 is correct. Two more pin that a run without experts is unchanged. Nine pass and 32 skip after. yapf 0.43.0 and flake8 pass on the changed files. This touches the two lines #8588 also rewrites, so whichever lands first, I will rebase the other. Executed on Apple M2 Pro CPU; no GPU or distributed run. Single rank, so the expert-parallel group is None and the helper runs as plain arithmetic.