feat(peft): add moe_shared_loras export layout for shared-outer MoE LoRA

docs/training/peft.md

@@ -179,6 +179,25 @@ finalize path takes over (Bridge's setup removes the hooks before the first step
 the wrapper's first call does). The fallback cannot cover fused weight-gradient accumulation
 and logs a warning in that case.
 
+##### Exporting shared-outer MoE adapters
+
+`AutoBridge.export_adapter_weights` emits shared-outer adapters in the legacy (SGLang)
+layout unless a serving layout is requested: the shared factor once as a `[1, ...]` tensor
+under the expert-agnostic name, plus per-expert factors under `experts.N.*` names. Three
+opt-in flags select a serving layout instead, listed from highest to lowest priority:
+
+| Option | Default | Export layout when enabled |
+|--------|---------|----------------------------|
+| `stack_3d_moe` | `False` | vLLM 3D-MoE layout: stacked tensors under `experts.base_layer` (gate/up) and `experts` (down) |
+| `moe_shared_loras` | `False` | vLLM 2D shared-MoE layout under `experts.w1/w2/w3`, with the shared factors kept at expert-dim 1 |
+| `expand_shared_outer` | `False` | vLLM 2D `pack_moe` layout: the shared factor replicated under per-expert `experts.N.*` names |
+
+The flags can be combined; when several are set, the highest-priority one decides the
+layout. With all three unset, existing callers keep the legacy layout. `moe_shared_loras=True`
+is the 2D serving format used when vLLM's `enable_moe_shared_loras` setting is enabled; it
+does not imply a 3D model. Callers can select `stack_3d_moe` from the model's MoE tensor
+layout and `moe_shared_loras` from the serving backend's adapter mode.
+
 #### LoRA+
 
 [LoRA+](https://arxiv.org/abs/2402.12354) trains the B matrix (`linear_out`) at a higher learning rate than A (`linear_in`) by a fixed ratio `lr_B = lora_plus_ratio * lr_A` (paper default 16). The ratio is set on the optimizer, not the `LoRA` adapter: when it differs from `1.0`, `get_lora_plus_config_overrides` (`megatron.bridge.peft.lora`) builds Megatron-Core per-group overrides giving the `*.linear_out.weight` group `max_lr`/`min_lr` scaled by the ratio. Pass `1.0` (or skip the call) to keep a single learning rate.

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

@@ -294,34 +294,29 @@ def _drop_readonly_config_properties(
     return {key: value for key, value in config_dict.items() if key not in readonly_properties}
 
 
-def _check_streaming_argument_support(stream_fn: Callable[..., Any], owner: object, argument: str) -> None:
-    """Fail early when a streaming override cannot accept an export argument."""
+def _check_stream_kwarg_support(stream_fn: Callable[..., Any], owner: object, kwarg: str) -> None:
+    """Fail up front when a bridge's streaming override cannot accept an opt-in keyword.
+
+    Opt-in keywords such as ``with_megatron_names`` or ``moe_shared_loras`` are only forwarded
+    when enabled. Bridges that override the streaming export with an explicit signature (rather
+    than ``*args, **kwargs``) have to forward them themselves; without this check the request
+    would only surface as an opaque ``TypeError`` once the generator is consumed.
+    """
     try:
         parameters = inspect.signature(stream_fn).parameters
     except (TypeError, ValueError):
         # Not introspectable (C callables); let the call itself decide.
         return
-    if argument in parameters or any(
+    if kwarg in parameters or any(
         parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()
     ):
         return
     raise TypeError(
         f"{type(owner).__name__}.{getattr(stream_fn, '__name__', 'stream_weights')} does not accept "
-        f"'{argument}'; a bridge that overrides the streaming export must forward this flag "
-        "before the requested export mode can be used."
+        f"'{kwarg}'; a bridge that overrides the streaming export must forward this flag before it can be enabled."
     )
 
 
-def _check_with_megatron_names_support(stream_fn: Callable[..., Any], owner: object) -> None:
-    """Fail up front when a streamer cannot accept ``with_megatron_names``."""
-    _check_streaming_argument_support(stream_fn, owner, "with_megatron_names")
-
-
-def _check_current_pp_stage_only_support(stream_fn: Callable[..., Any], owner: object) -> None:
-    """Fail early when a streamer cannot accept PP-local export."""
-    _check_streaming_argument_support(stream_fn, owner, "current_pp_stage_only")
-
-
 class AutoBridge(Generic[MegatronModelT]):
     """
     Automatically select and instantiate the appropriate bridge for a model.
@@ -869,9 +864,9 @@ def export_hf_weights(
 
         bridge = self._model_bridge
         if with_megatron_names:
-            _check_with_megatron_names_support(bridge.stream_weights_megatron_to_hf, bridge)
+            _check_stream_kwarg_support(bridge.stream_weights_megatron_to_hf, bridge, "with_megatron_names")
         if current_pp_stage_only:
-            _check_current_pp_stage_only_support(bridge.stream_weights_megatron_to_hf, bridge)
+            _check_stream_kwarg_support(bridge.stream_weights_megatron_to_hf, bridge, "current_pp_stage_only")
         return bridge.stream_weights_megatron_to_hf(
             model,
             self.hf_pretrained,
@@ -978,6 +973,7 @@ def export_adapter_weights(
         expand_shared_outer: bool = False,
         stack_3d_moe: bool = False,
         with_megatron_names: bool = False,
+        moe_shared_loras: bool = False,
     ) -> Iterable["HFWeightTuple"]:
         """
         Export only adapter weights from a Megatron model without merging them into base tensors.
@@ -999,6 +995,11 @@ def export_adapter_weights(
                 (``...experts.base_layer`` for gate_up_proj, bare ``...experts`` for
                 down_proj), instead of the per-expert 2D ``pack_moe`` layout.
                 Default ``False``; no effect for non-shared-outer adapters.
+            moe_shared_loras: Emit shared-outer routed-expert LoRA in the 2D ``experts.w1/w2/w3``
+                layout consumed by vLLM's ``enable_moe_shared_loras`` mode, keeping the shared
+                factors at expert-dim 1. When several layout flags are set, ``stack_3d_moe`` wins
+                over ``moe_shared_loras``, which wins over ``expand_shared_outer``.
+                Default ``False``; no effect for non-shared-outer adapters.
             with_megatron_names: Yield ``HFSourcedWeightTuple`` whose ``megatron_param_names``
                 holds the adapter's ``linear_in`` (lora_A) or ``linear_out`` (lora_B) Megatron
                 weight name. Default ``False`` keeps the two-field tuple.
@@ -1014,16 +1015,22 @@ def export_adapter_weights(
             first — :meth:`save_hf_adapter` already does.
         """
         bridge = self._model_bridge
+        # Only forward opt-in flags when set so bridges with a custom adapter streamer keep working by default.
+        opt_in_kwargs: dict[str, bool] = {}
         if with_megatron_names:
-            _check_with_megatron_names_support(bridge.stream_adapter_weights_megatron_to_hf, bridge)
+            opt_in_kwargs["with_megatron_names"] = True
+        if moe_shared_loras:
+            opt_in_kwargs["moe_shared_loras"] = True
+        for kwarg in opt_in_kwargs:
+            _check_stream_kwarg_support(bridge.stream_adapter_weights_megatron_to_hf, bridge, kwarg)
         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,
             stack_3d_moe=stack_3d_moe,
-            **({"with_megatron_names": True} if with_megatron_names else {}),
+            **opt_in_kwargs,
         )
 
     def save_hf_adapter(

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

@@ -2735,6 +2735,7 @@ def stream_adapter_weights_megatron_to_hf(
     exclude_adapter_base_prefixes: Optional[Iterable[str]] = None,
     expand_shared_outer: bool = False,
     stack_3d_moe: bool = False,
+    moe_shared_loras: bool = False,
 ) -> Iterable[HFWeightTuple]:
     """Bridge only adapter weights from Megatron to HuggingFace format."""
     ...
@@ -2831,6 +2832,7 @@ def _adapter_stream_registered_impl(
         exclude_adapter_base_prefixes: Optional[Iterable[str]] = None,
         expand_shared_outer: bool = False,
         stack_3d_moe: bool = False,
+        moe_shared_loras: bool = False,
     ) -> Iterable[HFWeightTuple]:
         bridge = bridge_class()
         return bridge.stream_adapter_weights_megatron_to_hf(
@@ -2840,6 +2842,8 @@ def _adapter_stream_registered_impl(
             exclude_adapter_base_prefixes=exclude_adapter_base_prefixes,
             expand_shared_outer=expand_shared_outer,
             stack_3d_moe=stack_3d_moe,
+            # Only forward the opt-in flag when set so bridges with a custom adapter streamer keep working by default.
+            **({"moe_shared_loras": True} if moe_shared_loras else {}),
         )
 
     # Set meaningful names for debugging

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

@@ -883,6 +883,7 @@ def stream_adapter_weights_megatron_to_hf(
         expand_shared_outer: bool = False,
         stack_3d_moe: bool = False,
         with_megatron_names: bool = False,
+        moe_shared_loras: bool = False,
     ) -> Iterable["HFWeightTuple"]:
         """Stream only adapter weights without merging them into base tensors.
 
@@ -896,6 +897,11 @@ def stream_adapter_weights_megatron_to_hf(
         2D ``pack_moe`` layout. It is the vLLM-3D-MoE analogue of ``expand_shared_outer``
         and takes precedence over it for shared-outer adapters.
 
+        ``moe_shared_loras`` emits the 2D shared-expert layout consumed by vLLM's
+        ``FusedMoEWithLoRA`` in ``enable_moe_shared_loras`` mode (``experts.w1/w2/w3``).
+        When several layout flags are set, ``stack_3d_moe`` wins over ``moe_shared_loras``,
+        which wins over ``expand_shared_outer``.
+
         ``with_megatron_names`` yields :class:`HFSourcedWeightTuple` values whose
         ``megatron_param_names`` is the one-element tuple holding the adapter's ``linear_in``
         (lora_A) or ``linear_out`` (lora_B) Megatron weight name; the default keeps the
@@ -940,6 +946,7 @@ def stream_adapter_weights_megatron_to_hf(
                     cpu,
                     expand_shared_outer=expand_shared_outer,
                     stack_3d_moe=stack_3d_moe,
+                    moe_shared_loras=moe_shared_loras,
                     source_names=source_names,
                     with_megatron_names=with_megatron_names,
                 )
@@ -1072,6 +1079,98 @@ def stream_adapter_weights_megatron_to_hf(
                 yield emit_in(linear_in_hf_names[0], current_linear_in_tensor)
                 yield emit_out(linear_out_hf_names[0], current_linear_out_tensor)
 
+    def _stream_shared_outer_adapter_weights_2d_moe(
+        self,
+        megatron_model: List[MegatronModel],
+        mapping_registry: "MegatronMappingRegistry",
+        adapter_task: AdapterWeightConversionTask,
+        linear_in_tensor: torch.Tensor,
+        linear_out_tensor: torch.Tensor,
+        num_moe_experts: int,
+        cpu: bool,
+        source_names: Optional[Tuple[str, str]] = None,
+        with_megatron_names: bool = False,
+    ) -> Iterable["HFWeightTuple"]:
+        """Reuse the ordinary shared-outer emission, then stack it for vLLM 2D MoE.
+
+        The unstacked path is the source of truth for mapping lookup, EP gather,
+        and fused gate/up splitting. This adapter only changes the final names and
+        expert axis to the ``experts.w1/w2/w3`` contract.
+        """
+        emitted = list(
+            self._stream_shared_outer_adapter_weights(
+                megatron_model,
+                mapping_registry,
+                adapter_task,
+                linear_in_tensor,
+                linear_out_tensor,
+                num_moe_experts,
+                cpu,
+                expand_shared_outer=False,
+                stack_3d_moe=False,
+                moe_shared_loras=False,
+                source_names=source_names,
+                with_megatron_names=with_megatron_names,
+            )
+        )
+        if not emitted:
+            return
+
+        in_name, out_name = source_names if source_names is not None else (None, None)
+        emit_in = functools.partial(self._make_hf_weight, in_name, with_megatron_names)
+        emit_out = functools.partial(self._make_hf_weight, out_name, with_megatron_names)
+
+        def _projection(name: str) -> str:
+            parts = name.split(".")
+            experts_idx = parts.index("experts")
+            projection_idx = experts_idx + 1
+            if projection_idx < len(parts) and parts[projection_idx].isdigit():
+                projection_idx += 1
+            return parts[projection_idx]
+
+        stem = emitted[0].param_name.split(".experts", 1)[0] + ".experts"
+        by_projection: Dict[str, Dict[str, List[Tuple[Optional[int], torch.Tensor]]]] = defaultdict(
+            lambda: {"A": [], "B": []}
+        )
+        for item in emitted:
+            suffix = "A" if item.param_name.endswith(".lora_A.weight") else "B"
+            by_projection[_projection(item.param_name)][suffix].append(
+                (self._infer_hf_expert_idx(item.param_name), item.weight)
+            )
+
+        def _shared(projection: str, suffix: str) -> torch.Tensor:
+            values = by_projection[projection][suffix]
+            if len(values) != 1 or values[0][0] is not None:
+                raise ValueError(f"Expected one shared {suffix} factor for {projection}")
+            return values[0][1]
+
+        def _per_expert(projection: str, suffix: str) -> torch.Tensor:
+            values = sorted(by_projection[projection][suffix], key=lambda value: value[0] or 0)
+            if len(values) != num_moe_experts or any(idx is None for idx, _ in values):
+                raise ValueError(f"Expected {num_moe_experts} per-expert {suffix} factors for {projection}")
+            return torch.stack([tensor for _, tensor in values], dim=0)
+
+        gate_projection = next((p for p in by_projection if p in {"gate_proj", "w1"}), None)
+        up_projection = next((p for p in by_projection if p in {"up_proj", "w3"}), None)
+        down_projection = next((p for p in by_projection if p in {"down_proj", "w2"}), None)
+        if gate_projection is not None and up_projection is not None:
+            shared_a = _shared(gate_projection, "A")
+            if not torch.equal(shared_a, _shared(up_projection, "A")):
+                raise ValueError("Shared gate/up LoRA-A factors differ")
+            yield emit_in(f"{stem}.w1.lora_A.weight", shared_a)
+            yield emit_out(f"{stem}.w1.lora_B.weight", _per_expert(gate_projection, "B"))
+            yield emit_in(f"{stem}.w3.lora_A.weight", shared_a)
+            yield emit_out(f"{stem}.w3.lora_B.weight", _per_expert(up_projection, "B"))
+            return
+
+        if down_projection is None:
+            raise ValueError(
+                "moe_shared_loras requires gated routed experts with gate/up (w1/w3) and down (w2) projections; "
+                f"found {sorted(by_projection)} under {stem}"
+            )
+        yield emit_in(f"{stem}.w2.lora_A.weight", _per_expert(down_projection, "A"))
+        yield emit_out(f"{stem}.w2.lora_B.weight", _shared(down_projection, "B"))
+
     def _stream_shared_outer_adapter_weights(
         self,
         megatron_model: List[MegatronModel],
@@ -1085,6 +1184,7 @@ def _stream_shared_outer_adapter_weights(
         stack_3d_moe: bool = False,
         source_names: Optional[Tuple[str, str]] = None,
         with_megatron_names: bool = False,
+        moe_shared_loras: bool = False,
     ) -> Iterable["HFWeightTuple"]:
         """Stream a shared-outer grouped-expert LoRA adapter (SGLang PR #21466).
 
@@ -1101,6 +1201,12 @@ def _stream_shared_outer_adapter_weights(
         vLLM-3D-MoE analogue of the 2D ``expand_shared_outer`` path and is required
         for 3D-MoE models (``is_3d_moe_weight=True``) that are not using vLLM's 2D
         mixed MoE LoRA format.
+
+        With ``moe_shared_loras``, the legacy shared-outer tensors are reused and
+        repacked as vLLM's 2D ``experts.w1/w2/w3`` layout.
+
+        The layout flags are resolved in priority order: ``stack_3d_moe``, then
+        ``moe_shared_loras``, then ``expand_shared_outer``; the first one set wins.
         """
 
         is_expert = is_expert_linear(adapter_task.global_base_prefix)
@@ -1121,6 +1227,20 @@ def _stream_shared_outer_adapter_weights(
             )
             return
 
+        if moe_shared_loras:
+            yield from self._stream_shared_outer_adapter_weights_2d_moe(
+                megatron_model,
+                mapping_registry,
+                adapter_task,
+                linear_in_tensor,
+                linear_out_tensor,
+                num_moe_experts,
+                cpu,
+                source_names=source_names,
+                with_megatron_names=with_megatron_names,
+            )
+            return
+
         for side_tensor, side_suffix, side_source in (
             (linear_in_tensor, ".linear_in.weight", in_name),
             (linear_out_tensor, ".linear_out.weight", out_name),

src/megatron/bridge/peft/utils.py

@@ -2634,13 +2634,15 @@ def __init__(
         init_method: Optional[Callable] = None,
         dtype: Optional[torch.dtype] = None,
         device: Optional[torch.device] = None,
+        pg_collection: ProcessGroupCollection | None = None,
     ):
         super().__init__()
         if not hasattr(torch, "_grouped_mm"):
             raise RuntimeError("PackedPerExpertLinear requires torch._grouped_mm (torch >= 2.9).")
         self.num_local_experts = num_local_experts
         self.in_features = in_features
         self.out_features = out_features
+        self.pg_collection = pg_collection
         weight = torch.empty(num_local_experts, out_features, in_features, dtype=dtype, device=device)
         if init_method is not None:
             for e in range(num_local_experts):
@@ -2670,7 +2672,11 @@ def sharded_state_dict(
         key = f"{prefix}weight"
         return {
             key: _make_grouped_expert_sharded_tensor(
-                self.weight.data, key, tp_axis=None, sharded_offsets=sharded_offsets
+                self.weight.data,
+                key,
+                tp_axis=None,
+                sharded_offsets=sharded_offsets,
+                pg_collection=self.pg_collection,
             )
         }
 
@@ -2770,6 +2776,7 @@ def __init__(
                 init_method=row_init,
                 device=params_device,
                 dtype=params_dtype,
+                pg_collection=self.pg_collection,
             )
         else:
             # Per-expert A (intermediate → rank); shared B (rank → hidden).
@@ -2780,6 +2787,7 @@ def __init__(
                 init_method=column_init,
                 device=params_device,
                 dtype=params_dtype,
+                pg_collection=self.pg_collection,
             )
             self.linear_out = RowParallelLinear(
                 dim,