From 9c1a6ce6c46ea9de072523b6914c7a2b63b4509f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=9E=E5=AD=9F?= Date: Wed, 12 Aug 2026 04:35:40 +0800 Subject: [PATCH] Add optional output normalization for latent MoE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add moe_latent_output_norm config that applies a TENorm over the latent dimension to the combined routed-expert output, after token combination and before the up-projection back to the transformer hidden size. Disabled by default; the normalization type and epsilon follow the existing normalization/layernorm_epsilon configs. Signed-off-by: 晞孟 --- megatron/core/transformer/moe/moe_layer.py | 12 +++++- .../core/transformer/transformer_config.py | 8 ++++ megatron/training/arguments.py | 6 +++ megatron/training/checkpointing.py | 1 + .../transformer/moe/test_latent_moe_layer.py | 41 ++++++++++++++++++- 5 files changed, 65 insertions(+), 3 deletions(-) diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 59684a34b0d..2f0ff95653a 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -63,9 +63,9 @@ HAVE_TRITON = False if HAVE_TE: - from megatron.core.extensions.transformer_engine import TELinear, te_checkpoint + from megatron.core.extensions.transformer_engine import TELinear, TENorm, te_checkpoint else: - TELinear, te_checkpoint = None, None + TELinear, TENorm, te_checkpoint = None, None, None class ExpertsInterface(Protocol): @@ -302,6 +302,12 @@ def __init__( is_expert=False, name=(name + ".fc2_latent_proj") if name is not None else None, ) + if self.config.moe_latent_output_norm: + self.routed_expert_norm = TENorm( + config=self.config, + hidden_size=self.config.moe_latent_size, + eps=self.config.layernorm_epsilon, + ) # Initialize token dispatcher if config.moe_token_dispatcher_type == "allgather": @@ -643,6 +649,8 @@ def postprocess(self, output: torch.Tensor, shared_expert_output: Optional[torch output = self.token_dispatcher.combine_postprocess(output) if self.config.moe_latent_size: + if self.config.moe_latent_output_norm: + output = self.routed_expert_norm(output) output, _ = self.fc2_latent_proj(output) if shared_expert_output is not None: diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 1f11d8d2d7b..5800d44f12b 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -999,6 +999,11 @@ class TransformerConfig(ModelParallelConfig): moe_latent_size: Optional[int] = None """Latent projection dimension for MoE. If None, MoE latent projections are not used.""" + moe_latent_output_norm: bool = False + """Apply normalization to the combined routed-expert output in the latent dimension before + projecting it back to the transformer hidden size. The normalization type and epsilon are + controlled by ``normalization`` and ``layernorm_epsilon``, respectively.""" + moe_flex_dispatcher_num_sms: Optional[int] = None """Number of SMs for the flex token dispatcher's dispatch/combine communication, for all backends (deepep, hybridep, ncclep). None lets each backend use its own default. Unifies the @@ -1928,6 +1933,9 @@ def __post_init__(self): if self.num_moe_experts is not None and self.num_moe_experts <= 0: raise ValueError("num_moe_experts must be non-negative.") + if self.moe_latent_output_norm and self.moe_latent_size is None: + raise ValueError("moe_latent_output_norm requires moe_latent_size to be set.") + if self.num_moe_experts is not None and self.moe_ffn_hidden_size is None: self.moe_ffn_hidden_size = self.ffn_hidden_size warnings.warn("moe_ffn_hidden_size is not set, using ffn_hidden_size instead.") diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index a168381eaa7..7e31c7cad14 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2143,6 +2143,11 @@ def validate_args(args, defaults={}): args.num_experts is not None ), "MoE latent projections are applicable only for MoE models." + if args.moe_latent_output_norm: + assert ( + args.moe_latent_size is not None + ), "--moe-latent-output-norm requires --moe-latent-size to be set." + # Print arguments. _print_args("arguments", args) @@ -2264,6 +2269,7 @@ def core_transformer_config_from_args(args, config_class=None): kw_args['quant_recipe'] = kitchen_quantization_recipe_config(args.kitchen_recipe_number) kw_args['moe_latent_size'] = args.moe_latent_size + kw_args['moe_latent_output_norm'] = args.moe_latent_output_norm if args.te_precision_config_file: assert not 'quant_recipe' in kw_args, "Quantization recipe already configured." diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index bfa3564c6cd..e8c9f75e674 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -1949,6 +1949,7 @@ def _set_arg(arg_name, old_arg_name=None, force=False): # MoE latent projection. _set_arg('moe_latent_size', force=True) + _set_arg('moe_latent_output_norm', force=True) # Tokenizer args. if args.use_tokenizer_model_from_checkpoint_args: diff --git a/tests/unit_tests/transformer/moe/test_latent_moe_layer.py b/tests/unit_tests/transformer/moe/test_latent_moe_layer.py index bb5ced291fc..9fc8d65e48c 100644 --- a/tests/unit_tests/transformer/moe/test_latent_moe_layer.py +++ b/tests/unit_tests/transformer/moe/test_latent_moe_layer.py @@ -1,5 +1,7 @@ # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +from types import SimpleNamespace + import pytest import torch @@ -26,7 +28,7 @@ def setup_method(self, method): @pytest.mark.parametrize("moe_token_dispatcher_type", ["allgather", "alltoall"]) @pytest.mark.parametrize("num_moe_experts", [4]) @pytest.mark.parametrize("use_te,grouped_gemm", [(True, True), (True, False), (False, False)]) - @pytest.mark.parametrize("moe_latent_size", [8, 16]) + @pytest.mark.parametrize("moe_latent_size", [64, 128]) def test_latent_moe_layer( self, num_moe_experts, moe_token_dispatcher_type, use_te, grouped_gemm, moe_latent_size ): @@ -48,6 +50,8 @@ def test_latent_moe_layer( gated_linear_unit=True, add_bias_linear=False, moe_latent_size=moe_latent_size, + moe_latent_output_norm=True, + normalization="RMSNorm", ) if use_te: transformer_layer_submodules = get_gpt_layer_with_transformer_engine_submodules( @@ -63,6 +67,8 @@ def test_latent_moe_layer( moe_layer.cuda() config = moe_layer.config + assert moe_layer.routed_expert_norm.weight.shape[0] == config.moe_latent_size + assert ( moe_layer.shared_experts.linear_fc1.weight.shape[1] == config.hidden_size ), "Shared expert computation has to happen in hidden dimension." @@ -99,3 +105,36 @@ def test_latent_moe_layer( assert output.shape[2] == config.hidden_size Utils.destroy_model_parallel() + + def test_latent_output_norm_requires_latent_size(self): + with pytest.raises(ValueError, match="requires moe_latent_size"): + TransformerConfig( + num_layers=1, hidden_size=32, num_attention_heads=4, moe_latent_output_norm=True + ) + + def test_latent_output_norm_is_applied_after_combine_and_before_up_projection(self): + class FakeDispatcher: + @staticmethod + def combine_postprocess(output): + return output + 1 + + class FakeNorm(torch.nn.Module): + def forward(self, output): + return output * 2 + + class FakeProjection(torch.nn.Module): + def forward(self, output): + return output + 3, None + + moe_layer = object.__new__(MoELayer) + torch.nn.Module.__init__(moe_layer) + moe_layer.config = SimpleNamespace(moe_latent_size=2, moe_latent_output_norm=True) + moe_layer.token_dispatcher = FakeDispatcher() + moe_layer.routed_expert_norm = FakeNorm() + moe_layer.fc2_latent_proj = FakeProjection() + moe_layer._latent_shared_expert_output = None + + output = moe_layer.postprocess(torch.zeros(1, 2), shared_expert_output=None) + + # combine (+1), normalize (*2), then project (+3) + torch.testing.assert_close(output, torch.full((1, 2), 5.0))