[BUG] get_global_grad_norm() returns None with fp32 at ZeRO stage 0

Same symptom as #7255 (open, no repro there). This adds a repro, the cause and a fix, so feel free to close one as a duplicate of the other.

Describe the bug With fp32 training at ZeRO stage 0 and gradient clipping enabled, engine.get_global_grad_norm() returns None after step(). The same config at stage 1 returns the norm.

The fp32 stage-0 path clips gradients in _take_model_step through clip_fp32_gradients(), which calls deepspeed.runtime.utils.clip_grad_norm_. That function returns the total norm, but the value is dropped (engine.py L3425-L3442). Afterwards the engine only copies optimizer._global_grad_norm. Only the fp16/bf16/ZeRO wrapper optimizers set that attribute. At fp32 stage 0 the basic optimizer is used directly (a torch.optim optimizer, or DeepSpeed's FusedAdam, CPUAdam, FusedLamb, ...), so it stays None. The same path is used with torch_autocast at stage 0.

This is what #7255 describes: HF Trainer maps max_grad_norm to gradient_clipping, reads the norm from get_global_grad_norm() (via Accelerate) when max_grad_norm > 0, and drops grad_norm from its logs when the value is None. Since #8068 made gradient_clipping default to 1.0, plain DeepSpeed configs hit it too. The docstring of get_global_grad_norm() only lists None for ZeRO-1/2 with gradient-norm computation disabled, so this looks unintended.

To Reproduce

1. Save as repro.py:

import sys

import torch
import deepspeed

stage = int(sys.argv[1]) if len(sys.argv) > 1 else 0
clip = float(sys.argv[2]) if len(sys.argv) > 2 else 1.0
torch.manual_seed(0)
model = torch.nn.Linear(8, 1)
config = {
    "train_micro_batch_size_per_gpu": 2,
    "gradient_clipping": clip,
    "zero_optimization": {"stage": stage},
    "optimizer": {"type": "Adam", "params": {"lr": 1e-3, "torch_adam": True}},
}
engine, *_ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config)
engine.backward(engine(torch.randn(2, 8, device=engine.device)).sum())
engine.step()
print(f"stage={stage} gradient_clipping={clip} get_global_grad_norm() = {engine.get_global_grad_norm()}")

1. Run with DS_ACCELERATOR=cpu RANK=0 LOCAL_RANK=0 WORLD_SIZE=1 MASTER_ADDR=127.0.0.1 MASTER_PORT=29500 python repro.py <stage>
2. Output on master (bc44b341):

stage=0 gradient_clipping=1.0 get_global_grad_norm() = None
stage=1 gradient_clipping=1.0 get_global_grad_norm() = 2.7553305625915527

Expected behavior Stage 0 reports the same pre-clip norm as stage 1 when clipping is enabled. The engine already computes that norm; it just doesn't keep it.

Proposed fix Return the norm from clip_fp32_gradients() and store it in _global_grad_norm in _take_model_step. It's stored as a tensor, like ZeRO does, so there's no extra host sync per step. Add a regression test that compares it against the pre-step gradient norm computed with plain torch. The test fails on master and passes with the fix, and with it stage 0 prints 2.7553305625915527, the same as stage 1. I'll open a PR.

One open point: with gradient_clipping: 0 nothing computes a norm on this path, so it would still return None, while stages 1/2 report it even then (compute_grad_norm defaults to true, #8331). HF doesn't read the norm when max_grad_norm is 0, so this isn't the #7255 case. Computing it anyway would add a norm pass and an all-reduce every step to a path that does neither today. I'd document None for that case in the docstring and leave parity for a follow-up, unless you'd rather have it in the same PR.

The stored value is the norm clipping already uses, so it inherits clip_grad_norm_'s behavior. The MoE expert-norm issue in #8475 is being addressed by #8478, which doesn't touch this code path.

ds_report output

DeepSpeed C++/CUDA extension op report
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
deepspeed_not_implemented  [NO] ....... [OKAY]
async_io ............... [NO] ....... [NO]
deepspeed_ccl_comm ..... [NO] ....... [OKAY]
deepspeed_shm_comm ..... [NO] ....... [NO]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
pin_memory ............. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['<venv>/site-packages/torch']
torch version .................... 2.8.0
deepspeed install path ........... ['<src>/DeepSpeed/deepspeed']
deepspeed info ................... 0.19.8+bc44b341, bc44b341, master
deepspeed wheel compiled w. ...... torch 0.0
shared memory (/dev/shm) size .... UNKNOWN

System info

• OS: macOS 26.4.1, Apple M5, CPU accelerator (DS_ACCELERATOR=cpu), single process, no GPU
• Python 3.13.9, torch 2.8.0, DeepSpeed 0.19.8+bc44b341 (master). The same code path exists in 0.16.5, the version in #7255.

Launcher context Plain python with RANK/WORLD_SIZE/MASTER_ADDR/MASTER_PORT set; no deepspeed launcher.

Docker context No.