feat(peft): expand_shared_outer export flag for shared-outer MoE LoRA (#5588)

Signed-off-by: Hollow Man <hollowman@opensuse.org>
Co-authored-by: Kamran Jafari <kjafarisadeg@nvidia.com>

Author: HollowMan6

examples/conversion/adapter/export_adapter.py

@@ -113,6 +113,15 @@ def parse_args() -> argparse.Namespace:
             "Can be specified multiple times; e.g. `mtp.layers` excludes MTP adapters."
         ),
     )
+    parser.add_argument(
+        "--expand-shared-outer",
+        action="store_true",
+        help=(
+            "For shared-outer MoE LoRA, replicate the shared factor across experts under "
+            "per-expert 2D names (vLLM `pack_moe`) instead of a single shared `[1, ...]` tensor "
+            "(SGLang). Multiplies the shared factor's on-disk size by the expert count."
+        ),
+    )
     parser.add_argument("--tp", type=int, default=1, help="Tensor parallel size for distributed GPU export.")
     parser.add_argument("--pp", type=int, default=1, help="Pipeline parallel size for distributed GPU export.")
     parser.add_argument("--ep", type=int, default=1, help="Expert parallel size for distributed GPU export.")
@@ -231,6 +240,7 @@ def _export_adapter_distributed(args: argparse.Namespace) -> None:
             peft_config=lora,
             base_model_name_or_path=args.hf_model_path,
             exclude_adapter_base_prefixes=tuple(args.exclude_adapter_base_prefix),
+            expand_shared_outer=args.expand_shared_outer,
         )
     finally:
         if parallel_state.is_initialized():
@@ -253,6 +263,7 @@ def main() -> None:
             peft_checkpoint=args.lora_checkpoint,
             output_path=args.output,
             exclude_adapter_base_prefixes=tuple(args.exclude_adapter_base_prefix),
+            expand_shared_outer=args.expand_shared_outer,
         )
 
 

src/megatron/bridge/models/conversion/auto_bridge.py

@@ -858,6 +858,7 @@ def export_adapter_weights(
         cpu: bool = True,
         show_progress: bool = True,
         exclude_adapter_base_prefixes: Iterable[str] | None = None,
+        expand_shared_outer: bool = False,
     ) -> Iterable["HFWeightTuple"]:
         """
         Export only adapter weights from a Megatron model without merging them into base tensors.
@@ -871,16 +872,26 @@ def export_adapter_weights(
             show_progress: Display progress bar during export
             exclude_adapter_base_prefixes: Megatron adapter base prefixes to
                 skip before resolving HuggingFace parameter mappings.
+            expand_shared_outer: Replicate the shared factor across experts under per-expert
+                names (vLLM 2D ``pack_moe``) instead of a shared ``[1, ...]`` tensor (SGLang).
+                Default ``False``; no effect for non-shared-outer adapters.
 
         Yields:
             HFWeightTuple: Named tuples of (param_name, weight_tensor) for adapter parameters
+
+        Note:
+            With ``expand_shared_outer``, the per-expert copies of the shared factor alias one
+            storage rather than being cloned. ``safetensors.torch.save_file`` rejects tensors
+            that share memory, so callers serializing these tensors directly must clone them
+            first — :meth:`save_hf_adapter` already does.
         """
         bridge = self._model_bridge
         return bridge.stream_adapter_weights_megatron_to_hf(
             model,
             cpu=cpu,
             show_progress=show_progress,
             exclude_adapter_base_prefixes=exclude_adapter_base_prefixes,
+            expand_shared_outer=expand_shared_outer,
         )
 
     def save_hf_adapter(
@@ -891,6 +902,7 @@ def save_hf_adapter(
         base_model_name_or_path: Optional[str] = None,
         show_progress: bool = True,
         exclude_adapter_base_prefixes: Iterable[str] | None = None,
+        expand_shared_outer: bool = False,
     ) -> None:
         """Save LoRA adapter weights as a HuggingFace PEFT-compatible directory.
 
@@ -909,6 +921,8 @@ def save_hf_adapter(
             show_progress: Display progress bar during export.
             exclude_adapter_base_prefixes: Megatron adapter base prefixes to
                 skip before resolving HuggingFace parameter mappings.
+            expand_shared_outer: Replicate the shared factor across experts under per-expert
+                names (vLLM 2D ``pack_moe``). Default ``False`` keeps the PEFT shared ``[1, ...]`` layout.
 
         Example:
             >>> bridge.save_hf_adapter(
@@ -949,6 +963,7 @@ def save_hf_adapter(
                 cpu=True,
                 show_progress=show_progress,
                 exclude_adapter_base_prefixes=exclude_adapter_base_prefixes,
+                expand_shared_outer=expand_shared_outer,
             )
         ]
         if not raw_adapter_weights:
@@ -1561,6 +1576,7 @@ def export_adapter_ckpt(
         output_path: str | Path,
         show_progress: bool = True,
         exclude_adapter_base_prefixes: Iterable[str] | None = None,
+        expand_shared_outer: bool = False,
     ) -> None:
         """Export LoRA adapter weights from a Megatron PEFT checkpoint to HuggingFace PEFT format.
 
@@ -1582,6 +1598,10 @@ def export_adapter_ckpt(
             show_progress: Display progress bar during export.
             exclude_adapter_base_prefixes: Megatron adapter base prefixes to
                 skip before resolving HuggingFace parameter mappings.
+            expand_shared_outer: Replicate the shared factor across experts under per-expert
+                names (vLLM 2D ``pack_moe``). Default ``False`` keeps the PEFT shared ``[1, ...]``
+                layout. Enabling it multiplies the shared factor's on-disk size by the expert
+                count; no effect for non-shared-outer adapters.
 
         Example:
             >>> bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B")
@@ -1689,6 +1709,7 @@ def _load_and_export_adapter(model):
                 base_model_name_or_path=base_model_name,
                 show_progress=show_progress,
                 exclude_adapter_base_prefixes=exclude_adapter_base_prefixes,
+                expand_shared_outer=expand_shared_outer,
             )
 
         model_context = (

src/megatron/bridge/models/conversion/model_bridge.py

@@ -2370,6 +2370,7 @@ def stream_adapter_weights_megatron_to_hf(
     cpu: bool = True,
     show_progress: bool = True,
     exclude_adapter_base_prefixes: Optional[Iterable[str]] = None,
+    expand_shared_outer: bool = False,
 ) -> Iterable[HFWeightTuple]:
     """Bridge only adapter weights from Megatron to HuggingFace format."""
     ...
@@ -2462,13 +2463,15 @@ def _adapter_stream_registered_impl(
         cpu: bool = True,
         show_progress: bool = True,
         exclude_adapter_base_prefixes: Optional[Iterable[str]] = None,
+        expand_shared_outer: bool = False,
     ) -> Iterable[HFWeightTuple]:
         bridge = bridge_class()
         return bridge.stream_adapter_weights_megatron_to_hf(
             megatron_model,
             cpu=cpu,
             show_progress=show_progress,
             exclude_adapter_base_prefixes=exclude_adapter_base_prefixes,
+            expand_shared_outer=expand_shared_outer,
         )
 
     # Set meaningful names for debugging

src/megatron/bridge/models/conversion/peft_bridge.py

@@ -853,6 +853,7 @@ def stream_adapter_weights_megatron_to_hf(
         cpu: bool = True,
         show_progress: bool = True,
         exclude_adapter_base_prefixes: Iterable[str] | None = None,
+        expand_shared_outer: bool = False,
     ) -> Iterable["HFWeightTuple"]:
         """Stream only adapter weights without merging them into base tensors.
 
@@ -895,6 +896,7 @@ def stream_adapter_weights_megatron_to_hf(
                     linear_out_tensor,
                     num_moe_experts,
                     cpu,
+                    expand_shared_outer=expand_shared_outer,
                 )
                 continue
 
@@ -1034,14 +1036,15 @@ def _stream_shared_outer_adapter_weights(
         linear_out_tensor: torch.Tensor,
         num_moe_experts: int,
         cpu: bool,
+        expand_shared_outer: bool,
     ) -> Iterable["HFWeightTuple"]:
         """Stream a shared-outer grouped-expert LoRA adapter (SGLang PR #21466).
 
-        One side is a 2D LoRA matrix replicated across local experts; the other
-        is a per-expert 3D pack. The shared side is emitted once as a ``[1, ...]``
-        tensor under the expert-agnostic HF name (so the serving loader takes its
-        3D-shared branch); the per-expert side is gathered across EP ranks and
-        emitted once per global expert.
+        One side is a 2D LoRA matrix shared across experts; the other is a
+        per-expert 3D pack. By default the shared side is emitted once as a
+        ``[1, ...]`` tensor under the expert-agnostic HF name. With
+        ``expand_shared_outer``, it is replicated under per-expert 2D names
+        (vLLM 2D ``pack_moe`` contract); the training-side parameter stays shared.
         """
 
         from megatron.bridge.models.conversion.model_bridge import HFWeightTuple
@@ -1051,7 +1054,7 @@ def _stream_shared_outer_adapter_weights(
             (linear_in_tensor, ".linear_in.weight"),
             (linear_out_tensor, ".linear_out.weight"),
         ):
-            if side_tensor.ndim == 2:
+            if side_tensor.ndim == 2 and not expand_shared_outer:
                 # Shared side: emit one [1, out, in] tensor. A shared linear_in
                 # feeding a fused gate/up FC1 maps to two HF names, so the same
                 # tensor is emitted for each projection.
@@ -1066,6 +1069,24 @@ def _stream_shared_outer_adapter_weights(
                     yield HFWeightTuple(hf_name, current)
                 continue
 
+            if side_tensor.ndim == 2 and expand_shared_outer:
+                # Expand the shared factor under per-expert 2D names (vLLM pack_moe).
+                # The tensor is reused across experts, not cloned.
+                shared_current = side_tensor.cpu() if cpu else side_tensor
+                for expert_idx in range(num_moe_experts):
+                    base_hf_weight_names = self._get_base_hf_param_names_for_adapter(
+                        mapping_registry,
+                        adapter_task.global_base_prefix,
+                        adapter_task.adapter_key,
+                        f".weight{expert_idx}",
+                    )
+                    for base_name in base_hf_weight_names:
+                        hf_name = self._make_lora_param_name(base_name, side_suffix)
+                        if hf_name is None:
+                            continue
+                        yield HFWeightTuple(hf_name, shared_current)
+                continue
+
             # Per-expert side: emit one slice per global expert. A fused FC1
             # linear_out (gate+up) is split per HF projection name; otherwise the
             # single projection is emitted directly.