Raise the per-parameter norms to norm_type in clip_grad_norm_

The p-norm over every gradient is (sum_i ||g_i||_p ** p) ** (1/p), so combining
the per-parameter norms means raising each to norm_type. clip_grad_norm_ squares
them instead and then takes the 1/norm_type root, which is only the same thing
when norm_type is 2:

    grads [3, -4] and [2], norm_type=1
    torch.nn.utils.clip_grad_norm_ ->  9.0
    deepspeed                      -> 53.0     (7**2 + 2**2)

    norm_type=3
    torch.nn.utils.clip_grad_norm_ ->  4.626
    deepspeed                      ->  2.894

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

This is a regression from #4915, which vectorized the accumulation. Before it the
loop read `total_norm += param_norm.item()**norm_type`; the rewrite replaced that
with `torch.stack(all_norms).square().sum()` and kept the 1/norm_type root. The
four other norm-combining sites in this file 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 at the top of the function, 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())

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