Raise the per-tensor norms to norm_type when combining them

[vineethsaivs]

Problem

The p-norm over a set of tensors is

||g||_p = ( sum_i ||g_i||_p ** p ) ** (1/p)

so combining per-tensor norms means raising each one to norm_type. Three functions take the 1/norm_type root but hardcode the exponent at 2, so they are correct only when norm_type == 2:

file	function	how it combines
runtime/utils.py	clip_grad_norm_	torch.stack(all_norms).square().sum()
runtime/zero/stage_1_and_2.py	get_grad_norm_direct	torch.stack(all_norms).square().sum()
runtime/zero/stage3.py	get_grad_norm_direct	.norm(2) per tensor, then torch.pow(..., 2)

stage3 is wrong twice over: it takes an L2 norm per tensor whatever norm_type says, sums the squares, and then takes the 1/norm_type root, so it never computes a p-norm at all.

norm_type is a documented argument on all three ("type of the used p-norm"). Measured against the p-norm of the concatenated gradients, on grads [3, -4] and [2]:

	p=1	p=2	p=3
ground truth	9.000	5.385	4.626
clip_grad_norm_	53.000	5.385	2.894
ZeRO 1/2	53.000	5.385	2.894
ZeRO 3	29.000	5.385	3.072

At p=1 clip_grad_norm_ reports 7**2 + 2**2 = 53 where the answer is 7 + 2 = 9; at p=3 all three under-clip. The returned norm is wrong and so are the resulting gradients, not just the reported number.

Scope, stated up front rather than buried. No DeepSpeed config reaches this. norm_type is not a config key (absent from runtime/config.py, constants.py and config_utils.py), and every in-tree call site passes the default 2: engine.py L3326 and L3339 for clip_grad_norm_, stage_1_and_2.py L2245 and stage3.py L2395 for get_grad_norm_direct. Both fp16 optimizers pin self.norm_type = 2, and BF16_Optimizer's norm_type default is not overridden; where a threaded self.norm_type does reach a norm helper it lands on get_global_norm_of_tensors and get_norm_with_moe_layers, which are already correct.

So p=2 is the only column a DeepSpeed run reaches today, and there pow(2.0) and .square() agree exactly. What this fixes is the documented norm_type argument for direct callers of runtime/utils.clip_grad_norm_ and the two get_grad_norm_direct methods, including code outside this repo, and it removes a defect that would become live the moment anything threads a non-default norm_type through.

Where it came from

clip_grad_norm_ is a regression from #4915, which vectorized the accumulation. Before that commit the loop read

param_norm = p.grad.data.float().norm(norm_type)
total_norm += param_norm.item()**norm_type

and the rewrite replaced the **norm_type with .square() while keeping .pow(1. / norm_type) below. The norm-combining sites that commit did not touch all still raise to norm_type, which is what makes the intended rule unambiguous rather than a matter of taste: get_flattened_grad_norm (L483), get_weight_norm (L579), get_global_norm_of_tensors (L927) and get_norm_with_moe_layers (L1144). scaled_global_norm in stage_1_and_2.py takes the other honest route and asserts norm_type == 2 outright.

Fix

Follow norm_type at each of the three sites. In stage3 that means the per-tensor .norm() as well as the exponent.

norm_type is float()-ed in all three functions, so the default becomes pow(2.0), which is bit-identical to square() on float32; checked over 20000 random vectors spanning a wide dynamic range, zero differing elements. Combined with the call-site audit above, the default path is provably unchanged.

Deliberately not touched: complete_grad_norm_calculation_for_cpu_offload in both ZeRO files takes the same 1/norm_type root over squares, but it accepts no norm_type argument and sets norm_type = 2.0 as a local, so it is self-consistent L2 and correct as written. A repo-wide sweep finds nine sites taking a 1/norm_type root: four already correct, those two self-consistent, and the three fixed here.

Test

• TestClipGradNormPNorm::test_matches_torch in tests/unit/runtime/test_runtime_utils.py, parametrized over norm_type in 1, 2, 3 against torch.nn.utils.clip_grad_norm_, asserting both the returned norm and the resulting gradients.
• TestZeroGradNormPNorm::test_matches_flat_norm in tests/unit/runtime/zero/test_zero_grad_clip.py, parametrized over zero_stage in 1, 2, 3 and norm_type in 1, 2, 3, asserting get_grad_norm_direct equals the p-norm of the concatenated gradients on a single rank.

norm_type=2 is the control throughout: it passes on both sides.

# new tests against the unpatched sources
8 failed, 8 passed, 6 skipped     (2 from clip_grad_norm_, 6 from ZeRO: p in {1,3} x stage in {1,2,3})

# after the fix
16 passed, 6 skipped

# clean tree, same two files, without the new tests
4 passed, 6 skipped

4 + 12 new = 16, so the net difference is exactly the twelve new cases and nothing else moved. The 6 skips are pre-existing and identical throughout.

Both new classes use world_size = 1 so they run wherever the suite runs; the defect is in the per-rank recombination and is independent of group size.

yapf --style .style.yapf -d and flake8 --config .flake8 are clean on all five files, with nothing reported on the clean-tree copies either.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f05a068e1a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

This reviewed commit is non-merge (single parent 84fd92a) and its commit message has no Signed-off-by trailer. DeepSpeed's commit requirements apply to every non-merge commit, so this needs to be recreated with --signoff or an equivalent trailer before it can satisfy the repo policy.

AGENTS.md reference: AGENTS.md:L6-L8

Useful? React with 👍 / 👎.

[vineethsaivs]

Widened this after finding the same defect in two more places, so it is now one rule fixed at all three sites rather than one of three.

get_grad_norm_direct in zero/stage_1_and_2.py combines with .square().sum() and then takes the 1/norm_type root, exactly like clip_grad_norm_ did. get_grad_norm_direct in zero/stage3.py is wrong twice over: it takes .norm(2) per tensor whatever norm_type says, sums the squares, and then takes the 1/norm_type root, so it never computes a p-norm at all.

Measured against the p-norm of the concatenated gradients, on grads [3, -4] and [2]:

              p=1     p=2      p=3
truth         9.000   5.385    4.626
clip_grad_   53.000   5.385    2.894
zero 1/2     53.000   5.385    2.894
zero 3       29.000   5.385    3.072

All three are reachable only through the norm_type argument; every in-tree caller uses the default, so the default path is untouched either way. Tests now cover all three: norm_type in 1, 2, 3 for clip_grad_norm_, and zero_stage in 1, 2, 3 crossed with norm_type in 1, 2, 3 for the optimizers, with norm_type=2 the control throughout. 8 failed / 8 passed before, 16 passed after, against 4 passed on a clean tree without the new tests.

Happy to split the ZeRO half back out into its own PR if you would rather review them separately.

[ebarkhordar]

Reachability check on norm_type at f10cb80, parsed rather than grepped since the severity claim rests on it: nothing inside DeepSpeed calls the three functions you are fixing with anything but the default 2.

clip_grad_norm_       engine.py:3326         (parameters, max_norm, mpu)
clip_grad_norm_       engine.py:3339         (parameters, max_norm, mpu)
get_grad_norm_direct  stage_1_and_2.py:2247  (2 positional args)
get_grad_norm_direct  stage3.py:2397         (2 positional args)

It is not a config key either: absent from runtime/config.py, constants.py and config_utils.py. Both fp16 optimizers pin self.norm_type = 2, and BF16_Optimizer's norm_type=2 default is not passed at engine.py:2347. The threaded self.norm_type reaches only get_global_norm_of_tensors and get_norm_with_moe_layers, already in your correct list.

So p=2 is the only column of your table a run reaches today, and there pow(2.0) and .square() agree. The change still looks right to me, and stage3 taking a per-tensor L2 is a real defect for direct callers of runtime/utils.clip_grad_norm_, which include code outside this repo. But "the clip coefficient is roughly 6x too small and the gradients are crushed" reads as a live training bug when no config reaches it. Scoping that to direct callers would make the PR harder to argue with.

Adjacent, not a request: complete_grad_norm_calculation_for_cpu_offload (stage_1_and_2.py:1608, stage3.py:1857) takes the same root over squares but sets norm_type = 2.0 as a local, so it is self consistent L2. If the goal is that ZeRO honors norm_type, that is where it stops.

I have no GPU here and did not reproduce your numbers.

[vineethsaivs]

Your reachability audit matches mine and the criticism of the wording is fair, so I have rewritten that part of the description.

"the clip coefficient is roughly 6x too small and the gradients are crushed" was accurate about the arithmetic and misleading about the blast radius. The body now says up front that no DeepSpeed config reaches this: norm_type is not a config key, every in-tree call site passes the default 2, both fp16 optimizers pin self.norm_type = 2, and the threaded self.norm_type only reaches get_global_norm_of_tensors and get_norm_with_moe_layers, which are already correct. What is left is the documented argument for direct callers of runtime/utils.clip_grad_norm_ and the two get_grad_norm_direct methods, including code outside this repo, plus removing a defect that goes live the moment anything threads a non-default norm_type. That is the claim I should have made.

On complete_grad_norm_calculation_for_cpu_offload: agreed, and I reached the same conclusion independently before your comment, which is why those two are untouched. They take no norm_type argument and set norm_type = 2.0 as a local, so the root over squares is self-consistent L2. I have added that to the description so the omission is explicit rather than looking like something I missed. A repo-wide sweep for 1/norm_type roots finds nine sites: four already correct, those two self-consistent, and the three fixed here.

You are also right that this is where "ZeRO honors norm_type" stops. Making complete_grad_norm_calculation_for_cpu_offload configurable would mean giving it the argument and threading it from the optimizer, which is a feature rather than a fix, so I have left it alone rather than widen this further.

No GPU needed for the numbers, in case it is useful: deepspeed.initialize with zero_optimization.stage 1, 2 and 3 runs on one CPU rank under the gloo backend, and get_grad_norm_direct is called directly on it. That is what the new TestZeroGradNormPNorm does, at world_size = 1.

[vineethsaivs]

Same false positive as before, now re-fired against f05a068e. The trailer is present on every commit in this PR, including the current head:

$ gh api repos/deepspeedai/DeepSpeed/pulls/8313/commits --jq '.[] | "\(.sha[0:8]) \(.commit.message | split("\n") | map(select(startswith("Signed-off-by")))[0])"'
f10cb801  Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>

The DCO check on this PR passes, which measures the same requirement independently. No change needed.

[ebarkhordar]

The rewritten scope section says what I measured, and calling out the two complete_grad_norm_calculation_for_cpu_offload sites as self-consistent rather than leaving them silent is the right call. Nothing further from me.