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

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,6 +872,9 @@ 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
@@ -881,6 +885,7 @@ def export_adapter_weights(
             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 +896,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 +915,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 +957,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:

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.

tests/unit_tests/models/test_adapter_export.py

@@ -1424,3 +1424,138 @@ def test_export_adapter_distributed_enables_legacy_shared_expert_adapter_loading
         mock_enable_legacy.assert_called_once_with([model_chunk], state_dicts[0], tmp_path)
         mock_load.assert_called_once_with(state_dicts[1], str(tmp_path), validate_access_integrity=False)
         model_chunk.load_state_dict.assert_called_once_with({"adapter": "weights"}, strict=False)
+
+
+class TestStreamSharedOuterAdapterWeights:
+    """Tests for ``MegatronPeftBridge._stream_shared_outer_adapter_weights``.
+
+    Verifies the two export contracts for the shared-outer grouped-expert LoRA
+    adapter: the default SGLang shared ``[1, ...]`` layout vs. the ``expand_shared_outer``
+    per-expert 2D (vLLM ``pack_moe``) layout, plus the unchanged per-expert side.
+    """
+
+    def _make_bridge(self):
+        return MegatronPeftBridge()
+
+    def _run(
+        self, bridge, linear_in, linear_out, num_experts, expand, cpu=False, mapping_registry=None, megatron_model=None
+    ):
+        # linear_in_task / linear_out_task are unused by this method.
+        task = SimpleNamespace(global_base_prefix="mlp.experts", adapter_key=None)
+        gen = bridge._stream_shared_outer_adapter_weights(
+            megatron_model=megatron_model,
+            mapping_registry=mapping_registry,
+            adapter_task=task,
+            linear_in_tensor=linear_in,
+            linear_out_tensor=linear_out,
+            num_moe_experts=num_experts,
+            cpu=cpu,
+            expand_shared_outer=expand,
+        )
+        return list(gen)
+
+    def _mapping(self):
+        # ``_get_base_hf_param_names_for_adapter`` looks up ``f"{prefix}{suffix}"``.
+        # Both the shared side (``.weight0``) and the per-expert side (``.weight{idx}``)
+        # map to per-expert HF names for the same grouped-expert linear (``gate_proj``);
+        # the shared-side branch strips the expert index via ``_strip_hf_expert_index``.
+        mapping = MagicMock()
+
+        def lookup(name):
+            idx = name.rsplit(".weight", 1)[1]
+            return SimpleNamespace(hf_param=f"model.layers.0.mlp.experts.{idx}.gate_proj.weight")
+
+        mapping.megatron_to_hf_lookup.side_effect = lookup
+        return mapping
+
+    @patch(
+        "megatron.bridge.models.conversion.peft_bridge.parallel_state.get_expert_model_parallel_world_size",
+        return_value=1,
+    )
+    def test_shared_side_default_emits_shared_1d_aggregate(self, _mock_ep):
+        """expand=False: shared 2D side is unsqueezed to [1, out, in] under expert-agnostic names."""
+        bridge = self._make_bridge()
+        linear_in = torch.randn(4, 8)  # [rank, hidden] shared across experts
+        linear_out = torch.randn(2, 2, 4)  # [num_experts, out, rank] per-expert side
+
+        with (
+            patch.object(bridge, "_gather_expert_adapter_weight", return_value=None),
+            patch.object(bridge, "_select_expert_adapter_weight", side_effect=lambda w, g, i, n: w[i]),
+            patch.object(bridge, "_get_fused_adapter_linear_out_slices", return_value=None),
+            patch("megatron.bridge.models.conversion.peft_bridge.is_expert_linear", return_value=True),
+        ):
+            out = self._run(
+                bridge, linear_in, linear_out, num_experts=2, expand=False, mapping_registry=self._mapping()
+            )
+
+        # Shared side (linear_in/A): one [1, rank, hidden] tensor under an expert-agnostic name.
+        shared = [t for t in out if t.param_name.endswith(".lora_A.weight")]
+        # One [1, rank, hidden] tensor under an expert-agnostic name (experts.gate_proj.lora_A).
+        assert len(shared) == 1
+        assert shared[0].weight.shape == (1, 4, 8)
+        assert shared[0].param_name == "model.layers.0.mlp.experts.gate_proj.lora_A.weight"
+
+    @patch(
+        "megatron.bridge.models.conversion.peft_bridge.parallel_state.get_expert_model_parallel_world_size",
+        return_value=1,
+    )
+    def test_shared_side_expanded_emits_per_expert_2d(self, _mock_ep):
+        """expand=True: shared side replicated under per-expert 2D names (vLLM pack_moe)."""
+        bridge = self._make_bridge()
+        linear_in = torch.randn(4, 8)
+        linear_out = torch.randn(2, 2, 4)
+
+        with (
+            patch.object(bridge, "_gather_expert_adapter_weight", return_value=None),
+            patch.object(bridge, "_select_expert_adapter_weight", side_effect=lambda w, g, i, n: w[i]),
+            patch.object(bridge, "_get_fused_adapter_linear_out_slices", return_value=None),
+            patch("megatron.bridge.models.conversion.peft_bridge.is_expert_linear", return_value=True),
+        ):
+            out = self._run(
+                bridge, linear_in, linear_out, num_experts=2, expand=True, mapping_registry=self._mapping()
+            )
+
+        shared = [t for t in out if t.param_name.endswith(".lora_A.weight")]
+        # One 2D tensor per expert, kept 2D (no unsqueeze), per-expert names.
+        assert len(shared) == 2
+        assert [t.param_name for t in shared] == [
+            "model.layers.0.mlp.experts.0.gate_proj.lora_A.weight",
+            "model.layers.0.mlp.experts.1.gate_proj.lora_A.weight",
+        ]
+        for t in shared:
+            assert t.weight.shape == (4, 8)
+        # Same underlying shared tensor reused across experts (not cloned).
+        assert shared[0].weight.data_ptr() == shared[1].weight.data_ptr()
+
+    @patch(
+        "megatron.bridge.models.conversion.peft_bridge.parallel_state.get_expert_model_parallel_world_size",
+        return_value=1,
+    )
+    def test_per_expert_side_unaffected_by_flag(self, _mock_ep):
+        """The per-expert 3D side emits one slice per expert regardless of expand."""
+        bridge = self._make_bridge()
+        linear_in = torch.randn(4, 8)
+        linear_out = torch.randn(2, 2, 4)  # per-expert side
+
+        results = {}
+        for expand in (False, True):
+            with (
+                patch.object(bridge, "_gather_expert_adapter_weight", return_value=None),
+                patch.object(bridge, "_select_expert_adapter_weight", side_effect=lambda w, g, i, n: w[i]),
+                patch.object(bridge, "_get_fused_adapter_linear_out_slices", return_value=None),
+                patch("megatron.bridge.models.conversion.peft_bridge.is_expert_linear", return_value=True),
+            ):
+                out = self._run(
+                    bridge, linear_in, linear_out, num_experts=2, expand=expand, mapping_registry=self._mapping()
+                )
+            results[expand] = [
+                (t.param_name, tuple(t.weight.shape)) for t in out if t.param_name.endswith(".lora_B.weight")
+            ]
+
+        assert results[False] == results[True]
+        assert [n for n, _ in results[False]] == [
+            "model.layers.0.mlp.experts.0.gate_proj.lora_B.weight",
+            "model.layers.0.mlp.experts.1.gate_proj.lora_B.weight",
+        ]
+        for _, shape in results[False]:
+            assert shape == (2, 4)