Add shared outer lora

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

@@ -341,6 +341,24 @@ def _infer_hf_expert_idx(self, hf_name: str) -> Optional[int]:
         except ValueError:
             return None
 
+    def _strip_hf_expert_index(self, hf_name: str) -> str:
+        """Drop the ``experts.<idx>`` index from an HF MoE weight name.
+
+        A shared-outer adapter's shared side is replicated across experts, so it
+        is exported under the expert-agnostic name (``experts.gate_proj`` rather
+        than ``experts.0.gate_proj``) that the serving loader keys its 3D-shared
+        branch on. Mirrors :meth:`_infer_hf_expert_idx`'s name parsing.
+        """
+
+        parts = hf_name.split(".")
+        try:
+            experts_idx = parts.index("experts")
+        except ValueError:
+            return hf_name
+        if experts_idx + 1 < len(parts) and parts[experts_idx + 1].isdigit():
+            del parts[experts_idx + 1]
+        return ".".join(parts)
+
     def _split_qkv_linear_out_weight(
         self,
         megatron_model: Union[MegatronModel, List[MegatronModel]],
@@ -823,9 +841,12 @@ def stream_adapter_weights_megatron_to_hf(
         cpu: bool = True,
         show_progress: bool = True,
     ) -> Iterable["HFWeightTuple"]:
-        """Stream only adapter weights without merging them into base tensors."""
+        """Stream only adapter weights without merging them into base tensors.
 
-        # Local import avoids circular dependency while ensuring runtime access.
+        Each adapter is classified into one export topology (default / packed-expert
+        / shared-outer) by :meth:`_select_adapter_emitter` and emitted by the
+        matching ``_emit_*_adapter`` method. The loop holds no per-topology logic.
+        """
         from megatron.bridge.models.conversion.model_bridge import HFWeightTuple
 
         if not isinstance(megatron_model, list):
@@ -847,6 +868,20 @@ def stream_adapter_weights_megatron_to_hf(
             linear_out_tensor = adapter_weight.linear_out_weight.weight
             is_expert = is_expert_linear(adapter_task.global_base_prefix)
             is_grouped_expert = is_expert and ".local_experts." not in adapter_task.global_base_prefix
+            is_shared_outer_lora = is_grouped_expert and linear_in_tensor.ndim != linear_out_tensor.ndim
+
+            if is_shared_outer_lora:
+                yield from self._stream_shared_outer_adapter_weights(
+                    megatron_model,
+                    mapping_registry,
+                    adapter_task,
+                    linear_in_tensor,
+                    linear_out_tensor,
+                    num_moe_experts,
+                    cpu,
+                )
+                continue
+
             expert_linear_in_gathered = None
             expert_linear_out_gathered = None
             if is_grouped_expert:
@@ -974,6 +1009,74 @@ def stream_adapter_weights_megatron_to_hf(
                 yield HFWeightTuple(linear_in_hf_names[0], current_linear_in_tensor)
                 yield HFWeightTuple(linear_out_hf_names[0], current_linear_out_tensor)
 
+    def _stream_shared_outer_adapter_weights(
+        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,
+    ) -> 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.
+        """
+
+        from megatron.bridge.models.conversion.model_bridge import HFWeightTuple
+
+        is_expert = is_expert_linear(adapter_task.global_base_prefix)
+        for side_tensor, side_suffix in (
+            (linear_in_tensor, ".linear_in.weight"),
+            (linear_out_tensor, ".linear_out.weight"),
+        ):
+            if side_tensor.ndim == 2:
+                # 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.
+                current = side_tensor.cpu() if cpu else side_tensor
+                current = current.unsqueeze(0)
+
+                base_hf_weight_names = self._get_base_hf_param_names_for_adapter(
+                    mapping_registry, adapter_task.global_base_prefix, adapter_task.adapter_key, ".weight0"
+                )
+                for base_name in base_hf_weight_names:
+                    hf_name = self._make_lora_param_name(self._strip_hf_expert_index(base_name), side_suffix)
+                    yield HFWeightTuple(hf_name, 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.
+            gathered = self._gather_expert_adapter_weight(side_tensor)
+            for expert_idx in range(num_moe_experts):
+                current = self._select_expert_adapter_weight(side_tensor, gathered, expert_idx, num_moe_experts)
+                if cpu:
+                    current = current.cpu()
+
+                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}"
+                )
+                side_hf_names = [self._make_lora_param_name(name, side_suffix) for name in base_hf_weight_names]
+
+                per_base = None
+                if side_suffix == ".linear_out.weight" and adapter_task.adapter_key is None:
+                    per_base = self._get_fused_adapter_linear_out_slices(
+                        megatron_model, base_hf_weight_names, current, is_expert=is_expert
+                    )
+                if per_base is None:
+                    yield HFWeightTuple(side_hf_names[0], current)
+                    continue
+                for index, base_name in enumerate(base_hf_weight_names):
+                    chunk = per_base.get(base_name)
+                    assert chunk is not None, f"unknown projection name: {base_name!r}"
+                    yield HFWeightTuple(side_hf_names[index], chunk)
+
     def _get_fused_adapter_linear_out_slices(
         self,
         megatron_model: List[MegatronModel],

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

@@ -533,7 +533,7 @@ def quantize_to_int4(
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Quantize bfloat16/float16 weights to Kimi INT4 packed format."""
     out_features, in_features = weight.shape
-    weight_shape = torch.tensor([out_features, in_features], dtype=torch.int32)
+    weight_shape = torch.tensor([out_features, in_features], dtype=torch.int32, device=weight.device)
 
     w = weight.float()
 

src/megatron/bridge/peft/lora.py

@@ -37,6 +37,7 @@
 from megatron.bridge.peft.utils import (
     GroupedExpertLinearAdapter,
     ParallelLinearAdapter,
+    SharedOuterGroupedExpertAdapter,
     align_expert_dim_for_tp,
     get_adapter_attributes_from_linear,
     get_effective_lora_dim,
@@ -92,6 +93,13 @@ class LoRA(PEFT, ModuleMatcher):
             so it is comparable to a dense model. Defaults to False.
         share_expert_adapters (bool): When True, grouped MoE expert linears share one adapter across all local
             experts on the EP rank. Set to False to create one adapter per local expert instead. Defaults to True.
+        experts_shared_outer_loras (bool): When True, grouped-expert LoRA
+            (``TE*ParallelGroupedLinear`` base modules) uses
+            :class:`SharedOuterGroupedExpertAdapter` — ``gate_up`` lora_A and
+            ``down`` lora_B are shared across experts (expert_dim=1), matching
+            SGLang's ``experts_shared_outer_loras=True`` serving contract (PR
+            #21466). Default False preserves the adapter layout selected by
+            ``share_expert_adapters``.
     """
 
     target_modules: List[str] = field(
@@ -107,6 +115,7 @@ class LoRA(PEFT, ModuleMatcher):
     lora_dtype: torch.dtype = None
     normalize_moe_lora: bool = False
     share_expert_adapters: bool = True
+    experts_shared_outer_loras: bool = False
 
     def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optional[str] = None) -> nn.Module:
         """
@@ -165,18 +174,28 @@ def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optio
                 is_expert=is_expert,
                 input_is_parallel=attrs.input_is_parallel,
             )
-            use_per_expert_adapter = is_grouped_expert_linear(full_name) and not self.share_expert_adapters
+            is_grouped_expert_name = is_grouped_expert_linear(full_name)
+            use_shared_outer_adapter = self.experts_shared_outer_loras and is_grouped_expert_name
+            use_per_expert_adapter = (
+                is_grouped_expert_name and not self.share_expert_adapters and not use_shared_outer_adapter
+            )
+            use_grouped_expert_adapter = use_shared_outer_adapter or use_per_expert_adapter
 
             enable_op_fuser = (
-                not use_per_expert_adapter
+                not use_grouped_expert_adapter
                 and not is_expert
                 and getattr(module.config, "use_transformer_engine_op_fuser", False)
                 # TP not yet supported
                 and parallel_state.get_tensor_model_parallel_world_size() == 1
             )
 
             logger.info(f"Adding lora to: {full_name}")
-            adapter_cls = GroupedExpertLinearAdapter if use_per_expert_adapter else ParallelLinearAdapter
+            if use_shared_outer_adapter:
+                adapter_cls = SharedOuterGroupedExpertAdapter
+            elif use_per_expert_adapter:
+                adapter_cls = GroupedExpertLinearAdapter
+            else:
+                adapter_cls = ParallelLinearAdapter
             adapter_kwargs = dict(
                 base_linear_name=full_name,
                 activation="identity",
@@ -189,7 +208,7 @@ def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optio
                 alpha=self.alpha,
                 base_linear_is_parallel=attrs.base_linear_is_parallel,
             )
-            if use_per_expert_adapter:
+            if use_grouped_expert_adapter:
                 first_param = next(module.parameters())
                 adapter_kwargs.update(
                     num_local_experts=module.num_gemms,