diff --git a/colossalai/moe/_operation.py b/colossalai/moe/_operation.py index 62904d90eef8..9360a583274f 100644 --- a/colossalai/moe/_operation.py +++ b/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, diff --git a/colossalai/shardformer/modeling/deepseek.py b/colossalai/shardformer/modeling/deepseek.py index 7bcdf6fc9892..2cd148bafe76 100644 --- a/colossalai/shardformer/modeling/deepseek.py +++ b/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,9 +173,7 @@ 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 = [] @@ -201,13 +181,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: 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) diff --git a/colossalai/shardformer/modeling/deepseek_v3.py b/colossalai/shardformer/modeling/deepseek_v3.py index 7515c2064dc2..c582ac9a3265 100644 --- a/colossalai/shardformer/modeling/deepseek_v3.py +++ b/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 diff --git a/colossalai/shardformer/modeling/mixtral.py b/colossalai/shardformer/modeling/mixtral.py index a739671ad81b..3d1d01ead731 100644 --- a/colossalai/shardformer/modeling/mixtral.py +++ b/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,10 +141,8 @@ 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 = [] @@ -170,14 +150,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: 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) diff --git a/tests/test_moe/test_moe_ep_zero_grad_scale.py b/tests/test_moe/test_moe_ep_zero_grad_scale.py new file mode 100644 index 000000000000..4f0317b54a70 --- /dev/null +++ b/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)