feat(moe+ssm): expert-LoRA GEMM stack and hybrid-recompute guard on the d12f6c8c base

megatron/core/extensions/transformer_engine.py

@@ -11,6 +11,7 @@
 import re
 import warnings
 from contextlib import contextmanager, nullcontext
+from itertools import accumulate
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, cast
 
 import torch
@@ -123,6 +124,65 @@ def _set_expert_parameter_attributes(
             param.partition_stride = 1
 
 
+def _run_torch_grouped_mm(x: Tensor, weight: Tensor, offsets: Tensor) -> Tensor:
+    """Run grouped matrix multiplication with a CPU reference for unit tests."""
+    if x.device.type == 'cuda':
+        return torch._grouped_mm(x, weight, offs=offsets)
+
+    outputs = []
+    start = 0
+    for group, end_tensor in enumerate(offsets):
+        end = int(end_tensor)
+        outputs.append(x[start:end].matmul(weight[group]))
+        start = end
+    if not outputs:
+        return x.new_empty((0, weight.shape[-1]))
+    return torch.cat(outputs, dim=0)
+
+
+class _TorchGroupedFusedLoRA(torch.autograd.Function):
+    """Differentiate a frozen grouped GEMM with a shared low-rank branch."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        augmented_weight: Tensor,
+        lora_a: Tensor,
+        lora_b: Tensor,
+        offsets: Tensor,
+        output_features: int,
+        scale: float,
+        counters: Dict[str, int],
+    ) -> Tensor:
+        """Forward pass."""
+        with torch.no_grad():
+            augmented_weight[:, output_features:, :].copy_(lora_a)
+        augmented_output = _run_torch_grouped_mm(x, augmented_weight.transpose(1, 2), offsets)
+        base_output = augmented_output[:, :output_features]
+        low_rank = augmented_output[:, output_features:]
+        output = base_output + scale * low_rank.matmul(lora_b.transpose(0, 1))
+
+        ctx.save_for_backward(x, lora_b, low_rank, offsets)
+        ctx.augmented_weight = augmented_weight
+        ctx.scale = scale
+        ctx.counters = counters
+        counters["forward_calls"] += 1
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: Tensor):
+        """Backward pass."""
+        x, lora_b, low_rank, offsets = ctx.saved_tensors
+        grad_low_rank = ctx.scale * grad_output.matmul(lora_b)
+        grad_augmented_output = torch.cat((grad_output, grad_low_rank), dim=-1)
+        grad_x = _run_torch_grouped_mm(grad_augmented_output, ctx.augmented_weight, offsets)
+        grad_lora_a = grad_low_rank.transpose(0, 1).matmul(x)
+        grad_lora_b = ctx.scale * grad_output.transpose(0, 1).matmul(low_rank)
+        ctx.counters["backward_calls"] += 1
+        return grad_x, None, grad_lora_a, grad_lora_b, None, None, None, None
+
+
 class TransformerEngineConfigType(enum.Enum):
     """Configuration object types in config dictionary"""
 
@@ -2358,6 +2418,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 +2693,282 @@ 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
+
+        def _torch_grouped_fused_lora_layout_is_current(self, rank: int) -> bool:
+            """Return whether expert parameters view the augmented grouped buffer."""
+            augmented_weight = self._buffers.get("_torch_grouped_fused_lora_weight")
+            if augmented_weight is None or augmented_weight.shape[0] != self.num_gemms:
+                return False
+
+            first_weight = self.weight0
+            expected_shape = (self.num_gemms, first_weight.shape[0] + rank, first_weight.shape[1])
+            if (
+                augmented_weight.device != first_weight.device
+                or augmented_weight.dtype != first_weight.dtype
+                or augmented_weight.shape != expected_shape
+            ):
+                return False
+
+            storage_pointer = augmented_weight.untyped_storage().data_ptr()
+            group_elements = augmented_weight.shape[1] * augmented_weight.shape[2]
+            for index in range(self.num_gemms):
+                weight = getattr(self, f"weight{index}")
+                weight_view = augmented_weight[index, : first_weight.shape[0], :]
+                if (
+                    weight.device != augmented_weight.device
+                    or weight.dtype != augmented_weight.dtype
+                    or weight.shape != weight_view.shape
+                    or weight.untyped_storage().data_ptr() != storage_pointer
+                    or weight.storage_offset() != index * group_elements
+                    or weight.stride() != 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'"
+                )
+            if not hasattr(torch, '_grouped_mm'):
+                raise RuntimeError("this PyTorch build does not provide 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]
+                self._buffers.pop("_torch_grouped_fused_lora_weight", None)
+                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()
+
+        @torch.no_grad()
+        def prepare_torch_grouped_mm_fused_lora(self, lora_a: Tensor) -> Dict[str, Any]:
+            """Prepare augmented frozen expert weights for shared LoRA-A fusion."""
+            if self.expert_gemm_backend != "torch":
+                raise RuntimeError("fused expert LoRA requires moe_expert_gemm_backend='torch'")
+            if not hasattr(torch, "_grouped_mm"):
+                raise RuntimeError("this PyTorch build does not provide torch._grouped_mm")
+            if self.use_bias:
+                raise RuntimeError("fused expert LoRA does not support expert bias")
+
+            weights = [getattr(self, f"weight{index}") for index in range(self.num_gemms)]
+            if any(weight.requires_grad for weight in weights):
+                raise RuntimeError("fused expert LoRA requires frozen base weights")
+            if any(weight.device.type != "cuda" for weight in weights):
+                raise RuntimeError("fused expert LoRA requires CUDA weights")
+            if any(weight.dtype != torch.bfloat16 for weight in weights):
+                raise RuntimeError("fused expert LoRA requires BF16 base weights")
+            if any(weight.shape != weights[0].shape for weight in weights[1:]):
+                raise RuntimeError("fused expert LoRA requires uniform expert shapes")
+            if lora_a.device != weights[0].device or lora_a.dtype != torch.bfloat16:
+                raise RuntimeError("fused expert LoRA requires a BF16 LoRA-A on the base device")
+            if lora_a.ndim != 2 or lora_a.shape[1] != weights[0].shape[1]:
+                raise RuntimeError(
+                    f"LoRA-A shape {tuple(lora_a.shape)} does not match base input "
+                    f"features {weights[0].shape[1]}"
+                )
+
+            memory_allocated_before = torch.cuda.memory_allocated(lora_a.device)
+            memory_allocated_at_coexistence = memory_allocated_before
+            peak_allocated_before = torch.cuda.max_memory_allocated(lora_a.device)
+            rank = lora_a.shape[0]
+            if not self._torch_grouped_fused_lora_layout_is_current(rank):
+                output_features, input_features = weights[0].shape
+                augmented_weight = weights[0].new_empty(
+                    (self.num_gemms, output_features + rank, input_features)
+                )
+                memory_allocated_at_coexistence = torch.cuda.memory_allocated(lora_a.device)
+                for index, weight in enumerate(weights):
+                    augmented_weight[index, :output_features, :].copy_(weight)
+                    weight.data = augmented_weight[index, :output_features, :]
+                augmented_weight[:, output_features:, :].copy_(lora_a)
+                self._buffers.pop("_torch_grouped_weight", None)
+                if "_torch_grouped_fused_lora_weight" in self._buffers:
+                    self._buffers["_torch_grouped_fused_lora_weight"] = augmented_weight
+                else:
+                    self.register_buffer(
+                        "_torch_grouped_fused_lora_weight", augmented_weight, persistent=False
+                    )
+
+            augmented_weight = self._buffers["_torch_grouped_fused_lora_weight"]
+            if "_torch_grouped_weight" in self._buffers:
+                raise RuntimeError("fused expert LoRA retained the separate grouped base buffer")
+            if not self._torch_grouped_fused_lora_layout_is_current(rank):
+                raise RuntimeError(
+                    "fused expert LoRA parameters do not share the augmented backing storage"
+                )
+            base_storage_bytes = sum(weight.numel() * weight.element_size() for weight in weights)
+            resident_storage_bytes = augmented_weight.numel() * augmented_weight.element_size()
+            memory_allocated_after = torch.cuda.memory_allocated(lora_a.device)
+            return {
+                "base_storage_bytes": base_storage_bytes,
+                "augmentation_storage_bytes": resident_storage_bytes - base_storage_bytes,
+                "resident_storage_bytes": resident_storage_bytes,
+                "duplicate_base_buffer_present": False,
+                "base_parameters_share_augmented_storage": True,
+                "cuda_memory_allocated_before": memory_allocated_before,
+                "cuda_memory_allocated_at_coexistence": memory_allocated_at_coexistence,
+                "cuda_relocation_peak_delta": (
+                    memory_allocated_at_coexistence - memory_allocated_before
+                ),
+                "cuda_memory_allocated_after": memory_allocated_after,
+                "cuda_memory_allocated_delta": memory_allocated_after - memory_allocated_before,
+                "cuda_peak_allocated_before": peak_allocated_before,
+                "cuda_peak_allocated_after": torch.cuda.max_memory_allocated(lora_a.device),
+            }
+
+        def torch_grouped_mm_fused_lora_forward(
+            self,
+            x: Tensor,
+            m_splits: List[int],
+            *,
+            lora_a: Tensor,
+            lora_b: Tensor,
+            scale: float,
+            adapter_enabled: bool,
+            counters: Dict[str, int],
+        ) -> Tuple[Tensor, None]:
+            """Run the frozen grouped base and shared LoRA branch in one grouped GEMM."""
+            if not self._torch_grouped_fused_lora_layout_is_current(lora_a.shape[0]):
+                self.prepare_torch_grouped_mm_fused_lora(lora_a)
+            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"
+                )
+
+            augmented_weight = self._buffers["_torch_grouped_fused_lora_weight"]
+            output_features = self.weight0.shape[0]
+            if lora_b.device != x.device or lora_b.dtype != torch.bfloat16:
+                raise RuntimeError("fused expert LoRA requires a BF16 LoRA-B on the input device")
+            if lora_b.shape != (output_features, lora_a.shape[0]):
+                raise RuntimeError(
+                    f"LoRA-B shape {tuple(lora_b.shape)} does not match "
+                    f"({output_features}, {lora_a.shape[0]})"
+                )
+
+            if input_rows == 0:
+                base_output = x_2d.matmul(self.weight0.transpose(0, 1))
+                if adapter_enabled:
+                    output = base_output + scale * x_2d.matmul(lora_a.transpose(0, 1)).matmul(
+                        lora_b.transpose(0, 1)
+                    )
+                    counters["forward_calls"] += 1
+                else:
+                    output = base_output
+                    counters["base_only_forward_calls"] += 1
+            else:
+                offsets = torch.tensor(
+                    list(accumulate(m_splits)), device=x.device, dtype=torch.int32
+                )
+                if adapter_enabled:
+                    output = _TorchGroupedFusedLoRA.apply(
+                        x_2d,
+                        augmented_weight,
+                        lora_a,
+                        lora_b,
+                        offsets,
+                        output_features,
+                        scale,
+                        counters,
+                    )
+                else:
+                    augmented_output = torch._grouped_mm(
+                        x_2d, augmented_weight.transpose(1, 2), offs=offsets
+                    )
+                    output = augmented_output[:, :output_features]
+                    counters["base_only_forward_calls"] += 1
+            return output.reshape(*x.shape[:-1], output.shape[-1]), None
+
+        @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()
+
+        def _torch_grouped_mm_forward(self, x, m_splits):
+            """Run the frozen expert base branch with ``torch._grouped_mm``."""
+            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
+                )
+                output = torch._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 +2984,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.