From e0d89a06e137a0829edddeaedbc5b773aa0bfc36 Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Mon, 21 Sep 2026 23:14:25 -0500 Subject: [PATCH 1/2] [checkpoint] unpad padded parameters when saving sharded MoE checkpoints `MoECheckpointIO._model_sharder` is a copy of `HybridParallelCheckpointIO._model_sharder` without the `to_unpadded_tensor` step, so `booster.save_model(..., shard=True)` with `MoeHybridParallelPlugin` writes padded parameters. With tensor parallelism, `VocabParallelEmbedding1D` pads the vocab to a multiple of `make_vocab_size_divisible_by * tp_size`, so the saved embedding has the padded vocab size and cannot be loaded by transformers: size mismatch for weight: copying a param with shape torch.Size([1024, 8]) from checkpoint, the shape in current model is torch.Size([1000, 8]). This affects any MoE model whose vocab is not a multiple of 64 * tp_size, e.g. Qwen3 (151936). The existing test uses Mixtral's default vocab (32000), which is never padded, and tp_size=1. Unpad the parameters like HybridParallelCheckpointIO does, and run the MoE checkpoint test with a padded vocab and tp_size=2 as well. --- colossalai/checkpoint_io/moe_checkpoint.py | 3 +++ tests/test_moe/test_moe_checkpoint.py | 9 +++++---- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/colossalai/checkpoint_io/moe_checkpoint.py b/colossalai/checkpoint_io/moe_checkpoint.py index 85e36f7c6336..fda3b3ce6953 100644 --- a/colossalai/checkpoint_io/moe_checkpoint.py +++ b/colossalai/checkpoint_io/moe_checkpoint.py @@ -34,6 +34,7 @@ ) from colossalai.interface import ModelWrapper, OptimizerWrapper from colossalai.tensor.moe_tensor.api import is_moe_tensor +from colossalai.tensor.padded_tensor import is_padded_tensor, to_unpadded_tensor try: from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX @@ -87,6 +88,8 @@ def _model_sharder( continue # Gather tensor pieces when using tensor parallel. param_ = gather_distributed_param(param, keep_vars=False) + if is_padded_tensor(param_): + param_ = to_unpadded_tensor(param_) block, block_size = state_dict_sharder.append_param(prefix + name, param_) if block is not None: yield block, block_size diff --git a/tests/test_moe/test_moe_checkpoint.py b/tests/test_moe/test_moe_checkpoint.py index f3f109192756..c408388d9a65 100644 --- a/tests/test_moe/test_moe_checkpoint.py +++ b/tests/test_moe/test_moe_checkpoint.py @@ -82,12 +82,15 @@ def check_optimizer_snapshot_equal(snapshot1, snapshot2, param2name, moe_dp_grou num_attention_heads=2, num_key_value_heads=2, num_hidden_layers=2, + # not a multiple of make_vocab_size_divisible_by * tp_size, so the embedding is padded with tp + vocab_size=1000, ), MixtralForCausalLM, ], ], ) -def check_moe_checkpoint(test_config): +@parameterize("plugin_config", [{"pp_size": 2, "ep_size": 2, "tp_size": 1}, {"pp_size": 2, "ep_size": 1, "tp_size": 2}]) +def check_moe_checkpoint(test_config, plugin_config): dtype, precision = torch.float16, "fp16" config, model_cls = test_config torch.cuda.set_device(dist.get_rank()) @@ -106,9 +109,7 @@ def check_moe_checkpoint(test_config): seed_all(10086) model = deepcopy(orig_model) optimizer = SGD(model.parameters(), lr=1e-3) - plugin = MoeHybridParallelPlugin( - pp_size=2, ep_size=2, tp_size=1, microbatch_size=1, zero_stage=1, precision=precision - ) + plugin = MoeHybridParallelPlugin(**plugin_config, microbatch_size=1, zero_stage=1, precision=precision) booster = Booster(plugin=plugin) model, optimizer, *_ = booster.boost(model=model, optimizer=optimizer) # initialize grads From 9c02e5d3f071c016809590a1d5f8415d19bf08f1 Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Mon, 21 Sep 2026 23:15:02 -0500 Subject: [PATCH 2/2] [shardformer] support Qwen3-MoE with expert parallelism Closes the Qwen part of #6180. Add a shardformer policy for `Qwen3MoeModel` / `Qwen3MoeForCausalLM` (transformers >= 4.51) that works with `MoeHybridParallelPlugin`: - expert parallelism: `EPQwen3MoeSparseMoeBlock`, adapted from the Mixtral EP block (all-to-all token dispatch, experts sharded across the EP group, optional TP inside each expert), with Qwen3's `norm_topk_prob`. Dense layers (`mlp_only_layers` / `decoder_sparse_step`) are kept as they are. - tensor parallelism for attention, router and dense MLP layers, vocab parallel embedding. - pipeline parallelism (1F1B and interleaved) with the router logits carried across stages for the load balancing loss; router logits of dense layers are skipped. - sequence parallelism (all_to_all), reusing the Qwen3 attention forward, since Qwen3MoeAttention is identical to Qwen3Attention. - ZeRO 1/2 through the plugin. Not supported yet (raise NotImplementedError): other SP modes, SP together with PP, and the zero bubble schedule. Tests: - tests/test_shardformer/test_model/test_shard_qwen3_moe.py: training (loss and weights) against a single-GPU reference for 11 combinations of EP / TP / PP / SP / ZeRO with a model that has dense and sparse layers and GQA, sharded checkpoint reloaded by `from_pretrained`, and the causal LM loss including the router aux loss for EP / PP / TP. - tests/test_moe/test_qwen3_moe_layer.py: fp32 layer-level check of the EP block (output, router logits and every gradient) against the transformers block, for 1 and 2 experts per rank and both `norm_topk_prob` values. The model-level test cannot catch wrong expert math, because the expert outputs are tiny compared to the residual stream at the default init; the layer test uses a larger init and catches e.g. a flipped `norm_topk_prob`. --- colossalai/shardformer/modeling/qwen3_moe.py | 458 ++++++++++++++++++ .../shardformer/policies/auto_policy.py | 7 + colossalai/shardformer/policies/qwen3_moe.py | 347 +++++++++++++ tests/test_moe/test_qwen3_moe_layer.py | 89 ++++ .../test_model/test_shard_qwen3_moe.py | 333 +++++++++++++ 5 files changed, 1234 insertions(+) create mode 100644 colossalai/shardformer/modeling/qwen3_moe.py create mode 100644 colossalai/shardformer/policies/qwen3_moe.py create mode 100644 tests/test_moe/test_qwen3_moe_layer.py create mode 100644 tests/test_shardformer/test_model/test_shard_qwen3_moe.py diff --git a/colossalai/shardformer/modeling/qwen3_moe.py b/colossalai/shardformer/modeling/qwen3_moe.py new file mode 100644 index 000000000000..656d41c3b62d --- /dev/null +++ b/colossalai/shardformer/modeling/qwen3_moe.py @@ -0,0 +1,458 @@ +# Modified from colossalai/shardformer/modeling/mixtral.py and colossalai/shardformer/modeling/qwen3.py +from typing import List, Optional, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.distributed import ProcessGroup +from transformers.modeling_attn_mask_utils import ( + _prepare_4d_causal_attention_mask, + _prepare_4d_causal_attention_mask_for_sdpa, +) +from transformers.models.qwen3_moe.modeling_qwen3_moe import ( + MoeCausalLMOutputWithPast, + MoeModelOutputWithPast, + Qwen3MoeForCausalLM, + Qwen3MoeModel, + Qwen3MoeSparseMoeBlock, + load_balancing_loss_func, +) +from transformers.utils import logging + +from colossalai.lazy import LazyInitContext +from colossalai.moe._operation import ( + DPGradScalerIn, + DPGradScalerOut, + 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 gather_sp_output, split_forward_gather_backward +from colossalai.shardformer.layer.linear import Linear1D_Col, Linear1D_Row, ParallelModule +from colossalai.shardformer.layer.utils import is_share_sp_tp +from colossalai.shardformer.shard import ShardConfig +from colossalai.shardformer.shard.utils import set_tensors_to_none +from colossalai.tensor.moe_tensor.api import set_moe_tensor_ep_group + +from ..layer import ColoAttention + + +class EPQwen3MoeSparseMoeBlock(ParallelModule): + def __init__(self, *args, **kwargs): + raise RuntimeError(f"Please use `from_native_module` to create an instance of {self.__class__.__name__}") + + def setup_process_groups( + self, + tp_group: ProcessGroup, + moe_dp_group: ProcessGroup, + ep_group: ProcessGroup, + fp8_communication: bool = False, + use_zbv: bool = False, + ): + assert tp_group is not None + assert moe_dp_group is not None + assert ep_group is not None + + # setup ep group + self.ep_size = dist.get_world_size(ep_group) + self.ep_rank = dist.get_rank(ep_group) + self.ep_group = ep_group + self.fp8_communication = fp8_communication + self.use_zbv = use_zbv + + if self.num_experts % self.ep_size != 0: + raise ValueError("The number of experts must be divisible by the number of expert parallel groups.") + + self.num_experts_per_ep = self.num_experts // self.ep_size + self.expert_start_idx = self.ep_rank * self.num_experts_per_ep + held_experts = self.experts[self.expert_start_idx : self.expert_start_idx + self.num_experts_per_ep] + + set_tensors_to_none(self.experts, exclude=set(held_experts)) + + # setup moe_dp group + self.moe_dp_group = moe_dp_group + self.moe_dp_size = moe_dp_group.size() + + # setup global tp group + self.tp_group = tp_group + if self.tp_group.size() > 1: + for expert in held_experts: + expert.gate_proj = Linear1D_Col.from_native_module( + expert.gate_proj, self.tp_group, fp8_communication=self.fp8_communication, use_zbv=self.use_zbv + ) + expert.up_proj = Linear1D_Col.from_native_module( + expert.up_proj, self.tp_group, fp8_communication=self.fp8_communication, use_zbv=self.use_zbv + ) + expert.down_proj = Linear1D_Row.from_native_module( + expert.down_proj, self.tp_group, fp8_communication=self.fp8_communication, use_zbv=self.use_zbv + ) + + for p in self.experts.parameters(): + set_moe_tensor_ep_group(p, ep_group) + + @staticmethod + def from_native_module( + module: Qwen3MoeSparseMoeBlock, + tp_group: ProcessGroup, + moe_dp_group: ProcessGroup, + ep_group: ProcessGroup, + *args, + **kwargs, + ) -> "EPQwen3MoeSparseMoeBlock": + LazyInitContext.materialize(module) + # layers in `mlp_only_layers` (or not on the `decoder_sparse_step`) keep a dense Qwen3MoeMLP + if not isinstance(module, Qwen3MoeSparseMoeBlock): + return module + module.__class__ = EPQwen3MoeSparseMoeBlock + fp8_communication = kwargs.get("fp8_communication", False) + use_zbv = kwargs.get("use_zbv", False) + module.setup_process_groups(tp_group, moe_dp_group, ep_group, fp8_communication, use_zbv) + return module + + def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + # router_logits: (batch * sequence_length, n_experts) + router_logits = self.gate(hidden_states) + + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + if self.norm_topk_prob: + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) + + selected_experts = selected_experts.t().reshape(-1) + selected_experts_idx = selected_experts.argsort() + dispatch_states = hidden_states.repeat(self.top_k, 1)[selected_experts_idx] + input_split_sizes = selected_experts.bincount(minlength=self.num_experts) + + output_split_sizes = torch.zeros_like(input_split_sizes) + + 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() + + output_states, _ = all_to_all_uneven( + dispatch_states, + input_split_list, + output_split_list, + self.ep_group, + fp8_communication=self.fp8_communication, + ) + # compute expert output + output_states = EPGradScalerIn.apply(output_states, self.ep_size) + if output_states.size(0) > 0: + 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(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(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) + dispatch_states, _ = all_to_all_uneven( + output_states, output_split_list, input_split_list, self.ep_group, fp8_communication=self.fp8_communication + ) + + recover_experts_idx = torch.empty_like(selected_experts_idx) + recover_experts_idx[selected_experts_idx] = torch.arange( + selected_experts_idx.size(0), device=selected_experts_idx.device + ) + dispatch_states = dispatch_states[recover_experts_idx] + k_hidden_states = dispatch_states.chunk(self.top_k) + output_states = k_hidden_states[0] * routing_weights[:, 0, None] + for i in range(1, self.top_k): + output_states += k_hidden_states[i] * routing_weights[:, i, None] + output_states = output_states.reshape(batch_size, sequence_length, hidden_dim) + return output_states, router_logits + + +class Qwen3MoePipelineForwards: + """ + This class serves as a micro library for forward function substitution of Qwen3-MoE models + under pipeline parallelism or sequence parallelism. + """ + + @staticmethod + def qwen3_moe_model_forward( + self: Qwen3MoeModel, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + past_router_logits: Optional[Tuple[torch.FloatTensor]] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + force_sp_output_gather: bool = True, + **kwargs, + ) -> Union[Tuple, MoeModelOutputWithPast, dict]: + """Used for pipeline parallelism (``stage_manager`` is set) or sequence parallelism (``stage_manager`` is + None, all the layers are run). The two are not supported together.""" + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + is_first_stage = stage_manager is None or stage_manager.is_first_stage() + is_last_stage = stage_manager is None or stage_manager.is_last_stage() + if stage_index is None: + stage_index = [0, len(self.layers)] + + # retrieve input_ids and inputs_embeds + if is_first_stage: + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + else: + batch_size, seq_length = hidden_states.shape[:-1] + device = hidden_states.device + + # TODO: kv cache, attentions and hidden states are not recorded, same as the other pipeline forwards + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + if use_cache: + logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") + use_cache = False + + sp_mode = shard_config.sequence_parallelism_mode if shard_config.enable_sequence_parallelism else None + sp_size = shard_config.sequence_parallel_size + sp_group = shard_config.sequence_parallel_process_group + + if position_ids is None: + position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0) + else: + position_ids = position_ids.view(-1, seq_length).long() + if cache_position is None: + cache_position = torch.arange(seq_length, device=device) + + if shard_config.enable_flash_attention or sp_mode is not None: + # the attention forward is replaced by `get_qwen3_flash_attention_forward` (see the policy), + # which takes ColoAttention kwargs or a 4d additive causal mask + if shard_config.enable_flash_attention: + attention_mask = ColoAttention.prepare_attn_kwargs( + (batch_size, 1, seq_length, seq_length), + hidden_states.dtype, + hidden_states.device, + q_padding_mask=attention_mask, + is_causal=True, + ) + else: + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + hidden_states, + 0, + sliding_window=self.config.sliding_window, + ) + elif self.config._attn_implementation == "flash_attention_2": + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + elif self.config._attn_implementation == "sdpa": + attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, (batch_size, seq_length), hidden_states, 0 + ) + else: + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + hidden_states, + 0, + sliding_window=self.config.sliding_window, + ) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + if is_first_stage and sp_mode is not None: + if is_share_sp_tp(sp_mode): + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication + ) + elif sp_mode == "all_to_all": + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication + ) + + all_router_logits = () if output_router_logits else None + start_idx, end_idx = stage_index[0], stage_index[1] + for decoder_layer in self.layers[start_idx:end_idx]: + layer_args = ( + hidden_states, + attention_mask, + position_ids, + None, # past_key_value + output_attentions, + output_router_logits, + use_cache, + cache_position, + position_embeddings, + ) + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func(decoder_layer.__call__, *layer_args) + else: + layer_outputs = decoder_layer(*layer_args) + hidden_states = layer_outputs[0] + + # dense layers (Qwen3MoeMLP) have no router logits + if output_router_logits and layer_outputs[-1] is not None: + all_router_logits += (layer_outputs[-1],) + + if output_router_logits and past_router_logits is not None: + all_router_logits = past_router_logits + all_router_logits + + if not is_last_stage: + out = {"hidden_states": hidden_states} + if output_router_logits: + out["past_router_logits"] = all_router_logits + return out + + hidden_states = self.norm(hidden_states) + if sp_mode is not None: + if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): + hidden_states = gather_sp_output(hidden_states, shard_config) + + if not return_dict: + return tuple(v for v in [hidden_states, all_router_logits] if v is not None) + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=None, + hidden_states=None, + attentions=None, + router_logits=all_router_logits, + ) + + @staticmethod + def qwen3_moe_for_causal_lm_forward( + self: Qwen3MoeForCausalLM, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + past_router_logits: Optional[Tuple[torch.FloatTensor]] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + **kwargs, + ): + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = Qwen3MoePipelineForwards.qwen3_moe_model_forward( + self.model, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_router_logits=output_router_logits, + cache_position=cache_position, + return_dict=True, + stage_manager=stage_manager, + hidden_states=hidden_states, + past_router_logits=past_router_logits, + stage_index=stage_index, + shard_config=shard_config, + ) + + if stage_manager is not None and not stage_manager.is_last_stage(): + return outputs + + logits = self.lm_head(outputs.last_hidden_state) + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.vocab_size) + + aux_loss = None + if output_router_logits: + aux_loss = load_balancing_loss_func( + outputs.router_logits, self.num_experts, self.num_experts_per_tok, attention_mask + ) + if labels is not None: + # make sure to reside in the same device + loss += self.router_aux_loss_coef * aux_loss.to(loss.device) + + if not return_dict: + output = (logits,) + if output_router_logits: + output = (aux_loss,) + output + (outputs.router_logits,) + return (loss,) + output if loss is not None else output + + return MoeCausalLMOutputWithPast( + loss=loss, + aux_loss=aux_loss, + logits=logits, + past_key_values=None, + hidden_states=None, + attentions=None, + router_logits=outputs.router_logits, + ) diff --git a/colossalai/shardformer/policies/auto_policy.py b/colossalai/shardformer/policies/auto_policy.py index 3d61af1e0aea..2e5dd759af8d 100644 --- a/colossalai/shardformer/policies/auto_policy.py +++ b/colossalai/shardformer/policies/auto_policy.py @@ -230,6 +230,13 @@ class PolicyLocation: "transformers.models.qwen3.modeling_qwen3.Qwen3ForSequenceClassification": PolicyLocation( file_name="qwen3", class_name="Qwen3ForSequenceClassificationPolicy" ), + # Qwen3-MoE + "transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeModel": PolicyLocation( + file_name="qwen3_moe", class_name="Qwen3MoeModelPolicy" + ), + "transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeForCausalLM": PolicyLocation( + file_name="qwen3_moe", class_name="Qwen3MoeForCausalLMPolicy" + ), # command "transformers.models.cohere.modeling_cohere.CohereModel": PolicyLocation( file_name="command", class_name="CommandModelPolicy" diff --git a/colossalai/shardformer/policies/qwen3_moe.py b/colossalai/shardformer/policies/qwen3_moe.py new file mode 100644 index 000000000000..27f6b3f97307 --- /dev/null +++ b/colossalai/shardformer/policies/qwen3_moe.py @@ -0,0 +1,347 @@ +# Modified from colossalai/shardformer/policies/mixtral.py and colossalai/shardformer/policies/qwen3.py +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module + +from colossalai.shardformer.layer import ( + FusedRMSNorm, + Linear1D_Col, + Linear1D_Row, + PaddingEmbedding, + RMSNorm, + VocabParallelEmbedding1D, +) +from colossalai.shardformer.modeling.qwen3 import get_qwen3_flash_attention_forward +from colossalai.shardformer.modeling.qwen3_moe import EPQwen3MoeSparseMoeBlock, Qwen3MoePipelineForwards +from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["Qwen3MoePolicy", "Qwen3MoeModelPolicy", "Qwen3MoeForCausalLMPolicy"] + + +class Qwen3MoePolicy(Policy): + def __init__(self) -> None: + super().__init__() + import transformers + from packaging.version import Version + + assert Version(transformers.__version__) >= Version( + "4.51.0" + ), "The Qwen3-MoE model should run on a transformers version of 4.51.0 or higher." + + def config_sanity_check(self): + pass + + def preprocess(self): + self.tie_weight = self.tie_weight_check() + self.origin_attn_implement = self.model.config._attn_implementation + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + from transformers.models.qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeAttention, + Qwen3MoeDecoderLayer, + Qwen3MoeModel, + ) + + policy = {} + + sp_mode = self.shard_config.sequence_parallelism_mode or None + sp_size = self.shard_config.sequence_parallel_size or None + sp_group = self.shard_config.sequence_parallel_process_group or None + tp_size = self.shard_config.tensor_parallel_size + if self.shard_config.enable_sequence_parallelism: + if sp_mode != "all_to_all": + raise NotImplementedError( + f"Sequence parallelism mode {sp_mode} is not supported for Qwen3-MoE yet, please use all_to_all." + ) + if self.pipeline_stage_manager is not None: + raise NotImplementedError("Sequence parallelism is not supported with pipeline parallelism.") + if self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv: + raise NotImplementedError("The zero bubble pipeline schedule is not supported for Qwen3-MoE yet.") + + norm_cls = FusedRMSNorm if self.shard_config.enable_fused_normalization else RMSNorm + + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = VocabParallelEmbedding1D + elif self.tie_weight: + embedding_cls = PaddingEmbedding + + # the number of heads held by each rank, used by the attention forward for sequence parallelism + num_q_heads = self.model.config.num_attention_heads + num_kv_heads = self.model.config.num_key_value_heads + if sp_mode == "all_to_all": + num_q_heads //= sp_size + num_kv_heads //= sp_size + decoder_attribute_replacement = {} + if self.shard_config.enable_tensor_parallelism: + assert ( + self.model.config.num_attention_heads % tp_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + assert ( + self.model.config.num_key_value_heads % tp_size == 0 + ), f"The number of key_value heads must be divisible by tensor parallel size." + num_q_heads //= tp_size + num_kv_heads //= tp_size + decoder_attribute_replacement["self_attn.hidden_size"] = self.model.config.hidden_size // tp_size + if sp_mode == "all_to_all" or self.shard_config.enable_tensor_parallelism: + decoder_attribute_replacement["self_attn.num_heads"] = num_q_heads + decoder_attribute_replacement["self_attn.num_key_value_heads"] = num_kv_heads + policy[Qwen3MoeDecoderLayer] = ModulePolicyDescription(attribute_replacement=decoder_attribute_replacement) + + if self.shard_config.enable_tensor_parallelism: + # tensor parallelism for the attention, the router and the dense mlp layers, + # the experts of the sparse layers are sharded by EPQwen3MoeSparseMoeBlock + fp8_communication = self.shard_config.fp8_communication + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=Linear1D_Col, + kwargs={"fp8_communication": fp8_communication}, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=Linear1D_Col, + kwargs={"fp8_communication": fp8_communication}, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=Linear1D_Col, + kwargs={"fp8_communication": fp8_communication}, + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=Linear1D_Row, + kwargs={"fp8_communication": fp8_communication}, + ), + # sparse layers + SubModuleReplacementDescription( + suffix="mlp.gate", + target_module=Linear1D_Col, + kwargs={"gather_output": True, "fp8_communication": fp8_communication}, + ignore_if_not_exist=True, + ), + # dense layers (`mlp_only_layers` or not on the `decoder_sparse_step`) + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=Linear1D_Col, + kwargs={"fp8_communication": fp8_communication}, + ignore_if_not_exist=True, + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=Linear1D_Col, + kwargs={"fp8_communication": fp8_communication}, + ignore_if_not_exist=True, + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=Linear1D_Row, + kwargs={"fp8_communication": fp8_communication}, + ignore_if_not_exist=True, + ), + ], + policy=policy, + target_key=Qwen3MoeDecoderLayer, + ) + + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=embedding_cls, + kwargs=( + { + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} + ), + ), + policy=policy, + target_key=Qwen3MoeModel, + ) + + if self.shard_config.ep_group: + # expert parallel, dense layers are kept as they are + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="mlp", + target_module=EPQwen3MoeSparseMoeBlock, + kwargs={ + "ep_group": self.shard_config.ep_group, + "tp_group": self.shard_config.tensor_parallel_process_group, + "moe_dp_group": self.shard_config.moe_dp_group, + "fp8_communication": self.shard_config.fp8_communication, + }, + ) + ], + policy=policy, + target_key=Qwen3MoeDecoderLayer, + ) + + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription(suffix="input_layernorm", target_module=norm_cls), + SubModuleReplacementDescription(suffix="post_attention_layernorm", target_module=norm_cls), + ], + policy=policy, + target_key=Qwen3MoeDecoderLayer, + ) + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription(suffix="norm", target_module=norm_cls), + policy=policy, + target_key=Qwen3MoeModel, + ) + + if self.shard_config.enable_flash_attention or self.shard_config.enable_sequence_parallelism: + # Qwen3MoeAttention is the same as Qwen3Attention + self.append_or_create_method_replacement( + description={ + "forward": get_qwen3_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), + }, + policy=policy, + target_key=Qwen3MoeAttention, + ) + if self.pipeline_stage_manager is None: + # the model forward prepares the attention mask and splits / gathers the sequence + self.append_or_create_method_replacement( + description={ + "forward": partial( + Qwen3MoePipelineForwards.qwen3_moe_model_forward, shard_config=self.shard_config + ), + }, + policy=policy, + target_key=Qwen3MoeModel, + ) + + return policy + + def postprocess(self): + return self.model + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if self.pipeline_stage_manager is None: + return + + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "Qwen3MoeModel": + module = self.model + else: + module = self.model.model + + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_interleave: + # stage_index is passed in by the interleaved schedule for each model chunk + stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial(new_forward, stage_manager=stage_manager, shard_config=self.shard_config) + } + else: + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial( + new_forward, stage_manager=stage_manager, stage_index=stage_index, shard_config=self.shard_config + ) + } + self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + if self.model.__class__.__name__ == "Qwen3MoeModel": + module = self.model + else: + module = self.model.model + stage_manager = self.pipeline_stage_manager + + held_layers = [] + held_layers.append(module.rotary_emb) + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + stage_indices = stage_manager.get_stage_index(layers_per_stage) + stage_manager.stage_indices = stage_indices + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embed_tokens) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.layers[start_idx:end_idx]) + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(module.norm) + else: + if stage_manager.is_first_stage(): + held_layers.append(module.embed_tokens) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.layers[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.norm) + return held_layers + + +class Qwen3MoeModelPolicy(Qwen3MoePolicy): + def module_policy(self): + from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeModel + + policy = super().module_policy() + self.set_pipeline_forward( + model_cls=Qwen3MoeModel, + new_forward=Qwen3MoePipelineForwards.qwen3_moe_model_forward, + policy=policy, + ) + return policy + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in Qwen3-MoE model""" + return [] + + +class Qwen3MoeForCausalLMPolicy(Qwen3MoePolicy): + def module_policy(self): + from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeForCausalLM + + policy = super().module_policy() + if self.shard_config.enable_tensor_parallelism: + policy[Qwen3MoeForCausalLM] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=Linear1D_Col, + kwargs=dict(gather_output=True, fp8_communication=self.shard_config.fp8_communication), + ) + ], + ) + self.set_pipeline_forward( + model_cls=Qwen3MoeForCausalLM, + new_forward=Qwen3MoePipelineForwards.qwen3_moe_for_causal_lm_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + held_layers = super().get_held_layers() + if self.pipeline_stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + qwen3_moe_model = self.model.model + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if id(qwen3_moe_model.embed_tokens.weight) == id(self.model.lm_head.weight): + # tie weights + return [ + { + 0: qwen3_moe_model.embed_tokens.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, + } + ] + return [] diff --git a/tests/test_moe/test_qwen3_moe_layer.py b/tests/test_moe/test_qwen3_moe_layer.py new file mode 100644 index 000000000000..802c10aa659c --- /dev/null +++ b/tests/test_moe/test_qwen3_moe_layer.py @@ -0,0 +1,89 @@ +from copy import deepcopy + +import pytest +import torch +import torch.distributed as dist +import transformers +from packaging.version import Version +from torch.testing import assert_close + +import colossalai +from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin +from colossalai.testing import parameterize +from colossalai.testing.utils import spawn + +tokens, n_experts = 7, 4 +hidden_size = 8 +top_k = 2 + + +@parameterize("ep_size", [2, 4]) +@parameterize("norm_topk_prob", [True, False]) +def check_qwen3_moe_layer(ep_size: int, norm_topk_prob: bool): + from transformers import Qwen3MoeConfig + from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock + + from colossalai.shardformer.modeling.qwen3_moe import EPQwen3MoeSparseMoeBlock + + torch.cuda.set_device(dist.get_rank()) + plugin = MoeHybridParallelPlugin( + precision="fp32", + tp_size=1, + pp_size=1, + zero_stage=1, + ep_size=ep_size, + ) + config = Qwen3MoeConfig( + hidden_size=hidden_size, + moe_intermediate_size=hidden_size * 2, + num_experts=n_experts, + num_experts_per_tok=top_k, + norm_topk_prob=norm_topk_prob, + ) + torch.manual_seed(0) + orig_model = Qwen3MoeSparseMoeBlock(config).cuda() + # a larger init than the default initializer_range so that the expert outputs are not negligible + for p in orig_model.parameters(): + torch.nn.init.normal_(p, std=0.5) + # as in training, the input requires grad, otherwise an ep rank whose experts get no tokens + # would skip the backward all-to-all + x = torch.rand(1, tokens, hidden_size, device="cuda", requires_grad=True) + ep_x = x.detach().clone().requires_grad_() + orig_output, orig_logits = orig_model(x) + model = deepcopy(orig_model) + model = EPQwen3MoeSparseMoeBlock.from_native_module( + model, + ep_group=plugin.ep_group, + tp_group=plugin.tp_group, + moe_dp_group=plugin.moe_dp_group, + ) + assert sum(e.gate_proj.weight is not None for e in model.experts) == n_experts // ep_size + + ep_output, ep_logits = model(ep_x) + assert_close(orig_logits, ep_logits) + assert_close(orig_output, ep_output) + + orig_output.pow(2).sum().backward() + ep_output.pow(2).sum().backward() + assert_close(x.grad, ep_x.grad) + name_to_p = {n: p for n, p in orig_model.named_parameters()} + for n, ep_p in model.named_parameters(): + if ep_p.grad is not None: + assert_close(name_to_p[n].grad, ep_p.grad, msg=lambda m: f"{n}: {m}") + + +def run_dist(rank: int, world_size: int, port: int): + colossalai.launch(rank, world_size, "localhost", port) + check_qwen3_moe_layer() + + +@pytest.mark.skipif( + Version(transformers.__version__) < Version("4.51.0"), reason="Requires transformers version 4.51.0 or later" +) +@pytest.mark.parametrize("world_size", [4]) +def test_qwen3_moe_layer(world_size: int): + spawn(run_dist, world_size) + + +if __name__ == "__main__": + test_qwen3_moe_layer(4) diff --git a/tests/test_shardformer/test_model/test_shard_qwen3_moe.py b/tests/test_shardformer/test_model/test_shard_qwen3_moe.py new file mode 100644 index 000000000000..b86f744bb00b --- /dev/null +++ b/tests/test_shardformer/test_model/test_shard_qwen3_moe.py @@ -0,0 +1,333 @@ +import os +import shutil +from copy import deepcopy +from typing import Tuple + +import pytest +import torch +import torch.distributed as dist +import transformers +from packaging.version import Version + +import colossalai +from colossalai.booster.booster import Booster +from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin +from colossalai.shardformer.layer.utils import Randomizer +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 assert_loose_close, check_model_equal + +NUM_BATCH = 8 +NUM_TOK_PER_BATCH, NUM_EXPERTS = 64, 4 +NUM_LAYERS = 4 +HIDDEN_SIZE_PER_HEAD = 4 +NUM_HEADS = 8 +NUM_KV_HEADS = 4 +TOP_K = 2 +VOCAB_SIZE = 128 + + +def make_config(): + from transformers import Qwen3MoeConfig + + return Qwen3MoeConfig( + hidden_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS, + head_dim=HIDDEN_SIZE_PER_HEAD, + intermediate_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS * 2, + moe_intermediate_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS, + num_hidden_layers=NUM_LAYERS, + num_attention_heads=NUM_HEADS, + num_key_value_heads=NUM_KV_HEADS, + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOP_K, + norm_topk_prob=True, + # layer 1 keeps a dense MLP, the others are sparse + mlp_only_layers=[1], + vocab_size=VOCAB_SIZE, + attn_implementation="sdpa", + ) + + +def make_plugin(stage, ep_size, pp_size, tp_size, sp_size): + return MoeHybridParallelPlugin( + pp_size=pp_size, + num_microbatches=pp_size, + tp_size=tp_size, + sp_size=sp_size, + ep_size=ep_size, + zero_stage=stage, + enable_sequence_parallelism=sp_size > 1, + sequence_parallelism_mode="all_to_all" if sp_size > 1 else None, + overlap_communication=False, + initial_scale=1, + precision="bf16", + find_unused_parameters=True, + ) + + +def run_qwen3_moe_commom(config: Tuple[int, ...]): + from transformers import Qwen3MoeModel + + Randomizer.reset_index() + stage, ep_size, pp_size, tp_size, sp_size = config + world_size = dist.get_world_size() + rank = dist.get_rank() + dtype = torch.bfloat16 + torch.cuda.set_device(dist.get_rank()) + + plugin = make_plugin(stage, ep_size, pp_size, tp_size, sp_size) + dp_size = plugin.dp_size + booster = Booster(plugin=plugin) + + assert pp_size <= NUM_LAYERS, "pp_size should be less than or equal to NUM_LAYERS" + model_config = make_config() + + # init model with the same seed + seed_all(10086) + + torch_model = Qwen3MoeModel(model_config).to(dtype).cuda() + torch_optimizer = torch.optim.SGD(torch_model.parameters(), lr=1) + + parallel_model = deepcopy(torch_model) + parallel_optimizer = torch.optim.SGD(parallel_model.parameters(), lr=1) + parallel_model, parallel_optimizer, _, _, _ = booster.boost(parallel_model, parallel_optimizer) + + # create different input along dp axis + seed_all(1453 + rank) + + torch_model.train() + parallel_model.train() + for _ in range(2): + # gen random input + input_embeddings = torch.rand( + NUM_BATCH, NUM_TOK_PER_BATCH, HIDDEN_SIZE_PER_HEAD * NUM_HEADS, requires_grad=True + ).cuda() + dist.all_reduce( + input_embeddings, group=plugin.pp_group + ) # pp inputs except the first stage doesn't matter, but need to be replicate for torch model check + + dist.all_reduce(input_embeddings, group=plugin.tp_group) # tp group duplicate input + dist.all_reduce(input_embeddings, group=plugin.sp_group) # sp group duplicate input + + # run the model with hybrid parallel + if booster.plugin.stage_manager is not None: + # for test with pp + data_iter = iter([{"inputs_embeds": input_embeddings}]) + sharded_output = booster.execute_pipeline( + data_iter, + parallel_model, + lambda x, y: x.last_hidden_state.mean(), + parallel_optimizer, + return_loss=True, + return_outputs=True, + ) + if booster.plugin.stage_manager.is_last_stage(): + parallel_output = sharded_output["loss"] + else: + parallel_output = torch.tensor(12345.0, device="cuda") + + # broadcast along pp axis + dist.broadcast( + parallel_output, src=dist.get_process_group_ranks(plugin.pp_group)[-1], group=plugin.pp_group + ) + else: + # for test without pp + parallel_output = parallel_model(inputs_embeds=input_embeddings.to(dtype)).last_hidden_state.mean() + parallel_optimizer.backward(parallel_output) + parallel_optimizer.step() + parallel_optimizer.zero_grad() + dist.all_reduce(parallel_output, group=plugin.mixed_dp_group) + + # =================================================================================== + # run normal model with all dp(different) inputs + all_inputs = [torch.empty_like(input_embeddings) for _ in range(dp_size)] + dist.all_gather(all_inputs, input_embeddings, group=plugin.mixed_dp_group) + torch_output_sum = 0 + for input_data_ in all_inputs: + torch_output = torch_model(inputs_embeds=input_data_.to(dtype)).last_hidden_state.mean() + torch_output.backward() + torch_output_sum += torch_output.detach() + # avg dp grads follows zero optimizer + for p in torch_model.parameters(): + if p.grad is not None: + p.grad /= dp_size + torch_optimizer.step() + torch_optimizer.zero_grad() + + assert_loose_close(parallel_output, torch_output_sum, dtype=dtype) + + # use checkpoint to load sharded zero model + model_dir = "./test_qwen3_moe" + if rank == world_size - 1: + os.makedirs(model_dir, exist_ok=True) + + dist.barrier() + booster.save_model(parallel_model, model_dir, shard=True) + dist.barrier() + + saved_model = Qwen3MoeModel.from_pretrained(model_dir).cuda().to(dtype) + check_model_equal(torch_model, saved_model, dtype=dtype) + dist.barrier() + + if rank == world_size - 1: + shutil.rmtree(model_dir) + + print(f"rank {dist.get_rank()} test passed") + + +def run_qwen3_moe_causal_lm_commom(config: Tuple[int, ...]): + # checks the causal lm head, the loss and the router aux loss, which is accumulated across pipeline stages + from transformers import Qwen3MoeForCausalLM + + Randomizer.reset_index() + stage, ep_size, pp_size, tp_size, sp_size = config + rank = dist.get_rank() + dtype = torch.bfloat16 + torch.cuda.set_device(dist.get_rank()) + + plugin = make_plugin(stage, ep_size, pp_size, tp_size, sp_size) + dp_size = plugin.dp_size + booster = Booster(plugin=plugin) + + model_config = make_config() + model_config.output_router_logits = True + model_config.router_aux_loss_coef = 0.1 + # the transformers forward also collects the router logits (None) of dense layers and then fails in + # load_balancing_loss_func, so only use sparse layers when computing the router aux loss + model_config.mlp_only_layers = [] + + seed_all(10086) + torch_model = Qwen3MoeForCausalLM(model_config).to(dtype).cuda() + parallel_model = deepcopy(torch_model) + parallel_optimizer = torch.optim.SGD(parallel_model.parameters(), lr=1) + parallel_model, parallel_optimizer, _, _, _ = booster.boost(parallel_model, parallel_optimizer) + + seed_all(1453 + rank) + input_ids = torch.randint(0, VOCAB_SIZE, (NUM_BATCH, NUM_TOK_PER_BATCH), device="cuda") + # replicate the input across every group but dp + for group in (plugin.pp_group, plugin.tp_group, plugin.sp_group): + dist.broadcast(input_ids, src=dist.get_process_group_ranks(group)[0], group=group) + data = {"input_ids": input_ids, "labels": input_ids.clone()} + + if booster.plugin.stage_manager is not None: + sharded_output = booster.execute_pipeline( + iter([data]), parallel_model, lambda x, y: x.loss, parallel_optimizer, return_loss=True + ) + if booster.plugin.stage_manager.is_last_stage(): + parallel_loss = sharded_output["loss"] + else: + parallel_loss = torch.tensor(12345.0, device="cuda") + dist.broadcast(parallel_loss, src=dist.get_process_group_ranks(plugin.pp_group)[-1], group=plugin.pp_group) + else: + parallel_loss = parallel_model(**data).loss + dist.all_reduce(parallel_loss, group=plugin.mixed_dp_group) + + all_input_ids = [torch.empty_like(input_ids) for _ in range(dp_size)] + dist.all_gather(all_input_ids, input_ids, group=plugin.mixed_dp_group) + # the pipeline schedule averages the loss over micro batches, the router aux loss is not linear in the batch, + # so compute the reference loss on the same micro batches + num_microbatches = pp_size if pp_size > 1 else 1 + with torch.no_grad(): + torch_loss = 0 + for ids in all_input_ids: + for mb in ids.chunk(num_microbatches): + torch_loss += torch_model(input_ids=mb, labels=mb).loss / num_microbatches + + assert_loose_close(parallel_loss, torch_loss, dtype=dtype) + print(f"rank {dist.get_rank()} causal lm test passed") + + +@parameterize( + "config", + [ + # DDP: ep == 1 since ep * moe_dp == dp == moe_dp; sp == 1 since sp * dp == moe_dp == dp + (0, 1, 4, 1, 1), + (0, 1, 1, 4, 1), + (0, 1, 2, 2, 1), + # zero 1 + (1, 4, 1, 1, 1), + (1, 1, 4, 1, 1), + (1, 1, 1, 4, 1), + (1, 2, 1, 1, 2), + # zero 2 + (2, 4, 1, 1, 1), + (2, 1, 4, 1, 1), + (2, 1, 1, 4, 1), + (2, 2, 1, 1, 2), + ], +) +def run_qwen3_moe_test(config: Tuple[int, ...]): + run_qwen3_moe_commom(config) + + +@parameterize( + "config", + [ + (1, 4, 1, 1, 1), + (1, 1, 4, 1, 1), + (1, 1, 1, 4, 1), + (1, 2, 2, 1, 1), + ], +) +def run_qwen3_moe_causal_lm_test(config: Tuple[int, ...]): + run_qwen3_moe_causal_lm_commom(config) + + +@parameterize( + "config", + [ + # DDP: ep == 1 since ep * moe_dp == dp == moe_dp; sp == 1 since sp * dp == moe_dp == dp + (0, 1, 2, 4, 1), + (0, 1, 4, 2, 1), + (0, 1, 1, 4, 1), + (0, 1, 4, 1, 1), + # zero 1: + (1, 2, 1, 1, 2), + (1, 2, 1, 4, 1), + (1, 1, 1, 2, 2), + (1, 2, 2, 2, 1), + # zero 2 + (2, 2, 1, 1, 2), + (2, 2, 1, 4, 1), + (2, 1, 1, 2, 2), + (2, 2, 2, 2, 1), + ], +) +def run_qwen3_moe_3d_test(config: Tuple[int, ...]): + print(f"{config=}") + run_qwen3_moe_commom(config) + + +def check_qwen3_moe(rank, world_size, port): + colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") + run_qwen3_moe_test() + run_qwen3_moe_causal_lm_test() + + +def check_qwen3_moe_3d(rank, world_size, port): + colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") + run_qwen3_moe_3d_test() + + +@pytest.mark.skipif( + Version(transformers.__version__) < Version("4.51.0"), reason="Requires transformers version 4.51.0 or later" +) +@pytest.mark.dist +@pytest.mark.parametrize("world_size", [4]) +@rerun_if_address_is_in_use() +def test_qwen3_moe(world_size): + spawn(check_qwen3_moe, world_size) + + +@pytest.mark.skipif( + Version(transformers.__version__) < Version("4.51.0"), reason="Requires transformers version 4.51.0 or later" +) +@pytest.mark.largedist +@pytest.mark.parametrize("world_size", [8]) +@rerun_if_address_is_in_use() +def test_qwen3_moe_3d(world_size): + spawn(check_qwen3_moe_3d, world_size) + + +if __name__ == "__main__": + test_qwen3_moe(world_size=4)