feat(peft): support multi-LoRA on grouped MoE expert linears

src/megatron/bridge/peft/multi_lora.py

@@ -29,14 +29,18 @@
 
 from megatron.bridge.peft.base import PEFT
 from megatron.bridge.peft.module_matcher import ModuleMatcher
-from megatron.bridge.peft.multi_lora_layers import MultiLoRALinear
-from megatron.bridge.peft.utils import is_expert_linear
+from megatron.bridge.peft.multi_lora_layers import (
+    MultiLoRAGroupedExpertLinear,
+    MultiLoRALinear,
+    install_moe_slot_routing,
+)
+from megatron.bridge.peft.utils import is_expert_linear, is_grouped_expert_linear
 
 
 logger = logging.getLogger(__name__)
 
-# One-shot flag so the "expert modules are skipped" warning fires once per
-# process rather than once per (many) expert linear.
+# One-shot flag so the "sequential expert modules are skipped" warning fires once
+# per process rather than once per (many) expert linear.
 _EXPERT_SKIP_WARNED = False
 
 
@@ -55,6 +59,9 @@ class MultiLoRA(PEFT, ModuleMatcher):
         lora_B_init_method: Initialisation method for the B matrix.
         a2a_experimental: Enable experimental all-to-all communication.
         lora_dtype: Data type for adapter weights.
+        normalize_moe_lora: Unsupported for multi-LoRA; see :meth:`__call__`.
+        share_expert_adapters: Unsupported for multi-LoRA; see :meth:`__call__`.
+        experts_shared_outer_loras: Unsupported for multi-LoRA; see :meth:`__call__`.
     """
 
     target_modules: List[str] = field(
@@ -69,6 +76,31 @@ class MultiLoRA(PEFT, ModuleMatcher):
     lora_B_init_method: str = "zero"
     a2a_experimental: bool = False
     lora_dtype: Optional[torch.dtype] = None
+    # Accepted (rather than rejected as unknown kwargs) so callers that share an
+    # argument surface with single-LoRA get an explicit error instead of a
+    # silently different adapter layout. Validated in __call__.
+    normalize_moe_lora: bool = False
+    share_expert_adapters: bool = False
+    experts_shared_outer_loras: bool = False
+
+    def __call__(self, model, training: bool = True):
+        """Apply multi-LoRA, then install MoE slot routing for wrapped expert linears."""
+        # Every slot shares one max-rank buffer and consumers slice all of an
+        # adapter's tensors to a single rank, so an expert-specific rank
+        # (normalize_moe_lora) or a layout that changes the exported tensor
+        # count per expert would break that contract rather than the layers.
+        for unsupported in ("normalize_moe_lora", "share_expert_adapters", "experts_shared_outer_loras"):
+            if getattr(self, unsupported):
+                raise NotImplementedError(
+                    f"MultiLoRA does not support {unsupported}=True; expert adapters use the "
+                    f"per-expert layout at the same max rank as every other target module."
+                )
+
+        model = super().__call__(model, training=training)
+        hooked = install_moe_slot_routing(model)
+        if hooked:
+            logger.info("Installed multi-lora MoE slot routing on %d MoE layer(s)", hooked)
+        return model
 
     def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optional[str] = None) -> nn.Module:
         if isinstance(module, MultiLoRALinear):
@@ -78,20 +110,37 @@ def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optio
             (match, full_name) = ans
 
             if is_expert_linear(full_name):
-                # MoE expert linears are not supported by the grouped-GEMM
-                # multi-LoRA layer; skipping keeps them from crashing, but on an
-                # MoE model that silently drops the experts (most of the
-                # trainable capacity) from LoRA. Warn once so it isn't invisible.
-                global _EXPERT_SKIP_WARNED
-                if not _EXPERT_SKIP_WARNED:
-                    logger.warning(
-                        "MultiLoRA does not support MoE expert linears; skipping all expert "
-                        "modules (e.g. %s). On MoE models only non-expert (attention/dense) "
-                        "layers will receive LoRA adapters.",
-                        full_name,
-                    )
-                    _EXPERT_SKIP_WARNED = True
-                return module
+                if not is_grouped_expert_linear(full_name):
+                    # SequentialMLP keeps one linear per expert
+                    # (mlp.experts.local_experts.N.linear_fc*), each seeing only
+                    # its own routed tokens in dispatcher order. Neither the
+                    # dense per-slot spans nor the grouped (slot, expert)
+                    # routing segments that, so skip with a one-shot warning.
+                    global _EXPERT_SKIP_WARNED
+                    if not _EXPERT_SKIP_WARNED:
+                        logger.warning(
+                            "MultiLoRA does not support sequential MoE expert linears; skipping "
+                            "them (e.g. %s). Use a grouped expert implementation "
+                            "(moe_grouped_gemm=True) to put adapters on experts.",
+                            full_name,
+                        )
+                        _EXPERT_SKIP_WARNED = True
+                    return module
+
+                logger.info(f"Adding multi-lora ({self.n_adapters} adapters) to expert: {full_name}")
+
+                return MultiLoRAGroupedExpertLinear(
+                    to_wrap=module,
+                    n_adapters=self.n_adapters,
+                    dim=self.dim,
+                    alpha=self.alpha,
+                    full_name=full_name,
+                    num_local_experts=module.num_gemms,
+                    column_init_method=self.lora_A_init_method,
+                    row_init_method=self.lora_B_init_method,
+                    dropout=self.dropout,
+                    dropout_position=self.dropout_position,
+                )
             if isinstance(module, TopKRouter):
                 return module