From 8c1228010ba426a63cbb86f5b537405d103080f6 Mon Sep 17 00:00:00 2001 From: zli03 Date: Mon, 15 Jun 2026 16:53:57 +0800 Subject: [PATCH 1/3] feat(moe): support sonic-moe for small model training --- .../core/transformer/moe/sonic_moe_layer.py | 930 ++++++++++++++++++ megatron/training/arguments.py | 56 ++ .../transformer/moe/test_sonic_moe_layer.py | 662 +++++++++++++ 3 files changed, 1648 insertions(+) create mode 100644 megatron/core/transformer/moe/sonic_moe_layer.py create mode 100644 tests/unit_tests/transformer/moe/test_sonic_moe_layer.py diff --git a/megatron/core/transformer/moe/sonic_moe_layer.py b/megatron/core/transformer/moe/sonic_moe_layer.py new file mode 100644 index 00000000000..7d8aede7f8c --- /dev/null +++ b/megatron/core/transformer/moe/sonic_moe_layer.py @@ -0,0 +1,930 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Pure SonicMoE module adapter for Megatron-Core model specs. + +This module wraps ``sonicmoe.MoE`` directly. SonicMoE owns the router, routing +metadata generation, and fused expert computation. The adapter owns Megatron +auxiliary-loss logging/scaling and the MLP/MoE call signature. +""" + +from __future__ import annotations + +from typing import Optional, Union + +import torch +import torch.nn.functional as F + +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.moe.moe_layer import BaseMoELayer, MoELayer, MoESubmodules +from megatron.core.transformer.moe.moe_utils import ( + MoEAuxLossAutoScaler, + compute_routing_scores_for_aux_loss, + get_default_pg_collection, + save_to_aux_losses_tracker, + switch_load_balancing_loss_func, + z_loss_func, +) +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.utils import ( + ensure_metadata_has_dp_cp_group, + make_sharded_tensors_for_checkpoint, +) + + +try: + from sonicmoe import KernelBackendMoE, MoE as _SonicMoE + from sonicmoe.enums import ActivationType +except ImportError as exc: + KernelBackendMoE = None + ActivationType = None + _SonicMoE = torch.nn.Module + _SONICMOE_IMPORT_ERROR = exc +else: + _SONICMOE_IMPORT_ERROR = None + + +def _require_sonicmoe() -> None: + if _SONICMOE_IMPORT_ERROR is not None: + raise ImportError( + "SonicMoELayer requires the optional sonic-moe package. " + "Install it on the training machine" + ) from _SONICMOE_IMPORT_ERROR + + +class _MegatronMoE(_SonicMoE): + """Sonic MoE variant that follows Megatron gradient-accumulation fusion.""" + + def __init__( + self, + num_experts: int, + num_experts_per_tok: int, + hidden_size: int, + intermediate_size: int, + activation_function, + add_bias: bool, + std: float, + router_score_function: str = "softmax", + router_score_over_topk: bool = True, + accumulate_wgrad_into_main_grad: bool = False, + ) -> None: + _require_sonicmoe() + super().__init__( + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + activation_function=activation_function, + add_bias=add_bias, + std=std, + router_score_function=router_score_function, + router_score_over_topk=router_score_over_topk, + ) + self.accumulate_wgrad_into_main_grad = accumulate_wgrad_into_main_grad + + +def _import_sonicmoe_functional(): + try: + from sonicmoe.enums import ActivationType, is_glu + from sonicmoe.functional import ( + TC_Softmax_Topk_Router_Function, + TC_topk_router_metadata_triton, + _DownProjection, + _UpProjection, + ) + except ImportError as exc: + raise ImportError( + "SonicMoELayer requires sonic-moe functional kernels from the optional " + "sonic-moe package." + ) from exc + return ( + ActivationType, + is_glu, + TC_Softmax_Topk_Router_Function, + TC_topk_router_metadata_triton, + _DownProjection, + _UpProjection, + ) + + +def _sonic_activation_type(config: TransformerConfig): + _require_sonicmoe() + if not config.gated_linear_unit: + raise ValueError("SonicMoELayer only supports gated MoE MLPs.") + activation_name = getattr(config.activation_func, "__name__", None) + if config.activation_func is F.silu or activation_name == "silu": + return ActivationType.SWIGLU + if config.activation_func is F.gelu or activation_name == "gelu": + return ActivationType.GEGLU + raise ValueError( + "SonicMoELayer only supports SonicMoE GLU activations SwiGLU and GEGLU. " + f"Got activation_func={config.activation_func}." + ) + + +def _as_list(value: Union[str, list]) -> list: + return value if isinstance(value, list) else [value] + + +def _loss_coeff(config: TransformerConfig, loss_type: str) -> float: + routing_types = _as_list(config.moe_router_load_balancing_type) + aux_coeffs = _as_list(config.moe_aux_loss_coeff) + if len(aux_coeffs) == 1 and len(routing_types) > 1: + aux_coeffs = aux_coeffs * len(routing_types) + for routing_type, coeff in zip(routing_types, aux_coeffs): + if routing_type == loss_type: + return float(coeff) + return 0.0 + + +def _has_positive_unsupported_aux_loss(config: TransformerConfig) -> bool: + routing_types = _as_list(config.moe_router_load_balancing_type) + aux_coeffs = _as_list(config.moe_aux_loss_coeff) + if len(aux_coeffs) == 1 and len(routing_types) > 1: + aux_coeffs = aux_coeffs * len(routing_types) + for routing_type, coeff in zip(routing_types, aux_coeffs): + if routing_type not in ("aux_loss", "seq_aux_loss", "global_aux_loss", "none") and float(coeff) > 0.0: + return True + return False + + +def _get_tokens_per_expert_and_token_count( + routing_map: torch.Tensor, + reduce_group: torch.distributed.ProcessGroup, + topk: int = None, + with_padding_mask: bool = False, +): + """Target-local copy of the newer MoE aux-loss token-count helper.""" + local_tokens_per_expert = routing_map.sum(dim=0) + global_tokens_per_expert = local_tokens_per_expert + group_size = reduce_group.size() if reduce_group is not None else 1 + if group_size > 1: + global_tokens_per_expert = local_tokens_per_expert.clone() + torch.distributed.all_reduce(global_tokens_per_expert, group=reduce_group) + + if with_padding_mask: + local_num_tokens = local_tokens_per_expert.sum() / topk + total_num_tokens = global_tokens_per_expert.sum() / topk + else: + local_num_tokens = routing_map.shape[0] + total_num_tokens = local_num_tokens * group_size + return global_tokens_per_expert, local_num_tokens, total_num_tokens + + +def _check_supported_config(config: TransformerConfig) -> None: + if config.num_moe_experts is None: + raise ValueError("SonicMoELayer requires config.num_moe_experts.") + if config.tensor_model_parallel_size != 1: + raise ValueError("SonicMoELayer does not integrate tensor parallelism yet.") + if config.expert_model_parallel_size != 1: + raise ValueError("SonicMoELayer does not integrate expert parallelism yet.") + if config.expert_tensor_parallel_size != 1: + raise ValueError("SonicMoELayer does not integrate expert tensor parallelism yet.") + if getattr(config, "moe_latent_size", None) is not None: + raise ValueError("SonicMoELayer does not support MoE latent projections.") + if config.moe_shared_expert_intermediate_size is not None: + raise ValueError("SonicMoELayer does not support shared experts.") + if config.overlap_moe_expert_parallel_comm: + raise ValueError("SonicMoELayer does not support EP communication overlap.") + if config.fp8 or config.fp4: + raise ValueError("SonicMoELayer does not support fp8/fp4 expert compute.") + if config.moe_expert_capacity_factor is not None: + raise ValueError("SonicMoELayer does not support token dropping or expert capacity.") + if config.moe_router_padding_for_quantization: + raise ValueError("SonicMoELayer does not support router padding for quantization.") + if config.moe_router_score_function not in ("softmax", "sigmoid"): + raise ValueError("SonicMoELayer pure Sonic router only supports softmax/sigmoid routing.") + if config.moe_router_num_groups is not None or config.moe_router_group_topk is not None: + raise ValueError("SonicMoELayer does not support group-limited routing.") + if config.moe_router_enable_expert_bias: + raise ValueError("SonicMoELayer does not support Megatron expert-bias routing.") + if config.moe_router_force_load_balancing or getattr( + config, "moe_router_force_biased", None + ) is not None: + raise ValueError("SonicMoELayer does not support forced benchmark routing.") + if config.moe_input_jitter_eps is not None: + raise ValueError("SonicMoELayer does not support Megatron router input jitter.") + if _has_positive_unsupported_aux_loss(config): + raise ValueError( + "SonicMoELayer only supports aux_loss, seq_aux_loss, global_aux_loss, or no aux loss." + ) + if getattr(config, "glu_linear_offset", 0.0) != 0.0: + raise ValueError("SonicMoELayer does not support nonzero glu_linear_offset.") + if getattr(config, "activation_func_clamp_value", None) is not None: + raise ValueError("SonicMoELayer does not support activation_func_clamp_value.") + _sonic_activation_type(config) + + +def _set_sonic_param_dtypes(module: torch.nn.Module, config: TransformerConfig) -> None: + module.router.to(dtype=torch.float32) + module.c_fc.to(dtype=config.params_dtype) + module.c_proj.to(dtype=config.params_dtype) + + +def _maybe_move_to_runtime_device(module: torch.nn.Module, config: TransformerConfig) -> None: + if not config.use_cpu_initialization and torch.cuda.is_available(): + module.to(device=torch.cuda.current_device()) + _set_sonic_param_dtypes(module, config) + + +class _SonicParamSync(torch.nn.Module): + """Expose SonicMoE's directly-read params to Megatron DDP pre-hooks. + + SonicMoE's fast path does not call the ``router``, ``c_fc``, or ``c_proj`` + modules; it reads their parameters directly. Megatron's overlapped + distributed optimizer waits for param all-gathers in module forward + pre-hooks, so this no-op module is called immediately before SonicMoE. + """ + + def __init__(self, sonic_moe: torch.nn.Module) -> None: + super().__init__() + self.register_parameter("router_weight", sonic_moe.router.weight) + self.register_parameter("c_fc_weight", sonic_moe.c_fc.weight) + self.register_parameter("c_proj_weight", sonic_moe.c_proj.weight) + if sonic_moe.c_fc.bias is not None: + self.register_parameter("c_fc_bias", sonic_moe.c_fc.bias) + if sonic_moe.c_proj.bias is not None: + self.register_parameter("c_proj_bias", sonic_moe.c_proj.bias) + + def forward(self) -> None: + return None + + def _save_to_state_dict(self, destination, prefix, keep_vars): + del destination, prefix, keep_vars + + # pylint: disable=arguments-differ + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) -> None: + del state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + del prefix, sharded_offsets, metadata + return {} + + +class SonicMoELayer(MoELayer): + """Megatron-compatible wrapper around pure ``sonicmoe.MoE``. + + This class subclasses ``MoELayer`` so ``TransformerLayer`` treats it as a + MoE block and forwards MoE-only kwargs. It deliberately bypasses Megatron + token dispatchers and therefore requires TP/EP/ETP sizes to be 1. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: Optional[MoESubmodules] = None, + layer_number: Optional[int] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + ) -> None: + del submodules + _check_supported_config(config) + if pg_collection is None: + pg_collection = get_default_pg_collection() + + BaseMoELayer.__init__( + self, config=config, layer_number=layer_number, pg_collection=pg_collection + ) + self.tp_group = pg_collection.tp + self.tp_cp_group = pg_collection.tp_cp + self.tp_dp_cp_group = pg_collection.tp_dp_cp + self.layer_number = layer_number + self.is_mtp = False + self._aux_loss_coeff = _loss_coeff(config, "aux_loss") + self._seq_aux_loss_coeff = _loss_coeff(config, "seq_aux_loss") + self._global_aux_loss_coeff = _loss_coeff(config, "global_aux_loss") + self._z_loss_coeff = config.moe_z_loss_coeff + + if self._global_aux_loss_coeff > 0.0: + device = torch.cuda.current_device() if torch.cuda.is_available() else None + self.register_buffer( + "global_tokens_per_expert", + torch.zeros(config.num_moe_experts, dtype=torch.float32, device=device), + persistent=False, + ) + self.register_buffer( + "ga_steps", + torch.tensor(0, dtype=torch.float32, device=device), + persistent=False, + ) + else: + self.global_tokens_per_expert = None + self.ga_steps = None + + _require_sonicmoe() + self.kernel_backend_moe = KernelBackendMoE.sonicmoe + self.sonic_moe = _MegatronMoE( + num_experts=config.num_moe_experts, + num_experts_per_tok=config.moe_router_topk, + hidden_size=config.hidden_size, + intermediate_size=config.moe_ffn_hidden_size, + activation_function=_sonic_activation_type(config), + add_bias=config.add_bias_linear, + std=config.init_method_std, + router_score_function=config.moe_router_score_function, + router_score_over_topk=self._score_over_topk(), + accumulate_wgrad_into_main_grad=config.gradient_accumulation_fusion, + ) + _maybe_move_to_runtime_device(self.sonic_moe, config) + self.sonic_param_sync = _SonicParamSync(self.sonic_moe) + + for param in self.sonic_moe.parameters(): + setattr(param, "allreduce", True) + + def set_layer_number(self, layer_number: int): + self.layer_number = layer_number + + def set_is_mtp(self): + self.is_mtp = True + + def reset_global_aux_loss_tracker(self): + if self.global_tokens_per_expert is not None: + self.global_tokens_per_expert.zero_() + self.ga_steps.zero_() + + def _num_layers_for_loss_tracker(self) -> int: + num_layers = self.config.num_layers + if self.config.mtp_num_layers is not None: + num_layers += self.config.mtp_num_layers + return num_layers + + def _attach_scaled_loss( + self, + output: torch.Tensor, + loss: torch.Tensor, + coeff: float, + name: str, + reduce_group: Optional[torch.distributed.ProcessGroup] = None, + reduce_group_has_dp: bool = False, + valid_token_count: Optional[Union[int, torch.Tensor]] = None, + ) -> torch.Tensor: + if coeff == 0.0: + return output + + save_to_aux_losses_tracker( + name, + loss / coeff, + self.layer_number, + self._num_layers_for_loss_tracker(), + reduce_group=reduce_group, + ) + if self.config.calculate_per_token_loss: + num_tokens = valid_token_count if valid_token_count is not None else output.shape[0] + loss = loss * num_tokens + return MoEAuxLossAutoScaler.apply(output, loss) + + def _needs_router_losses(self) -> bool: + return ( + self.training + and torch.is_grad_enabled() + and ( + self._aux_loss_coeff > 0.0 + or self._seq_aux_loss_coeff > 0.0 + or self._global_aux_loss_coeff > 0.0 + or (self._z_loss_coeff is not None and self._z_loss_coeff != 0.0) + ) + ) + + def _tokens_per_expert_and_count( + self, local_tokens_per_expert: torch.Tensor, reduce_group: torch.distributed.ProcessGroup + ): + global_tokens_per_expert = local_tokens_per_expert.float() + if reduce_group.size() > 1: + global_tokens_per_expert = global_tokens_per_expert.clone() + torch.distributed.all_reduce(global_tokens_per_expert, group=reduce_group) + local_num_tokens = local_tokens_per_expert.sum() / self.config.moe_router_topk + total_num_tokens = global_tokens_per_expert.sum() / self.config.moe_router_topk + return global_tokens_per_expert, local_num_tokens, total_num_tokens + + def _score_over_topk(self) -> bool: + return not self.config.moe_router_pre_softmax + + def _apply_sonic_router_topk( + self, + router_func, + router_logits: torch.Tensor, + num_experts: int, + ): + args = ( + router_logits, + num_experts, + self.config.moe_router_topk, + self._score_over_topk(), + False, + ) + if self.config.moe_router_score_function == "sigmoid": + return router_func.apply(*args, "sigmoid") + try: + return router_func.apply(*args, "softmax") + except TypeError: + return router_func.apply(*args) + + def _normalize_topk_scores(self, topk_scores: torch.Tensor) -> torch.Tensor: + if ( + self.config.moe_router_score_function == "sigmoid" + and self.config.moe_router_topk > 1 + ): + topk_scores = topk_scores / (topk_scores.sum(dim=-1, keepdim=True) + 1e-20) + if self.config.moe_router_topk_scaling_factor is not None: + topk_scores = topk_scores * self.config.moe_router_topk_scaling_factor + return topk_scores + + def _scores_for_aux_loss(self, router_logits: torch.Tensor) -> torch.Tensor: + if self.config.moe_router_score_function == "softmax": + return F.softmax(router_logits, dim=-1, dtype=torch.float32) + if self.config.moe_router_score_function == "sigmoid": + scores = torch.sigmoid(router_logits.float()) + return scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) + raise ValueError( + f"Invalid score_function: {self.config.moe_router_score_function}" + ) + + def _sonic_tc_forward( + self, + hidden_states: torch.Tensor, + is_inference_mode_enabled: bool = False, + ): + ( + ActivationType, + is_glu, + TC_Softmax_Topk_Router_Function, + TC_topk_router_metadata_triton, + _DownProjection, + _UpProjection, + ) = _import_sonicmoe_functional() + + original_shape = hidden_states.shape + x = hidden_states.view(-1, self.config.hidden_size) + router_logits = F.linear(x.float(), self.sonic_moe.router.weight) + num_experts = self.sonic_moe.router.weight.size(0) + topk_scores, topk_indices = self._apply_sonic_router_topk( + TC_Softmax_Topk_Router_Function, + router_logits, + num_experts, + ) + topk_scores = self._normalize_topk_scores(topk_scores) + + num_tokens, topk = topk_indices.size() + num_routed_tokens = num_tokens * topk + device = topk_indices.device + + s_scatter_idx = torch.empty(num_routed_tokens, dtype=torch.int32, device=device) + s_reverse_scatter_idx = torch.empty(num_routed_tokens, dtype=torch.int32, device=device) + expert_frequency = torch.empty(num_experts, dtype=torch.int32, device=device) + expert_frequency_offset = torch.empty(num_experts + 1, dtype=torch.int32, device=device) + x_gather_idx = torch.empty(num_routed_tokens, dtype=torch.int32, device=device) + + TC_topk_router_metadata_triton( + topk_indices, + num_experts, + expert_frequency, + expert_frequency_offset, + x_gather_idx, + s_scatter_idx, + s_reverse_scatter_idx, + ) + + activation_type = self.sonic_moe.activation_function + if type(activation_type) == str: + activation_type = ActivationType(activation_type) + + assert not torch.compiler.is_compiling() + assert is_glu(activation_type), "SonicMoELayer only supports GLU Sonic kernels." + + w1 = self.sonic_moe.c_fc.weight.permute(1, 2, 0) + w2 = self.sonic_moe.c_proj.weight.permute(1, 2, 0) + accumulate_wgrad_into_main_grad = getattr( + self.sonic_moe, "accumulate_wgrad_into_main_grad", False + ) + a, h = _UpProjection.apply( + x, + w1, + self.sonic_moe.c_fc.bias, + expert_frequency_offset, + num_routed_tokens, + topk, + x_gather_idx, + s_scatter_idx, + s_reverse_scatter_idx, + None, + False, + activation_type, + is_inference_mode_enabled, + True, + accumulate_wgrad_into_main_grad, + ) + + output = _DownProjection.apply( + a, + h, + w2, + self.sonic_moe.c_proj.bias, + topk_scores, + expert_frequency_offset, + num_tokens, + topk, + x_gather_idx, + s_scatter_idx, + s_reverse_scatter_idx, + None, + False, + activation_type, + accumulate_wgrad_into_main_grad, + ) + + return output.view(original_shape), router_logits, expert_frequency + + def _compute_seq_aux_inputs(self, router_logits: torch.Tensor): + routing_map, scores = compute_routing_scores_for_aux_loss( + router_logits, + self.config.moe_router_topk, + self.config.moe_router_score_function, + fused=False, + ) + return routing_map, scores + + def _apply_router_losses( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + tokens_per_expert: torch.Tensor, + ) -> torch.Tensor: + if not self.training or not torch.is_grad_enabled(): + return output + if ( + self._aux_loss_coeff == 0.0 + and self._seq_aux_loss_coeff == 0.0 + and self._global_aux_loss_coeff == 0.0 + and self._z_loss_coeff is None + ): + return output + + if self._z_loss_coeff is not None and self._z_loss_coeff != 0.0: + moe_z_loss_coeff = self._z_loss_coeff / self.tp_cp_group.size() + z_loss = z_loss_func(router_logits, moe_z_loss_coeff) + output = self._attach_scaled_loss( + output, + z_loss, + moe_z_loss_coeff, + "z_loss", + valid_token_count=router_logits.shape[0], + ) + + scores = None + if self._aux_loss_coeff > 0.0: + scores = self._scores_for_aux_loss(router_logits) + global_tokens_per_expert, local_num_tokens, total_num_tokens = ( + self._tokens_per_expert_and_count(tokens_per_expert, self.tp_cp_group) + ) + aux_loss = switch_load_balancing_loss_func( + probs=scores, + tokens_per_expert=global_tokens_per_expert, + total_num_tokens=total_num_tokens, + topk=self.config.moe_router_topk, + num_experts=self.config.num_moe_experts, + moe_aux_loss_coeff=self._aux_loss_coeff, + fused=False, + ) + output = self._attach_scaled_loss( + output, + aux_loss, + self._aux_loss_coeff, + "load_balancing_loss", + reduce_group=self.tp_cp_group, + valid_token_count=local_num_tokens, + ) + + if self._seq_aux_loss_coeff > 0.0: + if scores is None: + scores = self._scores_for_aux_loss(router_logits) + routing_map, _ = self._compute_seq_aux_inputs(router_logits) + if hidden_states.dim() >= 3: + seq_length, bsz = hidden_states.shape[0], hidden_states.shape[1] + else: + seq_length, bsz = hidden_states.shape[0], 1 + seq_scores = scores.reshape(seq_length, -1) + seq_routing_map = routing_map.reshape(seq_length, -1) + global_tokens_per_expert, local_num_tokens, total_num_tokens = ( + _get_tokens_per_expert_and_token_count( + routing_map=seq_routing_map, + reduce_group=self.tp_cp_group, + topk=self.config.moe_router_topk * bsz, + ) + ) + seq_aux_loss = ( + switch_load_balancing_loss_func( + probs=seq_scores, + tokens_per_expert=global_tokens_per_expert, + total_num_tokens=total_num_tokens, + topk=self.config.moe_router_topk, + num_experts=self.config.num_moe_experts, + moe_aux_loss_coeff=self._seq_aux_loss_coeff, + fused=False, + ) + / bsz + ) + output = self._attach_scaled_loss( + output, + seq_aux_loss, + self._seq_aux_loss_coeff, + "seq_load_balancing_loss", + reduce_group=self.tp_cp_group, + valid_token_count=local_num_tokens, + ) + + if self._global_aux_loss_coeff > 0.0: + if scores is None: + scores = self._scores_for_aux_loss(router_logits) + global_tokens_per_expert, local_num_tokens, total_num_tokens = ( + self._tokens_per_expert_and_count(tokens_per_expert, self.tp_dp_cp_group) + ) + self.global_tokens_per_expert += global_tokens_per_expert + self.ga_steps += 1 + averaged_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps + global_aux_loss = switch_load_balancing_loss_func( + probs=scores, + tokens_per_expert=averaged_tokens_per_expert, + total_num_tokens=total_num_tokens, + topk=self.config.moe_router_topk, + num_experts=self.config.num_moe_experts, + moe_aux_loss_coeff=self._global_aux_loss_coeff, + fused=False, + ) + output = self._attach_scaled_loss( + output, + global_aux_loss, + self._global_aux_loss_coeff, + "global_load_balancing_loss", + reduce_group=self.tp_dp_cp_group, + reduce_group_has_dp=True, + valid_token_count=local_num_tokens, + ) + + return output + + def forward( + self, + hidden_states: torch.Tensor, + intermediate_tensors=None, + padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + ): + del input_ids + if intermediate_tensors is not None: + raise ValueError("SonicMoELayer does not support partial MoE CUDA graph replay.") + if padding_mask is not None and torch.any(padding_mask): + raise ValueError("Pure SonicMoE routing does not support padding masks.") + + self.sonic_param_sync() + need_router_losses = self._needs_router_losses() + output, router_logits, tokens_per_expert = self._sonic_tc_forward( + hidden_states, + is_inference_mode_enabled=(not self.training), + ) + if need_router_losses: + output = self._apply_router_losses( + output, hidden_states, router_logits, tokens_per_expert + ) + return output, None + + def backward_dw(self, routed_experts: bool = True, shared_experts: bool = False): + del routed_experts, shared_experts + + def set_for_recompute_pre_mlp_layernorm(self): + raise ValueError( + "SonicMoELayer does not support fp8/fp4 pre-MLP layernorm recompute." + ) + + def _state_tensor(self, tensor: torch.Tensor, keep_vars: bool) -> torch.Tensor: + return tensor if keep_vars else tensor.detach() + + def _origin_grouped_fc1_weight(self, keep_vars: bool) -> torch.Tensor: + weight = self._state_tensor(self.sonic_moe.c_fc.weight, keep_vars) + return ( + weight.transpose(1, 2) + .contiguous() + .view( + self.config.hidden_size, + self.config.num_moe_experts * 2 * self.config.moe_ffn_hidden_size, + ) + ) + + def _origin_grouped_fc2_weight(self, keep_vars: bool) -> torch.Tensor: + weight = self._state_tensor(self.sonic_moe.c_proj.weight, keep_vars) + return ( + weight.transpose(1, 2) + .contiguous() + .view( + self.config.num_moe_experts * self.config.moe_ffn_hidden_size, + self.config.hidden_size, + ) + ) + + def _origin_state_dict_entries(self, prefix: str, keep_vars: bool) -> dict: + entries = { + f"{prefix}router.weight": self._state_tensor( + self.sonic_moe.router.weight, keep_vars + ), + f"{prefix}experts.weight1": self._origin_grouped_fc1_weight(keep_vars), + f"{prefix}experts.weight2": self._origin_grouped_fc2_weight(keep_vars), + } + if self.sonic_moe.c_fc.bias is not None: + entries[f"{prefix}experts.linear_fc1.bias"] = self._state_tensor( + self.sonic_moe.c_fc.bias, keep_vars + ) + if self.sonic_moe.c_proj.bias is not None: + entries[f"{prefix}experts.linear_fc2.bias"] = self._state_tensor( + self.sonic_moe.c_proj.bias, keep_vars + ) + return entries + + def _remove_sonic_state_dict_entries(self, state_dict, prefix: str) -> None: + for key in ( + "sonic_moe.router.weight", + "sonic_moe.c_fc.weight", + "sonic_moe.c_fc.bias", + "sonic_moe.c_proj.weight", + "sonic_moe.c_proj.bias", + "sonic_param_sync.router_weight", + "sonic_param_sync.c_fc_weight", + "sonic_param_sync.c_fc_bias", + "sonic_param_sync.c_proj_weight", + "sonic_param_sync.c_proj_bias", + ): + state_dict.pop(f"{prefix}{key}", None) + + # pylint: disable=arguments-differ + def state_dict(self, *args, destination=None, prefix="", keep_vars=False): + state_dict = super().state_dict( + *args, destination=destination, prefix=prefix, keep_vars=keep_vars + ) + self._remove_sonic_state_dict_entries(state_dict, prefix) + state_dict.update(self._origin_state_dict_entries(prefix, keep_vars)) + return state_dict + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + metadata = ensure_metadata_has_dp_cp_group(metadata) + return make_sharded_tensors_for_checkpoint( + self._origin_state_dict_entries("", keep_vars=True), + prefix, + {}, + sharded_offsets=sharded_offsets, + tp_group=self.tp_group, + dp_cp_group=metadata["dp_cp_group"], + ) + + def _move_if_present(self, state_dict, src_key: str, dst_key: str) -> None: + if dst_key not in state_dict and src_key in state_dict: + state_dict[dst_key] = state_dict.pop(src_key) + + def _stack_per_expert_if_present( + self, state_dict, prefix: str, src_prefix: str, dst_key: str + ) -> None: + if dst_key in state_dict: + return + keys = [ + f"{prefix}{src_prefix}{expert_idx}" + for expert_idx in range(self.config.num_moe_experts) + ] + if all(key in state_dict for key in keys): + state_dict[dst_key] = torch.stack([state_dict.pop(key) for key in keys], dim=0) + + def _merge_legacy_grouped_mlp_state_dict(self, state_dict, prefix: str) -> None: + weight1_key = f"{prefix}experts.weight1" + weight2_key = f"{prefix}experts.weight2" + if f"{prefix}sonic_moe.c_fc.weight" not in state_dict and weight1_key in state_dict: + weight1 = state_dict.pop(weight1_key) + state_dict[f"{prefix}sonic_moe.c_fc.weight"] = ( + weight1.view( + self.config.num_moe_experts, + self.config.hidden_size, + 2 * self.config.moe_ffn_hidden_size, + ) + .transpose(1, 2) + .contiguous() + ) + if f"{prefix}sonic_moe.c_proj.weight" not in state_dict and weight2_key in state_dict: + weight2 = state_dict.pop(weight2_key) + state_dict[f"{prefix}sonic_moe.c_proj.weight"] = ( + weight2.view( + self.config.num_moe_experts, + self.config.moe_ffn_hidden_size, + self.config.hidden_size, + ) + .transpose(1, 2) + .contiguous() + ) + + # pylint: disable=arguments-differ + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) -> None: + self._move_if_present( + state_dict, f"{prefix}router.weight", f"{prefix}sonic_moe.router.weight" + ) + state_dict.pop(f"{prefix}router.bias", None) + + self._move_if_present( + state_dict, f"{prefix}experts.linear_fc1.weight", f"{prefix}sonic_moe.c_fc.weight" + ) + self._move_if_present( + state_dict, f"{prefix}experts.linear_fc1.bias", f"{prefix}sonic_moe.c_fc.bias" + ) + self._move_if_present( + state_dict, f"{prefix}experts.linear_fc2.weight", f"{prefix}sonic_moe.c_proj.weight" + ) + self._move_if_present( + state_dict, f"{prefix}experts.linear_fc2.bias", f"{prefix}sonic_moe.c_proj.bias" + ) + self._stack_per_expert_if_present( + state_dict, + prefix, + "experts.linear_fc1.weight", + f"{prefix}sonic_moe.c_fc.weight", + ) + self._stack_per_expert_if_present( + state_dict, + prefix, + "experts.linear_fc1.bias", + f"{prefix}sonic_moe.c_fc.bias", + ) + self._stack_per_expert_if_present( + state_dict, + prefix, + "experts.linear_fc2.weight", + f"{prefix}sonic_moe.c_proj.weight", + ) + self._stack_per_expert_if_present( + state_dict, + prefix, + "experts.linear_fc2.bias", + f"{prefix}sonic_moe.c_proj.bias", + ) + self._merge_legacy_grouped_mlp_state_dict(state_dict, prefix) + + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + +def get_sonic_moe_module_spec() -> ModuleSpec: + """Return an MLP/MoE ModuleSpec for use in a TransformerLayer spec.""" + return ModuleSpec( + module=SonicMoELayer, + metainfo={"fuse_pre_mlp_layernorm": False}, + ) + + +def replace_moe_layer_specs_with_sonic_moe(transformer_layer_spec) -> int: + """Replace Megatron MoELayer MLP specs with SonicMoELayer specs in-place. + + Accepts a TransformerBlockSubmodules object, a TransformerLayer ModuleSpec, + or a list/tuple of layer specs. Returns the number of MoE MLP specs that + use SonicMoELayer after the replacement. + """ + if transformer_layer_spec is None: + return 0 + if isinstance(transformer_layer_spec, (list, tuple)): + return sum( + replace_moe_layer_specs_with_sonic_moe(spec) for spec in transformer_layer_spec + ) + + layer_specs = getattr(transformer_layer_spec, "layer_specs", None) + if layer_specs is not None: + return replace_moe_layer_specs_with_sonic_moe(layer_specs) + + submodules = getattr(transformer_layer_spec, "submodules", None) + if submodules is None: + return 0 + + mlp_spec = getattr(submodules, "mlp", None) + module = getattr(mlp_spec, "module", None) + if module is SonicMoELayer: + return 1 + if module is not MoELayer: + return 0 + + submodules.mlp = get_sonic_moe_module_spec() + return 1 diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 7d45337029a..70d9a97643f 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -69,6 +69,7 @@ def add_megatron_arguments(parser: argparse.ArgumentParser): parser = _add_biencoder_args(parser) parser = _add_vision_args(parser) parser = _add_moe_args(parser) + parser = _add_sonic_moe_args(parser) parser = _add_mla_args(parser) parser = _add_experimental_attention_variant_args(parser) parser = _add_heterogeneous_args(parser) @@ -1092,6 +1093,53 @@ def validate_args(args, defaults={}): if isinstance(args.moe_aux_loss_coeff, list) and len(args.moe_aux_loss_coeff) == 1: args.moe_aux_loss_coeff = args.moe_aux_loss_coeff[0] + if args.use_sonic_moe: + assert args.num_experts is not None, "--use-sonic-moe requires --num-experts." + assert not args.use_legacy_models, "--use-sonic-moe is only supported for mcore models." + assert ( + args.tensor_model_parallel_size == 1 + ), "--use-sonic-moe does not support tensor model parallelism yet." + assert ( + args.expert_model_parallel_size == 1 + ), "--use-sonic-moe does not support expert model parallelism yet." + assert ( + args.expert_tensor_parallel_size == 1 + ), "--use-sonic-moe does not support expert tensor parallelism yet." + assert not getattr(args, "overlap_moe_expert_parallel_comm", False), ( + "--use-sonic-moe does not support --overlap-moe-expert-parallel-comm." + ) + assert getattr(args, "moe_shared_expert_intermediate_size", None) is None, ( + "--use-sonic-moe does not support shared experts." + ) + assert getattr(args, "moe_latent_size", None) is None, ( + "--use-sonic-moe does not support MoE latent projections." + ) + assert getattr(args, "moe_expert_capacity_factor", None) is None, ( + "--use-sonic-moe does not support token dropping or expert capacity." + ) + assert not getattr(args, "moe_router_padding_for_quantization", False), ( + "--use-sonic-moe does not support router padding for quantization." + ) + assert getattr(args, "moe_router_score_function", "softmax") in ( + "softmax", + "sigmoid", + ), "--use-sonic-moe only supports softmax or sigmoid routing." + assert getattr(args, "moe_router_num_groups", None) is None, ( + "--use-sonic-moe does not support group-limited routing." + ) + assert getattr(args, "moe_router_group_topk", None) is None, ( + "--use-sonic-moe does not support group-limited routing." + ) + assert not getattr(args, "moe_router_enable_expert_bias", False), ( + "--use-sonic-moe does not support expert-bias routing." + ) + assert getattr(args, "moe_input_jitter_eps", None) is None, ( + "--use-sonic-moe does not support router input jitter." + ) + assert getattr(args, "fp8", None) is None and not getattr(args, "fp4", False), ( + "--use-sonic-moe does not support fp8/fp4 expert compute." + ) + # Distributed checkpointing checks if args.use_dist_ckpt and args.use_legacy_models: raise RuntimeError('--use-dist-ckpt is not supported in legacy models.') @@ -3321,6 +3369,14 @@ def _add_moe_args(parser): " max that an expert could see during inference so no tokens are actually dropped.") return parser +def _add_sonic_moe_args(parser): + group = parser.add_argument_group(title="sonic moe") + group.add_argument('--use-sonic-moe', action='store_true', default=False, + help='Replace mcore MoE MLP specs with SonicMoELayer. ' + 'This requires --num-experts and currently supports only TP=EP=ETP=1; ' + 'checkpoint load/save remains in the standard Megatron MoE format.') + return parser + def _add_mla_args(parser): group = parser.add_argument_group(title="mla") group.add_argument('--q-lora-rank', type=int, default=None, diff --git a/tests/unit_tests/transformer/moe/test_sonic_moe_layer.py b/tests/unit_tests/transformer/moe/test_sonic_moe_layer.py new file mode 100644 index 00000000000..0598007b033 --- /dev/null +++ b/tests/unit_tests/transformer/moe/test_sonic_moe_layer.py @@ -0,0 +1,662 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import dataclasses + +import pytest +import torch +import torch.nn.functional as F + +from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec +from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec +from megatron.core.transformer.mlp import MLP +from megatron.core.transformer.moe.moe_utils import ( + clear_aux_losses_tracker, + get_default_pg_collection, + get_moe_layer_wise_logging_tracker, +) +from megatron.core.transformer.moe.moe_layer import MoELayer +from megatron.core.transformer.moe.router import TopKRouter +from megatron.core.transformer.moe.sonic_moe_layer import ( + SonicMoELayer, + replace_moe_layer_specs_with_sonic_moe, +) +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.training.initialize import _set_random_seed +from tests.unit_tests.test_utilities import Utils + + +pytestmark = [ + pytest.mark.internal, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available"), +] + +EXACT_RTOL = 0.0 +EXACT_ATOL = 0.0 +ROUTER_KERNEL_RTOL = 1.0e-4 +ROUTER_KERNEL_ATOL = 3.0e-5 +FORWARD_RTOL = 1.0e-2 +FORWARD_ATOL = 1.0e-5 + + +def _require_sonicmoe(): + pytest.importorskip("sonicmoe") + + +class TestSonicMoELayerRouterLoss: + def setup_method(self, method): + _require_sonicmoe() + Utils.initialize_model_parallel(1, 1) + _set_random_seed(seed_=123, data_parallel_random_init=False) + clear_aux_losses_tracker() + + self.default_config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + ffn_hidden_size=32, + num_moe_experts=8, + moe_ffn_hidden_size=8, + moe_router_topk=2, + moe_router_load_balancing_type="global_aux_loss", + moe_aux_loss_coeff=0.7, + moe_z_loss_coeff=0.3, + moe_router_score_function="softmax", + moe_router_dtype="fp32", + add_bias_linear=False, + gated_linear_unit=True, + activation_func=F.silu, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + ) + + def teardown_method(self, method): + clear_aux_losses_tracker() + Utils.destroy_model_parallel() + + def _new_topk_router(self, config): + router = TopKRouter(config=config, pg_collection=get_default_pg_collection()).cuda() + router.set_layer_number(0) + router.weight.data = router.weight.data.float() + return router + + def _new_sonic_layer(self, config): + layer = SonicMoELayer(config=config, pg_collection=get_default_pg_collection()).cuda() + layer.set_layer_number(0) + return layer + + def _new_original_sequential_moe_layer(self, config): + spec = get_moe_module_spec( + use_te=False, + num_experts=config.num_moe_experts, + moe_grouped_gemm=False, + ) + layer = MoELayer( + config=config, + submodules=spec.submodules, + pg_collection=get_default_pg_collection(), + ).cuda() + layer.set_layer_number(0) + return layer + + def test_replace_moe_layer_specs_with_sonic_moe(self): + config = dataclasses.replace( + self.default_config, + num_layers=2, + moe_layer_freq=[1, 0], + ) + + block_spec = get_gpt_decoder_block_spec( + config=config, + use_transformer_engine=False, + normalization=config.normalization, + ) + + assert block_spec.layer_specs[0].submodules.mlp.module is MoELayer + assert block_spec.layer_specs[1].submodules.mlp.module is MLP + assert replace_moe_layer_specs_with_sonic_moe(block_spec) == 1 + assert block_spec.layer_specs[0].submodules.mlp.module is SonicMoELayer + assert block_spec.layer_specs[1].submodules.mlp.module is MLP + assert replace_moe_layer_specs_with_sonic_moe(block_spec) == 1 + + def test_sonic_param_and_main_grad_dtypes(self): + layer = self._new_sonic_layer(self.default_config) + + assert type(layer.sonic_moe).__name__ == "_MegatronMoE" + assert ( + type(layer.sonic_moe).__module__ + == "megatron.core.transformer.moe.sonic_moe_layer" + ) + assert ( + layer.sonic_moe.accumulate_wgrad_into_main_grad + is self.default_config.gradient_accumulation_fusion + ) + assert layer.sonic_moe.router.weight.dtype is torch.float32 + assert layer.sonic_moe.c_fc.weight.dtype is torch.bfloat16 + assert layer.sonic_moe.c_proj.weight.dtype is torch.bfloat16 + if layer.sonic_moe.c_fc.bias is not None: + assert layer.sonic_moe.c_fc.bias.dtype is torch.bfloat16 + if layer.sonic_moe.c_proj.bias is not None: + assert layer.sonic_moe.c_proj.bias.dtype is torch.bfloat16 + + ddp = DistributedDataParallel( + self.default_config, + DistributedDataParallelConfig(grad_reduce_in_fp32=True), + module=layer, + disable_bucketing=True, + ) + for name, param in ddp.module.named_parameters(): + assert param.main_grad.dtype is torch.float32, name + + def test_megatron_moe_accumulates_expert_wgrad_into_main_grad(self): + config = dataclasses.replace( + self.default_config, + moe_z_loss_coeff=None, + moe_aux_loss_coeff=0.0, + gradient_accumulation_fusion=True, + ) + layer = self._new_sonic_layer(config) + ddp = DistributedDataParallel( + config, + DistributedDataParallelConfig(grad_reduce_in_fp32=True), + module=layer, + disable_bucketing=True, + ) + ddp.zero_grad_buffer() + hidden_state = torch.randn( + (8192, 1, config.hidden_size), + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + output, bias = ddp(hidden_state) + assert bias is None + output.float().square().mean().backward() + + for name, param in ( + ("c_fc", ddp.module.sonic_moe.c_fc.weight), + ("c_proj", ddp.module.sonic_moe.c_proj.weight), + ): + assert param.grad_added_to_main_grad is True, name + assert param.grad is None, name + assert param.main_grad.dtype is torch.float32, name + assert torch.count_nonzero(param.main_grad).item() > 0, name + + @pytest.mark.parametrize("moe_router_pre_softmax", [False, True]) + def test_softmax_router_flags_follow_megatron_config(self, moe_router_pre_softmax): + config = dataclasses.replace( + self.default_config, + moe_router_pre_softmax=moe_router_pre_softmax, + moe_router_topk_scaling_factor=1.5 if moe_router_pre_softmax else None, + moe_z_loss_coeff=None, + moe_aux_loss_coeff=0.0, + ) + layer = self._new_sonic_layer(config) + + assert layer._score_over_topk() is (not moe_router_pre_softmax) + assert layer.sonic_moe.router_score_function == "softmax" + assert layer.sonic_moe.router_score_over_topk is (not moe_router_pre_softmax) + + @pytest.mark.parametrize("moe_router_pre_softmax", [False, True]) + def test_sigmoid_router_flags_follow_megatron_config(self, moe_router_pre_softmax): + config = dataclasses.replace( + self.default_config, + moe_router_score_function="sigmoid", + moe_router_pre_softmax=moe_router_pre_softmax, + moe_router_topk_scaling_factor=1.5, + moe_z_loss_coeff=None, + moe_aux_loss_coeff=0.0, + ) + layer = self._new_sonic_layer(config) + + assert layer.sonic_moe.router_score_function == "sigmoid" + assert layer.sonic_moe.router_score_over_topk is (not moe_router_pre_softmax) + assert layer._score_over_topk() is (not moe_router_pre_softmax) + + @pytest.mark.parametrize( + "score_over_topk", + [False, True], + ids=["topk_over_sigmoid", "sigmoid_over_topk"], + ) + @pytest.mark.parametrize("num_experts, topk", [(8, 2), (128, 8)]) + def test_sigmoid_router_kernel_matches_reference_pre_and_post_topk( + self, score_over_topk, num_experts, topk + ): + from sonicmoe.functional import TC_Softmax_Topk_Router_Function + + config = dataclasses.replace( + self.default_config, + num_moe_experts=num_experts, + moe_router_topk=topk, + moe_router_score_function="sigmoid", + moe_router_pre_softmax=not score_over_topk, + moe_router_topk_scaling_factor=None, + ) + router_logits = self._make_margin_separated_router_logits( + 8192, config.num_moe_experts + ) + + topk_scores, topk_indices = TC_Softmax_Topk_Router_Function.apply( + router_logits, + config.num_moe_experts, + config.moe_router_topk, + score_over_topk, + False, + "sigmoid", + ) + + ref_logits = router_logits.detach().clone().requires_grad_(True) + if score_over_topk: + ref_topk = ref_logits.topk(config.moe_router_topk, dim=-1) + ref_scores = ref_topk.values.sigmoid() + else: + ref_topk = ref_logits.sigmoid().topk(config.moe_router_topk, dim=-1) + ref_scores = ref_topk.values + + assert torch.equal(topk_indices, ref_topk.indices.to(dtype=topk_indices.dtype)) + torch.testing.assert_close( + topk_scores, ref_scores, rtol=ROUTER_KERNEL_RTOL, atol=ROUTER_KERNEL_ATOL + ) + + grad = torch.randn_like(topk_scores) + topk_scores.backward(grad) + ref_scores.backward(grad) + torch.testing.assert_close( + router_logits.grad, + ref_logits.grad, + rtol=ROUTER_KERNEL_RTOL, + atol=ROUTER_KERNEL_ATOL, + ) + + @staticmethod + def _make_margin_separated_router_logits(num_tokens: int, num_experts: int): + expert_scores = torch.linspace(-2.0, 2.0, num_experts, device="cuda") + row_offsets = torch.arange(num_tokens, device="cuda").unsqueeze(1) + expert_offsets = torch.arange(num_experts, device="cuda").unsqueeze(0) + rotated_experts = (expert_offsets + row_offsets) % num_experts + logits = expert_scores[rotated_experts] + logits = logits + 1.0e-4 * torch.randn_like(logits) + return logits.requires_grad_(True) + + @pytest.mark.parametrize("moe_router_score_function", ["softmax", "sigmoid"]) + @pytest.mark.parametrize("moe_router_pre_softmax", [False, True]) + @pytest.mark.parametrize("num_accumulation_steps", [1, 2, 4, 8]) + def test_z_loss_and_global_aux_loss_match_original_router( + self, num_accumulation_steps, moe_router_pre_softmax, moe_router_score_function + ): + config = dataclasses.replace( + self.default_config, + moe_router_score_function=moe_router_score_function, + moe_router_pre_softmax=moe_router_pre_softmax, + moe_router_topk_scaling_factor=1.5 if moe_router_pre_softmax else None, + ) + router = self._new_topk_router(config) + sonic_layer = self._new_sonic_layer(config) + with torch.no_grad(): + sonic_layer.sonic_moe.router.weight.copy_(router.weight) + + hidden_states = [ + torch.randn((8192, 1, config.hidden_size), device="cuda", dtype=torch.bfloat16) + for _ in range(num_accumulation_steps) + ] + + for hidden_state in hidden_states: + ref = self._run_original_router(router, hidden_state) + sonic = self._run_sonic_loss_path(sonic_layer, hidden_state) + + torch.testing.assert_close( + sonic["z_loss"], ref["z_loss"], rtol=EXACT_RTOL, atol=EXACT_ATOL + ) + torch.testing.assert_close( + sonic["global_aux_loss"], + ref["global_aux_loss"], + rtol=EXACT_RTOL, + atol=EXACT_ATOL, + ) + assert sonic["router_weight_grad"].dtype is torch.float32 + torch.testing.assert_close( + sonic["router_weight_grad"], + ref["router_weight_grad"], + rtol=EXACT_RTOL, + atol=EXACT_ATOL, + ) + torch.testing.assert_close( + sonic["hidden_state_grad"], + ref["hidden_state_grad"], + rtol=EXACT_RTOL, + atol=EXACT_ATOL, + ) + torch.testing.assert_close( + sonic_layer.global_tokens_per_expert, + router.global_tokens_per_expert, + rtol=EXACT_RTOL, + atol=EXACT_ATOL, + ) + torch.testing.assert_close( + sonic_layer.ga_steps, router.ga_steps, rtol=EXACT_RTOL, atol=EXACT_ATOL + ) + + sonic_layer.reset_global_aux_loss_tracker() + router.reset_global_aux_loss_tracker() + torch.testing.assert_close( + sonic_layer.global_tokens_per_expert, + router.global_tokens_per_expert, + rtol=EXACT_RTOL, + atol=EXACT_ATOL, + ) + torch.testing.assert_close( + sonic_layer.ga_steps, router.ga_steps, rtol=EXACT_RTOL, atol=EXACT_ATOL + ) + + def test_sigmoid_end_to_end_forward_backward(self): + config = dataclasses.replace( + self.default_config, + moe_router_score_function="sigmoid", + moe_router_pre_softmax=False, + moe_router_topk_scaling_factor=None, + ) + layer = self._new_sonic_layer(config) + hidden_state = torch.randn( + (8192, 1, config.hidden_size), + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + output, bias = layer(hidden_state) + assert bias is None + assert output.shape == hidden_state.shape + assert output.dtype is torch.bfloat16 + assert torch.isfinite(output).all() + + output.float().square().mean().backward() + + tracker = get_moe_layer_wise_logging_tracker() + assert "z_loss" in tracker + assert "global_load_balancing_loss" in tracker + assert hidden_state.grad.dtype is torch.bfloat16 + assert torch.isfinite(hidden_state.grad).all() + assert layer.sonic_moe.router.weight.grad.dtype is torch.float32 + assert torch.isfinite(layer.sonic_moe.router.weight.grad).all() + assert layer.sonic_moe.c_fc.weight.grad.dtype is torch.bfloat16 + assert torch.isfinite(layer.sonic_moe.c_fc.weight.grad).all() + assert layer.sonic_moe.c_proj.weight.grad.dtype is torch.bfloat16 + assert torch.isfinite(layer.sonic_moe.c_proj.weight.grad).all() + + @pytest.mark.parametrize("moe_router_score_function", ["softmax", "sigmoid"]) + @pytest.mark.parametrize("moe_router_pre_softmax", [False, True]) + def test_forward_matches_original_moe_layer( + self, moe_router_pre_softmax, moe_router_score_function + ): + config = dataclasses.replace( + self.default_config, + moe_router_load_balancing_type="none", + moe_aux_loss_coeff=0.0, + moe_z_loss_coeff=None, + moe_router_score_function=moe_router_score_function, + moe_router_pre_softmax=moe_router_pre_softmax, + moe_router_topk_scaling_factor=1.5 if moe_router_pre_softmax else None, + moe_token_dispatcher_type="allgather", + bias_activation_fusion=False, + ) + original_layer = self._new_original_sequential_moe_layer(config) + sonic_layer = self._new_sonic_layer(config) + origin_state = self._origin_grouped_moe_state_dict( + config, expert_weight_scale=config.init_method_std + ) + self._set_margin_separated_router_weight(origin_state["router.weight"], config) + self._load_origin_grouped_state_into_sequential_moe( + original_layer, origin_state, config + ) + load_result = sonic_layer.load_state_dict(origin_state, strict=True) + assert load_result.missing_keys == [] + assert load_result.unexpected_keys == [] + + original_layer.eval() + sonic_layer.eval() + hidden_state = self._make_forward_hidden_state(config) + + with torch.no_grad(): + original_output, original_bias = original_layer(hidden_state) + sonic_output, sonic_bias = sonic_layer(hidden_state) + + assert original_bias is None + assert sonic_bias is None + assert sonic_output.shape == original_output.shape == hidden_state.shape + assert sonic_output.dtype is torch.bfloat16 + torch.testing.assert_close( + sonic_output, + original_output, + rtol=FORWARD_RTOL, + atol=FORWARD_ATOL, + ) + + def test_load_origin_grouped_moe_state_and_save_origin_format(self): + config = dataclasses.replace( + self.default_config, + moe_z_loss_coeff=None, + moe_aux_loss_coeff=0.0, + ) + origin_state = self._origin_grouped_moe_state_dict(config) + layer = self._new_sonic_layer(config) + + load_result = layer.load_state_dict(origin_state, strict=True) + assert load_result.missing_keys == [] + assert load_result.unexpected_keys == [] + + expected_fc1 = origin_state["experts.weight1"].view( + config.num_moe_experts, + config.hidden_size, + 2 * config.moe_ffn_hidden_size, + ).transpose(1, 2).contiguous() + expected_fc2 = origin_state["experts.weight2"].view( + config.num_moe_experts, + config.moe_ffn_hidden_size, + config.hidden_size, + ).transpose(1, 2).contiguous() + torch.testing.assert_close( + layer.sonic_moe.router.weight, + origin_state["router.weight"], + rtol=EXACT_RTOL, + atol=EXACT_ATOL, + ) + torch.testing.assert_close( + layer.sonic_moe.c_fc.weight, expected_fc1, rtol=EXACT_RTOL, atol=EXACT_ATOL + ) + torch.testing.assert_close( + layer.sonic_moe.c_proj.weight, expected_fc2, rtol=EXACT_RTOL, atol=EXACT_ATOL + ) + + saved_state = layer.state_dict() + assert set(saved_state.keys()) == { + "router.weight", + "experts.weight1", + "experts.weight2", + } + assert not any(key.startswith("sonic_moe.") for key in saved_state) + assert not any(key.startswith("sonic_param_sync.") for key in saved_state) + torch.testing.assert_close( + saved_state["router.weight"], + origin_state["router.weight"], + rtol=EXACT_RTOL, + atol=EXACT_ATOL, + ) + torch.testing.assert_close( + saved_state["experts.weight1"], + origin_state["experts.weight1"], + rtol=EXACT_RTOL, + atol=EXACT_ATOL, + ) + torch.testing.assert_close( + saved_state["experts.weight2"], + origin_state["experts.weight2"], + rtol=EXACT_RTOL, + atol=EXACT_ATOL, + ) + + prefixed_state = layer.state_dict(prefix="mlp.") + assert set(prefixed_state.keys()) == { + "mlp.router.weight", + "mlp.experts.weight1", + "mlp.experts.weight2", + } + torch.testing.assert_close( + prefixed_state["mlp.experts.weight1"], + origin_state["experts.weight1"], + rtol=EXACT_RTOL, + atol=EXACT_ATOL, + ) + + sharded_state = layer.sharded_state_dict(prefix="mlp.") + assert set(sharded_state.keys()) == { + "mlp.router.weight", + "mlp.experts.weight1", + "mlp.experts.weight2", + } + + roundtrip_layer = self._new_sonic_layer(config) + roundtrip_load_result = roundtrip_layer.load_state_dict(saved_state, strict=True) + assert roundtrip_load_result.missing_keys == [] + assert roundtrip_load_result.unexpected_keys == [] + roundtrip_state = roundtrip_layer.state_dict() + assert set(roundtrip_state.keys()) == set(origin_state.keys()) + for key, tensor in saved_state.items(): + assert roundtrip_state[key].shape == tensor.shape + assert roundtrip_state[key].dtype == tensor.dtype + torch.testing.assert_close( + roundtrip_state[key], tensor, rtol=EXACT_RTOL, atol=EXACT_ATOL + ) + + @staticmethod + def _origin_grouped_moe_state_dict(config, expert_weight_scale: float = 1.0): + weight1_shape = ( + config.hidden_size, + config.num_moe_experts * 2 * config.moe_ffn_hidden_size, + ) + weight2_shape = ( + config.num_moe_experts * config.moe_ffn_hidden_size, + config.hidden_size, + ) + return { + "router.weight": torch.randn( + config.num_moe_experts, + config.hidden_size, + device="cuda", + dtype=torch.float32, + ), + "experts.weight1": torch.randn( + weight1_shape, device="cuda", dtype=config.params_dtype + ) + * expert_weight_scale, + "experts.weight2": torch.randn( + weight2_shape, device="cuda", dtype=config.params_dtype + ) + * expert_weight_scale, + } + + @staticmethod + def _set_margin_separated_router_weight(router_weight, config): + assert config.hidden_size >= config.num_moe_experts + router_weight.zero_() + router_weight[:, : config.num_moe_experts].copy_( + torch.eye( + config.num_moe_experts, + device=router_weight.device, + dtype=router_weight.dtype, + ) + ) + + @staticmethod + def _load_origin_grouped_state_into_sequential_moe(layer, origin_state, config): + with torch.no_grad(): + layer.router.weight.data = origin_state["router.weight"].detach().clone() + fc1_by_expert = origin_state["experts.weight1"].view( + config.num_moe_experts, + config.hidden_size, + 2 * config.moe_ffn_hidden_size, + ) + fc2_by_expert = origin_state["experts.weight2"].view( + config.num_moe_experts, + config.moe_ffn_hidden_size, + config.hidden_size, + ) + for expert_idx, expert in enumerate(layer.experts.local_experts): + expert.linear_fc1.weight.copy_( + fc1_by_expert[expert_idx].transpose(0, 1).contiguous() + ) + expert.linear_fc2.weight.copy_( + fc2_by_expert[expert_idx].transpose(0, 1).contiguous() + ) + + def _make_forward_hidden_state(self, config): + hidden_state = torch.randn( + (8192, 1, config.hidden_size), + device="cuda", + dtype=torch.bfloat16, + ) * 0.1 + router_logits = self._make_margin_separated_router_logits( + hidden_state.shape[0], config.num_moe_experts + ).detach() + hidden_state[:, 0, : config.num_moe_experts] = router_logits.to( + dtype=hidden_state.dtype + ) + return hidden_state + + def _run_original_router(self, router, hidden_state): + clear_aux_losses_tracker() + router.weight.grad = None + hidden_state = hidden_state.detach().clone().requires_grad_(True) + + scores, _ = router(hidden_state) + scores.backward(torch.zeros_like(scores)) + + tracker = get_moe_layer_wise_logging_tracker() + return { + "z_loss": tracker["z_loss"]["values"][0].detach().clone(), + "global_aux_loss": tracker["global_load_balancing_loss"]["values"][0] + .detach() + .clone(), + "router_weight_grad": router.weight.grad.detach().clone(), + "hidden_state_grad": hidden_state.grad.detach().clone(), + } + + def _run_sonic_loss_path(self, sonic_layer, hidden_state): + clear_aux_losses_tracker() + sonic_layer.sonic_moe.router.weight.grad = None + hidden_state = hidden_state.detach().clone().requires_grad_(True) + + router_logits = F.linear(hidden_state.float(), sonic_layer.sonic_moe.router.weight).view( + -1, sonic_layer.config.num_moe_experts + ) + tokens_per_expert = self._tokens_per_expert(router_logits, sonic_layer.config) + output = torch.zeros_like(hidden_state, requires_grad=True) + + output = sonic_layer._apply_router_losses( + output, hidden_state, router_logits, tokens_per_expert + ) + output.backward(torch.zeros_like(output)) + + tracker = get_moe_layer_wise_logging_tracker() + return { + "z_loss": tracker["z_loss"]["values"][0].detach().clone(), + "global_aux_loss": tracker["global_load_balancing_loss"]["values"][0] + .detach() + .clone(), + "router_weight_grad": sonic_layer.sonic_moe.router.weight.grad.detach().clone(), + "hidden_state_grad": hidden_state.grad.detach().clone(), + } + + @staticmethod + def _tokens_per_expert(router_logits, config): + if config.moe_router_score_function == "sigmoid": + routing_scores = torch.sigmoid(router_logits.float()) + elif config.moe_router_pre_softmax: + routing_scores = torch.softmax(router_logits, dim=-1, dtype=torch.float32) + else: + routing_scores = router_logits + selected_experts = routing_scores.topk(config.moe_router_topk, dim=-1).indices + return selected_experts.flatten().bincount( + minlength=config.num_moe_experts + ).to(dtype=torch.int32) From f11962b8feb3eec07208628ac4e5c3643a8669a4 Mon Sep 17 00:00:00 2001 From: zli03 Date: Mon, 22 Jun 2026 14:23:32 +0800 Subject: [PATCH 2/3] move to real b300 test --- .../core/transformer/moe/sonic_moe_layer.py | 208 +++++++++++++++++- .../transformer/moe/test_sonic_moe_layer.py | 198 ++++++++++++++++- 2 files changed, 393 insertions(+), 13 deletions(-) diff --git a/megatron/core/transformer/moe/sonic_moe_layer.py b/megatron/core/transformer/moe/sonic_moe_layer.py index 7d8aede7f8c..8b09aa88405 100644 --- a/megatron/core/transformer/moe/sonic_moe_layer.py +++ b/megatron/core/transformer/moe/sonic_moe_layer.py @@ -14,12 +14,16 @@ import torch import torch.nn.functional as F +from megatron.core.dist_checkpointing import ShardedTensor +from megatron.core.dist_checkpointing.mapping import ShardedTensorFactory from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.utils import get_pg_rank, get_pg_size, make_sharded_tensor_for_checkpoint from megatron.core.transformer.moe.moe_layer import BaseMoELayer, MoELayer, MoESubmodules from megatron.core.transformer.moe.moe_utils import ( MoEAuxLossAutoScaler, compute_routing_scores_for_aux_loss, get_default_pg_collection, + router_gating_linear, save_to_aux_losses_tracker, switch_load_balancing_loss_func, z_loss_func, @@ -28,7 +32,6 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.utils import ( ensure_metadata_has_dp_cp_group, - make_sharded_tensors_for_checkpoint, ) @@ -66,6 +69,7 @@ def __init__( std: float, router_score_function: str = "softmax", router_score_over_topk: bool = True, + router_dtype: Optional[str] = None, accumulate_wgrad_into_main_grad: bool = False, ) -> None: _require_sonicmoe() @@ -80,7 +84,24 @@ def __init__( router_score_function=router_score_function, router_score_over_topk=router_score_over_topk, ) + self.megatron_router_dtype = router_dtype self.accumulate_wgrad_into_main_grad = accumulate_wgrad_into_main_grad + self._maintain_router_param_dtype() + + def _target_router_param_dtype(self) -> torch.dtype: + if self.megatron_router_dtype == "fp32": + return torch.float32 + return self.c_fc.weight.dtype + + def _maintain_router_param_dtype(self) -> None: + target_dtype = self._target_router_param_dtype() + if self.router.weight.dtype != target_dtype: + self.router.to(dtype=target_dtype) + + def _apply(self, fn): + module = super()._apply(fn) + self._maintain_router_param_dtype() + return module def _import_sonicmoe_functional(): @@ -216,9 +237,11 @@ def _check_supported_config(config: TransformerConfig) -> None: def _set_sonic_param_dtypes(module: torch.nn.Module, config: TransformerConfig) -> None: - module.router.to(dtype=torch.float32) module.c_fc.to(dtype=config.params_dtype) module.c_proj.to(dtype=config.params_dtype) + module.router.to( + dtype=torch.float32 if config.moe_router_dtype == "fp32" else config.params_dtype + ) def _maybe_move_to_runtime_device(module: torch.nn.Module, config: TransformerConfig) -> None: @@ -331,6 +354,7 @@ def __init__( std=config.init_method_std, router_score_function=config.moe_router_score_function, router_score_over_topk=self._score_over_topk(), + router_dtype=config.moe_router_dtype, accumulate_wgrad_into_main_grad=config.gradient_accumulation_fusion, ) _maybe_move_to_runtime_device(self.sonic_moe, config) @@ -339,6 +363,11 @@ def __init__( for param in self.sonic_moe.parameters(): setattr(param, "allreduce", True) + def _apply(self, fn): + module = super()._apply(fn) + self.sonic_moe._maintain_router_param_dtype() + return module + def set_layer_number(self, layer_number: int): self.layer_number = layer_number @@ -447,6 +476,13 @@ def _scores_for_aux_loss(self, router_logits: torch.Tensor) -> torch.Tensor: f"Invalid score_function: {self.config.moe_router_score_function}" ) + def _router_dtype(self, input: torch.Tensor) -> torch.dtype: + if self.config.moe_router_dtype == "fp32": + return torch.float32 + if self.config.moe_router_dtype == "fp64": + return torch.float64 + return input.dtype + def _sonic_tc_forward( self, hidden_states: torch.Tensor, @@ -463,7 +499,12 @@ def _sonic_tc_forward( original_shape = hidden_states.shape x = hidden_states.view(-1, self.config.hidden_size) - router_logits = F.linear(x.float(), self.sonic_moe.router.weight) + router_logits = router_gating_linear( + x, + self.sonic_moe.router.weight, + self.sonic_moe.router.bias, + self._router_dtype(x), + ) num_experts = self.sonic_moe.router.weight.size(0) topk_scores, topk_indices = self._apply_sonic_router_topk( TC_Softmax_Topk_Router_Function, @@ -773,14 +814,147 @@ def state_dict(self, *args, destination=None, prefix="", keep_vars=False): def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): metadata = ensure_metadata_has_dp_cp_group(metadata) - return make_sharded_tensors_for_checkpoint( - self._origin_state_dict_entries("", keep_vars=True), - prefix, - {}, - sharded_offsets=sharded_offsets, - tp_group=self.tp_group, - dp_cp_group=metadata["dp_cp_group"], - ) + dp_cp_group = metadata["dp_cp_group"] + prepend_axis_num = len(sharded_offsets) + ep_rank = get_pg_rank(self.ep_group) + ep_size = get_pg_size(self.ep_group) + tp_rank = get_pg_rank(self.tp_group) + tp_size = get_pg_size(self.tp_group) + dp_rank = get_pg_rank(dp_cp_group) + replica_id = (0, 0, dp_rank) + + def fc1_build_fn(key, tensor, replica_id, flattened_range): + if flattened_range is not None: + raise ValueError("SonicMoELayer does not support flattened-range MoE fc1 checkpointing.") + tensor = ( + tensor.view( + self.config.num_moe_experts, + self.config.hidden_size, + 2 * self.config.moe_ffn_hidden_size, + ) + .transpose(1, 2) + .contiguous() + ) + gate_weight, up_weight = torch.chunk(tensor, 2, dim=1) + return [ + ShardedTensor.from_rank_offsets( + key, + gate_weight.contiguous(), + *sharded_offsets, + (prepend_axis_num, ep_rank, ep_size), + (prepend_axis_num + 1, tp_rank, tp_size * 2), + replica_id=replica_id, + prepend_axis_num=prepend_axis_num, + ), + ShardedTensor.from_rank_offsets( + key, + up_weight.contiguous(), + *sharded_offsets, + (prepend_axis_num, ep_rank, ep_size), + (prepend_axis_num + 1, tp_size + tp_rank, tp_size * 2), + replica_id=replica_id, + prepend_axis_num=prepend_axis_num, + ), + ] + + def fc1_merge_fn(sub_state_dict): + if isinstance(sub_state_dict, dict): + assert sub_state_dict["singleton_local_shards"] + sub_state_dict = torch.cat( + ( + torch.stack(sub_state_dict["data"]["w"]), + torch.stack(sub_state_dict["data"]["v"]), + ), + dim=-2, + ) + else: + sub_state_dict = torch.cat(sub_state_dict, dim=-2) + return ( + sub_state_dict.transpose(1, 2) + .contiguous() + .view( + self.config.hidden_size, + self.config.num_moe_experts * 2 * self.config.moe_ffn_hidden_size, + ) + ) + + def fc2_build_fn(key, tensor, replica_id, flattened_range): + if flattened_range is not None: + raise ValueError("SonicMoELayer does not support flattened-range MoE fc2 checkpointing.") + tensor = ( + tensor.view( + self.config.num_moe_experts, + self.config.moe_ffn_hidden_size, + self.config.hidden_size, + ) + .transpose(1, 2) + .contiguous() + ) + return ShardedTensor.from_rank_offsets( + key, + tensor.contiguous(), + *sharded_offsets, + (prepend_axis_num, ep_rank, ep_size), + (prepend_axis_num + 2, tp_rank, tp_size), + replica_id=replica_id, + prepend_axis_num=prepend_axis_num, + ) + + def fc2_merge_fn(sub_state_dict): + if isinstance(sub_state_dict, dict): + assert sub_state_dict["singleton_local_shards"] + sub_state_dict = torch.stack(sub_state_dict["data"]) + return ( + sub_state_dict.transpose(1, 2) + .contiguous() + .view( + self.config.num_moe_experts * self.config.moe_ffn_hidden_size, + self.config.hidden_size, + ) + ) + + sharded_state_dict = { + f"{prefix}router.weight": make_sharded_tensor_for_checkpoint( + self.sonic_moe.router.weight, + f"{prefix}router.weight", + prepend_offsets=sharded_offsets, + tp_group=self.tp_group, + dp_cp_group=dp_cp_group, + ), + f"{prefix}experts.weight1": ShardedTensorFactory( + f"{prefix}experts.experts.linear_fc1.weight", + self._origin_grouped_fc1_weight(keep_vars=True), + fc1_build_fn, + fc1_merge_fn, + replica_id, + ), + f"{prefix}experts.weight2": ShardedTensorFactory( + f"{prefix}experts.experts.linear_fc2.weight", + self._origin_grouped_fc2_weight(keep_vars=True), + fc2_build_fn, + fc2_merge_fn, + replica_id, + ), + } + + if self.sonic_moe.c_fc.bias is not None: + sharded_state_dict[f"{prefix}experts.linear_fc1.bias"] = make_sharded_tensor_for_checkpoint( + self.sonic_moe.c_fc.bias, + f"{prefix}experts.experts.linear_fc1.bias", + prepend_offsets=sharded_offsets, + tp_group=self.tp_group, + dp_cp_group=dp_cp_group, + ) + if self.sonic_moe.c_proj.bias is not None: + sharded_state_dict[f"{prefix}experts.linear_fc2.bias"] = make_sharded_tensor_for_checkpoint( + self.sonic_moe.c_proj.bias, + f"{prefix}experts.experts.linear_fc2.bias", + prepend_offsets=sharded_offsets, + tp_group=self.tp_group, + dp_cp_group=dp_cp_group, + ) + + return sharded_state_dict def _move_if_present(self, state_dict, src_key: str, dst_key: str) -> None: if dst_key not in state_dict and src_key in state_dict: @@ -843,15 +1017,27 @@ def _load_from_state_dict( self._move_if_present( state_dict, f"{prefix}experts.linear_fc1.weight", f"{prefix}sonic_moe.c_fc.weight" ) + self._move_if_present( + state_dict, f"{prefix}experts.experts.linear_fc1.weight", f"{prefix}sonic_moe.c_fc.weight" + ) self._move_if_present( state_dict, f"{prefix}experts.linear_fc1.bias", f"{prefix}sonic_moe.c_fc.bias" ) + self._move_if_present( + state_dict, f"{prefix}experts.experts.linear_fc1.bias", f"{prefix}sonic_moe.c_fc.bias" + ) self._move_if_present( state_dict, f"{prefix}experts.linear_fc2.weight", f"{prefix}sonic_moe.c_proj.weight" ) + self._move_if_present( + state_dict, f"{prefix}experts.experts.linear_fc2.weight", f"{prefix}sonic_moe.c_proj.weight" + ) self._move_if_present( state_dict, f"{prefix}experts.linear_fc2.bias", f"{prefix}sonic_moe.c_proj.bias" ) + self._move_if_present( + state_dict, f"{prefix}experts.experts.linear_fc2.bias", f"{prefix}sonic_moe.c_proj.bias" + ) self._stack_per_expert_if_present( state_dict, prefix, diff --git a/tests/unit_tests/transformer/moe/test_sonic_moe_layer.py b/tests/unit_tests/transformer/moe/test_sonic_moe_layer.py index 0598007b033..94c3107a52a 100644 --- a/tests/unit_tests/transformer/moe/test_sonic_moe_layer.py +++ b/tests/unit_tests/transformer/moe/test_sonic_moe_layer.py @@ -1,11 +1,16 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import dataclasses +import os +import shutil import pytest import torch +import torch.distributed as dist import torch.nn.functional as F +from megatron.core import dist_checkpointing +from megatron.core.dist_checkpointing import ShardedTensor from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec @@ -14,6 +19,7 @@ clear_aux_losses_tracker, get_default_pg_collection, get_moe_layer_wise_logging_tracker, + router_gating_linear, ) from megatron.core.transformer.moe.moe_layer import MoELayer from megatron.core.transformer.moe.router import TopKRouter @@ -21,6 +27,7 @@ SonicMoELayer, replace_moe_layer_specs_with_sonic_moe, ) +from megatron.core.transformer.module import Float16Module from megatron.core.transformer.transformer_config import TransformerConfig from megatron.training.initialize import _set_random_seed from tests.unit_tests.test_utilities import Utils @@ -37,12 +44,81 @@ ROUTER_KERNEL_ATOL = 3.0e-5 FORWARD_RTOL = 1.0e-2 FORWARD_ATOL = 1.0e-5 +EP2_DCP_PREFIX = "decoder.layers.0.mlp." def _require_sonicmoe(): pytest.importorskip("sonicmoe") +def _ep2_dcp_expected_tensors(config, device): + router = ( + torch.arange( + config.num_moe_experts * config.hidden_size, + dtype=torch.float32, + device=device, + ).view(config.num_moe_experts, config.hidden_size) + / 100.0 + ) + fc1 = torch.arange( + config.num_moe_experts * 2 * config.moe_ffn_hidden_size * config.hidden_size, + dtype=torch.float32, + device=device, + ).view(config.num_moe_experts, 2 * config.moe_ffn_hidden_size, config.hidden_size) + fc2 = torch.arange( + config.num_moe_experts * config.hidden_size * config.moe_ffn_hidden_size, + dtype=torch.float32, + device=device, + ).view(config.num_moe_experts, config.hidden_size, config.moe_ffn_hidden_size) + return router, fc1.to(dtype=config.params_dtype), fc2.to(dtype=config.params_dtype) + + +def _ep2_dcp_sharded_state(config, rank, world_size, device, with_data): + router, fc1, fc2 = _ep2_dcp_expected_tensors(config, device) + if not with_data: + router = torch.empty_like(router) + fc1 = torch.empty_like(fc1) + fc2 = torch.empty_like(fc2) + + local_experts = config.num_moe_experts // world_size + start = rank * local_experts + end = start + local_experts + fc1_gate, fc1_up = torch.chunk(fc1[start:end], 2, dim=1) + + return { + f"{EP2_DCP_PREFIX}router.weight": ShardedTensor.from_rank_offsets( + f"{EP2_DCP_PREFIX}router.weight", + router, + replica_id=(0, 0, rank), + ), + f"{EP2_DCP_PREFIX}experts.experts.linear_fc1.weight.gate": ( + ShardedTensor.from_rank_offsets( + f"{EP2_DCP_PREFIX}experts.experts.linear_fc1.weight", + fc1_gate.contiguous(), + (0, rank, world_size), + (1, 0, 2), + replica_id=(0, 0, 0), + ) + ), + f"{EP2_DCP_PREFIX}experts.experts.linear_fc1.weight.up": ( + ShardedTensor.from_rank_offsets( + f"{EP2_DCP_PREFIX}experts.experts.linear_fc1.weight", + fc1_up.contiguous(), + (0, rank, world_size), + (1, 1, 2), + replica_id=(0, 0, 0), + ) + ), + f"{EP2_DCP_PREFIX}experts.experts.linear_fc2.weight": ShardedTensor.from_rank_offsets( + f"{EP2_DCP_PREFIX}experts.experts.linear_fc2.weight", + fc2[start:end].contiguous(), + (0, rank, world_size), + (2, 0, 1), + replica_id=(0, 0, 0), + ), + } + + class TestSonicMoELayerRouterLoss: def setup_method(self, method): _require_sonicmoe() @@ -140,6 +216,11 @@ def test_sonic_param_and_main_grad_dtypes(self): if layer.sonic_moe.c_proj.bias is not None: assert layer.sonic_moe.c_proj.bias.dtype is torch.bfloat16 + wrapped_layer = Float16Module(self.default_config, self._new_sonic_layer(self.default_config)) + assert wrapped_layer.module.sonic_moe.router.weight.dtype is torch.float32 + assert wrapped_layer.module.sonic_moe.c_fc.weight.dtype is torch.bfloat16 + assert wrapped_layer.module.sonic_moe.c_proj.weight.dtype is torch.bfloat16 + ddp = DistributedDataParallel( self.default_config, DistributedDataParallelConfig(grad_reduce_in_fp32=True), @@ -529,6 +610,113 @@ def test_load_origin_grouped_moe_state_and_save_origin_format(self): roundtrip_state[key], tensor, rtol=EXACT_RTOL, atol=EXACT_ATOL ) + def test_load_ep2_distributed_checkpoint(self, tmp_path): + if Utils.world_size != 2: + pytest.skip("Run with torchrun --nproc_per_node=2 to exercise EP=2 DCP.") + if torch.cuda.device_count() < 2: + pytest.skip("EP=2 checkpoint test requires at least 2 CUDA devices.") + + config = TransformerConfig( + num_layers=1, + hidden_size=8, + num_attention_heads=1, + ffn_hidden_size=16, + num_moe_experts=4, + moe_ffn_hidden_size=4, + moe_router_topk=2, + moe_router_load_balancing_type="none", + moe_aux_loss_coeff=0.0, + moe_z_loss_coeff=None, + moe_router_score_function="sigmoid", + moe_router_dtype="fp32", + add_bias_linear=False, + gated_linear_unit=True, + activation_func=F.silu, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + ) + rank = Utils.rank + world_size = Utils.world_size + device = torch.device("cuda", torch.cuda.current_device()) + shared_tmp_path = [str(tmp_path) if rank == 0 else None] + dist.broadcast_object_list(shared_tmp_path, src=0) + src_ckpt_dir = os.path.join(shared_tmp_path[0], "source_ep2_dcp") + sonic_ckpt_dir = os.path.join(shared_tmp_path[0], "sonic_saved_dcp") + + if rank == 0: + for ckpt_dir in (src_ckpt_dir, sonic_ckpt_dir): + if os.path.exists(ckpt_dir): + shutil.rmtree(ckpt_dir) + os.makedirs(ckpt_dir, exist_ok=True) + dist.barrier() + + src_ep2_state = _ep2_dcp_sharded_state( + config, rank, world_size, device, with_data=True + ) + dist_checkpointing.save( + src_ep2_state, src_ckpt_dir, validate_access_integrity=False + ) + dist.barrier() + + layer = SonicMoELayer(config=config, pg_collection=get_default_pg_collection()).cuda() + layer.set_layer_number(0) + loaded = dist_checkpointing.load( + layer.sharded_state_dict(prefix=EP2_DCP_PREFIX), + src_ckpt_dir, + validate_access_integrity=False, + strict="raise_all", + ) + load_result = layer.load_state_dict( + { + key.removeprefix(EP2_DCP_PREFIX): value + for key, value in loaded.items() + if key.startswith(EP2_DCP_PREFIX) + }, + strict=True, + ) + assert load_result.missing_keys == [] + assert load_result.unexpected_keys == [] + + router, fc1, fc2 = _ep2_dcp_expected_tensors( + config, torch.device("cuda", torch.cuda.current_device()) + ) + torch.testing.assert_close( + layer.sonic_moe.router.weight, router, rtol=EXACT_RTOL, atol=EXACT_ATOL + ) + torch.testing.assert_close( + layer.sonic_moe.c_fc.weight, fc1, rtol=EXACT_RTOL, atol=EXACT_ATOL + ) + torch.testing.assert_close( + layer.sonic_moe.c_proj.weight, fc2, rtol=EXACT_RTOL, atol=EXACT_ATOL + ) + assert layer.sonic_moe.router.weight.dtype is torch.float32 + + dist_checkpointing.save( + layer.sharded_state_dict(prefix=EP2_DCP_PREFIX), + sonic_ckpt_dir, + validate_access_integrity=False, + ) + dist.barrier() + + src_ep2_loaded = dist_checkpointing.load( + _ep2_dcp_sharded_state(config, rank, world_size, device, with_data=False), + src_ckpt_dir, + validate_access_integrity=False, + strict="raise_all", + ) + sonic_ep2_loaded = dist_checkpointing.load( + _ep2_dcp_sharded_state(config, rank, world_size, device, with_data=False), + sonic_ckpt_dir, + validate_access_integrity=False, + strict="raise_all", + ) + assert set(src_ep2_loaded.keys()) == set(sonic_ep2_loaded.keys()) + for key, tensor in src_ep2_loaded.items(): + torch.testing.assert_close( + sonic_ep2_loaded[key], tensor, rtol=EXACT_RTOL, atol=EXACT_ATOL + ) + @staticmethod def _origin_grouped_moe_state_dict(config, expert_weight_scale: float = 1.0): weight1_shape = ( @@ -627,8 +815,14 @@ def _run_sonic_loss_path(self, sonic_layer, hidden_state): sonic_layer.sonic_moe.router.weight.grad = None hidden_state = hidden_state.detach().clone().requires_grad_(True) - router_logits = F.linear(hidden_state.float(), sonic_layer.sonic_moe.router.weight).view( - -1, sonic_layer.config.num_moe_experts + router_logits = router_gating_linear( + hidden_state, + sonic_layer.sonic_moe.router.weight, + sonic_layer.sonic_moe.router.bias, + sonic_layer._router_dtype(hidden_state), + ).view( + -1, + sonic_layer.config.num_moe_experts, ) tokens_per_expert = self._tokens_per_expert(router_logits, sonic_layer.config) output = torch.zeros_like(hidden_state, requires_grad=True) From b0a8270629ec929f235a0cec0cd18419321ace65 Mon Sep 17 00:00:00 2001 From: zli03 Date: Tue, 23 Jun 2026 14:33:56 +0800 Subject: [PATCH 3/3] fix router dtype --- .../core/transformer/moe/sonic_moe_layer.py | 6 +--- .../transformer/moe/test_sonic_moe_layer.py | 33 ++++++++++++++----- 2 files changed, 26 insertions(+), 13 deletions(-) diff --git a/megatron/core/transformer/moe/sonic_moe_layer.py b/megatron/core/transformer/moe/sonic_moe_layer.py index 8b09aa88405..9dae8abfda2 100644 --- a/megatron/core/transformer/moe/sonic_moe_layer.py +++ b/megatron/core/transformer/moe/sonic_moe_layer.py @@ -89,8 +89,6 @@ def __init__( self._maintain_router_param_dtype() def _target_router_param_dtype(self) -> torch.dtype: - if self.megatron_router_dtype == "fp32": - return torch.float32 return self.c_fc.weight.dtype def _maintain_router_param_dtype(self) -> None: @@ -239,9 +237,7 @@ def _check_supported_config(config: TransformerConfig) -> None: def _set_sonic_param_dtypes(module: torch.nn.Module, config: TransformerConfig) -> None: module.c_fc.to(dtype=config.params_dtype) module.c_proj.to(dtype=config.params_dtype) - module.router.to( - dtype=torch.float32 if config.moe_router_dtype == "fp32" else config.params_dtype - ) + module.router.to(dtype=config.params_dtype) def _maybe_move_to_runtime_device(module: torch.nn.Module, config: TransformerConfig) -> None: diff --git a/tests/unit_tests/transformer/moe/test_sonic_moe_layer.py b/tests/unit_tests/transformer/moe/test_sonic_moe_layer.py index 94c3107a52a..624112745b5 100644 --- a/tests/unit_tests/transformer/moe/test_sonic_moe_layer.py +++ b/tests/unit_tests/transformer/moe/test_sonic_moe_layer.py @@ -59,7 +59,7 @@ def _ep2_dcp_expected_tensors(config, device): device=device, ).view(config.num_moe_experts, config.hidden_size) / 100.0 - ) + ).to(dtype=config.params_dtype) fc1 = torch.arange( config.num_moe_experts * 2 * config.moe_ffn_hidden_size * config.hidden_size, dtype=torch.float32, @@ -154,7 +154,6 @@ def teardown_method(self, method): def _new_topk_router(self, config): router = TopKRouter(config=config, pg_collection=get_default_pg_collection()).cuda() router.set_layer_number(0) - router.weight.data = router.weight.data.float() return router def _new_sonic_layer(self, config): @@ -208,7 +207,21 @@ def test_sonic_param_and_main_grad_dtypes(self): layer.sonic_moe.accumulate_wgrad_into_main_grad is self.default_config.gradient_accumulation_fusion ) - assert layer.sonic_moe.router.weight.dtype is torch.float32 + assert layer.sonic_moe.router.weight.dtype is self.default_config.params_dtype + assert ( + layer._router_dtype(torch.empty(1, device="cuda", dtype=torch.bfloat16)) + is torch.float32 + ) + router_input = torch.empty( + 1, 1, self.default_config.hidden_size, device="cuda", dtype=torch.bfloat16 + ) + router_logits = router_gating_linear( + router_input, + layer.sonic_moe.router.weight, + layer.sonic_moe.router.bias, + layer._router_dtype(router_input), + ) + assert router_logits.dtype is torch.float32 assert layer.sonic_moe.c_fc.weight.dtype is torch.bfloat16 assert layer.sonic_moe.c_proj.weight.dtype is torch.bfloat16 if layer.sonic_moe.c_fc.bias is not None: @@ -217,7 +230,11 @@ def test_sonic_param_and_main_grad_dtypes(self): assert layer.sonic_moe.c_proj.bias.dtype is torch.bfloat16 wrapped_layer = Float16Module(self.default_config, self._new_sonic_layer(self.default_config)) - assert wrapped_layer.module.sonic_moe.router.weight.dtype is torch.float32 + assert wrapped_layer.module.sonic_moe.router.weight.dtype is self.default_config.params_dtype + assert ( + wrapped_layer.module._router_dtype(torch.empty(1, device="cuda", dtype=torch.bfloat16)) + is torch.float32 + ) assert wrapped_layer.module.sonic_moe.c_fc.weight.dtype is torch.bfloat16 assert wrapped_layer.module.sonic_moe.c_proj.weight.dtype is torch.bfloat16 @@ -396,7 +413,7 @@ def test_z_loss_and_global_aux_loss_match_original_router( rtol=EXACT_RTOL, atol=EXACT_ATOL, ) - assert sonic["router_weight_grad"].dtype is torch.float32 + assert sonic["router_weight_grad"].dtype is config.params_dtype torch.testing.assert_close( sonic["router_weight_grad"], ref["router_weight_grad"], @@ -459,7 +476,7 @@ def test_sigmoid_end_to_end_forward_backward(self): assert "global_load_balancing_loss" in tracker assert hidden_state.grad.dtype is torch.bfloat16 assert torch.isfinite(hidden_state.grad).all() - assert layer.sonic_moe.router.weight.grad.dtype is torch.float32 + assert layer.sonic_moe.router.weight.grad.dtype is config.params_dtype assert torch.isfinite(layer.sonic_moe.router.weight.grad).all() assert layer.sonic_moe.c_fc.weight.grad.dtype is torch.bfloat16 assert torch.isfinite(layer.sonic_moe.c_fc.weight.grad).all() @@ -690,7 +707,7 @@ def test_load_ep2_distributed_checkpoint(self, tmp_path): torch.testing.assert_close( layer.sonic_moe.c_proj.weight, fc2, rtol=EXACT_RTOL, atol=EXACT_ATOL ) - assert layer.sonic_moe.router.weight.dtype is torch.float32 + assert layer.sonic_moe.router.weight.dtype is config.params_dtype dist_checkpointing.save( layer.sharded_state_dict(prefix=EP2_DCP_PREFIX), @@ -732,7 +749,7 @@ def _origin_grouped_moe_state_dict(config, expert_weight_scale: float = 1.0): config.num_moe_experts, config.hidden_size, device="cuda", - dtype=torch.float32, + dtype=config.params_dtype, ), "experts.weight1": torch.randn( weight1_shape, device="cuda", dtype=config.params_dtype