Report the unscaled gradient norm from FP16_UnfusedOptimizer

[vineethsaivs]

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.

[vineethsaivs]

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]

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.

[vineethsaivs]

The adjacent MoE finding from this thread is now #8613: step drops the expert half of the split and step_fused_lamb spends it on the overflow check only, so the unfused clip norm is built from the shared parameters alone. It touches the two lines this PR rewrites, so whichever lands first I will rebase the other.

[ebarkhordar]

I ran this against e0068d7 in a clean python:3.11-slim container.

_update_scale runs a few lines above both new divisions, so on a step where dynamic scaling raises the scale, cur_scale is already the NEXT scale while the gradients in hand were produced under the previous one. The reported norm then comes out half the true value. Both functions already bind prev_scale just above the overflow check, and zero/stage_1_and_2.py:2533 divides by prev_scale for exactly this reason.

Dynamic scale, scale_window=2, gradients set so the true norm is 1.0:

iter 0: scale  128.0 ->  128.0   reported_norm=1.000000
iter 1: scale  128.0 ->  128.0   reported_norm=1.000000
iter 2: scale  128.0 ->  256.0   reported_norm=0.500000   WRONG
iter 3: scale  256.0 ->  256.0   reported_norm=1.000000
iter 4: scale  256.0 ->  512.0   reported_norm=0.500000   WRONG

Identical on both paths, fused_lamb_legacy False and True. Changing both sites to divide by prev_scale gives 1.000000 on all five iterations, and the clip still receives the scaled norm.

Your new test pins static_loss_scale=128, where cur_scale never moves, so it cannot tell the two divisors apart.

[vineethsaivs]

You are right, fixed in 16b9c87: both sites now divide by prev_scale. Added a dynamic-scale case that reproduced your numbers exactly (step 2 reported 0.5 on both paths) and passes on all six steps now.

[ebarkhordar]

Confirmed at 16b9c87. I re-ran the same script in a clean container and both paths report the true norm on every step now, including the two where the scale doubles:

  iter 2: scale  128.0 ->  256.0   reported_norm=1.000000 (true 1.0)
  iter 4: scale  256.0 ->  512.0   reported_norm=1.000000 (true 1.0)

The any(before != after ...) assertion in the new test is worth keeping. It means a future change that stops the scale moving fails the test rather than passing a version of it that proves nothing. The file is 4 passed here.

One thing I hit while checking, and it is not a request to change this PR. The same pattern looks live in fused_optimizer.py, which binds prev_scale at :166 and :264 and then divides by cur_scale at :180 and :333. Same script, same container, that file unmodified at your head:

--- step_fused_adam (line 180) ---
  iter 2: scale  128.0 ->  256.0   reported_norm=0.500000 (true 1.0)   WRONG
--- step (line 333) ---
  iter 2: scale  128.0 ->  256.0   reported_norm=0.500000 (true 1.0)   WRONG

One caveat if anyone picks it up: prev_scale is bound inside if self.loss_scale_config.use_grad_scaling: there, so it is unbound under bf16 where that block is skipped, and a straight swap would need to handle that. I have not tried it against a real training run, only this harness.

[vineethsaivs]

One thing I found while checking this: the same ordering affects the gradients themselves, not only the reported norm. step() calls _update_scale before unscale_and_clip_grads, and that function divides by self.loss_scale_config.cur_scale, so on a scale-increase step the gradients are divided by twice the scale backward() used. Measured with FP16_Optimizer + SGD(lr=1) on CPU, init_scale=4, scale_window=1, clip_grad=0:

step 0: backward scale=4  scale after update=4   true grad=1  grad actually used=1
step 1: backward scale=4  scale after update=8   true grad=1  grad actually used=0.5
step 2: backward scale=8  scale after update=16  true grad=1  grad actually used=0.5

It cancels whenever the clip factor is above 1, since combined_scale = clip * cur_scale puts cur_scale on both sides, so it only bites when clipping is not binding. zero/stage_1_and_2.py has the same split: line 2536 divides the reported norm by prev_scale while unscale_and_clip_grads nine lines later uses the updated self.loss_scale. I have not reproduced the ZeRO sites, only read them, since fp16 deepspeed.initialize will not build on CPU here.

Out of scope for this PR. Happy to send it separately if you think it is worth fixing.