Keep empty gradient-norm ranks in collective reductions

deepspeed/runtime/utils.py

@@ -913,6 +913,27 @@ def combine_grad_norm_groups(norm_groups, norm_type=2):
     return torch.where(invalid, torch.full_like(combined, float("inf")), combined)
 
 
+def _get_local_norm_contribution(input_tensors, norm_type, device):
+    """Compute the FP32 input to the global MAX or SUM reduction."""
+    all_norms = []
+    for tensor in input_tensors:
+        if norm_type == inf:
+            if tensor.numel() == 0:
+                continue
+            tensor_norm = tensor.detach().abs().max().float()
+        else:
+            tensor_norm = tensor.detach().float().norm(norm_type)
+        all_norms.append(tensor_norm.to(device))
+
+    if not all_norms and norm_type > 0:
+        return torch.zeros((), dtype=torch.float32, device=device)
+
+    local_norms = torch.stack(all_norms)
+    if norm_type == inf:
+        return local_norms.max()
+    return local_norms.pow(norm_type).sum()
+
+
 def get_global_norm_of_tensors(input_tensors, norm_type=2, mpu=None, moe_ep_group=None):
     """Get norm of an iterable of tensors.
 
@@ -928,50 +949,28 @@ def get_global_norm_of_tensors(input_tensors, norm_type=2, mpu=None, moe_ep_grou
         Total norm of the tensors (viewed as a single vector), as a detached float32 scalar tensor.
     """
     assert isinstance(input_tensors, Iterable), f'expected Iterable type not {type(input_tensors)}'
+    input_tensors = list(input_tensors)
     assert all([torch.is_tensor(t) for t in input_tensors]), 'expected list of only tensors'
 
     norm_type = float(norm_type)
-    all_norms = []
-    if norm_type == inf:
-        for t in input_tensors:
-            all_norms.append(t.data.abs().max().float())
-        total_norm = torch.stack(all_norms).max()
-        device_total_norm = total_norm.to(get_accelerator().current_device_name())
-        # Max across model parallel
-        if mpu is not None:
-            # For MoE grads, max over model parallel only if MoE-TP is enabled
-            if moe_ep_group is None or groups._get_expert_model_parallel_world_size() > 1:
-                dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group())
-            # If MoE grads and MoE-TP disabled, max over pipeline parallel
-            elif bwc_pipeline_parallel_world_size(mpu) > 1:
-                dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=bwc_pipeline_parallel_group(mpu))
-
-        # MoE grads: max across expert parallel group
-        if pg_utils.is_parallel(moe_ep_group):
-            dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=moe_ep_group)
-        total_norm = device_total_norm.to(input_tensors[0].device)
-    else:
-
-        device = get_accelerator().current_device_name()
-        for tensor in input_tensors:
-            tensor_norm = tensor.detach().float().norm(norm_type)
-            all_norms.append(tensor_norm.to(device))
-        norm_powers = torch.stack(all_norms).pow(norm_type)
-        device_total_norm = norm_powers.sum()
+    device = get_accelerator().current_device_name()
+    return_device = input_tensors[0].device if input_tensors else device
+    device_total_norm = _get_local_norm_contribution(input_tensors, norm_type, device)
+    reduce_op = dist.ReduceOp.MAX if norm_type == inf else dist.ReduceOp.SUM
 
-        # Sum across model parallel
-        if mpu is not None:
-            # For MoE grads, sum over model parallel only if MoE-TP is enabled
-            if moe_ep_group is None or groups._get_expert_model_parallel_world_size() > 1:
-                dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group())
-            # If MoE grads and MoE-TP disabled, sum over pipeline parallel
-            elif bwc_pipeline_parallel_world_size(mpu) > 1:
-                dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=bwc_pipeline_parallel_group(mpu))
-
-        # MoE grads: sum across expert parallel group
-        if pg_utils.is_parallel(moe_ep_group):
-            dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=moe_ep_group)
-        total_norm = device_total_norm.to(input_tensors[0].device).pow(1. / norm_type)
+    if mpu is not None:
+        # For MoE grads, reduce over model parallel only if MoE-TP is enabled.
+        if moe_ep_group is None or groups._get_expert_model_parallel_world_size() > 1:
+            dist.all_reduce(device_total_norm, op=reduce_op, group=mpu.get_model_parallel_group())
+        # If MoE grads and MoE-TP disabled, reduce over pipeline parallel.
+        elif bwc_pipeline_parallel_world_size(mpu) > 1:
+            dist.all_reduce(device_total_norm, op=reduce_op, group=bwc_pipeline_parallel_group(mpu))
+
+    if pg_utils.is_parallel(moe_ep_group):
+        dist.all_reduce(device_total_norm, op=reduce_op, group=moe_ep_group)
+    total_norm = device_total_norm.to(return_device)
+    if norm_type != inf:
+        total_norm = total_norm.pow(1. / norm_type)
 
     mask_nan_or_inf_with_val_inplace(total_norm, device=total_norm.device)
 

tests/unit/runtime/test_runtime_utils.py

@@ -4,6 +4,7 @@
 # DeepSpeed Team
 
 import torch
+import deepspeed
 from torch._utils import _flatten_dense_tensors
 import deepspeed.comm as dist
 import pytest
@@ -12,6 +13,7 @@
 import deepspeed.runtime.utils as ds_utils
 import deepspeed.utils.groups as groups
 from deepspeed.accelerator import get_accelerator
+from deepspeed.runtime.bf16_optimizer import BF16_Optimizer
 
 from unit.common import DistributedTest
 
@@ -128,6 +130,132 @@ def test(self, norm_type, kind):
         assert norm.item() == pytest.approx(expected)
 
 
+class _NormModelParallelUnit:
+
+    def __init__(self, model_parallel_group, data_parallel_group=None):
+        self.model_parallel_group = model_parallel_group
+        self.data_parallel_group = data_parallel_group
+
+    def get_model_parallel_group(self):
+        return self.model_parallel_group
+
+    def get_model_parallel_rank(self):
+        return dist.get_rank(self.model_parallel_group)
+
+    def get_model_parallel_world_size(self):
+        return dist.get_world_size(self.model_parallel_group)
+
+    def get_data_parallel_group(self):
+        return self.data_parallel_group
+
+    def get_data_parallel_rank(self):
+        return dist.get_rank(self.data_parallel_group)
+
+    def get_data_parallel_world_size(self):
+        return dist.get_world_size(self.data_parallel_group)
+
+
+class TestGlobalNormEmptyInput(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.parametrize("parallelism", ["model", "expert", "none"])
+    @pytest.mark.parametrize("input_case",
+                             ["nonempty", "one_empty", "all_empty", "one_zero_tensor", "all_zero_tensors", "mixed"])
+    def test(self, parallelism, input_case):
+        rank = dist.get_rank()
+        parallel_group = dist.new_group([0, 1])
+        kwargs = {}
+        if parallelism == "model":
+            kwargs["mpu"] = _NormModelParallelUnit(parallel_group)
+        elif parallelism == "expert":
+            kwargs["moe_ep_group"] = parallel_group
+
+        local_values = [3.0, -4.0] if rank == 0 else [12.0]
+        if input_case in ("all_empty", "all_zero_tensors") or (input_case in ("one_empty", "one_zero_tensor")
+                                                               and rank == 1):
+            local_values = []
+        tensors = [torch.tensor(local_values, device=get_accelerator().current_device_name())] if local_values else []
+        if input_case in ("one_zero_tensor", "all_zero_tensors", "mixed"):
+            tensors.insert(0, torch.empty(0, device=get_accelerator().current_device_name()))
+        expected_norms = [0.0, 0.0, 0.0, 0.0]
+        if local_values:
+            expected_norms = [7.0, 5.0, 91.0**(1.0 / 3.0), 4.0] if rank == 0 else [12.0] * 4
+        if parallelism != "none" and input_case not in ("all_empty", "all_zero_tensors"):
+            expected_norms = ([19.0, 13.0, 1819.0**(1.0 / 3.0), 12.0]
+                              if input_case in ("nonempty", "mixed") else [7.0, 5.0, 91.0**(1.0 / 3.0), 4.0])
+
+        for norm_type, expected in zip([1, 2, 3, float("inf")], expected_norms):
+            norm = ds_utils.get_global_norm_of_tensors(tensors, norm_type=norm_type, **kwargs)
+            assert norm.item() == pytest.approx(expected)
+            assert norm.dtype == torch.float32
+            assert norm.ndim == 0
+            assert norm.device == torch.device(get_accelerator().current_device_name())
+            assert not norm.requires_grad
+
+
+class TestGlobalNormIterable(DistributedTest):
+    world_size = 1
+
+    @pytest.mark.parametrize("input_case", ["nonempty", "empty", "zero_tensor", "mixed"])
+    def test(self, input_case):
+        for device in ["cpu", get_accelerator().current_device_name()]:
+            tensors = [torch.tensor([3.0, -4.0], device=device, requires_grad=True)
+                       ] if input_case in ("nonempty", "mixed") else []
+            if input_case in ("zero_tensor", "mixed"):
+                tensors.insert(0, torch.empty(0, device=device))
+            for norm_type, nonempty_expected in [(1, 7.0), (2, 5.0), (3, 91.0**(1.0 / 3.0)), (float("inf"), 4.0)]:
+                norm = ds_utils.get_global_norm_of_tensors((tensor for tensor in tensors), norm_type=norm_type)
+                assert norm.item() == pytest.approx(nonempty_expected if input_case in ("nonempty", "mixed") else 0.0)
+                expected_device = tensors[0].device if tensors else torch.device(
+                    get_accelerator().current_device_name())
+                assert norm.device == expected_device
+                assert norm.dtype == torch.float32
+                assert norm.ndim == 0
+                assert not norm.requires_grad
+
+
+class TestBF16EmptyNormRank(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        if torch.bfloat16 not in get_accelerator().supported_dtypes():
+            pytest.skip("requires bfloat16 support")
+        model_parallel_group = dist.new_group([0, 1])
+        data_parallel_group = [dist.new_group([rank]) for rank in range(self.world_size)][dist.get_rank()]
+        mpu = _NormModelParallelUnit(model_parallel_group, data_parallel_group)
+        model = torch.nn.Linear(2, 1, bias=False)
+        torch.nn.init.ones_(model.weight)
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+        config = {
+            "train_micro_batch_size_per_gpu": 1,
+            "bf16": {
+                "enabled": True
+            },
+            "data_types": {
+                "grad_accum_dtype": "fp32"
+            },
+            "zero_optimization": {
+                "stage": 1
+            },
+            "zero_allow_untested_optimizer": True,
+            "gradient_clipping": 1.0,
+        }
+        engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config, mpu=mpu)
+        try:
+            assert isinstance(engine.optimizer, BF16_Optimizer)
+            inputs = torch.tensor([[3.0, 4.0]], device=engine.device, dtype=torch.bfloat16)
+            engine.backward(engine(inputs).sum())
+            # Replicated parameters contribute to the norm only on tensor-parallel rank zero.
+            norm_grads, _ = engine.optimizer.get_grads_for_norm()
+            assert len(norm_grads) == (1 if dist.get_rank() == 0 else 0)
+            engine.step()
+            assert engine.get_global_grad_norm().item() == pytest.approx(5.0)
+            expected = torch.tensor([[0.94, 0.92]], device=engine.device, dtype=torch.bfloat16)
+            torch.testing.assert_close(engine.module.weight, expected, rtol=0, atol=0)
+        finally:
+            engine.destroy()
+
+
 @pytest.mark.parametrize("check_using_norm", [(False), (True)])
 class TestCheckOverflow(DistributedTest):
     world_size = 2