Raise the per-tensor norms to norm_type when combining them

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"),