diff --git a/src/megatron/bridge/data/vlm_datasets/__init__.py b/src/megatron/bridge/data/vlm_datasets/__init__.py index 476ece1f8f..897bdf021c 100644 --- a/src/megatron/bridge/data/vlm_datasets/__init__.py +++ b/src/megatron/bridge/data/vlm_datasets/__init__.py @@ -21,7 +21,10 @@ - Collate fns: model-specific batch builders """ -from megatron.bridge.data.energon.energon_provider import EnergonProvider +try: + from megatron.bridge.data.energon.energon_provider import EnergonProvider +except ImportError: + EnergonProvider = None # megatron.energon not installed from megatron.bridge.data.vlm_datasets.collate import ( COLLATE_FNS, default_collate_fn, diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index b70af56bf6..f11f108118 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/src/megatron/bridge/models/conversion/auto_bridge.py @@ -483,8 +483,8 @@ def save_hf_adapter( dist.barrier() adapter_state: dict[str, torch.Tensor] = {} - for name, tensor in self.export_adapter_weights(model, cpu=True, show_progress=show_progress): - adapter_state[f"base_model.model.{name}"] = tensor.clone().float() + for item in self.export_adapter_weights(model, cpu=True, show_progress=show_progress): + adapter_state[f"base_model.model.{item.param_name}"] = item.weight.clone().float() if not adapter_state: raise RuntimeError( diff --git a/src/megatron/bridge/models/conversion/model_bridge.py b/src/megatron/bridge/models/conversion/model_bridge.py index a99250564b..ddd3d31f4c 100644 --- a/src/megatron/bridge/models/conversion/model_bridge.py +++ b/src/megatron/bridge/models/conversion/model_bridge.py @@ -86,6 +86,7 @@ class HFWeightTuple(NamedTuple): param_name: str weight: torch.Tensor + megatron_param_name: Optional[str] = None @dataclass(frozen=True) @@ -1007,13 +1008,13 @@ def stream_weights_megatron_to_hf( # TODO(yuya): fix this hard coded naming if embeddings_are_tied and hf_name == "model.embed_tokens.weight": # Yield the embedding weight - yield HFWeightTuple(hf_name, final_tensor) + yield HFWeightTuple(hf_name, final_tensor, task.param_name) # Also yield as lm_head.weight if it's expected if hasattr(hf_pretrained, "state") and hasattr(hf_pretrained.state, "source"): expected_keys = hf_pretrained.state.source.get_all_keys() if "lm_head.weight" in expected_keys: - yield HFWeightTuple("lm_head.weight", final_tensor.clone().detach()) + yield HFWeightTuple("lm_head.weight", final_tensor.clone().detach(), task.param_name) elif embeddings_are_tied and hf_name == "lm_head.weight": # This should not happen when embeddings are tied - assert error raise ValueError( @@ -1021,7 +1022,7 @@ def stream_weights_megatron_to_hf( ) else: # Regular case - yield the tensor normally - yield HFWeightTuple(hf_name, final_tensor) + yield HFWeightTuple(hf_name, final_tensor, task.param_name) def dtype_from_hf(self, config, default=None): """Extract torch dtype from a HuggingFace config. diff --git a/src/megatron/bridge/models/conversion/param_mapping.py b/src/megatron/bridge/models/conversion/param_mapping.py index 0258b422da..a04212733a 100644 --- a/src/megatron/bridge/models/conversion/param_mapping.py +++ b/src/megatron/bridge/models/conversion/param_mapping.py @@ -821,6 +821,14 @@ def hf_to_megatron( ) hf_weights = hf_weights.to(target_param.dtype) + actual_dim0_size = hf_weights.shape[0] + expect_dim0_size = target_param.shape[0] * self.tp_size + if actual_dim0_size != expect_dim0_size: + assert self.megatron_param in {"embedding.word_embeddings.weight", "output_layer.weight"}, ( + f"{hf_weights.shape=} {target_param.shape=} {self.tp_size=} {self.megatron_param=} {self.hf_param=}" + ) + hf_weights = _pad_right_dim0(hf_weights, pad_size=expect_dim0_size - actual_dim0_size) + # For bias (1D), we still split along dim 0 # For weight (2D), we split along dim 0 (output dimension) full_size = hf_weights.shape[0] @@ -867,6 +875,11 @@ def megatron_to_hf( return {str(self.hf_param): full_weights} +def _pad_right_dim0(x: torch.Tensor, pad_size: int) -> torch.Tensor: + padder = torch.zeros((pad_size, *x.shape[1:]), dtype=x.dtype, device=x.device) + return torch.cat([x, padder], dim=0) + + class RowParallelMapping(MegatronParamMapping[torch.Tensor]): """Mapping for **row-parallel** linear weights. diff --git a/src/megatron/bridge/models/conversion/peft_bridge.py b/src/megatron/bridge/models/conversion/peft_bridge.py index 61e8d440c4..d101224f39 100644 --- a/src/megatron/bridge/models/conversion/peft_bridge.py +++ b/src/megatron/bridge/models/conversion/peft_bridge.py @@ -97,7 +97,7 @@ def _select_hf_base_param_name(base_mapping, adapter_key: Optional[str], expecte hf_param = base_mapping.hf_param if isinstance(hf_param, str): - return hf_param if hf_param.endswith(expected_suffix) else None + return hf_param if isinstance(hf_param, dict): if adapter_key: @@ -169,10 +169,12 @@ def _resolve_hf_adapter_param_name( # Strip expert layers numbering base_suffix = base_suffix.rstrip(digits) hf_base_name = _select_hf_base_param_name(base_mapping, adapter_key, base_suffix) - if hf_base_name is None or not hf_base_name.endswith(base_suffix): + if hf_base_name is None: return None - return hf_base_name[: -len(base_suffix)] + hf_suffix + if hf_base_name.endswith(base_suffix): + return hf_base_name[: -len(base_suffix)] + hf_suffix + return hf_base_name + hf_suffix def _get_base_hf_param_names_for_adapter( self, @@ -203,14 +205,31 @@ def _get_base_hf_param_names_for_adapter( def _make_lora_param_name(self, base_name: str, megatron_adapter_suffix: str) -> Optional[str]: """Translate a base HF weight name into its LoRA-specific counterpart.""" - if not base_name.endswith(".weight"): - return None - hf_suffix = MEGATRON_TO_HF_LORA_SUFFIX.get(megatron_adapter_suffix) if hf_suffix is None: return None - return base_name[: -len(".weight")] + hf_suffix + if base_name.endswith(".weight"): + return base_name[: -len(".weight")] + hf_suffix + return base_name + hf_suffix + + def _get_grouped_expert_base_suffixes(self, num_moe_experts: int) -> list[str]: + """Return base suffixes used for grouped expert adapter export. + + Default behavior iterates per-expert. Override for models whose grouped + expert adapters share a single 2D weight across all experts. + """ + return [f".weight{i}" for i in range(num_moe_experts)] + + def _prepare_expert_adapter_for_hf( + self, tensor: torch.Tensor, is_grouped_expert: bool + ) -> torch.Tensor: + """Adjust adapter weight shape before yielding to HF format. + + Default is identity (no-op). Override for models that need shape + adjustment for grouped expert adapters. + """ + return tensor def _is_fused_qkv(self, hf_weight_names: Iterable[str]) -> bool: """Check whether the provided HF names correspond to a fused QKV weight.""" @@ -624,6 +643,8 @@ def stream_adapter_weights_megatron_to_hf( linear_in_tensor = adapter_weight.linear_in_weight.weight linear_out_tensor = adapter_weight.linear_out_weight.weight + megatron_linear_in_name = adapter_weight.linear_in_weight.param_name + megatron_linear_out_name = adapter_weight.linear_out_weight.param_name 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 expert_linear_in_gathered = None @@ -638,12 +659,12 @@ def stream_adapter_weights_megatron_to_hf( base_suffixes = [".weight"] if is_grouped_expert: - base_suffixes = [f".weight{expert_num}" for expert_num in range(num_moe_experts)] + base_suffixes = self._get_grouped_expert_base_suffixes(num_moe_experts) for base_suffix in base_suffixes: current_linear_in_tensor = linear_in_tensor current_linear_out_tensor = linear_out_tensor - if is_grouped_expert: + if is_grouped_expert and len(base_suffixes) > 1: expert_idx = int(base_suffix[len(".weight") :]) current_linear_in_tensor = self._select_expert_adapter_weight( linear_in_tensor, @@ -688,12 +709,16 @@ def stream_adapter_weights_megatron_to_hf( current_linear_out_tensor = per_base_linear_out.get(base_name) assert current_linear_out_tensor is not None, "unknown projection name" - yield HFWeightTuple(linear_in_hf_names[index], current_linear_in_tensor) - yield HFWeightTuple(linear_out_hf_names[index], current_linear_out_tensor) + lin_in = self._prepare_expert_adapter_for_hf(current_linear_in_tensor, is_grouped_expert) + lin_out = self._prepare_expert_adapter_for_hf(current_linear_out_tensor, is_grouped_expert) + yield HFWeightTuple(linear_in_hf_names[index], lin_in, megatron_linear_in_name) + yield HFWeightTuple(linear_out_hf_names[index], lin_out, megatron_linear_out_name) continue - yield HFWeightTuple(linear_in_hf_names[0], current_linear_in_tensor) - yield HFWeightTuple(linear_out_hf_names[0], current_linear_out_tensor) + lin_in = self._prepare_expert_adapter_for_hf(current_linear_in_tensor, is_grouped_expert) + lin_out = self._prepare_expert_adapter_for_hf(current_linear_out_tensor, is_grouped_expert) + yield HFWeightTuple(linear_in_hf_names[0], lin_in, megatron_linear_in_name) + yield HFWeightTuple(linear_out_hf_names[0], lin_out, megatron_linear_out_name) def _get_fused_adapter_linear_out_slices( self, diff --git a/src/megatron/bridge/models/conversion/utils.py b/src/megatron/bridge/models/conversion/utils.py index 25c5b360a9..fb2dea5d55 100644 --- a/src/megatron/bridge/models/conversion/utils.py +++ b/src/megatron/bridge/models/conversion/utils.py @@ -44,7 +44,8 @@ def weights_verification_table(bridge, megatron_model) -> Table: table.add_column("Matches Original", justify="center") # Check each weight against the original HF-model - for name, param in bridge.export_hf_weights(megatron_model, show_progress=True): + for item in bridge.export_hf_weights(megatron_model, show_progress=True): + name, param = item.param_name, item.weight original_param = bridge.hf_pretrained.state[name] table.add_row( name, diff --git a/src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py b/src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py index 9dcc0c341e..cb75000f48 100644 --- a/src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py +++ b/src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py @@ -108,6 +108,21 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> GPTModelProvider return provider + def _get_grouped_expert_base_suffixes(self, num_moe_experts: int) -> list[str]: + """GPT-OSS grouped expert adapters are 2D (shared across all experts). + + A single .weight0 lookup is sufficient for mapping resolution. + """ + return [".weight0"] + + def _prepare_expert_adapter_for_hf( + self, tensor: torch.Tensor, is_grouped_expert: bool + ) -> torch.Tensor: + """Restore 3D shape expected by HF for grouped expert adapters.""" + if is_grouped_expert: + return tensor.unsqueeze(0) + return tensor + def maybe_modify_loaded_hf_weight( self, hf_param: str | dict[str, str], hf_state_dict: Mapping[str, torch.Tensor] ) -> torch.Tensor: diff --git a/src/megatron/bridge/models/hf_pretrained/state.py b/src/megatron/bridge/models/hf_pretrained/state.py index 0b79feb2a6..3cd9b6c36d 100644 --- a/src/megatron/bridge/models/hf_pretrained/state.py +++ b/src/megatron/bridge/models/hf_pretrained/state.py @@ -732,8 +732,11 @@ def save_generator( key_to_filename_map = self.key_to_filename_map all_expected_keys = set(key_to_filename_map.keys()) + buffered_tensors = {} if not key_to_filename_map: - buffered_tensors = dict(generator) + for item in generator: + name, tensor = item.param_name, item.weight + buffered_tensors[name] = tensor if buffered_tensors: save_file(buffered_tensors, output_path / "model.safetensors") return @@ -743,11 +746,11 @@ def save_generator( filename_to_keys_map[filename].add(key) files_to_save = dict(filename_to_keys_map) - buffered_tensors = {} all_yielded_keys = set() all_saved_keys = set() - for name, tensor in generator: + for item in generator: + name, tensor = item.param_name, item.weight all_yielded_keys.add(name) if name not in all_expected_keys: if strict: @@ -939,7 +942,8 @@ def _save_generator_distributed( buffered_tensors: Dict[str, torch.Tensor] = {} actually_saved_keys: Set[str] = set() - for name, tensor in generator: + for item in generator: + name, tensor = item.param_name, item.weight all_yielded_keys.add(name) if name not in all_expected_keys: diff --git a/src/megatron/bridge/models/qwen_vl/qwen35_vl_bridge.py b/src/megatron/bridge/models/qwen_vl/qwen35_vl_bridge.py index 6e586e07ea..2fe95f932f 100644 --- a/src/megatron/bridge/models/qwen_vl/qwen35_vl_bridge.py +++ b/src/megatron/bridge/models/qwen_vl/qwen35_vl_bridge.py @@ -64,6 +64,7 @@ _QWEN3_5_DENSE_HF_CLASS_NAME = "Qwen3_5ForConditionalGeneration" _QWEN3_5_MOE_HF_CLASS_NAME = "Qwen3_5MoeForConditionalGeneration" +_MTP_LAYER_MODULE_NAMES = ("transformer_layer", "mtp_model_layer") @MegatronModelBridge.register_bridge( @@ -444,48 +445,56 @@ def mapping_registry(self) -> MegatronMappingRegistry: "language_model.mtp.layers.0.enorm.weight": "mtp.pre_fc_norm_embedding.weight", "language_model.mtp.layers.0.hnorm.weight": "mtp.pre_fc_norm_hidden.weight", "language_model.mtp.layers.0.final_layernorm.weight": "mtp.norm.weight", - "language_model.mtp.layers.0.mtp_model_layer.mlp.router.weight": "mtp.layers.0.mlp.gate.weight", - "language_model.mtp.layers.0.mtp_model_layer.pre_mlp_layernorm.weight": "mtp.layers.0.post_attention_layernorm.weight", - "language_model.mtp.layers.0.mtp_model_layer.self_attention.linear_qkv.layer_norm_weight": "mtp.layers.0.input_layernorm.weight", - "language_model.mtp.layers.0.mtp_model_layer.self_attention.q_layernorm.weight": "mtp.layers.0.self_attn.q_norm.weight", - "language_model.mtp.layers.0.mtp_model_layer.self_attention.k_layernorm.weight": "mtp.layers.0.self_attn.k_norm.weight", - "language_model.mtp.layers.0.mtp_model_layer.self_attention.linear_proj.weight": "mtp.layers.0.self_attn.o_proj.weight", } for megatron_param, hf_param in mtp_param_mappings.items(): mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) - mapping_list.extend( - [ - QKVMapping( - megatron_param="language_model.mtp.layers.*.mtp_model_layer.self_attention.linear_qkv.weight", - q="mtp.layers.*.self_attn.q_proj.weight", - k="mtp.layers.*.self_attn.k_proj.weight", - v="mtp.layers.*.self_attn.v_proj.weight", - ), - GatedMLPMapping( - megatron_param="language_model.mtp.layers.*.mtp_model_layer.mlp.experts.linear_fc1.weight*", - gate="mtp.layers.*.mlp.experts.*.gate_proj.weight", - up="mtp.layers.*.mlp.experts.*.up_proj.weight", - ), - AutoMapping( - megatron_param="language_model.mtp.layers.*.mtp_model_layer.mlp.experts.linear_fc2.weight*", - hf_param="mtp.layers.*.mlp.experts.*.down_proj.weight", - ), - GatedMLPMapping( - megatron_param="language_model.mtp.layers.*.mtp_model_layer.mlp.shared_experts.linear_fc1.weight", - gate="mtp.layers.*.mlp.shared_expert.gate_proj.weight", - up="mtp.layers.*.mlp.shared_expert.up_proj.weight", - ), - AutoMapping( - megatron_param="language_model.mtp.layers.*.mtp_model_layer.mlp.shared_experts.linear_fc2.weight", - hf_param="mtp.layers.*.mlp.shared_expert.down_proj.weight", - ), - ReplicatedMapping( - megatron_param="language_model.mtp.layers.0.mtp_model_layer.mlp.shared_experts.gate_weight", - hf_param="mtp.layers.0.mlp.shared_expert_gate.weight", - ), - ] - ) + mtp_layer_param_mappings = { + "language_model.mtp.layers.0.{mtp_layer}.mlp.router.weight": "mtp.layers.0.mlp.gate.weight", + "language_model.mtp.layers.0.{mtp_layer}.pre_mlp_layernorm.weight": "mtp.layers.0.post_attention_layernorm.weight", + "language_model.mtp.layers.0.{mtp_layer}.self_attention.linear_qkv.layer_norm_weight": "mtp.layers.0.input_layernorm.weight", + "language_model.mtp.layers.0.{mtp_layer}.self_attention.q_layernorm.weight": "mtp.layers.0.self_attn.q_norm.weight", + "language_model.mtp.layers.0.{mtp_layer}.self_attention.k_layernorm.weight": "mtp.layers.0.self_attn.k_norm.weight", + "language_model.mtp.layers.0.{mtp_layer}.self_attention.linear_proj.weight": "mtp.layers.0.self_attn.o_proj.weight", + } + for mtp_layer in _MTP_LAYER_MODULE_NAMES: + for megatron_param, hf_param in mtp_layer_param_mappings.items(): + mapping_list.append( + AutoMapping(megatron_param=megatron_param.format(mtp_layer=mtp_layer), hf_param=hf_param) + ) + + mapping_list.extend( + [ + QKVMapping( + megatron_param=f"language_model.mtp.layers.*.{mtp_layer}.self_attention.linear_qkv.weight", + q="mtp.layers.*.self_attn.q_proj.weight", + k="mtp.layers.*.self_attn.k_proj.weight", + v="mtp.layers.*.self_attn.v_proj.weight", + ), + GatedMLPMapping( + megatron_param=f"language_model.mtp.layers.*.{mtp_layer}.mlp.experts.linear_fc1.weight*", + gate="mtp.layers.*.mlp.experts.*.gate_proj.weight", + up="mtp.layers.*.mlp.experts.*.up_proj.weight", + ), + AutoMapping( + megatron_param=f"language_model.mtp.layers.*.{mtp_layer}.mlp.experts.linear_fc2.weight*", + hf_param="mtp.layers.*.mlp.experts.*.down_proj.weight", + ), + GatedMLPMapping( + megatron_param=f"language_model.mtp.layers.*.{mtp_layer}.mlp.shared_experts.linear_fc1.weight", + gate="mtp.layers.*.mlp.shared_expert.gate_proj.weight", + up="mtp.layers.*.mlp.shared_expert.up_proj.weight", + ), + AutoMapping( + megatron_param=f"language_model.mtp.layers.*.{mtp_layer}.mlp.shared_experts.linear_fc2.weight", + hf_param="mtp.layers.*.mlp.shared_expert.down_proj.weight", + ), + ReplicatedMapping( + megatron_param=f"language_model.mtp.layers.0.{mtp_layer}.mlp.shared_experts.gate_weight", + hf_param="mtp.layers.0.mlp.shared_expert_gate.weight", + ), + ] + ) return MegatronMappingRegistry(*mapping_list) @@ -721,30 +730,38 @@ def mapping_registry(self) -> MegatronMappingRegistry: "language_model.mtp.layers.0.enorm.weight": "mtp.pre_fc_norm_embedding.weight", "language_model.mtp.layers.0.hnorm.weight": "mtp.pre_fc_norm_hidden.weight", "language_model.mtp.layers.0.final_layernorm.weight": "mtp.norm.weight", - "language_model.mtp.layers.0.mtp_model_layer.mlp.linear_fc1.layer_norm_weight": "mtp.layers.0.post_attention_layernorm.weight", - "language_model.mtp.layers.0.mtp_model_layer.mlp.linear_fc2.weight": "mtp.layers.0.mlp.down_proj.weight", - "language_model.mtp.layers.0.mtp_model_layer.self_attention.linear_qkv.layer_norm_weight": "mtp.layers.0.input_layernorm.weight", - "language_model.mtp.layers.0.mtp_model_layer.self_attention.q_layernorm.weight": "mtp.layers.0.self_attn.q_norm.weight", - "language_model.mtp.layers.0.mtp_model_layer.self_attention.k_layernorm.weight": "mtp.layers.0.self_attn.k_norm.weight", - "language_model.mtp.layers.0.mtp_model_layer.self_attention.linear_proj.weight": "mtp.layers.0.self_attn.o_proj.weight", } for megatron_param, hf_param in mtp_param_mappings.items(): mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) - mapping_list.extend( - [ - QKVMapping( - megatron_param="language_model.mtp.layers.*.mtp_model_layer.self_attention.linear_qkv.weight", - q="mtp.layers.*.self_attn.q_proj.weight", - k="mtp.layers.*.self_attn.k_proj.weight", - v="mtp.layers.*.self_attn.v_proj.weight", - ), - GatedMLPMapping( - megatron_param="language_model.mtp.layers.*.mtp_model_layer.mlp.linear_fc1.weight", - gate="mtp.layers.*.mlp.gate_proj.weight", - up="mtp.layers.*.mlp.up_proj.weight", - ), - ] - ) + mtp_layer_param_mappings = { + "language_model.mtp.layers.0.{mtp_layer}.mlp.linear_fc1.layer_norm_weight": "mtp.layers.0.post_attention_layernorm.weight", + "language_model.mtp.layers.0.{mtp_layer}.mlp.linear_fc2.weight": "mtp.layers.0.mlp.down_proj.weight", + "language_model.mtp.layers.0.{mtp_layer}.self_attention.linear_qkv.layer_norm_weight": "mtp.layers.0.input_layernorm.weight", + "language_model.mtp.layers.0.{mtp_layer}.self_attention.q_layernorm.weight": "mtp.layers.0.self_attn.q_norm.weight", + "language_model.mtp.layers.0.{mtp_layer}.self_attention.k_layernorm.weight": "mtp.layers.0.self_attn.k_norm.weight", + "language_model.mtp.layers.0.{mtp_layer}.self_attention.linear_proj.weight": "mtp.layers.0.self_attn.o_proj.weight", + } + for mtp_layer in _MTP_LAYER_MODULE_NAMES: + for megatron_param, hf_param in mtp_layer_param_mappings.items(): + mapping_list.append( + AutoMapping(megatron_param=megatron_param.format(mtp_layer=mtp_layer), hf_param=hf_param) + ) + + mapping_list.extend( + [ + QKVMapping( + megatron_param=f"language_model.mtp.layers.*.{mtp_layer}.self_attention.linear_qkv.weight", + q="mtp.layers.*.self_attn.q_proj.weight", + k="mtp.layers.*.self_attn.k_proj.weight", + v="mtp.layers.*.self_attn.v_proj.weight", + ), + GatedMLPMapping( + megatron_param=f"language_model.mtp.layers.*.{mtp_layer}.mlp.linear_fc1.weight", + gate="mtp.layers.*.mlp.gate_proj.weight", + up="mtp.layers.*.mlp.up_proj.weight", + ), + ] + ) return MegatronMappingRegistry(*mapping_list) diff --git a/tests/unit_tests/models/test_model_bridge_lora.py b/tests/unit_tests/models/test_model_bridge_lora.py index 6641f72c4c..c00e179824 100644 --- a/tests/unit_tests/models/test_model_bridge_lora.py +++ b/tests/unit_tests/models/test_model_bridge_lora.py @@ -723,6 +723,64 @@ def _raise_on_build(*_args, **_kwargs): assert weights[0].param_name in {"hf.weight", "hf.base_layer.weight"} +def test_stream_adapter_weights_megatron_to_hf_grouped_expert(monkeypatch): + """Verify grouped MoE expert adapters are exported with correct names and shapes.""" + bridge = DummyBridge() + + bridge._get_grouped_expert_base_suffixes = lambda num_experts: [".weight0"] + bridge._prepare_expert_adapter_for_hf = ( + lambda tensor, is_grouped: tensor.unsqueeze(0) if is_grouped else tensor + ) + + 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(), + ), + ) + + 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", torch.ones(2, 2), vp_stage=0), + linear_out_weight=MegatronWeightTuple("local_out", 2 * torch.ones(2, 2), vp_stage=0), + ) + + 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", + lambda *_args, **_kwargs: ["model.layers.0.mlp.experts.gate_up_proj"], + ) + monkeypatch.setattr(bridge, "_gather_expert_adapter_weight", lambda self_or_tensor, *a: self_or_tensor) + + megatron_model = [SimpleNamespace(config=SimpleNamespace(num_moe_experts=8))] + weights = list(bridge.stream_adapter_weights_megatron_to_hf(megatron_model, cpu=False, show_progress=False)) + + assert len(weights) == 2 + assert weights[0].param_name == "model.layers.0.mlp.experts.gate_up_proj.lora_A.weight" + assert weights[1].param_name == "model.layers.0.mlp.experts.gate_up_proj.lora_B.weight" + assert weights[0].weight.shape == (1, 2, 2) + assert weights[1].weight.shape == (1, 2, 2) + + def test_column_parallel_mapping_skips_ep_gather_for_adapters(monkeypatch): mapping = ColumnParallelMapping( "decoder.layers.0.mlp.experts.linear_fc1.adapter.linear_in.weight",