perf(moe): add frozen grouped GEMM for adapter training

megatron/core/extensions/transformer_engine.py

@@ -11,6 +11,8 @@
 import re
 import warnings
 from contextlib import contextmanager, nullcontext
+from itertools import accumulate
+from numbers import Integral
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, cast
 
 import torch
@@ -2358,6 +2360,7 @@ def __init__(
                 name (str | None): module instance name passed top-down from its paranet module
             """
             self.config = config
+            self.expert_gemm_backend = config.moe_expert_gemm_backend
 
             # TE returns a zero length Tensor when bias=False and
             # return_bias=True, but we prefer None.  So in that case we
@@ -2632,6 +2635,141 @@ def _split_grouped_checkpoint_tensor(
                 f"into {self.num_gemms} GEMM shards."
             )
 
+        def _torch_grouped_weight_layout_is_current(self) -> bool:
+            """Return whether each expert parameter still views the contiguous buffer."""
+            grouped_weight = self._buffers.get('_torch_grouped_weight')
+            if grouped_weight is None or grouped_weight.shape[0] != self.num_gemms:
+                return False
+
+            first_weight = self.weight0
+            if (
+                grouped_weight.device != first_weight.device
+                or grouped_weight.dtype != first_weight.dtype
+                or grouped_weight.shape[1:] != first_weight.shape
+            ):
+                return False
+
+            storage_pointer = grouped_weight.untyped_storage().data_ptr()
+            expert_elements = first_weight.numel()
+            for index in range(self.num_gemms):
+                weight = getattr(self, f'weight{index}')
+                grouped_weight_view = grouped_weight[index]
+                if (
+                    weight.device != grouped_weight.device
+                    or weight.dtype != grouped_weight.dtype
+                    or weight.shape != grouped_weight_view.shape
+                    or weight.untyped_storage().data_ptr() != storage_pointer
+                    or weight.storage_offset() != index * expert_elements
+                    or weight.stride() != grouped_weight_view.stride()
+                ):
+                    return False
+            return True
+
+        @torch.no_grad()
+        def prepare_torch_grouped_mm(self) -> int:
+            """Relocate frozen BF16 expert weights into one contiguous allocation.
+
+            The existing per-expert ``Parameter`` objects become views of the new allocation, so
+            parameter names and checkpoint structure remain unchanged. The backing allocation is a
+            non-persistent buffer and therefore does not add a checkpoint entry.
+
+            Returns:
+                Number of bytes in the contiguous backing allocation.
+            """
+            if self.expert_gemm_backend != 'torch':
+                raise RuntimeError(
+                    "prepare_torch_grouped_mm requires moe_expert_gemm_backend='torch'"
+                )
+            self._resolve_torch_grouped_mm()
+            if self.use_bias:
+                raise RuntimeError("torch grouped expert GEMM does not support bias")
+            if getattr(self, 'single_grouped_weight', False):
+                raise RuntimeError(
+                    "torch grouped expert GEMM expects per-GEMM weight parameters, but "
+                    "moe_single_grouped_weight=True makes TE hold one grouped weight tensor"
+                )
+
+            weights = [getattr(self, f'weight{index}') for index in range(self.num_gemms)]
+            if any(weight.requires_grad for weight in weights):
+                raise RuntimeError("torch grouped expert GEMM requires frozen base weights")
+            if any(weight.device.type != 'cuda' for weight in weights):
+                raise RuntimeError("torch grouped expert GEMM requires CUDA weights")
+            if any(weight.dtype != torch.bfloat16 for weight in weights):
+                raise RuntimeError("torch grouped expert GEMM requires BF16 weights")
+            if any(weight.shape != weights[0].shape for weight in weights[1:]):
+                raise RuntimeError("torch grouped expert GEMM requires uniform expert shapes")
+
+            if not self._torch_grouped_weight_layout_is_current():
+                grouped_weight = torch.stack([weight.detach() for weight in weights]).contiguous()
+                for index, weight in enumerate(weights):
+                    weight.data = grouped_weight[index]
+                if '_torch_grouped_weight' in self._buffers:
+                    self._buffers['_torch_grouped_weight'] = grouped_weight
+                else:
+                    self.register_buffer('_torch_grouped_weight', grouped_weight, persistent=False)
+
+            grouped_weight = self._buffers['_torch_grouped_weight']
+            return grouped_weight.numel() * grouped_weight.element_size()
+
+        @property
+        def torch_grouped_mm_prepared(self) -> bool:
+            """Return whether the torch grouped GEMM backing allocation is ready."""
+            return self._torch_grouped_weight_layout_is_current()
+
+        @staticmethod
+        def _resolve_torch_grouped_mm():
+            """Return the public torch grouped GEMM, with the legacy private op as fallback."""
+            grouped_mm = getattr(F, 'grouped_mm', None)
+            if grouped_mm is None:
+                grouped_mm = getattr(torch, '_grouped_mm', None)
+            if grouped_mm is None:
+                raise RuntimeError(
+                    "this PyTorch build does not provide torch.nn.functional.grouped_mm "
+                    "or torch._grouped_mm"
+                )
+            return grouped_mm
+
+        def _validate_torch_grouped_mm_splits(self, m_splits) -> list[int]:
+            """Validate and normalize the token count for each local expert."""
+            try:
+                split_count = len(m_splits)
+            except TypeError as error:
+                raise RuntimeError(
+                    f"expected {self.num_gemms} expert splits, but expert splits are not a sequence"
+                ) from error
+            if split_count != self.num_gemms:
+                raise RuntimeError(f"expected {self.num_gemms} expert splits, got {split_count}")
+            if any(
+                isinstance(split, bool) or not isinstance(split, Integral) or split < 0
+                for split in m_splits
+            ):
+                raise RuntimeError("expert splits must be non-negative integers")
+            return [int(split) for split in m_splits]
+
+        def _torch_grouped_mm_forward(self, x, m_splits):
+            """Run the frozen expert base branch with torch grouped GEMM."""
+            m_splits = self._validate_torch_grouped_mm_splits(m_splits)
+            if not self._torch_grouped_weight_layout_is_current():
+                self.prepare_torch_grouped_mm()
+
+            x_2d = x.reshape(-1, x.shape[-1])
+            input_rows = sum(m_splits)
+            if input_rows != x_2d.shape[0]:
+                raise RuntimeError(
+                    f"expert splits sum to {input_rows}, but input has {x_2d.shape[0]} rows"
+                )
+
+            grouped_weight = self._buffers['_torch_grouped_weight']
+            if input_rows == 0:
+                output = x_2d.matmul(grouped_weight[0].transpose(0, 1))
+            else:
+                offsets = torch.tensor(
+                    list(accumulate(m_splits)), device=x.device, dtype=torch.int32
+                )
+                grouped_mm = self._resolve_torch_grouped_mm()
+                output = grouped_mm(x_2d, grouped_weight.transpose(1, 2), offs=offsets)
+            return output.reshape(*x.shape[:-1], output.shape[-1]), None
+
         def finish_init(self, quantization_config: QuantizationConfig):
             """Post-init of quantization override"""
             if quantization_config is None:
@@ -2647,6 +2785,9 @@ def will_execute_quantized(self, is_context_quantized: bool) -> bool:
 
         def forward(self, x, m_splits):
             """Forward."""
+            if self.expert_gemm_backend == 'torch':
+                return self._torch_grouped_mm_forward(x, m_splits)
+
             _is_first_microbatch = (
                 None if self.disable_parameter_transpose_cache else self.is_first_microbatch
             )

megatron/core/recompute.py

@@ -154,6 +154,13 @@ def chunk_runner(start: int, end: int, use_checkpoint: bool):
             if (start + layer_offset) in extract_layer_indices:
                 intermediate_hidden_states.append(hidden_states)
 
+    if not hidden_states.requires_grad:
+        # Re-entrant checkpointing only attaches a grad_fn when some tensor input
+        # requires grad. With a frozen embedding (adapter-only training) every chunk
+        # output would otherwise carry no grad_fn and the adapters inside the chunks
+        # would receive no gradient.
+        hidden_states = hidden_states.detach().requires_grad_(True)
+
     if self.config.recompute_method == 'uniform':
         # Uniformly divide the total number of layers and checkpoint
         # the input activation of each divided chunk.

megatron/core/transformer/transformer_config.py

@@ -843,6 +843,14 @@ class TransformerConfig(ModelParallelConfig):
     parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True``
     and ``add_bias_linear=True``."""
 
+    moe_expert_gemm_backend: Literal['transformer_engine', 'torch'] = 'transformer_engine'
+    """Backend for grouped expert linear layers during training.
+
+    ``torch`` uses ``torch._grouped_mm`` for frozen BF16, bias-free expert weights. It is intended
+    for parameter-efficient fine-tuning where only the adapter branch is trainable. The default
+    ``transformer_engine`` backend supports trainable expert weights and other precisions.
+    """
+
     moe_aux_loss_coeff: Union[float, List[float]] = 0.0
     """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended.
     If a list of load balancing types is provided for `moe_router_load_balancing_type`,
@@ -1551,6 +1559,60 @@ 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_expert_gemm_backend not in ('transformer_engine', 'torch'):
+            raise ValueError(
+                "moe_expert_gemm_backend must be 'transformer_engine' or 'torch', "
+                f"got {self.moe_expert_gemm_backend!r}"
+            )
+        if self.moe_expert_gemm_backend == 'torch':
+            if self.num_moe_experts is None:
+                raise ValueError("moe_expert_gemm_backend='torch' requires num_moe_experts")
+            if not self.moe_grouped_gemm:
+                raise ValueError("moe_expert_gemm_backend='torch' requires moe_grouped_gemm=True")
+            if not self.bf16 or self.params_dtype != torch.bfloat16:
+                raise ValueError("moe_expert_gemm_backend='torch' requires BF16 parameters")
+            if self.add_bias_linear:
+                raise ValueError("moe_expert_gemm_backend='torch' does not support expert bias")
+            if self.fp8 or self.fp4:
+                raise ValueError(
+                    "moe_expert_gemm_backend='torch' does not support FP8 or FP4 experts"
+                )
+            if self.moe_single_grouped_weight:
+                raise ValueError(
+                    "moe_expert_gemm_backend='torch' requires per-GEMM expert weight parameters "
+                    "and is incompatible with moe_single_grouped_weight=True"
+                )
+            if self.use_transformer_engine_op_fuser:
+                raise ValueError(
+                    "moe_expert_gemm_backend='torch' is incompatible with "
+                    "use_transformer_engine_op_fuser=True"
+                )
+            if self.delay_wgrad_compute:
+                raise ValueError(
+                    "moe_expert_gemm_backend='torch' is incompatible with "
+                    "delay_wgrad_compute=True"
+                )
+            if self.overlap_dispatch_backward_with_experts_wgrad:
+                raise ValueError(
+                    "moe_expert_gemm_backend='torch' is incompatible with "
+                    "overlap_dispatch_backward_with_experts_wgrad=True"
+                )
+            if self.transformer_impl != 'transformer_engine':
+                raise ValueError(
+                    "moe_expert_gemm_backend='torch' requires "
+                    "transformer_impl='transformer_engine' to construct TEGroupedMLP"
+                )
+
+            # Import lazily to avoid the transformer_engine -> TransformerConfig import cycle and
+            # prevent TEGroupedMLP from silently falling back to SequentialMLP.
+            from megatron.core.extensions.transformer_engine import TEColumnParallelGroupedLinear
+
+            if TEColumnParallelGroupedLinear is None:
+                raise ValueError(
+                    "moe_expert_gemm_backend='torch' requires Transformer Engine >= 1.9.0.dev0 "
+                    "with GroupedLinear support"
+                )
+
         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.")

tests/unit_tests/ssm/test_hybrid_block.py

@@ -228,6 +228,58 @@ def run(block, hs, am):
             gb, gr = base_grads[name], rec_grads[name]
             assert torch.equal(gr, gb), f"Grad should be bitwise matched for {name}"
 
+    @pytest.mark.parametrize(
+        ("recompute_method", "num_layers", "recompute_num_layers"),
+        [pytest.param("uniform", 2, 1, id="uniform"), pytest.param("block", 3, 2, id="block")],
+    )
+    def test_full_recompute_with_frozen_input_trains_adapter(
+        self, recompute_method: str, num_layers: int, recompute_num_layers: int
+    ):
+        """Full re-entrant recompute preserves gradients for adapter-only training."""
+        block = self.get_hybrid_block(
+            Symbols.MLP * num_layers,
+            add_bias_linear=False,
+            recompute_granularity="full",
+            recompute_method=recompute_method,
+            recompute_num_layers=recompute_num_layers,
+        ).cuda()
+        block.requires_grad_(False)
+        block.train()
+
+        def register_adapter(layer):
+            adapter = torch.nn.Sequential(
+                torch.nn.Linear(block.config.hidden_size, 4, bias=False),
+                torch.nn.Linear(4, block.config.hidden_size, bias=False),
+            ).cuda()
+            layer.add_module("adapter", adapter)
+
+            def apply_adapter(_module, _args, kwargs, output):
+                hidden_states, context = output
+                return hidden_states + adapter(kwargs["hidden_states"]), context
+
+            layer.register_forward_hook(apply_adapter, with_kwargs=True)
+            return adapter
+
+        # Block recompute checkpoints layer 0 but not the final layer in this configuration.
+        adapters = [register_adapter(block.layers[index]) for index in (0, num_layers - 1)]
+
+        sequence_length, micro_batch_size = 4, 1
+        hidden_states = torch.randn(
+            sequence_length, micro_batch_size, block.config.hidden_size, device="cuda"
+        )
+        attention_mask = torch.ones(
+            (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool, device="cuda"
+        )
+
+        assert not hidden_states.requires_grad
+        output = block(hidden_states, attention_mask=attention_mask)
+        output.float().square().mean().backward()
+
+        for adapter in adapters:
+            for parameter in adapter.parameters():
+                assert parameter.grad is not None
+                assert torch.count_nonzero(parameter.grad) > 0
+
     def test_layer_types(self):
         """
         Make sure that the layer types specified with layer_pattern