diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py old mode 100755 new mode 100644 index e22a49c1d5d6..82bfb445110e --- a/deepspeed/runtime/utils.py +++ b/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()) diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index bac260c297c4..deac254dd91a 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -2283,7 +2283,7 @@ 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: @@ -2291,7 +2291,9 @@ def get_grad_norm_direct(self, gradients, params, norm_type=2): 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) diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py old mode 100755 new mode 100644 index ca2493792ca5..582aebdfac9b --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/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. diff --git a/tests/unit/runtime/test_runtime_utils.py b/tests/unit/runtime/test_runtime_utils.py index 8aebbe500eb7..814fce830f72 100644 --- a/tests/unit/runtime/test_runtime_utils.py +++ b/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 diff --git a/tests/unit/runtime/zero/test_zero_grad_clip.py b/tests/unit/runtime/zero/test_zero_grad_clip.py index 9ce50c2ca9ce..253393cab215 100644 --- a/tests/unit/runtime/zero/test_zero_grad_clip.py +++ b/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"),