Count each expert once in the fp32 gradient-clipping norm

deepspeed/runtime/utils.py

@@ -376,20 +376,81 @@ def clip_grad_norm_(parameters, max_norm, norm_type=2, mpu=None):
     """
     if isinstance(parameters, torch.Tensor):
         parameters = [parameters]
-    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    parameters = list(parameters)
     norm_type = float(norm_type)
+
+    # Expert parameters are not replicated across the data parallel group, so the
+    # rank-averaging below cannot reconstruct a global norm from them: each rank's
+    # norm covers a different set of experts. Combine those the way the bf16 and
+    # fp16 optimizers already do, over the expert parallel group.
+    #
+    # Read ownership from the parameters, before dropping the ones with no gradient.
+    # The two branches below issue different collectives, so this decision has to
+    # come out the same on every rank -- and a rank whose experts received no tokens
+    # this step has expert parameters but no expert gradients.
+    # The registry of expert groups is built identically on every rank, so an empty one
+    # means no MoE is configured anywhere and the choice is already made: take the
+    # original path, with no collective and no host sync. This is a per-step call from
+    # `_take_model_step`, so the common case must not pay for the MoE one.
+    registry = sorted(getattr(groups, "_EXPERT_PARALLEL_GROUP", None) or {})
+    expert_group_names = []
+    if registry:
+        # A rank can legitimately own no expert at all -- under pipeline parallelism, a
+        # stage holding only dense layers -- so reading ownership locally cannot guarantee
+        # the same choice everywhere. One scalar settles it, and it has to be reduced over
+        # the world rather than the data parallel group: under pipeline parallelism that
+        # group is one stage wide (`PipelineParallelGrid.get_data_parallel_group`), while
+        # the expert norm below reduces over the pipeline group, which spans stages. A
+        # per-stage vote let an expert stage and a dense stage issue different collectives
+        # on that group. The second slot carries "an expert here has no group name", which
+        # has no expert parallel group to reduce over, so such a call stays on the
+        # original path rather than guessing which group the parameter belongs to.
+        vote = torch.tensor([
+            float(any(is_moe_param(p) and getattr(p, "group_name", None) is not None for p in parameters)),
+            float(any(is_moe_param(p) and getattr(p, "group_name", None) is None for p in parameters)),
+        ],
+                            device=get_accelerator().current_device_name(),
+                            dtype=torch.float)
+        dist.all_reduce(vote, group=groups._clone_world_group())
+        if vote[0].item() > 0 and vote[1].item() == 0:
+            expert_group_names = registry
+
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+
+    expert_tensors: Dict[str, List[torch.Tensor]] = {}
+    if expert_group_names:
+        # `get_norm_with_moe_layers` issues one collective per entry, so every rank has to
+        # walk the groups in the same order. Start from the registry's order rather than
+        # the order this rank's parameters happen to come in.
+        expert_tensors = {group_name: [] for group_name in expert_group_names}
+        for p in parameters:
+            if is_moe_param(p):
+                expert_tensors.setdefault(p.group_name, []).append(p.grad.data)
+        # Every rank has to reduce over every group, whether or not its own experts
+        # were used this step; a zero contributes nothing to the sum of p-th powers
+        # and nothing to a max over absolute values.
+        zero = torch.zeros(1, device=get_accelerator().current_device_name())
+        for tensors in expert_tensors.values():
+            if not tensors:
+                tensors.append(zero)
+
+    norm_parameters = [p for p in parameters if not is_moe_param(p)] if expert_tensors else parameters
+
     all_norms = []
     if norm_type == inf:
-        for p in parameters:
+        for p in norm_parameters:
             all_norms.append(p.grad.data.abs().max().float())
-        total_norm = torch.stack(all_norms).max()
+        # Everything may be an expert, leaving nothing replicated to take a max over.
+        # Zero is the identity here: these are absolute values.
+        total_norm = torch.stack(all_norms).max() if all_norms \
+            else torch.zeros((), device=get_accelerator().current_device_name())
         total_norm = total_norm.to(get_accelerator().current_device_name())
         # Take max across all GPUs.
         if mpu is not None:
             dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group())
     else:
         total_norm = 0
-        for p in parameters:
+        for p in norm_parameters:
             if mpu is not None:
                 if (mpu.get_model_parallel_rank() == 0) or is_model_parallel_parameter(p):
                     param_norm = p.grad.data.detach().float().norm(norm_type)
@@ -410,14 +471,26 @@ def clip_grad_norm_(parameters, max_norm, norm_type=2, mpu=None):
             dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group())
         total_norm = total_norm.pow(1. / norm_type)
 
-    # Need to average total_norm across different GPUs due to the presence of moe params
-    pg = groups._get_data_parallel_group()
-    scaled_norm = total_norm * 1.0 / float(dist.get_world_size(group=pg))
-    scaled_norm_tensor = scaled_norm
-
-    dist.all_reduce(scaled_norm_tensor, group=pg)
-    total_norm = scaled_norm_tensor
-    total_norm = total_norm.to(parameters[0].device)
+    if expert_tensors:
+        # `total_norm` now covers only the replicated parameters, so it is already the
+        # same on every rank. Fold the experts in over their own group, counting each
+        # expert once, rather than averaging norms that describe different parameters.
+        moe_norm = get_norm_with_moe_layers(total_norm, mpu=mpu, expert_tensors=expert_tensors, norm_type=norm_type)
+        # That helper reports a non-finite norm as -1. Left as-is it would make
+        # `clip_coef` negative and flip the sign of every gradient; inf keeps the
+        # behaviour the averaging path already had, which is to scale them to zero.
+        if moe_norm == -1:
+            moe_norm = float('inf')
+        total_norm = torch.tensor([float(moe_norm)], device=total_norm.device, dtype=torch.float)
+    else:
+        # Need to average total_norm across different GPUs due to the presence of moe params
+        pg = groups._get_data_parallel_group()
+        scaled_norm = total_norm * 1.0 / float(dist.get_world_size(group=pg))
+        scaled_norm_tensor = scaled_norm
+
+        dist.all_reduce(scaled_norm_tensor, group=pg)
+        total_norm = scaled_norm_tensor
+        total_norm = total_norm.to(parameters[0].device)
 
     max_norm = torch.tensor([float(max_norm)], device=total_norm.device)
     clip_coef = max_norm / (total_norm + 1e-6)

tests/unit/runtime/test_runtime_utils.py

@@ -9,9 +9,12 @@
 import pytest
 from typing import Dict
 
+import deepspeed
 import deepspeed.runtime.utils as ds_utils
 import deepspeed.utils.groups as groups
 from deepspeed.accelerator import get_accelerator
+from deepspeed.moe.layer import MoE
+from deepspeed.moe.utils import is_moe_param, split_params_into_different_moe_groups_for_optimizer
 
 from unit.common import DistributedTest
 
@@ -51,6 +54,129 @@ def test_gather(self):
 
         assert gathered_norm[0] == gathered_norm[1], "norm at rank 0 does not match the norm at rank 1"
 
+    def test_sharded_experts_count_once(self):
+        """The global norm must not depend on how experts are spread over ranks.
+
+        Rank-averaging the per-rank norms only reconstructs a global norm when every
+        rank holds the same parameters. Under expert parallelism they do not, so the
+        four experts below gave sqrt(20) and sqrt(100) averaged to 7.236 instead of
+        the sqrt(120) = 10.954 that counting each expert once produces (#8469).
+        """
+        groups._create_expert_and_data_parallel(2)
+        rank = dist.get_rank()
+        device = get_accelerator().device_name(rank)
+
+        experts = [torch.full((4, ), float(v)) for v in (1.0, 2.0, 3.0, 4.0)]
+        expected = torch.cat(experts).norm(2).item()
+
+        owned = experts[:2] if rank == 0 else experts[2:]
+        params = []
+        for grad in owned:
+            param = torch.nn.Parameter(torch.zeros(4, device=device))
+            param.grad = grad.clone().to(device)
+            param.allreduce = False
+            param.group_name = "ep_size_2"
+            params.append(param)
+
+        # A max_norm far above the norm leaves the gradients alone, so this reads the
+        # computed norm rather than the clipped result.
+        norm = ds_utils.clip_grad_norm_(params, max_norm=1e9)
+
+        assert abs(float(norm) -
+                   expected) < 1e-4, (f"global norm {float(norm)} should be {expected} regardless of expert placement")
+
+    def _expert(self, value, device, with_grad=True):
+        param = torch.nn.Parameter(torch.zeros(4, device=device))
+        if with_grad:
+            param.grad = torch.full((4, ), value, device=device)
+        param.allreduce = False
+        param.group_name = "ep_size_2"
+        return param
+
+    def _plain(self, value, device):
+        param = torch.nn.Parameter(torch.zeros(4, device=device))
+        param.grad = torch.full((4, ), value, device=device)
+        return param
+
+    def test_experts_idle_on_one_rank_do_not_deadlock(self):
+        """A rank whose experts got no tokens must take the same collectives (#8469).
+
+        Ownership has to be read from the parameters, not the gradients: `p.grad is
+        not None` filtering drops an unused expert, and branching on what is left made
+        one rank reduce over the expert group while its peer reduced over the data
+        parallel group. They then waited on each other.
+        """
+        groups._create_expert_and_data_parallel(2)
+        rank = dist.get_rank()
+        device = get_accelerator().device_name(rank)
+
+        expert = self._expert(2.0, device, with_grad=(rank == 0))
+        norm = ds_utils.clip_grad_norm_([self._plain(1.0, device), expert], max_norm=1e9)
+
+        # non-expert 2.0 replicated, one expert of norm 4.0 owned by rank 0:
+        # sqrt(2**2 + 4**2).
+        assert abs(float(norm) - 20**0.5) < 1e-4
+
+    def test_a_rank_owning_no_expert_does_not_deadlock(self):
+        """Under pipeline parallelism a rank can hold no expert at all.
+
+        Local ownership then differs in kind rather than in degree, so the branch is
+        settled by one all-reduced flag instead of by what this rank happens to hold.
+        """
+        groups._create_expert_and_data_parallel(2)
+        rank = dist.get_rank()
+        device = get_accelerator().device_name(rank)
+
+        params = [self._plain(1.0, device)]
+        if rank == 0:
+            params.append(self._expert(2.0, device))
+
+        norm = ds_utils.clip_grad_norm_(params, max_norm=1e9)
+        assert abs(float(norm) - 20**0.5) < 1e-4
+
+    def test_expert_only_inf_norm(self):
+        """Nothing replicated leaves no tensor to take a max over; zero is the identity."""
+        groups._create_expert_and_data_parallel(2)
+        rank = dist.get_rank()
+        device = get_accelerator().device_name(rank)
+
+        norm = ds_utils.clip_grad_norm_([self._expert(float(rank + 1), device)], max_norm=1e9, norm_type=float('inf'))
+
+        assert abs(float(norm) - 2.0) < 1e-4
+
+    def test_a_rank_with_no_gradient_at_all_still_gets_the_norm(self):
+        """Every local gradient can be absent, expert ones included, and the rank still has to answer."""
+        groups._create_expert_and_data_parallel(2)
+        rank = dist.get_rank()
+        device = get_accelerator().device_name(rank)
+
+        norm = ds_utils.clip_grad_norm_([self._expert(2.0, device, with_grad=(rank == 1))], max_norm=1e9)
+
+        # Only rank 1's expert has a gradient, of norm 4.0, and both ranks report it.
+        assert abs(float(norm) - 4.0) < 1e-4
+
+    def test_expert_groups_are_visited_in_registry_order(self, monkeypatch):
+        """Ranks owning different expert groups must still walk the groups in one order."""
+        groups._create_expert_and_data_parallel(2)
+        groups._create_expert_and_data_parallel(1)
+        rank = dist.get_rank()
+        device = get_accelerator().device_name(rank)
+
+        visited = []
+        real = ds_utils.get_norm_with_moe_layers
+
+        def spy(non_expert_norm, mpu, expert_tensors, norm_type=2):
+            visited.append(list(expert_tensors))
+            return real(non_expert_norm, mpu=mpu, expert_tensors=expert_tensors, norm_type=norm_type)
+
+        monkeypatch.setattr(ds_utils, "get_norm_with_moe_layers", spy)
+
+        expert = self._expert(2.0, device)
+        expert.group_name = "ep_size_2" if rank == 0 else "ep_size_1"
+        ds_utils.clip_grad_norm_([expert], max_norm=1e9)
+
+        assert visited == [["ep_size_1", "ep_size_2"]]
+
     def test_clipped_val(self):
         max_norm = 0.1
 
@@ -73,6 +199,84 @@ def test_params():
         assert torch.equal(params_expected[1].grad, params_actual[1].grad)
 
 
+class _OneMoELayer(torch.nn.Module):
+    """A dense layer, then one expert per rank behind a gate."""
+
+    def __init__(self, hidden_dim):
+        super().__init__()
+        self.dense = torch.nn.Linear(hidden_dim, hidden_dim)
+        expert = torch.nn.Linear(hidden_dim, hidden_dim)
+        # k=2 sends every token to both experts, so each rank's expert gets a gradient.
+        self.moe = MoE(hidden_size=hidden_dim, expert=expert, num_experts=2, ep_size=2, k=2, capacity_factor=2.0)
+
+    def forward(self, x):
+        output, _, _ = self.moe(self.dense(x))
+        return output.pow(2).mean()
+
+
+class TestClipFp32GradientsThroughEngine(DistributedTest):
+    """The clip `DeepSpeedEngine.step()` applies in an fp32 run without ZeRO, with experts on two ranks.
+
+    `clip_grad_norm_` returns its norm and the engine drops it, so the norm is read from its effect. With
+    plain SGD the update is `lr * clip_coef * grad`, and the coefficient has to come from the norm of every
+    gradient with each expert counted once, not from the ranks' own norms averaged.
+    """
+
+    world_size = 2
+
+    def test_the_update_is_clipped_by_the_global_norm(self):
+        hidden_dim, max_norm, lr = 8, 1e-3, 1.0
+        torch.manual_seed(0)
+        model = _OneMoELayer(hidden_dim)
+        param_group = {'params': list(model.parameters()), 'name': 'all'}
+        optimizer = torch.optim.SGD(split_params_into_different_moe_groups_for_optimizer(param_group), lr=lr)
+        config = {
+            "train_micro_batch_size_per_gpu": 2,
+            "gradient_clipping": max_norm,
+            "zero_optimization": {
+                "stage": 0
+            }
+        }
+        engine, _, _, _ = deepspeed.initialize(config=config, model=model, optimizer=optimizer)
+
+        # A different loss scale on each rank gives the two experts gradients of different size.
+        rank = dist.get_rank()
+        generator = torch.Generator().manual_seed(100 + rank)
+        x = torch.randn(2, 4, hidden_dim, generator=generator).to(engine.device)
+        engine.backward(engine(x) * (1 + 3 * rank))
+
+        params = list(engine.module.parameters())
+        dense = [p for p in params if not is_moe_param(p)]
+        experts = [p for p in params if is_moe_param(p)]
+        assert dense and experts
+        dense_sq = torch.stack([p.grad.float().pow(2).sum() for p in dense]).sum()
+        expert_sq = torch.stack([p.grad.float().pow(2).sum() for p in experts]).sum()
+
+        # The oracle counts a dense gradient once because it is the same on both ranks, and an expert once
+        # because each rank owns a different one.
+        dense_on_every_rank = dense_sq.clone()
+        dist.all_reduce(dense_on_every_rank, op=dist.ReduceOp.MIN)
+        assert torch.equal(dense_on_every_rank, dense_sq), "dense gradients should match across ranks"
+        all_experts_sq = expert_sq.clone()
+        dist.all_reduce(all_experts_sq)
+        true_norm = (dense_sq + all_experts_sq).sqrt().item()
+
+        # Averaging the ranks' own norms is the answer this test exists to rule out, so it has to differ.
+        own_norm = (dense_sq + expert_sq).sqrt()
+        dist.all_reduce(own_norm)
+        averaged_norm = own_norm.item() / dist.get_world_size()
+        assert true_norm > 5 * max_norm, "the clip must be active for this to say anything"
+        assert abs(averaged_norm - true_norm) > 0.05 * true_norm, "the two norms must be distinguishable"
+
+        clip_coef = max_norm / (true_norm + 1e-6)
+        before = [p.detach().clone() for p in params]
+        grads = [p.grad.detach().clone() for p in params]
+        engine.step()
+
+        for p, start, grad in zip(params, before, grads):
+            torch.testing.assert_close(p.detach(), start - lr * clip_coef * grad, rtol=1e-4, atol=1e-6)
+
+
 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.