diff --git a/src/megatron/bridge/models/conversion/peft_bridge.py b/src/megatron/bridge/models/conversion/peft_bridge.py index 7a873972b1..cf4ccc31da 100644 --- a/src/megatron/bridge/models/conversion/peft_bridge.py +++ b/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.`` 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], diff --git a/src/megatron/bridge/models/conversion/quantization_utils.py b/src/megatron/bridge/models/conversion/quantization_utils.py index fe3f7d5a57..c9488467d1 100644 --- a/src/megatron/bridge/models/conversion/quantization_utils.py +++ b/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() diff --git a/src/megatron/bridge/peft/lora.py b/src/megatron/bridge/peft/lora.py index fddc30420a..aeb5b53b10 100644 --- a/src/megatron/bridge/peft/lora.py +++ b/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,10 +174,15 @@ 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 @@ -176,7 +190,12 @@ def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optio ) 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, diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py index 996c51e8c1..de13432586 100644 --- a/src/megatron/bridge/peft/utils.py +++ b/src/megatron/bridge/peft/utils.py @@ -24,7 +24,7 @@ import packaging import torch import torch.nn as nn -from megatron.core import ModelParallelConfig, dist_checkpointing +from megatron.core import ModelParallelConfig, dist_checkpointing, parallel_state from megatron.core.dist_checkpointing.mapping import ShardedStateDict, ShardedTensor, ShardedTensorFactory from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear, set_tensor_model_parallel_attributes @@ -2118,3 +2118,292 @@ def sharded_state_dict( sharded_state_dict.update(linear_in_sd) sharded_state_dict.update(linear_out_sd) return sharded_state_dict + + +def _make_cross_ep_replicated(weight: nn.Parameter) -> None: + """Mark a weight as logically replicated across the intra-PP-stage group. + + Megatron's DDP routes ``is_expert=True`` parameters through the expert + data-parallel group only, which does not span the EP axis. A weight + that must stay bit-identical across all EP ranks (e.g., the shared + side of :class:`SharedOuterGroupedExpertAdapter`, which a serving + engine consumes as a single global LoRA tensor) is otherwise left + unsynced. This helper closes that gap with two primitives: + + * a one-shot broadcast from group rank 0 so every rank starts with + bit-identical values despite per-rank RNG forks; + * a backward hook that SUM all-reduces the gradient across the group + so the optimizer step on every rank applies the same update. + + SUM is the correct reduction: each rank's local gradient is the partial + loss gradient over its (token, expert) subset, and the total gradient + is the sum of those partials. AVG would train at 1/N the intended rate. + + The intra-PP-stage group is ``tensor_and_data_parallel_group`` with + context parallel included, which by Megatron's construction equals + ETP × EP × EDP — all ranks within the current pipeline stage. + + Args: + weight: The parameter to keep replicated across the group. Must + be a leaf parameter so the backward hook fires when its + gradient is computed. + """ + + if not (torch.distributed.is_available() and torch.distributed.is_initialized()): + return + try: + group = parallel_state.get_tensor_and_data_parallel_group(with_context_parallel=True) + except AssertionError: + return + if torch.distributed.get_world_size(group=group) <= 1: + return + + if weight.is_cuda: + # NCCL requires CUDA tensors; pre-GPU construction relies on + # deterministic init matching across ranks. + src_rank = torch.distributed.get_global_rank(group, 0) + with torch.no_grad(): + torch.distributed.broadcast(weight.data, src=src_rank, group=group) + + def _all_reduce_grad(grad: torch.Tensor) -> torch.Tensor: + grad = grad.contiguous() + torch.distributed.all_reduce(grad, op=torch.distributed.ReduceOp.SUM, group=group) + return grad + + weight.register_hook(_all_reduce_grad) + + +class PackedPerExpertLinear(nn.Module): + """Per-expert linear with a packed 3D weight ``[N_local, out, in]``. + + Used as the per-expert side of :class:`SharedOuterGroupedExpertAdapter`. + Stores one ``nn.Parameter`` (3D) so Bridge's adapter export sees a single + ``.weight`` per side, matching the ``linear_in.weight`` / ``linear_out.weight`` + convention in :mod:`megatron.bridge.models.conversion.peft_bridge`. Forward + dispatches to :func:`torch._grouped_mm` (the same grouped GEMM kernel TE's + :class:`te.pytorch.GroupedLinear` calls) via a single fused op with native + autograd, which keeps rank kernel launch counts in lockstep so CP's ring + P2P does not deadlock. + """ + + def __init__( + self, + num_local_experts: int, + in_features: int, + out_features: int, + *, + init_method: Optional[Callable] = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = 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 + 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): + init_method(weight[e]) + else: + nn.init.zeros_(weight) + self.weight = nn.Parameter(weight) + # DDP routes ``is_expert`` weights through the EDP group; the cross-EP + # axis is naturally distinct here (different experts on each EP rank). + setattr(self.weight, "allreduce", False) + + def forward(self, x: torch.Tensor, m_splits) -> Tuple[torch.Tensor, None]: + # torch._grouped_mm expects mat2 as [num_groups, K, N]; our weight is + # [N_local, out, in] so transpose the last two dims. + if isinstance(m_splits, torch.Tensor): + m_splits_i32 = m_splits.to(device=x.device, dtype=torch.int32) + else: + m_splits_i32 = torch.tensor(m_splits, device=x.device, dtype=torch.int32) + offs = torch.cumsum(m_splits_i32, dim=0, dtype=torch.int32) + out = torch._grouped_mm(x, self.weight.transpose(1, 2), offs=offs) + return out, None + + def sharded_state_dict( + self, prefix: str = "", sharded_offsets: Tuple = (), metadata: Optional[Dict] = None + ) -> ShardedStateDict: + """Shard the packed 3D weight along dim 0 (experts) across EP ranks.""" + key = f"{prefix}weight" + return { + key: _make_grouped_expert_sharded_tensor( + self.weight.data, key, tp_axis=None, sharded_offsets=sharded_offsets + ) + } + + +class SharedOuterGroupedExpertAdapter(nn.Module): + """LoRA adapter for grouped expert MLP with shared-outer semantics. + + Matches SGLang PR #21466's ``experts_shared_outer_loras=True`` contract: + + * fc1 (gate_up): linear_in = SHARED (hidden -> rank) + linear_out = PER-EXPERT (rank -> 2*intermediate) + * fc2 (down): linear_in = PER-EXPERT (intermediate -> rank) + linear_out = SHARED (rank -> hidden) + + The shared side is an ``is_expert=True`` ``ColumnParallelLinear`` (fc1) + or ``RowParallelLinear`` (fc2): the TP group is ETP (ETP=1 → local + forward), DDP routes the weight through the EDP group, and the + logically-replicated cross-EP axis is covered by + :func:`_make_cross_ep_replicated`. + + The per-expert side is :class:`PackedPerExpertLinear` (packed 3D weight + + :func:`torch._grouped_mm`) — kept as a single ``.weight`` Parameter so + Bridge's adapter-export materializer (which reads ``linear_in.weight`` / + ``linear_out.weight``) sees a standard single-weight linear per side. + + Differs from ``ParallelLinearAdapter`` in ``__init__`` and ``forward``; + ``sharded_state_dict`` is specialized for the packed 3D per-expert side. + """ + + def __init__( + self, + in_features: int, + out_features: int, + dim: int, + *, + num_local_experts: int, + base_linear_name: str, + activation: str = "swish", + column_init_method: str = "xavier", + row_init_method: str = "zero", + input_is_parallel: bool = False, + dropout: float = 0.0, + model_parallel_config: Optional[ModelParallelConfig] = None, + alpha: Optional[float] = None, + dropout_position: str = "pre", + base_linear_is_parallel: bool = True, + params_device: Optional[torch.device] = None, + params_dtype: Optional[torch.dtype] = None, + ) -> None: + """Initialize shared-outer LoRA weights with one shared and one per-expert side.""" + + super().__init__() + self.base_linear_name = base_linear_name + self.activation = ParallelLinearAdapter._get_activation_fn(self, activation) + self.dim = dim + self.alpha = alpha if alpha is not None else self.dim + self.dropout_position = dropout_position + self.num_local_experts = num_local_experts + self.base_linear_is_parallel = base_linear_is_parallel + # ``is_expert=True`` is observed by param_mapping.py and by inherited + # checkpoint helpers; the per-expert side's grad routing is set on its + # 3D weight directly inside :class:`PackedPerExpertLinear`. + self.is_expert = True + + if model_parallel_config is None: + model_parallel_config = ModelParallelConfig() + model_parallel_config.perform_initialization = True + self.config = model_parallel_config + + # ``input_is_parallel`` selects fc1 (column-parallel base) vs fc2 + # (row-parallel base). Mirrors :class:`ParallelLinearAdapter` and + # :class:`GroupedExpertLinearAdapter`. + self._is_fc1 = not input_is_parallel + + column_init = ParallelLinearAdapter._get_init_fn(self, column_init_method) + row_init = ParallelLinearAdapter._get_init_fn(self, row_init_method) + if self._is_fc1: + # Shared A (hidden → rank); per-expert B (rank → 2*intermediate). + self.linear_in = ColumnParallelLinear( + in_features, + dim, + config=model_parallel_config, + bias=False, + gather_output=True, + init_method=column_init, + is_expert=True, + ) + self.linear_out = PackedPerExpertLinear( + num_local_experts, + dim, + out_features, + init_method=row_init, + device=params_device, + dtype=params_dtype, + ) + else: + # Per-expert A (intermediate → rank); shared B (rank → hidden). + self.linear_in = PackedPerExpertLinear( + num_local_experts, + in_features, + dim, + init_method=column_init, + device=params_device, + dtype=params_dtype, + ) + self.linear_out = RowParallelLinear( + dim, + out_features, + config=model_parallel_config, + bias=False, + input_is_parallel=True, + skip_bias_add=True, + init_method=row_init, + is_expert=True, + ) + + self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() + + if model_parallel_config.bf16: + self.bfloat16() + elif model_parallel_config.fp16: + self.half() + + # The shared weight is logically replicated across EP; close the gap + # that Megatron's expert-DDP routing leaves open. + shared_weight = self.linear_in.weight if self._is_fc1 else self.linear_out.weight + _make_cross_ep_replicated(shared_weight) + + def forward(self, x: torch.Tensor, m_splits=None) -> torch.Tensor: + """Forward. ``m_splits`` is the tokens-per-expert split passed through + from the base TEGroupedLinear; required for the per-expert side. + """ + if self.dropout_position == "pre": + x = self.dropout(x) + + if self._is_fc1: + # Shared A → activation → per-expert B. + x, _ = self.linear_in(x) + x = self.activation(x) + x, _ = self.linear_out(x, m_splits) + else: + # Per-expert A → activation → shared B. + x, _ = self.linear_in(x, m_splits) + x = self.activation(x) + x, _ = self.linear_out(x) + + if self.dropout_position == "post": + x = self.dropout(x) + + return x * (self.alpha / self.dim) + + def sharded_state_dict( + self, + prefix: str = "", + sharded_offsets: Tuple = (), + metadata: Optional[Dict] = None, + ) -> ShardedStateDict: + """Create sharded state dictionary for mixed shared/per-expert adapter weights.""" + + linear_in_sd = self.linear_in.sharded_state_dict(f"{prefix}linear_in.", sharded_offsets, metadata) + linear_out_sd = self.linear_out.sharded_state_dict(f"{prefix}linear_out.", sharded_offsets, metadata) + + if self._is_fc1: + singleton_local_shards = (metadata or {}).get("singleton_local_shards", False) + linear_out_key = f"{prefix}linear_out.weight" + linear_out_sd[linear_out_key] = _apply_grouped_expert_swiglu_sharded_factory( + linear_out_sd[linear_out_key], + sharded_offsets, + singleton_local_shards, + ) + + sharded_state_dict = {} + sharded_state_dict.update(linear_in_sd) + sharded_state_dict.update(linear_out_sd) + return sharded_state_dict diff --git a/tests/unit_tests/models/test_model_bridge_lora.py b/tests/unit_tests/models/test_model_bridge_lora.py index 2bbadc431d..1f2d4640f5 100644 --- a/tests/unit_tests/models/test_model_bridge_lora.py +++ b/tests/unit_tests/models/test_model_bridge_lora.py @@ -1405,6 +1405,172 @@ def test_stream_adapter_weights_megatron_to_hf_grouped_expert_exports_per_expert assert weights[1].param_name == "model.layers.0.mlp.experts.0.down_proj.lora_B.weight" +def test_stream_adapter_weights_megatron_to_hf_shared_outer_fc1_gate_up(monkeypatch): + # Shared-outer FC1 (SGLang PR #21466): linear_in (lora_A) is a single 2D matrix + # shared across experts, while linear_out (lora_B) is a per-expert 3D pack of the + # fused gate/up projection. The shared side is emitted once under an + # expert-agnostic name; the per-expert side is split into gate/up per expert. + bridge = DummyBridge() + + adapter_task = AdapterWeightConversionTask( + global_base_prefix="decoder.layers.0.mlp.experts.linear_fc1", + adapter_key=None, + alpha=2, + dim=4, + linear_in_task=WeightConversionTask( + param_name="local_in", + global_param_name="decoder.layers.0.mlp.experts.linear_fc1.adapter.linear_in.weight", + mapping=Mock(), + ), + linear_out_task=WeightConversionTask( + param_name="local_out", + global_param_name="decoder.layers.0.mlp.experts.linear_fc1.adapter.linear_out.weight", + mapping=Mock(), + ), + ) + + # Shared lora_A: [rank=2, hidden=3]. Per-expert lora_B: [num_experts=2, 2*inter=4, rank=2], + # gate = first 2 rows, up = last 2 rows, with distinct values per expert/projection. + shared_lora_a = torch.ones(2, 3) + expert0_lora_b = torch.cat([torch.full((2, 2), 10.0), torch.full((2, 2), 20.0)], dim=0) + expert1_lora_b = torch.cat([torch.full((2, 2), 30.0), torch.full((2, 2), 40.0)], dim=0) + per_expert_lora_b = torch.stack([expert0_lora_b, expert1_lora_b], dim=0) + + adapter_weight = AdapterWeight( + global_base_prefix="decoder.layers.0.mlp.experts.linear_fc1", + adapter_key=None, + alpha=2, + dim=4, + linear_in_weight=MegatronWeightTuple("local_in", shared_lora_a, vp_stage=0), + linear_out_weight=MegatronWeightTuple("local_out", per_expert_lora_b, vp_stage=0), + ) + + def fake_base_names(_registry, _prefix, _adapter_key, base_suffix): + # base_suffix is ".weight0"/".weight1"/...; reflect the expert index into the + # HF names so the per-expert side keeps experts. and the shared side strips it. + idx = base_suffix[len(".weight") :] + return [ + f"model.layers.0.mlp.experts.{idx}.gate_proj.weight", + f"model.layers.0.mlp.experts.{idx}.up_proj.weight", + ] + + monkeypatch.setattr( + bridge, + "build_adapter_conversion_tasks", + lambda *_: {"decoder.layers.0.mlp.experts.linear_fc1": [adapter_task]}, + ) + monkeypatch.setattr(bridge, "materialize_adapter_weights", lambda *_: [adapter_weight]) + monkeypatch.setattr(bridge, "_get_base_hf_param_names_for_adapter", fake_base_names) + monkeypatch.setattr( + "megatron.bridge.models.conversion.peft_bridge.parallel_state.get_expert_model_parallel_world_size", + lambda: 1, + ) + + weights = list( + bridge.stream_adapter_weights_megatron_to_hf( + [SimpleNamespace(config=SimpleNamespace(num_moe_experts=2))], + cpu=False, + show_progress=False, + ) + ) + + # Shared lora_A emitted once per fused projection under the expert-agnostic name, + # then per-expert lora_B split into gate/up for each expert. + assert [w.param_name for w in weights] == [ + "model.layers.0.mlp.experts.gate_proj.lora_A.weight", + "model.layers.0.mlp.experts.up_proj.lora_A.weight", + "model.layers.0.mlp.experts.0.gate_proj.lora_B.weight", + "model.layers.0.mlp.experts.0.up_proj.lora_B.weight", + "model.layers.0.mlp.experts.1.gate_proj.lora_B.weight", + "model.layers.0.mlp.experts.1.up_proj.lora_B.weight", + ] + + # Shared side is unsqueezed to [1, ...] and identical for gate and up. + assert weights[0].weight.shape == (1, 2, 3) + torch.testing.assert_close(weights[0].weight, shared_lora_a.unsqueeze(0)) + torch.testing.assert_close(weights[1].weight, shared_lora_a.unsqueeze(0)) + # Per-expert side carries each expert's gate/up halves. + torch.testing.assert_close(weights[2].weight, torch.full((2, 2), 10.0)) + torch.testing.assert_close(weights[3].weight, torch.full((2, 2), 20.0)) + torch.testing.assert_close(weights[4].weight, torch.full((2, 2), 30.0)) + torch.testing.assert_close(weights[5].weight, torch.full((2, 2), 40.0)) + + +def test_stream_adapter_weights_megatron_to_hf_shared_outer_fc2_down(monkeypatch): + # Shared-outer FC2 is the mirror of FC1: linear_in (lora_A) is the per-expert 3D + # pack and linear_out (lora_B) is the shared 2D matrix. down_proj is not fused, so + # the per-expert side emits one weight per expert and the shared side emits once. + bridge = DummyBridge() + + adapter_task = AdapterWeightConversionTask( + global_base_prefix="decoder.layers.0.mlp.experts.linear_fc2", + adapter_key=None, + alpha=2, + dim=4, + linear_in_task=WeightConversionTask( + param_name="local_in", + global_param_name="decoder.layers.0.mlp.experts.linear_fc2.adapter.linear_in.weight", + mapping=Mock(), + ), + linear_out_task=WeightConversionTask( + param_name="local_out", + global_param_name="decoder.layers.0.mlp.experts.linear_fc2.adapter.linear_out.weight", + mapping=Mock(), + ), + ) + + # Per-expert lora_A: [num_experts=2, rank=2, inter=3]. Shared lora_B: [hidden=3, rank=2]. + per_expert_lora_a = torch.stack([torch.full((2, 3), 5.0), torch.full((2, 3), 6.0)], dim=0) + shared_lora_b = torch.full((3, 2), 7.0) + + adapter_weight = AdapterWeight( + global_base_prefix="decoder.layers.0.mlp.experts.linear_fc2", + adapter_key=None, + alpha=2, + dim=4, + linear_in_weight=MegatronWeightTuple("local_in", per_expert_lora_a, vp_stage=0), + linear_out_weight=MegatronWeightTuple("local_out", shared_lora_b, vp_stage=0), + ) + + def fake_base_names(_registry, _prefix, _adapter_key, base_suffix): + idx = base_suffix[len(".weight") :] + return [f"model.layers.0.mlp.experts.{idx}.down_proj.weight"] + + monkeypatch.setattr( + bridge, + "build_adapter_conversion_tasks", + lambda *_: {"decoder.layers.0.mlp.experts.linear_fc2": [adapter_task]}, + ) + monkeypatch.setattr(bridge, "materialize_adapter_weights", lambda *_: [adapter_weight]) + monkeypatch.setattr(bridge, "_get_base_hf_param_names_for_adapter", fake_base_names) + monkeypatch.setattr( + "megatron.bridge.models.conversion.peft_bridge.parallel_state.get_expert_model_parallel_world_size", + lambda: 1, + ) + + weights = list( + bridge.stream_adapter_weights_megatron_to_hf( + [SimpleNamespace(config=SimpleNamespace(num_moe_experts=2))], + cpu=False, + show_progress=False, + ) + ) + + # Per-expert lora_A emitted once per expert (experts.), then the shared lora_B + # emitted once under the expert-agnostic name. + assert [w.param_name for w in weights] == [ + "model.layers.0.mlp.experts.0.down_proj.lora_A.weight", + "model.layers.0.mlp.experts.1.down_proj.lora_A.weight", + "model.layers.0.mlp.experts.down_proj.lora_B.weight", + ] + + torch.testing.assert_close(weights[0].weight, torch.full((2, 3), 5.0)) + torch.testing.assert_close(weights[1].weight, torch.full((2, 3), 6.0)) + # Shared side is unsqueezed to [1, hidden, rank]. + assert weights[2].weight.shape == (1, 3, 2) + torch.testing.assert_close(weights[2].weight, shared_lora_b.unsqueeze(0)) + + def test_split_gdn_in_proj_linear_out_weight_roundtrip(monkeypatch): bridge = DummyBridge() config = SimpleNamespace(