Skip to content

Report the unscaled gradient norm from FP16_UnfusedOptimizer - #8588

Open
vineethsaivs wants to merge 1 commit into
deepspeedai:masterfrom
vineethsaivs:fix-unfused-grad-norm-scale-20260918
Open

vineethsaivs wants to merge 1 commit into
deepspeedai:masterfrom
vineethsaivs:fix-unfused-grad-norm-scale-20260918

Conversation

@vineethsaivs

Copy link
Copy Markdown
Contributor

DeepSpeedEngine.get_global_grad_norm() is documented as the 2-norm of all gradients, but on the fp16 non-fused path it returns that norm still multiplied by the loss scale.

step and step_fused_lamb stash the norm built from the fp16 gradients, which still carry the loss scale. The fused optimizer divides it out, and this class's own unscale_and_clip_grads does too. Keep the scaled norm in a local for clipping and stash the unscaled one, matching the fused optimizer. The clip decision is unchanged.

Test: TORCHDYNAMO_DISABLE=1 DS_ACCELERATOR=cpu PYTHONPATH=. python -m pytest tests/unit/runtime/half_precision/test_unfused_optimizer.py -q. Two regressions fail before, reporting 128.0 for a true norm of 1.0 at a static loss scale of 128, on both step paths. They pass after, and pin that clipping still gets the scaled norm. Changed-file pre-commit passes. Executed on Apple M2 Pro CPU; no CUDA or distributed training run.

Keep the scaled norm in a local for clipping and stash the unscaled one, matching the fused optimizer. The clip decision is unchanged.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
@vineethsaivs

Copy link
Copy Markdown
Contributor Author

Could you approve the CI run? Two regressions fail before, reporting 128.0 for a true norm of 1.0 at a static loss scale of 128, on both step paths. They pass after, and pin that clipping still gets the scaled norm. Changed-file pre-commit passes.

@vineethsaivs

Copy link
Copy Markdown
Contributor Author

One adjacent thing I found while checking the siblings, deliberately left out of this diff.

The two classes also disagree about MoE. fused_optimizer.py folds the expert gradients into the norm it clips with (get_norm_with_moe_layers, line 322). The unfused step computes expert_norm_groups but spends it only on the overflow check, and _global_grad_norm and unscale_and_clip_grads see norm_groups alone; step_fused_lamb discards the expert grads outright at line 209. So on the unfused fp16 path with MoE layers the clip norm is built from the shared parameters only.

That is a behaviour change to clipping rather than to reporting, which is what this PR is about, and this PR states the clip decision is unchanged. Happy to open it separately, or fold it in here if you would rather have both at once.

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