Changes to support latent MoEs

Signed-off-by: Deepak Narayanan <[email protected]>

Author: venmugil

megatron/core/extensions/transformer_engine.py

@@ -312,7 +312,7 @@ def __init__(
             )
 
         if is_te_min_version("0.8.0"):
-            if self.config.tp_comm_overlap:
+            if self.config.tp_comm_overlap and parallel_mode != "duplicated":
                 if is_te_min_version("1.5.0"):
                     # Use old overlap flags if they were supplied instead
                     extra_kwargs["ub_overlap_ag"] = (

megatron/core/transformer/mlp.py

@@ -104,9 +104,12 @@ def __init__(
         if self.config.gated_linear_unit:
             ffn_hidden_size *= 2
 
+        # Use moe_latent_size only for routed experts. 'is_expert' is false for shared_experts
+        use_latent_size = (self.config.moe_latent_size is not None) and is_expert
+
         self.linear_fc1 = build_module(
             submodules.linear_fc1,
-            self.input_size,
+            self.input_size if not use_latent_size else self.config.moe_latent_size,
             ffn_hidden_size,
             config=self.config,
             init_method=self.config.init_method,
@@ -126,7 +129,7 @@ def __init__(
         self.linear_fc2 = build_module(
             submodules.linear_fc2,
             self.config.ffn_hidden_size,
-            self.config.hidden_size,
+            self.config.hidden_size if not use_latent_size else self.config.moe_latent_size,
             config=self.config,
             init_method=self.config.output_layer_init_method,
             bias=self.config.add_bias_linear,

megatron/core/transformer/moe/experts.py

@@ -118,6 +118,9 @@ def __init__(
         assert (
             config.add_bias_linear == False
         ), "bias not supported in Grouped GEMM yet, please set '--disable-bias-linear' instead."
+        assert (
+            config.moe_latent_size is None
+        ), "MoE latent projection not supported in GroupedMLP yet."
 
         self.expert_parallel = config.expert_model_parallel_size > 1
         if self.config.gated_linear_unit:
@@ -778,7 +781,7 @@ def __init__(
         self.linear_fc1 = build_module(
             submodules.linear_fc1,
             self.num_local_experts,
-            self.input_size,
+            self.input_size if self.config.moe_latent_size is None else self.config.moe_latent_size,
             ffn_hidden_size,
             config=self.config,
             init_method=self.config.init_method,
@@ -799,7 +802,11 @@ def __init__(
             submodules.linear_fc2,
             self.num_local_experts,
             self.config.moe_ffn_hidden_size,
-            self.config.hidden_size,
+            (
+                self.config.hidden_size
+                if self.config.moe_latent_size is None
+                else self.config.moe_latent_size
+            ),
             config=self.config,
             init_method=self.config.output_layer_init_method,
             bias=self.config.add_bias_linear,

megatron/core/transformer/moe/moe_layer.py

@@ -23,7 +23,7 @@
 try:
     import transformer_engine as te  # pylint: disable=unused-import
 
-    from megatron.core.extensions.transformer_engine import te_checkpoint
+    from megatron.core.extensions.transformer_engine import TELinear, te_checkpoint
 
     HAVE_TE = True
 except ImportError:
@@ -123,6 +123,32 @@ def __init__(
         # Initialize router
         self.router = TopKRouter(config=self.config, pg_collection=pg_collection)
 
+        # Initialize latent projections
+        if self.config.moe_latent_size:
+            assert HAVE_TE
+            self.fc1_latent_proj = TELinear(
+                self.config.hidden_size,
+                self.config.moe_latent_size,
+                parallel_mode="duplicated",
+                config=self.config,
+                init_method=self.config.init_method,
+                bias=self.config.add_bias_linear,
+                skip_bias_add=False,
+                skip_weight_param_allocation=False,
+                is_expert=False,
+            )
+            self.fc2_latent_proj = TELinear(
+                self.config.moe_latent_size,
+                self.config.hidden_size,
+                parallel_mode="duplicated",
+                config=self.config,
+                init_method=self.config.output_layer_init_method,
+                bias=self.config.add_bias_linear,
+                skip_bias_add=True,
+                skip_weight_param_allocation=False,
+                is_expert=False,
+            )
+
         # Initialize token dispatcher
         if config.moe_token_dispatcher_type == "allgather":
             self.token_dispatcher = MoEAllGatherTokenDispatcher(
@@ -272,8 +298,23 @@ def forward(self, hidden_states: torch.Tensor):
         def custom_forward(hidden_states):
             shared_expert_output = self.shared_experts_compute(hidden_states)
             hidden_states, probs, residual = self.router_and_preprocess(hidden_states)
+
+            # Project the hidden_states from hidden dimension down to latent dimenion.
+            if self.config.moe_latent_size:
+                assert (
+                    not self.shared_expert_overlap
+                ), "Shared expert overlap not supported when MoE latent projections are used."
+                hidden_states, _ = self.fc1_latent_proj(hidden_states)
+
             dispatched_input, probs = self.dispatch(hidden_states, probs)
             output, mlp_bias = self.routed_experts_compute(dispatched_input, probs, residual)
+
+            # Project the output back from latent dimension to hidden dimension
+            if self.config.moe_latent_size:
+                if mlp_bias is not None:
+                    output = output + mlp_bias
+                output, mlp_bias = self.fc2_latent_proj(output)
+
             output = self.combine(output, shared_expert_output)
             return output, mlp_bias
 

megatron/core/transformer/transformer_config.py

@@ -216,6 +216,9 @@ class TransformerConfig(ModelParallelConfig):
     """Number of SMs to use for HybridEP. In pure NVL scenarios, 
     16 SMs can generally achieve good bandwidth."""
 
+    moe_latent_size: Optional[int] = None
+    """Latent projection dimension for MoE. If None, MoE latent projections are not used."""
+
     ####################
     # initialization
     ####################

megatron/training/arguments.py

@@ -1203,6 +1203,13 @@ def validate_args(args, defaults={}):
             args.recompute_granularity != 'full'
         ), 'recompute_granularity must not be full when CUDA Graphs are enabled.'
 
+    # MoE latent projections
+    if args.moe_latent_size is not None:
+        assert args.moe_latent_size > 0, "MoE latent projection dimension has to be greater than zero."
+        assert args.num_experts is not None, "MoE latent projections are applicable only for MoE models."
+        assert not args.use_legacy_models, "MoE latent projections are only supported for mcore models."
+        assert not args.moe_use_legacy_grouped_gemm, "MoE latent projection is not supported yet with legacy grouped GEMM."
+
     # Print arguments.
     _print_args("arguments", args)
 
@@ -1302,6 +1309,7 @@ def core_transformer_config_from_args(args, config_class=None):
         kw_args['use_kitchen'] = True
         kw_args['quant_recipe'] = kitchen_quantization_recipe_config(args.kitchen_recipe_number)
 
+    kw_args['moe_latent_size'] = args.moe_latent_size
 
     # Return config.
     return config_class(**kw_args)
@@ -1671,6 +1679,8 @@ def _add_network_size_args(parser):
                        'We compute the average of the MTP losses across all depths, '
                        'and multiply it the scaling factor to obtain the overall MTP loss, '
                        'which serves as an additional training objective.')
+    group.add_argument('--moe-latent-size', type=int, default=None,
+                       help='Latent projection dimension for MoE. If None, MoE latent projections are not used.')
     return parser
 
 

megatron/training/checkpointing.py

@@ -1335,6 +1335,9 @@ def _set_arg(arg_name, old_arg_name=None, force=False):
     _set_arg('heterogeneous_layers_config_path', force=True)
     _set_arg('heterogeneous_layers_config_encoded_json', force=True)
 
+    # MoE latent projection
+    _set_arg('moe_latent_size', force=True)
+
     # Tokenizer args.
     _set_arg('tokenizer_type', force=True)
     # Using checkpoint version might not always be safe (e.g., if running on different cluster).