From 330dc04237df7d5d79845b580a5690bf4e746a77 Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Sat, 15 Aug 2026 00:47:50 +0300 Subject: [PATCH] feat(peft): expand_shared_outer export flag for shared-outer MoE LoRA Shared-outer MoE LoRA exports the shared adapter side as a single `[1, ...]` tensor under an expert-agnostic name (SGLang contract). vLLM's 2D `pack_moe` path for text MoE models expects per-expert 2D tensors, so it can't find the shared factor and crashes at `lora.optimize()` (`lora_b is None`). So this PR adds `expand_shared_outer` to the adapter export path. When set, the shared 2D factor is replicated under per-expert 2D names (reused, not cloned); training keeps the shared parameter. Default `False` preserves the existing layout. Signed-off-by: Hollow Man --- .../bridge/models/conversion/auto_bridge.py | 9 ++ .../bridge/models/conversion/model_bridge.py | 3 + .../bridge/models/conversion/peft_bridge.py | 33 ++++- .../unit_tests/models/test_adapter_export.py | 135 ++++++++++++++++++ 4 files changed, 174 insertions(+), 6 deletions(-) diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index b6f89a16af..f7df0041f1 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/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: diff --git a/src/megatron/bridge/models/conversion/model_bridge.py b/src/megatron/bridge/models/conversion/model_bridge.py index 181896b883..69c4fdacd5 100644 --- a/src/megatron/bridge/models/conversion/model_bridge.py +++ b/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,6 +2463,7 @@ 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( @@ -2469,6 +2471,7 @@ def _adapter_stream_registered_impl( 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 diff --git a/src/megatron/bridge/models/conversion/peft_bridge.py b/src/megatron/bridge/models/conversion/peft_bridge.py index 84d8992c7a..ec29dbd49e 100644 --- a/src/megatron/bridge/models/conversion/peft_bridge.py +++ b/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. diff --git a/tests/unit_tests/models/test_adapter_export.py b/tests/unit_tests/models/test_adapter_export.py index cb3021df54..827deea708 100644 --- a/tests/unit_tests/models/test_adapter_export.py +++ b/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)