diff --git a/docs/training/peft.md b/docs/training/peft.md index 5fc33d0a50..9707ccccf2 100644 --- a/docs/training/peft.md +++ b/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. diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index d0c8ec5e0a..ea3643a999 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/src/megatron/bridge/models/conversion/auto_bridge.py @@ -303,34 +303,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. @@ -864,9 +859,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, @@ -973,6 +968,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. @@ -994,6 +990,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. @@ -1009,8 +1010,14 @@ 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, @@ -1018,7 +1025,7 @@ def export_adapter_weights( 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( diff --git a/src/megatron/bridge/models/conversion/model_bridge.py b/src/megatron/bridge/models/conversion/model_bridge.py index 87ac3c177b..b50d24c056 100644 --- a/src/megatron/bridge/models/conversion/model_bridge.py +++ b/src/megatron/bridge/models/conversion/model_bridge.py @@ -2750,6 +2750,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.""" ... @@ -2846,6 +2847,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( @@ -2855,6 +2857,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 diff --git a/src/megatron/bridge/models/conversion/peft_bridge.py b/src/megatron/bridge/models/conversion/peft_bridge.py index 45c4277065..a7eed7c126 100644 --- a/src/megatron/bridge/models/conversion/peft_bridge.py +++ b/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), diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py index ca33919be8..a46df168f3 100644 --- a/src/megatron/bridge/peft/utils.py +++ b/src/megatron/bridge/peft/utils.py @@ -2634,6 +2634,7 @@ 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"): @@ -2641,6 +2642,7 @@ def __init__( 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, diff --git a/tests/unit_tests/models/test_adapter_export.py b/tests/unit_tests/models/test_adapter_export.py index dc62827311..53cec5c266 100644 --- a/tests/unit_tests/models/test_adapter_export.py +++ b/tests/unit_tests/models/test_adapter_export.py @@ -2085,3 +2085,132 @@ def test_stack_3d_moe_matches_vllm_consumer_layout(self, _mock_ep): torch.testing.assert_close(gu[f"{stem}.base_layer.lora_B.weight"], exp_base_b) torch.testing.assert_close(dn[f"{stem}.lora_A.weight"], exp_down_a) torch.testing.assert_close(dn[f"{stem}.lora_B.weight"], exp_down_b) + + # ------------------------------------------------------------------ + # moe_shared_loras: vLLM 2D shared-MoE layout (experts.w1/w2/w3) + # ------------------------------------------------------------------ + + def _run_serving_layout( + self, bridge, linear_in, linear_out, num_experts, mapping_registry, *, gate_up_split=None, **flags + ): + """Stream one shared-outer adapter at EP=1 and return ``{hf_name: tensor}``. + + ``gate_up_split`` is the gate-half size of a fused FC1 ``linear_out``; ``None`` means + there is no fused projection to split. + """ + + def split(_model, names, tensor, is_expert=False): + return {names[0]: tensor[:gate_up_split], names[1]: tensor[gate_up_split:]} + + task = SimpleNamespace(global_base_prefix="mlp.experts", adapter_key=None) + with ( + patch( + "megatron.bridge.models.conversion.peft_bridge.parallel_state.get_expert_model_parallel_world_size", + return_value=1, + ), + 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", + side_effect=split if gate_up_split is not None else None, + return_value=None, + ), + patch("megatron.bridge.models.conversion.peft_bridge.is_expert_linear", return_value=True), + ): + items = bridge._stream_shared_outer_adapter_weights( + 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=False, + **flags, + ) + return {item.param_name: item.weight for item in items} + + def test_moe_shared_loras_repacks_legacy_emission_as_w1_w2_w3(self): + """Shared factors keep expert-dim 1, per-expert factors stack in expert order, w1=gate and w3=up.""" + bridge = self._make_bridge() + E, r, hid, mi = 3, 2, 6, 4 + stem = "model.layers.0.mlp.experts" + gate_up_a, gate_up_b = torch.randn(r, hid), torch.randn(E, 2 * mi, r) + down_a, down_b = torch.randn(E, r, mi), torch.randn(hid, r) + + gate_up = self._run_serving_layout( + bridge, + gate_up_a, + gate_up_b, + E, + self._mapping_gate_up(), + gate_up_split=mi, + expand_shared_outer=False, + moe_shared_loras=True, + ) + assert set(gate_up) == { + f"{stem}.w1.lora_A.weight", + f"{stem}.w1.lora_B.weight", + f"{stem}.w3.lora_A.weight", + f"{stem}.w3.lora_B.weight", + } + torch.testing.assert_close(gate_up[f"{stem}.w1.lora_A.weight"], gate_up_a.unsqueeze(0)) + torch.testing.assert_close(gate_up[f"{stem}.w3.lora_A.weight"], gate_up_a.unsqueeze(0)) + torch.testing.assert_close(gate_up[f"{stem}.w1.lora_B.weight"], gate_up_b[:, :mi]) + torch.testing.assert_close(gate_up[f"{stem}.w3.lora_B.weight"], gate_up_b[:, mi:]) + + down = self._run_serving_layout( + bridge, down_a, down_b, E, self._mapping_down(), expand_shared_outer=False, moe_shared_loras=True + ) + assert set(down) == {f"{stem}.w2.lora_A.weight", f"{stem}.w2.lora_B.weight"} + torch.testing.assert_close(down[f"{stem}.w2.lora_A.weight"], down_a) + torch.testing.assert_close(down[f"{stem}.w2.lora_B.weight"], down_b.unsqueeze(0)) + + @pytest.mark.parametrize( + ("flags", "expected_modules"), + [ + ({"expand_shared_outer": True, "stack_3d_moe": True, "moe_shared_loras": True}, ["base_layer"]), + ({"expand_shared_outer": False, "stack_3d_moe": True, "moe_shared_loras": True}, ["base_layer"]), + ({"expand_shared_outer": True, "stack_3d_moe": False, "moe_shared_loras": True}, ["w1", "w3"]), + ], + ) + def test_serving_layout_flag_priority(self, flags, expected_modules): + """stack_3d_moe wins over moe_shared_loras, which wins over expand_shared_outer.""" + bridge = self._make_bridge() + E, r, hid, mi = 3, 2, 6, 4 + + out = self._run_serving_layout( + bridge, + torch.randn(r, hid), + torch.randn(E, 2 * mi, r), + E, + self._mapping_gate_up(), + gate_up_split=mi, + **flags, + ) + + assert set(out) == { + f"model.layers.0.mlp.experts.{module}.lora_{side}.weight" + for module in expected_modules + for side in ("A", "B") + } + + def test_moe_shared_loras_rejects_non_gated_experts(self): + """Non-gated experts have no w1/w3 pair, so the shared-MoE layout cannot represent them.""" + bridge = self._make_bridge() + E, r, hid, mi = 3, 2, 6, 4 + mapping = MagicMock() + mapping.megatron_to_hf_lookup.side_effect = lambda name: SimpleNamespace( + hf_param=f"model.layers.0.mlp.experts.{name.rsplit('.weight', 1)[1]}.up_proj.weight" + ) + + with pytest.raises(ValueError, match=r"requires gated routed experts.*\['up_proj'\]"): + self._run_serving_layout( + bridge, + torch.randn(r, hid), + torch.randn(E, mi, r), + E, + mapping, + expand_shared_outer=False, + moe_shared_loras=True, + ) diff --git a/tests/unit_tests/models/test_auto_bridge.py b/tests/unit_tests/models/test_auto_bridge.py index 49815c3ce4..e04f907b0c 100644 --- a/tests/unit_tests/models/test_auto_bridge.py +++ b/tests/unit_tests/models/test_auto_bridge.py @@ -1831,7 +1831,7 @@ def test_export_hf_weights(self): ) def test_export_with_megatron_names_rejects_streamers_without_the_flag(self): - """A bridge whose streaming overrides lack ``with_megatron_names`` fails before streaming.""" + """A bridge whose streaming overrides lack an opt-in flag fails before streaming.""" mock_hf_model = Mock(spec=PreTrainedCausalLM) mock_hf_model.config = Mock() mock_hf_model.config.architectures = ["LlamaForCausalLM"] @@ -1881,6 +1881,8 @@ def stream_adapter_weights_megatron_to_hf( bridge.export_hf_weights([object()], cpu=True, current_pp_stage_only=True) with pytest.raises(TypeError, match="stream_adapter_weights_megatron_to_hf does not accept"): bridge.export_adapter_weights([object()], with_megatron_names=True) + with pytest.raises(TypeError, match="does not accept 'moe_shared_loras'"): + bridge.export_adapter_weights([object()], moe_shared_loras=True) def test_export_adapter_weights(self): """Test exporting adapter weights from Megatron to HF format.""" @@ -1919,8 +1921,9 @@ def test_export_adapter_weights(self): stack_3d_moe=False, ) - def test_export_adapter_weights_forwards_expand_shared_outer(self): - """A non-default expand_shared_outer must reach the model bridge.""" + @pytest.mark.parametrize("flag", ["expand_shared_outer", "moe_shared_loras"]) + def test_export_adapter_weights_forwards_layout_flag(self, flag): + """A non-default adapter layout flag must reach the model bridge.""" mock_hf_model = Mock(spec=PreTrainedCausalLM) mock_hf_model.config = Mock() mock_hf_model.config.architectures = ["LlamaForCausalLM"] @@ -1941,10 +1944,10 @@ def test_export_adapter_weights_forwards_expand_shared_outer(self): with patch.object(AutoBridge, "_causal_lm_architecture", new_callable=PropertyMock) as mock_prop: mock_prop.return_value = mock_arch_class - list(bridge.export_adapter_weights(mock_megatron_model, expand_shared_outer=True)) + list(bridge.export_adapter_weights(mock_megatron_model, **{flag: True})) call_kwargs = mock_model_bridge.stream_adapter_weights_megatron_to_hf.call_args.kwargs - assert call_kwargs["expand_shared_outer"] is True + assert call_kwargs[flag] is True def test_get_causal_lm_architecture(self): """Test getting the CausalLM architecture class."""