From ec5f1ed4ef734211b79a9f67a944b9d0eb54f7ce Mon Sep 17 00:00:00 2001 From: JazyJiang Date: Thu, 15 Oct 2026 20:12:24 +0800 Subject: [PATCH] [moe] sync and shard expert grads over moe_dp * sp with all_to_all sequence parallelism With all_to_all sequence parallelism, the ranks of an sp group hold the same experts but run the MoE blocks on different parts of the sequence. Dense grads are synced over (moe_dp, ep, sp), but expert grads were put in a ZeRO store on the moe_dp group only, so each sp rank kept the grad of its own tokens and the expert copies drifted apart (0.8-1.0 relative error against a single-GPU reference with ep=2, sp=2). Add MoeHybridParallelPlugin.mixed_moe_dp_group, a group along the (moe_dp, sp) axes when all_to_all SP is on and moe_dp_group otherwise, and use it for the expert ZeRO store and MoECheckpointIO. The output gather already scales grads by sp_size, so ZeRO's division by the group size gives the correct average. Expert optimizer states are now sharded over sp too. When loading a sharded optimizer, expert states were padded by the global dp size but split by the moe_dp size, while ZeRO pads each param by the size of its own group; pad by the moe_dp size so they agree. Add test_moe_ep_sp_grad_sync.py: ep=2, sp=2, ZeRO-1/2, every grad compared with an fp32 single-GPU reference and across the sp group, plus a sharded model/optimizer checkpoint round trip with sizes that exercise the padding. Co-Authored-By: Claude Opus 5.5 --- .../plugin/moe_hybrid_parallel_plugin.py | 7 +- colossalai/checkpoint_io/moe_checkpoint.py | 5 +- tests/test_moe/test_moe_ep_sp_grad_sync.py | 163 ++++++++++++++++++ 3 files changed, 170 insertions(+), 5 deletions(-) create mode 100644 tests/test_moe/test_moe_ep_sp_grad_sync.py diff --git a/colossalai/booster/plugin/moe_hybrid_parallel_plugin.py b/colossalai/booster/plugin/moe_hybrid_parallel_plugin.py index a733fc5f5261..d360fdc6249a 100644 --- a/colossalai/booster/plugin/moe_hybrid_parallel_plugin.py +++ b/colossalai/booster/plugin/moe_hybrid_parallel_plugin.py @@ -355,8 +355,11 @@ def __init__( if self.enable_sequence_parallelism and self.sequence_parallelism_mode == "all_to_all": self.mixed_dp_group = self.pg_mesh.create_group_along_axis([self.moe_dp_axis, self.ep_axis, self.sp_axis]) self.dp_size = dist.get_world_size(self.mixed_dp_group) + # sp ranks hold the same experts but see different tokens, so expert grads are synced over moe_dp * sp + self.mixed_moe_dp_group = self.pg_mesh.create_group_along_axis([self.moe_dp_axis, self.sp_axis]) else: self.mixed_dp_group = self.dp_group + self.mixed_moe_dp_group = self.moe_dp_group self.use_fp8 = use_fp8 @@ -416,7 +419,7 @@ def get_checkpoint_io(self) -> MoECheckpointIO: self.tp_group, self.sp_group, self.ep_group, - self.moe_dp_group, + self.mixed_moe_dp_group, self.zero_stage, ) @@ -506,7 +509,7 @@ def configure( dp_process_group=self.mixed_dp_group, tp_process_group=self.tp_group, pp_process_group=self.pp_group, - moe_dp_group=self.moe_dp_group, + moe_dp_group=self.mixed_moe_dp_group, verbose=True, clip_grad_norm=self.max_norm, **self.zero_config, diff --git a/colossalai/checkpoint_io/moe_checkpoint.py b/colossalai/checkpoint_io/moe_checkpoint.py index 85e36f7c6336..2989146bcedb 100644 --- a/colossalai/checkpoint_io/moe_checkpoint.py +++ b/colossalai/checkpoint_io/moe_checkpoint.py @@ -664,9 +664,8 @@ def shard_from_complete_optimizer_state( v = v.split(slice_size, dim=0)[self.global_dp_rank] elif self.use_zero and is_moe_param and self.moe_dp_size > 1: - # LowLevelZeRO pads by global dp size for now. - # TODO: update both to use moe dp size - padding_size = (self.global_dp_size - v.numel() % self.global_dp_size) % self.global_dp_size + # LowLevelZeRO pads each param by the size of the group it is sharded over + padding_size = (self.moe_dp_size - v.numel() % self.moe_dp_size) % self.moe_dp_size with torch.no_grad(): v = v.flatten() if padding_size > 0: diff --git a/tests/test_moe/test_moe_ep_sp_grad_sync.py b/tests/test_moe/test_moe_ep_sp_grad_sync.py new file mode 100644 index 000000000000..d6bc7f6549f4 --- /dev/null +++ b/tests/test_moe/test_moe_ep_sp_grad_sync.py @@ -0,0 +1,163 @@ +import os +import tempfile +from contextlib import nullcontext +from copy import deepcopy + +import pytest +import torch +import torch.distributed as dist +from torch.optim import Adam +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 +from tests.test_moe.moe_utils import check_model_equal + +# Relative L2 error allowed between ZeRO grads (bf16) and an fp32 single-GPU reference. bf16 noise stays below 2e-2 +# here, while without the sp reduction the expert grads are off by 0.8-1.0. +GRAD_RTOL = 3e-2 + +EP_SIZE, SP_SIZE = 2, 2 +# Sizes chosen so that expert params (5 x 6 = 30 elements) are not divisible by the dp size (4), which exercises the +# ZeRO padding when optimizer states are re-sharded on loading. head_dim is set separately because RoPE needs it even. +HIDDEN_SIZE, NUM_HEADS, HEAD_DIM, INTERMEDIATE_SIZE = 6, 2, 4, 5 +BATCH_SIZE, SEQ_LEN = 2, 8 + + +def build_plugin(zero_stage: int) -> MoeHybridParallelPlugin: + return MoeHybridParallelPlugin( + pp_size=1, + tp_size=1, + ep_size=EP_SIZE, + sp_size=SP_SIZE, + enable_sequence_parallelism=True, + sequence_parallelism_mode="all_to_all", + zero_stage=zero_stage, + precision="bf16", + overlap_communication=False, + ) + + +def build_model() -> MixtralModel: + seed_all(10086) + config = MixtralConfig( + hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, + num_hidden_layers=2, + num_attention_heads=NUM_HEADS, + num_key_value_heads=NUM_HEADS, + head_dim=HEAD_DIM, + num_local_experts=4, + num_experts_per_tok=2, + attn_implementation="sdpa", + ) + return MixtralModel(config).to(torch.bfloat16).cuda() + + +def replica_batch(plugin: MoeHybridParallelPlugin, index: int = None) -> torch.Tensor: + """The batch of a data replica. The ranks of one sp group share it; sp splits its sequence.""" + if index is None: + index = dist.get_rank(plugin.dp_group) + generator = torch.Generator(device="cuda").manual_seed(1453 + index) + return torch.rand(BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE, generator=generator, device="cuda") + + +@parameterize("zero_stage", [1, 2]) +def run_grad_sync(zero_stage: int): + """Expert grads must be reduced over the sp group too, like the dense ones. + + Each sp rank runs the MoE layers on its own part of the sequence. The reference is a single-GPU fp32 copy of the + model run on every data replica's batch, whose loss is the global mean loss. + """ + plugin = build_plugin(zero_stage) + booster = Booster(plugin=plugin) + torch_model = build_model() + 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) + + zero_output = zero_model(inputs_embeds=replica_batch(plugin).bfloat16()).last_hidden_state.mean() + zero_optimizer.backward(zero_output) + num_replicas = dist.get_world_size(plugin.dp_group) + single_loss = sum( + single_model(inputs_embeds=replica_batch(plugin, i)).last_hidden_state.mean() for i in range(num_replicas) + ) + (single_loss / num_replicas).backward() + + # get_param_grad is collective: finish the loop on every rank before asserting + errors = [] + single_params = dict(single_model.named_parameters()) + for name, param in zero_model.named_parameters(): + name = name.removeprefix("module.") + grad = zero_optimizer.get_param_grad(param) + grad = torch.zeros_like(param, dtype=torch.float) if grad is None else grad.float() + ref = single_params[name].grad + ref = torch.zeros_like(grad) if ref is None else ref.float() + err, ref_norm = (grad - ref).norm().item(), ref.norm().item() + if err > GRAD_RTOL * ref_norm + 1e-6: + errors.append(f"{name}: |zero - single| / |single| = {err / max(ref_norm, 1e-12):.3f}") + peers = [torch.empty_like(grad) for _ in range(SP_SIZE)] + dist.all_gather(peers, grad, group=plugin.sp_group) + if not all(torch.equal(peer, peers[0]) for peer in peers[1:]): + errors.append(f"{name}: grads differ across the sp group") + assert not errors, f"rank {dist.get_rank()} zero_stage={zero_stage}: " + "; ".join(errors) + + +def run_checkpoint(): + """Sharded model and optimizer checkpoints must round-trip when expert states are sharded over moe_dp * sp.""" + plugin = build_plugin(zero_stage=1) + booster = Booster(plugin=plugin) + model = build_model() + optimizer = Adam(model.parameters(), lr=1e-3) + model, optimizer, _, _, _ = booster.boost(model, optimizer) + optimizer.backward(model(inputs_embeds=replica_batch(plugin).bfloat16()).last_hidden_state.mean()) + optimizer.step() + + context = tempfile.TemporaryDirectory() if dist.get_rank() == 0 else nullcontext() + with context as tmpdir: + tmpdir = [tmpdir] + dist.broadcast_object_list(tmpdir, src=0) + model_dir, optim_dir = os.path.join(tmpdir[0], "model"), os.path.join(tmpdir[0], "optim") + booster.save_model(model, model_dir, shard=True) + booster.save_optimizer(optimizer, optim_dir, shard=True) + dist.barrier() + + new_model = build_model() + new_optimizer = Adam(new_model.parameters(), lr=1e-3) + new_model, new_optimizer, _, _, _ = booster.boost(new_model, new_optimizer) + booster.load_model(new_model, model_dir) + check_model_equal(model, new_model, dtype=torch.bfloat16) + + snapshot = {id(p): {k: v.clone() for k, v in s.items()} for p, s in optimizer.unwrap().state.items()} + for state in optimizer.unwrap().state.values(): + for value in state.values(): + if isinstance(value, torch.Tensor): + value.zero_() + booster.load_optimizer(optimizer, optim_dir) + for param, state in optimizer.unwrap().state.items(): + for key, value in state.items(): + assert torch.equal(value, snapshot[id(param)][key]), f"rank {dist.get_rank()}: optimizer state differs" + dist.barrier() + + +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_sync() + run_checkpoint() + + +@pytest.mark.dist +@pytest.mark.parametrize("world_size", [4]) +@rerun_if_address_is_in_use() +def test_moe_ep_sp_grad_sync(world_size): + spawn(run_dist, world_size) + + +if __name__ == "__main__": + test_moe_ep_sp_grad_sync(world_size=4)