Raise the per-tensor norms to norm_type when combining them

The p-norm over a set of tensors is (sum_i ||g_i||_p ** p) ** (1/p), so combining
per-tensor norms means raising each to norm_type. Three sites take the
1/norm_type root but hardcode the exponent at 2, so they are only correct for
norm_type == 2:

  runtime/utils.py            clip_grad_norm_          .square().sum()
  zero/stage_1_and_2.py       get_grad_norm_direct     .square().sum()
  zero/stage3.py              get_grad_norm_direct     .norm(2), then pow(..., 2)

stage3 is doubly wrong: it takes an L2 norm per tensor whatever norm_type says,
then sums the squares, then takes the 1/norm_type root.

Measured against the p-norm of the concatenated gradients, for 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

The returned norm is wrong, and so is the clip coefficient derived from it, so
the gradients get scaled by the wrong factor.

clip_grad_norm_ is a regression from #4915, which vectorized the accumulation:
the loop there read `total_norm += param_norm.item()**norm_type` before it, and
the rewrite replaced that with `.square()` while keeping the 1/norm_type root.
The four norm-combining sites that were not touched by that commit all still
raise to norm_type: get_flattened_grad_norm, get_weight_norm,
get_global_norm_of_tensors and get_norm_with_moe_layers.

norm_type is float()'d in every one of these functions, so the default path
becomes pow(2.0), which is bit-identical to square() on float32.

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

Author: vineethsaivs

deepspeed/runtime/utils.py

@@ -401,7 +401,10 @@ def clip_grad_norm_(parameters, max_norm, norm_type=2, mpu=None):
                 param_norm = p.grad.data.detach().float().norm(norm_type)
                 all_norms.append(param_norm)
         if len(all_norms) > 0:
-            total_norm = torch.stack(all_norms).square().sum().float()
+            # The p-norm over every gradient is (sum_i ||g_i||_p ** p) ** (1/p), and the
+            # 1/norm_type root is taken below, so each per-parameter norm has to be raised
+            # to norm_type here. Squaring only matches that for norm_type == 2.
+            total_norm = torch.stack(all_norms).pow(norm_type).sum().float()
         else:
             total_norm = get_accelerator().FloatTensor([0.0])
         total_norm = total_norm.to(get_accelerator().current_device_name())

deepspeed/runtime/zero/stage3.py

@@ -2283,15 +2283,17 @@ def get_grad_norm_direct(self, gradients, params, norm_type=2):
             for g, p in zip(gradients, params):
                 if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0):
                     grad_norms.append(
-                        g.to(get_accelerator().device_name(), non_blocking=True).to(get_norm_dtype()).norm(2))
+                        g.to(get_accelerator().device_name(), non_blocking=True).to(get_norm_dtype()).norm(norm_type))
 
             # Sum across all model parallel GPUs.
             if len(grad_norms) == 0:
                 # FIX https://github.com/deepspeedai/DeepSpeed/issues/3564
                 total_norm_cuda = torch.tensor(0, dtype=gradients[0].dtype).to(get_accelerator().device_name()).to(
                     get_norm_dtype())
             else:
-                total_norm_cuda = torch.sum(torch.pow(torch.stack(grad_norms), 2))
+                # Each entry is ||g||_norm_type and the 1/norm_type root is taken below, so
+                # both the per-tensor norm above and this exponent follow norm_type.
+                total_norm_cuda = torch.sum(torch.pow(torch.stack(grad_norms), norm_type))
 
             dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=process_group)
 

deepspeed/runtime/zero/stage_1_and_2.py

@@ -2090,7 +2090,9 @@ def get_grad_norm_direct(self, gradients, params, norm_type=2):
                         torch.linalg.vector_norm(g.data.to(get_norm_dtype()).detach(),
                                                  ord=norm_type).to(get_accelerator().current_device_name()))
             if len(all_norms) > 0:
-                total_norm = torch.stack(all_norms).square().sum().float()
+                # vector_norm above already gives each ||g||_norm_type, and the 1/norm_type
+                # root is taken below, so the exponent here has to be norm_type too.
+                total_norm = torch.stack(all_norms).pow(norm_type).sum().float()
             else:
                 total_norm = torch.tensor(0.0, dtype=torch.float32).to(self.device)
             # Sum across all model parallel Device.

tests/unit/runtime/test_runtime_utils.py

@@ -73,6 +73,34 @@ def test_params():
         assert torch.equal(params_expected[1].grad, params_actual[1].grad)
 
 
+class TestClipGradNormPNorm(DistributedTest):
+    # world_size 1 so this runs wherever the suite runs; the bug is in the per-rank
+    # recombination of the norms, which is independent of the group size.
+    world_size = 1
+
+    @pytest.mark.parametrize("norm_type", [1, 2, 3])
+    def test_matches_torch(self, norm_type):
+        # The p-norm over all gradients is (sum_i ||g_i||_p ** p) ** (1/p). Squaring the
+        # per-parameter norms computes that only for p == 2, which is the control here.
+        def test_params():
+            param1 = torch.nn.Parameter(torch.zeros(2))
+            param1.grad = torch.Tensor([3.0, -4.0])
+            param2 = torch.nn.Parameter(torch.zeros(1))
+            param2.grad = torch.Tensor([2.0])
+            return [param1, param2]
+
+        max_norm = 1.0
+        params_expected = test_params()
+        expected_norm = torch.nn.utils.clip_grad_norm_(params_expected, max_norm, norm_type=norm_type)
+
+        params_actual = test_params()
+        actual_norm = ds_utils.clip_grad_norm_(params_actual, max_norm=max_norm, norm_type=norm_type)
+
+        assert torch.allclose(actual_norm.float().cpu(), expected_norm.float().cpu())
+        for expected, actual in zip(params_expected, params_actual):
+            assert torch.allclose(actual.grad, expected.grad)
+
+
 @pytest.mark.parametrize("check_using_norm", [(False), (True)])
 class TestCheckOverflow(DistributedTest):
     world_size = 2

tests/unit/runtime/zero/test_zero_grad_clip.py

@@ -10,6 +10,7 @@
 from deepspeed.utils import safe_get_local_grad, safe_set_local_grad
 from deepspeed.accelerator import get_accelerator
 from unit.simple_model import SimpleModel
+from unit.common import DistributedTest
 import os
 
 
@@ -48,6 +49,38 @@ def get_config(precision, clip_value, offload_device="cpu"):
     return config
 
 
+@pytest.mark.parametrize("zero_stage", [1, 2, 3])
+@pytest.mark.parametrize("norm_type", [1, 2, 3])
+class TestZeroGradNormPNorm(DistributedTest):
+    world_size = 1
+
+    def test_matches_flat_norm(self, zero_stage, norm_type):
+        # get_grad_norm_direct returns the norm of the gradients viewed as a single vector,
+        # so on one rank with no model parallelism it must equal the p-norm of the
+        # concatenation. norm_type 2 is the control: it is right on both sides.
+        config = {
+            "train_batch_size": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-4
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage
+            },
+        }
+        model = SimpleModel(hidden_dim=4, nlayers=2)
+        engine, optimizer, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config)
+
+        gradients = [torch.Tensor([3.0, -4.0]), torch.Tensor([2.0])]
+        params = list(model.parameters())[:len(gradients)]
+        expected = torch.cat([g.reshape(-1) for g in gradients]).norm(float(norm_type))
+
+        actual = optimizer.get_grad_norm_direct(gradients, params, norm_type=norm_type)
+        assert torch.allclose(torch.as_tensor(actual).float().cpu(), expected.float().cpu())
+
+
 @pytest.mark.parametrize("precision,clip_value,offload_device", [
     ("fp16", 0.5, "cpu"),
     ("bf16", 0.05, "cpu"),