Skip to content

Include the expert gradients in the unfused fp16 clip norm - #8613

Open
vineethsaivs wants to merge 1 commit into
deepspeedai:masterfrom
vineethsaivs:fix-moe-unfused-clip-norm-20260920
Open

vineethsaivs wants to merge 1 commit into
deepspeedai:masterfrom
vineethsaivs:fix-moe-unfused-clip-norm-20260920

Conversation

@vineethsaivs

Copy link
Copy Markdown
Contributor

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.

step calls split_params_grads_into_shared_and_expert_params and discards the expert half; step_fused_lamb spends its expert_norm_groups on 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 that get_weight_norm cannot do. FP16_Optimizer folds them in with get_norm_with_moe_layers, and this class already carries the norm_type that call wants but never uses it. Fold them in the same way, and pass the engine's has_moe_layers through 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.

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>

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