[zero] compute the grad norm over all gradient stores at once

colossalai/booster/plugin/hybrid_parallel_plugin.py

@@ -692,6 +692,7 @@ def __init__(
         forced_dtype: Optional[torch.dtype] = None,
         overlap_allgather: bool = False,
         fp8_communication: bool = False,
+        grad_norm_process_group: Optional[ProcessGroup] = None,
     ):
         self.model = model
         self.param_info = param_info
@@ -723,6 +724,7 @@ def __init__(
             overlap_allgather=overlap_allgather,
             fp8_communication=fp8_communication,
             backward_context=model._hook_context,
+            grad_norm_process_group=grad_norm_process_group,
         )
 
     def sync_dp_grads(self):
@@ -840,90 +842,29 @@ def backward_by_grad(self, tensor, grad, inputs: Tensor = None, retain_graph: bo
             # If gradient synchronization is is not required, return.
             return
 
-    def _compute_grad_norm(self, dp_pg, gradients: List[Tensor], norm_type: int = 2) -> float:
-        r"""
-        Compute and return the gradient norm for gradient clipping.
-
-        Args:
-            gradients (List[Tensor]): A list of tensors containing gradients.
-            norm_type (int, optional): Type of the p-norm to be computed. Defaults to 2.
-
-        Returns:
-            float: The computed gradient norm.
-        """
-
-        # Check if the list of gradients is empty
-        if len(gradients) == 0:
-            return 0.0
-
-        dp_size = get_world_size(dp_pg) if dp_pg is not None else 1
-        tp_size = get_world_size(self.tp_pg) if self.tp_pg is not None else 1
-        pp_size = get_world_size(self.pp_pg) if self.pp_pg is not None else 1
-        norm_type = float(norm_type)
-
-        if norm_type == inf:
-            # The parent class calculates the norm of 'dp' gradients,
-            # so we only need to calculate the norm 'tp' of 'pp' gradients.
-            total_norm = super()._compute_grad_norm(gradients, norm_type)
-
-            total_norm_cuda = torch.tensor(
-                [float(total_norm)], device=get_accelerator().get_current_device(), dtype=torch.float32
-            )
-
-            if tp_size > 1:
-                dist.all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.MAX, group=self.tp_pg)
-            if pp_size > 1:
-                dist.all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.MAX, group=self.pp_pg)
-
-            total_norm = total_norm_cuda.item()
-        else:
-            total_norm_exponentiated = 0.0
-            for grad in gradients:
-                grad_norm_exponentiated = grad.data.double().norm(norm_type) ** norm_type
-
-                # If 'tp_size' is greater than 1 and the parameter for the gradient is not a distributed tensor,
-                # it indicates that the parameter is not distributed across devices of the 'tp_group'.
-                # Consequently, there is no need to perform an 'all_reduce' operation for 'grad_norm'.
-                # However, we still perform the 'all_reduce' operation for the sake of good coding practices.
-                # To ensure mathematical equivalence, we divide the 'grad_norm' by 'tp_size.'
-                if tp_size > 1:
-                    param_id_for_grad = self.get_param_id_for_grad(grad)
-                    param_for_grad = ctypes.cast(param_id_for_grad, ctypes.py_object).value
-
-                    if not is_distributed_tensor(param_for_grad):
-                        grad_norm_exponentiated /= tp_size
-
-                # If 'pp_size' is greater than 1 and the gradient belongs to shared parameters,
-                # it means that this parameter is used in two different pipeline stages.
-                # To avoid redundant norm calculations, we divide the exponent of this norm by
-                # the number of shared stages.
-                if pp_size > 1:
-                    for shared_param in self.shared_params:
-                        if self.stage_manager.stage in shared_param:
-                            stage_shared_param = shared_param[self.stage_manager.stage]
-                            working_grad = self.get_working_grad_by_param_id(id(stage_shared_param))
-                            if grad is working_grad:
-                                grad_norm_exponentiated /= len(shared_param)
-
-                total_norm_exponentiated += grad_norm_exponentiated
-
-            total_norm_exponentiated_cuda = torch.tensor(
-                [float(total_norm_exponentiated)], device=get_accelerator().get_current_device(), dtype=torch.float32
-            )
-            if dp_size > 1:
-                # compute norm in dp process group
-                dist.all_reduce(tensor=total_norm_exponentiated_cuda, op=dist.ReduceOp.SUM, group=dp_pg)
-            if tp_size > 1:
-                # compute norm in tp process group
-                dist.all_reduce(tensor=total_norm_exponentiated_cuda, op=dist.ReduceOp.SUM, group=self.tp_pg)
-            if pp_size > 1:
-                # compute norm in pp process group
-                dist.all_reduce(tensor=total_norm_exponentiated_cuda, op=dist.ReduceOp.SUM, group=self.pp_pg)
-
-            # Compute the 'total_norm' from 'total_norm_exponentiated'
-            total_norm = total_norm_exponentiated_cuda.item() ** (1.0 / norm_type)
-
-        return total_norm
+    def _grad_norm_process_groups(self) -> List[ProcessGroup]:
+        # gradients are also split across tp and pp ranks, so the norm is summed over them too
+        groups = super()._grad_norm_process_groups()
+        for group in (self.tp_pg, self.pp_pg):
+            if group is not None and get_world_size(group) > 1:
+                groups.append(group)
+        return groups
+
+    def _grad_norm_weight(self, grad: Tensor) -> float:
+        weight = 1.0
+        # a parameter not split by tp is replicated on every tp rank: each copy counts 1 / tp_size
+        if self.tp_pg is not None and get_world_size(self.tp_pg) > 1:
+            param_for_grad = ctypes.cast(self.get_param_id_for_grad(grad), ctypes.py_object).value
+            if not is_distributed_tensor(param_for_grad):
+                weight /= get_world_size(self.tp_pg)
+        # a parameter shared by several pipeline stages (e.g. tied embeddings) counts 1 / num_stages per stage
+        if self.pp_pg is not None and get_world_size(self.pp_pg) > 1:
+            for shared_param in self.shared_params:
+                if self.stage_manager.stage in shared_param:
+                    stage_shared_param = shared_param[self.stage_manager.stage]
+                    if grad is self.get_working_grad_by_param_id(id(stage_shared_param)):
+                        weight /= len(shared_param)
+        return weight
 
 
 class HybridParallelPlugin(PipelinePluginBase):

colossalai/booster/plugin/moe_hybrid_parallel_plugin.py

@@ -101,6 +101,8 @@ def __init__(
             forced_dtype=forced_dtype,
             pg_to_param_list=pg_param_list,
             overlap_allgather=overlap_allgather,
+            # the dp group spans the moe_dp and ep axes, so it covers every shard of dense and expert grads
+            grad_norm_process_group=dp_process_group,
         )
 
 

colossalai/zero/low_level/low_level_optim.py

@@ -99,6 +99,7 @@ def __init__(
         overlap_allgather: bool = False,
         fp8_communication: bool = False,
         backward_context=None,
+        grad_norm_process_group: Optional[Union[ProcessGroup, Tuple[ProcessGroup, ...]]] = None,
     ):
         super(LowLevelZeroOptimizer, self).__init__(optim=optimizer)
 
@@ -124,6 +125,15 @@ def __init__(
                 pg_to_param_list[unique_dp_group].extend(group["params"])
 
         self.pg_to_param_list = pg_to_param_list
+        # The grad norm is reduced once over this group (see _compute_grad_norm), so it must contain every rank that
+        # owns a distinct gradient shard of any store. With a single store this is the store's own group.
+        if grad_norm_process_group is None:
+            if len(pg_to_param_list) > 1:
+                raise ValueError(
+                    "grad_norm_process_group should be provided when parameters are split across several process groups."
+                )
+            grad_norm_process_group = next(iter(pg_to_param_list))
+        self._grad_norm_process_group = grad_norm_process_group
         param_to_pg = {}
         for grp, param_list in pg_to_param_list.items():
             for p in param_list:
@@ -550,7 +560,6 @@ def step(self, closure=None):
             working_params = self._working_param_groups[group_id]
             real_working_params[group_id] = []
             real_master_params[group_id] = []
-            working_grads = []
             for working_param, master_param in zip(working_params, master_params):
                 # if a working param requires grad and has no grad
                 # it is not 'really' working, e.g. the droped layer
@@ -568,13 +577,8 @@ def step(self, closure=None):
                     grad_partition_groups.append(grad)
                     real_master_params[group_id].append(master_param)
 
-            # compute norm
-            norm_group = 0
-            for grad_store in self.pg_to_grad_store.values():
-                working_grads = grad_store.get_working_grads_by_group_id(group_id)
-                norm_group += self._compute_grad_norm(dp_pg=grad_store.torch_pg, gradients=working_grads)
-
-            norm_groups.append(norm_group)
+            # one norm over the grads of all stores, so that experts and dense params are combined correctly
+            norm_groups.append(self._compute_grad_norm(self.get_working_grads_by_group_id(group_id)))
 
             # update the params in the optimizer
             self.optim.param_groups[group_id]["params"] = real_master_params[group_id]
@@ -634,66 +638,45 @@ def step(self, closure=None):
                 if not tensor_bucket.is_empty():
                     tensor_bucket.all_gather(pg, fp8_communication=self._fp8_communication)
 
-    def _compute_grad_norm(
-        self, dp_pg: Union[ProcessGroup, Tuple[ProcessGroup, ...]], gradients: List[Tensor], norm_type: int = 2
-    ) -> float:
+    def _compute_grad_norm(self, gradients: List[Tensor], norm_type: int = 2) -> float:
         r"""
-        Compute and return the gradient norm for gradient clipping.
+        Compute the norm of the whole gradient, viewed as a single vector, for gradient clipping.
+
+        Each rank passes the gradient shards it owns, from all gradient stores. Their p-th power sum (or max for the
+        infinity norm) is reduced once over the groups of ``_grad_norm_process_groups`` and the root is taken once,
+        so every element is counted once and all ranks get the same value. Copies of the same element held by several
+        ranks are handled by ``_grad_norm_weight``. Every rank must call this, even without gradients.
 
         Args:
-            gradients (List[Tensor]): The gradients to compute norm
+            gradients (List[Tensor]): The gradient shards owned by this rank
             norm_type (int, optional): type of the used p-norm, Can be ``'inf'`` for infinity norm. Defaults to 2.
 
         Returns:
-            float: The total norm of given gradients
+            float: The total norm of the gradients
         """
-
-        if len(gradients) == 0:
-            return 0.0
-
         norm_type = float(norm_type)
         if norm_type == inf:
-            total_norm = max(grad.data.abs().max() for grad in gradients)
-            total_norm_cuda = torch.tensor(
-                [float(total_norm)],
-                device=get_accelerator().get_current_device(),
-                dtype=torch.float,
-            )
-            if isinstance(dp_pg, tuple):
-                for grp in dp_pg:
-                    dist.all_reduce(total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=grp)
-            else:
-                dist.all_reduce(total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=dp_pg)
-            total_norm = total_norm_cuda.item()
-
+            local = max((float(grad.data.abs().max()) for grad in gradients), default=0.0)
+            op = dist.ReduceOp.MAX
         else:
-            total_norm_exponentiated = 0.0
+            local = 0.0
             for grad in gradients:
-                grad_norm_exponentiated = grad.data.double().norm(norm_type) ** norm_type
-                total_norm_exponentiated += grad_norm_exponentiated
-
-            # Sum across all model parallel GPUs.
-            total_norm_exponentiated_cuda = torch.tensor(
-                [float(total_norm_exponentiated)],
-                device=get_accelerator().get_current_device(),
-                dtype=torch.float,
-            )
-            if isinstance(dp_pg, tuple):
-                for grp in dp_pg:
-                    dist.all_reduce(
-                        total_norm_exponentiated_cuda,
-                        op=torch.distributed.ReduceOp.SUM,
-                        group=grp,
-                    )
-            else:
-                torch.distributed.all_reduce(
-                    total_norm_exponentiated_cuda,
-                    op=torch.distributed.ReduceOp.SUM,
-                    group=dp_pg,
-                )
-            total_norm = total_norm_exponentiated_cuda.item() ** (1.0 / norm_type)
-
-        return total_norm
+                local += self._grad_norm_weight(grad) * grad.data.double().norm(norm_type) ** norm_type
+            op = dist.ReduceOp.SUM
+
+        total = torch.tensor([float(local)], device=get_accelerator().get_current_device(), dtype=torch.float)
+        for group in self._grad_norm_process_groups():
+            dist.all_reduce(total, op=op, group=group)
+        return total.item() if norm_type == inf else total.item() ** (1.0 / norm_type)
+
+    def _grad_norm_process_groups(self) -> List[ProcessGroup]:
+        """The groups the local grad norm is reduced over, one after another."""
+        group = self._grad_norm_process_group
+        return list(group) if isinstance(group, tuple) else [group]
+
+    def _grad_norm_weight(self, grad: Tensor) -> float:
+        """How much this rank's copy of a gradient counts towards the norm; copies of one gradient sum to 1."""
+        return 1.0
 
     #############################
     # Mixed Precision Utilities #

tests/test_moe/test_moe_ep_zero_grad_norm.py

@@ -0,0 +1,103 @@
+from copy import deepcopy
+
+import pytest
+import torch
+import torch.distributed as dist
+from transformers.models.mixtral.configuration_mixtral import MixtralConfig
+from transformers.models.mixtral.modeling_mixtral import MixtralModel
+
+import colossalai
+from colossalai.booster.booster import Booster
+from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin
+from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.testing.random import seed_all
+
+# Relative error allowed between the ZeRO grad norm (bf16 grads) and the norm of an fp32 single-GPU reference. With
+# the norm combined correctly it is below 1.2e-3 here; the old per-store norms were off by 0.6% to 67%.
+NORM_RTOL = 5e-3
+
+HIDDEN_SIZE = 8
+NUM_EXPERTS = 8
+
+
+@parameterize("ep_size", [2, 4])
+@parameterize("moe_only", [True, False])
+def run_grad_norm(ep_size: int, moe_only: bool):
+    """The ZeRO grad norm must be the norm of the whole gradient and equal on every rank.
+
+    Dense and expert grads live in two ZeRO stores, and every EP rank holds different experts, so the norm has to
+    combine both stores and all EP ranks. Training only the MoE blocks makes experts carry most of the norm, as in a
+    real MoE model; training everything mixes dense and expert grads. Routing is random, so every expert is used on
+    every moe_dp replica and the moe_dp scaling of #6464 cannot interfere.
+    """
+    if dist.get_world_size() % ep_size != 0:
+        return
+    dtype = torch.bfloat16
+    rank = dist.get_rank()
+
+    plugin = MoeHybridParallelPlugin(
+        pp_size=1,
+        tp_size=1,
+        ep_size=ep_size,
+        zero_stage=1,
+        precision="bf16",
+        overlap_communication=False,
+    )
+    booster = Booster(plugin=plugin)
+
+    seed_all(10086)
+    config = MixtralConfig(
+        hidden_size=HIDDEN_SIZE,
+        intermediate_size=HIDDEN_SIZE * 64,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        num_local_experts=NUM_EXPERTS,
+        num_experts_per_tok=2,
+    )
+    torch_model = MixtralModel(config).to(dtype).cuda()
+    for name, param in torch_model.named_parameters():
+        param.requires_grad_(not moe_only or "block_sparse_moe" in name)
+
+    # single-GPU reference: same weights in fp32 on the concatenation of every rank's batch (= global mean loss)
+    single_model = deepcopy(torch_model).float()
+    zero_model = deepcopy(torch_model)
+    zero_optimizer = torch.optim.SGD(zero_model.parameters(), lr=0.0)
+    zero_model, zero_optimizer, _, _, _ = booster.boost(zero_model, zero_optimizer)
+
+    seed_all(1453 + rank)
+    input_data = torch.rand(8, 16, HIDDEN_SIZE).cuda().to(dtype)
+    zero_optimizer.backward(zero_model(inputs_embeds=input_data).last_hidden_state.mean())
+    zero_optimizer.step()  # lr = 0: nothing moves, but this computes the norm used for clipping
+    zero_norm = zero_optimizer.get_grad_norm()
+
+    all_inputs = [torch.empty_like(input_data) for _ in range(dist.get_world_size())]
+    dist.all_gather(all_inputs, input_data)
+    single_model(inputs_embeds=torch.cat(all_inputs).float()).last_hidden_state.mean().backward()
+    grads = [p.grad.double() for p in single_model.parameters() if p.grad is not None]
+    ref_norm = sum(g.pow(2).sum() for g in grads).sqrt().item()
+
+    norms = [None] * dist.get_world_size()
+    dist.all_gather_object(norms, zero_norm)
+    assert len(set(norms)) == 1, f"ep_size={ep_size} moe_only={moe_only}: grad norm differs across ranks: {norms}"
+    assert abs(zero_norm - ref_norm) <= NORM_RTOL * ref_norm, (
+        f"ep_size={ep_size} moe_only={moe_only}: ZeRO grad norm {zero_norm:.6g}, reference {ref_norm:.6g}, "
+        f"ratio {zero_norm / ref_norm:.4f}"
+    )
+
+
+def run_dist(rank, world_size, port):
+    colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
+    torch.cuda.set_device(rank)
+    run_grad_norm()
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize("world_size", [4])
+@rerun_if_address_is_in_use()
+def test_moe_ep_zero_grad_norm(world_size):
+    spawn(run_dist, world_size)
+
+
+if __name__ == "__main__":
+    test_moe_ep_zero_grad_norm(world_size=4)