[MoE] Add optional output normalization for latent MoE

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:

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.")

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."

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:

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))