[moe] remove DPGradScaler that over-scales expert grads under moe_dp

colossalai/moe/_operation.py

@@ -332,50 +332,6 @@ def backward(ctx: Any, *grad_outputs: Tensor) -> Tuple[Tensor, None]:
         return grad, None
 
 
-class DPGradScalerIn(torch.autograd.Function):
-    """
-    Scale the gradient back by the number of experts
-    because the batch size increases in the moe stage
-    """
-
-    @staticmethod
-    def forward(ctx: Any, inputs: Tensor, moe_dp_size: int, activated_experts: int) -> Tensor:
-        assert activated_experts != 0, f"shouldn't be called when no expert is activated"
-        ctx.moe_dp_size = moe_dp_size
-        ctx.activated_experts = activated_experts
-        return inputs
-
-    @staticmethod
-    def backward(ctx: Any, *grad_outputs: Tensor) -> Tuple[Tensor, None, None]:
-        assert len(grad_outputs) == 1
-        grad = grad_outputs[0]
-        if ctx.moe_dp_size != ctx.activated_experts:
-            grad.mul_(ctx.activated_experts / ctx.moe_dp_size)
-        return grad, None, None
-
-
-class DPGradScalerOut(torch.autograd.Function):
-    """
-    Scale the gradient by the number of experts
-    because the batch size increases in the moe stage
-    """
-
-    @staticmethod
-    def forward(ctx: Any, inputs: Tensor, moe_dp_size: int, activated_experts: int) -> Tensor:
-        assert activated_experts != 0, f"shouldn't be called when no expert is activated"
-        ctx.moe_dp_size = moe_dp_size
-        ctx.activated_experts = activated_experts
-        return inputs
-
-    @staticmethod
-    def backward(ctx: Any, *grad_outputs: Tensor) -> Tuple[Tensor, None, None]:
-        assert len(grad_outputs) == 1
-        grad = grad_outputs[0]
-        if ctx.moe_dp_size != ctx.activated_experts:
-            grad.mul_(ctx.moe_dp_size / ctx.activated_experts)
-        return grad, None, None
-
-
 def _all_to_all(
     inputs: torch.Tensor,
     input_split_sizes: Optional[List[int]] = None,

colossalai/shardformer/modeling/deepseek.py

@@ -16,15 +16,8 @@
 from transformers.utils import is_flash_attn_2_available, logging
 
 from colossalai.lazy import LazyInitContext
-from colossalai.moe._operation import (
-    DPGradScalerIn,
-    DPGradScalerOut,
-    EPGradScalerIn,
-    EPGradScalerOut,
-    all_to_all_uneven,
-)
+from colossalai.moe._operation import EPGradScalerIn, EPGradScalerOut, all_to_all_uneven
 from colossalai.pipeline.stage_manager import PipelineStageManager
-from colossalai.quantization.fp8 import all_reduce_fp8
 from colossalai.shardformer.layer._operation import (
     all_to_all_comm,
     gather_forward_split_backward,
@@ -166,17 +159,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             group=self.ep_group,
         )
 
-        with torch.no_grad():
-            activate_experts = output_split_sizes[: self.num_experts_per_ep].clone()
-            for i in range(1, self.ep_size):
-                activate_experts += output_split_sizes[i * self.num_experts_per_ep : (i + 1) * self.num_experts_per_ep]
-            activate_experts = (activate_experts > 0).float()
-
-        if self.fp8_communication:
-            all_reduce_fp8(activate_experts, group=self.moe_dp_group)
-        else:
-            dist.all_reduce(activate_experts, group=self.moe_dp_group)
-
         input_split_list = input_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist()
         output_split_list = output_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist()
         output_states, _ = all_to_all_uneven(
@@ -191,23 +173,15 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         if output_states.size(0) > 0:
             if self.num_experts_per_ep == 1:
                 expert = self.experts[self.expert_start_idx]
-                output_states = DPGradScalerIn.apply(output_states, self.moe_dp_size, activate_experts[0])
                 output_states = expert(output_states)
-                output_states = DPGradScalerOut.apply(output_states, self.moe_dp_size, activate_experts[0])
             else:
                 output_states_splits = output_states.split(output_split_sizes.tolist())
                 output_states_list = []
                 for i, split_states in enumerate(output_states_splits):
                     if split_states.size(0) == 0:  # no token routed to this experts
                         continue
                     expert = self.experts[self.expert_start_idx + i % self.num_experts_per_ep]
-                    split_states = DPGradScalerIn.apply(
-                        split_states, self.moe_dp_size, activate_experts[i % self.num_experts_per_ep]
-                    )
                     split_states = expert(split_states)
-                    split_states = DPGradScalerOut.apply(
-                        split_states, self.moe_dp_size, activate_experts[i % self.num_experts_per_ep]
-                    )
                     output_states_list.append(split_states)
                 output_states = torch.cat(output_states_list)
         output_states = EPGradScalerOut.apply(output_states, self.ep_size)

colossalai/shardformer/modeling/deepseek_v3.py

@@ -10,13 +10,7 @@
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 
 from colossalai.lazy import LazyInitContext
-from colossalai.moe._operation import (
-    DPGradScalerIn,
-    DPGradScalerOut,
-    EPGradScalerIn,
-    EPGradScalerOut,
-    all_to_all_uneven,
-)
+from colossalai.moe._operation import EPGradScalerIn, EPGradScalerOut, all_to_all_uneven
 from colossalai.pipeline.stage_manager import PipelineStageManager
 from colossalai.shardformer.layer.linear import ParallelModule
 from colossalai.shardformer.shard.utils import set_tensors_to_none
@@ -108,11 +102,6 @@ def moe_forward(self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torc
             sorted_tokens = gathered_tokens[gatherd_idxs]
             tokens_per_expert = tokens_per_expert_post_gather
 
-            # moe-dp related code
-            activate_experts = tokens_per_expert_post_gather > 0
-            activate_experts = activate_experts.int()
-            dist.all_reduce(activate_experts, group=self.moe_dp_group)
-
             # ep related code
             sorted_tokens = EPGradScalerIn.apply(sorted_tokens, self.ep_size)
 
@@ -126,11 +115,7 @@ def moe_forward(self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torc
                 continue
             expert = self.experts[i + self.ep_rank * self.experts_per_rank]
             tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
-            # moe-dp related code
-            tokens_for_this_expert = DPGradScalerIn.apply(tokens_for_this_expert, self.moe_dp_size, activate_experts[i])
             expert_out = expert(tokens_for_this_expert)
-            # moe-dp related code
-            expert_out = DPGradScalerOut.apply(expert_out, self.moe_dp_size, activate_experts[i])
             outputs.append(expert_out)
             start_idx = end_idx
 

colossalai/shardformer/modeling/mixtral.py

@@ -22,15 +22,8 @@
 from transformers.utils import is_flash_attn_2_available, logging
 
 from colossalai.lazy import LazyInitContext
-from colossalai.moe._operation import (
-    DPGradScalerIn,
-    DPGradScalerOut,
-    EPGradScalerIn,
-    EPGradScalerOut,
-    all_to_all_uneven,
-)
+from colossalai.moe._operation import EPGradScalerIn, EPGradScalerOut, all_to_all_uneven
 from colossalai.pipeline.stage_manager import PipelineStageManager
-from colossalai.quantization.fp8 import all_reduce_fp8
 from colossalai.shardformer.layer._operation import (
     all_to_all_comm,
     gather_forward_split_backward,
@@ -132,17 +125,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
         dist.all_to_all_single(output_split_sizes, input_split_sizes, group=self.ep_group)
 
-        with torch.no_grad():
-            activate_experts = output_split_sizes[: self.num_experts_per_ep].clone()
-            for i in range(1, self.ep_size):
-                activate_experts += output_split_sizes[i * self.num_experts_per_ep : (i + 1) * self.num_experts_per_ep]
-            activate_experts = (activate_experts > 0).float()
-
-        if self.fp8_communication:
-            all_reduce_fp8(activate_experts, group=self.moe_dp_group)
-        else:
-            dist.all_reduce(activate_experts, group=self.moe_dp_group)
-
         input_split_list = input_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist()
         output_split_list = output_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist()
 
@@ -159,25 +141,17 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             if self.num_experts_per_ep == 1:
                 # no need to split
                 expert = self.experts[self.expert_start_idx]
-                output_states = DPGradScalerIn.apply(output_states, self.moe_dp_size, activate_experts[0])
                 output_states = expert.act_fn(expert.w1(output_states)) * expert.w3(output_states)
                 output_states = expert.w2(output_states)
-                output_states = DPGradScalerOut.apply(output_states, self.moe_dp_size, activate_experts[0])
             else:
                 output_states_splits = output_states.split(output_split_sizes.tolist())
                 output_states_list = []
                 for i, split_states in enumerate(output_states_splits):
                     if split_states.size(0) == 0:
                         continue
                     expert = self.experts[self.expert_start_idx + i % self.num_experts_per_ep]
-                    split_states = DPGradScalerIn.apply(
-                        split_states, self.moe_dp_size, activate_experts[i % self.num_experts_per_ep]
-                    )
                     split_states = expert.act_fn(expert.w1(split_states)) * expert.w3(split_states)
                     split_states = expert.w2(split_states)
-                    split_states = DPGradScalerOut.apply(
-                        split_states, self.moe_dp_size, activate_experts[i % self.num_experts_per_ep]
-                    )
                     output_states_list.append(split_states)
                 output_states = torch.cat(output_states_list)
 

tests/test_moe/test_moe_ep_zero_grad_scale.py

@@ -0,0 +1,127 @@
+from copy import deepcopy
+
+import pytest
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+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 rerun_if_address_is_in_use, spawn
+from colossalai.testing.random import seed_all
+from tests.test_moe.moe_utils import assert_loose_close
+
+# Relative L2 error allowed against the reference grads. bf16 noise stays below 1e-2 here, while the old moe_dp
+# scaling was off by 0.33x-3x on the experts. The atol of assert_loose_close (4e-3) exceeds these tiny grads.
+GRAD_RTOL = 2e-2
+
+EP_SIZE = 2
+HIDDEN_SIZE = 8
+
+
+def run_forced_routing():
+    """Expert grads must match the references when an expert is only used by some moe_dp replicas.
+
+    The gates send tokens with hidden[0] > 0 to expert 0 and the others to expert 1. moe_dp replica 0 feeds
+    hidden[0] = +5 and the other replicas feed -5, so each of the two experts gets tokens from only part of the
+    moe_dp group. With random inputs every expert is hit on every replica and this case is never exercised.
+
+    Two references: DDP over the whole dp group, and a single-GPU fp32 copy run on the concatenation of every
+    rank's batch. The latter has no parallel logic at all and its loss is the global mean loss by definition, so it
+    also checks that DDP (which averages in the zero grads of replicas that skip an expert) is the right target.
+    """
+    dtype = torch.bfloat16
+    rank = dist.get_rank()
+    torch.cuda.set_device(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 * 2,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        num_local_experts=4,
+        num_experts_per_tok=1,
+    )
+    torch_model = MixtralModel(config).to(dtype).cuda()
+    with torch.no_grad():
+        for layer in torch_model.layers:
+            gate = layer.block_sparse_moe.gate.weight
+            gate.zero_()
+            gate[0, 0], gate[1, 0] = 10.0, -10.0
+
+    single_model = deepcopy(torch_model).float()
+    zero_model = deepcopy(torch_model)
+    zero_optimizer = torch.optim.SGD(zero_model.parameters(), lr=1)
+    zero_model, zero_optimizer, _, _, _ = booster.boost(zero_model, zero_optimizer)
+    ddp_model = DDP(torch_model, process_group=plugin.dp_group, find_unused_parameters=True)
+
+    seed_all(1453 + rank)
+    input_data = torch.rand(4, 7, HIDDEN_SIZE).cuda().to(dtype)
+    input_data[..., 0] = 5.0 if rank // EP_SIZE == 0 else -5.0
+
+    zero_output = zero_model(inputs_embeds=input_data).last_hidden_state.mean()
+    zero_optimizer.backward(zero_output)
+    ddp_output = ddp_model(inputs_embeds=input_data).last_hidden_state.mean()
+    ddp_output.backward()
+    assert_loose_close(zero_output, ddp_output, dtype=dtype)
+
+    # every rank has the same batch size, so the mean over the concatenation equals the mean of per-rank losses
+    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()
+
+    # get_param_grad is collective: finish the loop on every rank before asserting, otherwise a mismatch on one rank
+    # leaves the other ranks waiting in all-gather until the NCCL timeout.
+    errors = []
+    ddp_params = dict(ddp_model.module.named_parameters())
+    single_params = dict(single_model.named_parameters())
+    for n, p in zero_model.named_parameters():
+        n = n.removeprefix("module.")
+        grads = {
+            "zero": zero_optimizer.get_param_grad(p),
+            "ddp": ddp_params[n].grad,
+            "single": single_params[n].grad,
+        }
+        grads = {k: torch.zeros_like(p, dtype=torch.float) if g is None else g.float() for k, g in grads.items()}
+        for name, ref_name in (("zero", "ddp"), ("zero", "single"), ("ddp", "single")):
+            grad, ref_grad = grads[name], grads[ref_name]
+            err, ref = (grad - ref_grad).norm().item(), ref_grad.norm().item()
+            if err > GRAD_RTOL * ref + 1e-6:
+                errors.append(
+                    f"{n}: |{name} - {ref_name}| / |{ref_name}| = {err / max(ref, 1e-12):.3f}, "
+                    f"|{name}| / |{ref_name}| = {grad.norm().item() / max(ref, 1e-12):.3f}"
+                )
+    assert not errors, f"rank {rank}: " + "; ".join(errors)
+
+
+def run_dist(rank, world_size, port):
+    colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
+    run_forced_routing()
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize("world_size", [4, 8])
+@rerun_if_address_is_in_use()
+def test_moe_ep_zero_grad_scale(world_size):
+    if world_size > torch.cuda.device_count():
+        pytest.skip(f"needs {world_size} GPUs")
+    spawn(run_dist, world_size)
+
+
+if __name__ == "__main__":
+    test_moe_ep_zero_grad_scale(world_size=4)