From 196265af1ff7ff669657c71b5df0730d419eb983 Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Tue, 7 Apr 2026 02:49:03 +0000 Subject: [PATCH 1/8] make it align with miles megatron-core Guard EnergonProvider import with try/except so megatron-bridge works when megatron.energon is not installed. Co-Authored-By: Yusheng Su Co-Authored-By: gongyisheng Co-Authored-By: Claude --- src/megatron/bridge/data/vlm_datasets/__init__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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, From 3b48c9a30f4f5d0d9434cbdcca3fb68f1fd5c2ea Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Tue, 7 Apr 2026 02:50:14 +0000 Subject: [PATCH 2/8] allow megatron-bridge to return param_name - rl require feat Add megatron_param_name field to HFWeightTuple so downstream consumers (e.g. RL pipelines) can map exported HF weights back to their original Megatron parameter names. Co-Authored-By: Yusheng Su Co-Authored-By: gongyisheng Co-Authored-By: Claude --- src/megatron/bridge/models/conversion/model_bridge.py | 7 ++++--- src/megatron/bridge/models/conversion/peft_bridge.py | 8 ++++---- 2 files changed, 8 insertions(+), 7 deletions(-) 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/peft_bridge.py b/src/megatron/bridge/models/conversion/peft_bridge.py index 61e8d440c4..0033ea0c72 100644 --- a/src/megatron/bridge/models/conversion/peft_bridge.py +++ b/src/megatron/bridge/models/conversion/peft_bridge.py @@ -688,12 +688,12 @@ 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) + yield HFWeightTuple(linear_in_hf_names[index], current_linear_in_tensor, None) + yield HFWeightTuple(linear_out_hf_names[index], current_linear_out_tensor, None) continue - yield HFWeightTuple(linear_in_hf_names[0], current_linear_in_tensor) - yield HFWeightTuple(linear_out_hf_names[0], current_linear_out_tensor) + yield HFWeightTuple(linear_in_hf_names[0], current_linear_in_tensor, None) + yield HFWeightTuple(linear_out_hf_names[0], current_linear_out_tensor, None) def _get_fused_adapter_linear_out_slices( self, From 73e4ec8cd607640b6858c202dd4ad4558b53f1e4 Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Tue, 7 Apr 2026 02:50:56 +0000 Subject: [PATCH 3/8] add - return megatron_param_name Propagate actual megatron linear_in/linear_out parameter names through HFWeightTuple in peft_bridge adapter weight streaming instead of None. Co-Authored-By: Yusheng Su Co-Authored-By: gongyisheng Co-Authored-By: Claude --- src/megatron/bridge/models/conversion/peft_bridge.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/megatron/bridge/models/conversion/peft_bridge.py b/src/megatron/bridge/models/conversion/peft_bridge.py index 0033ea0c72..793b70e911 100644 --- a/src/megatron/bridge/models/conversion/peft_bridge.py +++ b/src/megatron/bridge/models/conversion/peft_bridge.py @@ -624,6 +624,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 @@ -688,12 +690,12 @@ 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, None) - yield HFWeightTuple(linear_out_hf_names[index], current_linear_out_tensor, None) + yield HFWeightTuple(linear_in_hf_names[index], current_linear_in_tensor, megatron_linear_in_name) + yield HFWeightTuple(linear_out_hf_names[index], current_linear_out_tensor, megatron_linear_out_name) continue - yield HFWeightTuple(linear_in_hf_names[0], current_linear_in_tensor, None) - yield HFWeightTuple(linear_out_hf_names[0], current_linear_out_tensor, None) + yield HFWeightTuple(linear_in_hf_names[0], current_linear_in_tensor, megatron_linear_in_name) + yield HFWeightTuple(linear_out_hf_names[0], current_linear_out_tensor, megatron_linear_out_name) def _get_fused_adapter_linear_out_slices( self, From d2ee05178d382414bec006fb94dc415483ec6cda Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Tue, 7 Apr 2026 02:55:18 +0000 Subject: [PATCH 4/8] merged Tom's patch - Add _pad_right_dim0 padding for vocab size mismatch in ColumnParallelMapping (embedding/output_layer weights) - Fix HFWeightTuple unpacking in state.py, utils.py, and auto_bridge.py to handle the new 3-field NamedTuple (param_name, weight, megatron_param_name) Co-Authored-By: Yusheng Su Co-Authored-By: gongyisheng Co-Authored-By: Claude --- .../bridge/models/conversion/auto_bridge.py | 4 ++-- .../bridge/models/conversion/param_mapping.py | 13 +++++++++++++ src/megatron/bridge/models/conversion/utils.py | 3 ++- src/megatron/bridge/models/hf_pretrained/state.py | 12 ++++++++---- 4 files changed, 25 insertions(+), 7 deletions(-) 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/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/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/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: From d1232659282474c70e5233fbb6f29a5527f22bb0 Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Sat, 18 Apr 2026 07:05:44 +0000 Subject: [PATCH 5/8] fix: support MoE expert LoRA adapter conversion in peft_bridge Handle MoE expert layers (e.g. GPT-OSS gate_up_proj) whose HF param names do not end with the conventional ".weight" suffix: - _select_hf_base_param_name: return hf_param even when it lacks the expected suffix, so MoE expert mappings are no longer silently dropped. - _megatron_to_hf_adapter_name: append hf_suffix directly when hf_base_name does not end with base_suffix instead of returning None. - _make_lora_param_name: remove hard ".weight" suffix requirement; gracefully handle both standard and MoE-style param names. For grouped MoE experts: - Use a single ".weight0" lookup instead of per-expert iteration, since grouped expert adapters share 2D weights across all experts. - unsqueeze(0) on adapter tensors to restore the expected 3D shape when yielding HFWeightTuples for grouped experts. Made-with: Cursor --- .../bridge/models/conversion/peft_bridge.py | 39 ++++++++++++------- 1 file changed, 26 insertions(+), 13 deletions(-) diff --git a/src/megatron/bridge/models/conversion/peft_bridge.py b/src/megatron/bridge/models/conversion/peft_bridge.py index 793b70e911..6f0ebe4384 100644 --- a/src/megatron/bridge/models/conversion/peft_bridge.py +++ b/src/megatron/bridge/models/conversion/peft_bridge.py @@ -97,7 +97,11 @@ 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 + if hf_param.endswith(expected_suffix): + return hf_param + else: + # MoE expert layers (e.g. GPT-OSS gate_up_proj) don't end with .weight + return hf_param if isinstance(hf_param, dict): if adapter_key: @@ -169,10 +173,13 @@ 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 + else: + return hf_base_name + hf_suffix def _get_base_hf_param_names_for_adapter( self, @@ -203,14 +210,14 @@ 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 + else: + return base_name + hf_suffix def _is_fused_qkv(self, hf_weight_names: Iterable[str]) -> bool: """Check whether the provided HF names correspond to a fused QKV weight.""" @@ -640,12 +647,14 @@ 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)] + # MoE grouped expert adapters: weights are shared (2D) across experts. + # Yield once with .weight0 for mapping lookup instead of per-expert iteration. + base_suffixes = [".weight0"] 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, @@ -690,12 +699,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, megatron_linear_in_name) - yield HFWeightTuple(linear_out_hf_names[index], current_linear_out_tensor, megatron_linear_out_name) + lin_in = current_linear_in_tensor.unsqueeze(0) if is_grouped_expert else current_linear_in_tensor + lin_out = current_linear_out_tensor.unsqueeze(0) if is_grouped_expert else current_linear_out_tensor + 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, megatron_linear_in_name) - yield HFWeightTuple(linear_out_hf_names[0], current_linear_out_tensor, megatron_linear_out_name) + lin_in = current_linear_in_tensor.unsqueeze(0) if is_grouped_expert else current_linear_in_tensor + lin_out = current_linear_out_tensor.unsqueeze(0) if is_grouped_expert else current_linear_out_tensor + 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, From dc22ea0938726514e042bfb0e3381498792b9768 Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Mon, 20 Apr 2026 01:58:54 +0000 Subject: [PATCH 6/8] Revert "fix: support MoE expert LoRA adapter conversion in peft_bridge" This reverts commit d1232659282474c70e5233fbb6f29a5527f22bb0. --- .../bridge/models/conversion/peft_bridge.py | 39 +++++++------------ 1 file changed, 13 insertions(+), 26 deletions(-) diff --git a/src/megatron/bridge/models/conversion/peft_bridge.py b/src/megatron/bridge/models/conversion/peft_bridge.py index 6f0ebe4384..793b70e911 100644 --- a/src/megatron/bridge/models/conversion/peft_bridge.py +++ b/src/megatron/bridge/models/conversion/peft_bridge.py @@ -97,11 +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): - if hf_param.endswith(expected_suffix): - return hf_param - else: - # MoE expert layers (e.g. GPT-OSS gate_up_proj) don't end with .weight - return hf_param + return hf_param if hf_param.endswith(expected_suffix) else None if isinstance(hf_param, dict): if adapter_key: @@ -173,13 +169,10 @@ 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: + if hf_base_name is None or not hf_base_name.endswith(base_suffix): return None - if hf_base_name.endswith(base_suffix): - return hf_base_name[: -len(base_suffix)] + hf_suffix - else: - return hf_base_name + hf_suffix + return hf_base_name[: -len(base_suffix)] + hf_suffix def _get_base_hf_param_names_for_adapter( self, @@ -210,14 +203,14 @@ 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 - if base_name.endswith(".weight"): - return base_name[: -len(".weight")] + hf_suffix - else: - return base_name + hf_suffix + return base_name[: -len(".weight")] + hf_suffix def _is_fused_qkv(self, hf_weight_names: Iterable[str]) -> bool: """Check whether the provided HF names correspond to a fused QKV weight.""" @@ -647,14 +640,12 @@ def stream_adapter_weights_megatron_to_hf( base_suffixes = [".weight"] if is_grouped_expert: - # MoE grouped expert adapters: weights are shared (2D) across experts. - # Yield once with .weight0 for mapping lookup instead of per-expert iteration. - base_suffixes = [".weight0"] + base_suffixes = [f".weight{expert_num}" for expert_num in range(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 and len(base_suffixes) > 1: + if is_grouped_expert: expert_idx = int(base_suffix[len(".weight") :]) current_linear_in_tensor = self._select_expert_adapter_weight( linear_in_tensor, @@ -699,16 +690,12 @@ 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" - lin_in = current_linear_in_tensor.unsqueeze(0) if is_grouped_expert else current_linear_in_tensor - lin_out = current_linear_out_tensor.unsqueeze(0) if is_grouped_expert else current_linear_out_tensor - 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) + yield HFWeightTuple(linear_in_hf_names[index], current_linear_in_tensor, megatron_linear_in_name) + yield HFWeightTuple(linear_out_hf_names[index], current_linear_out_tensor, megatron_linear_out_name) continue - lin_in = current_linear_in_tensor.unsqueeze(0) if is_grouped_expert else current_linear_in_tensor - lin_out = current_linear_out_tensor.unsqueeze(0) if is_grouped_expert else current_linear_out_tensor - 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) + yield HFWeightTuple(linear_in_hf_names[0], current_linear_in_tensor, megatron_linear_in_name) + yield HFWeightTuple(linear_out_hf_names[0], current_linear_out_tensor, megatron_linear_out_name) def _get_fused_adapter_linear_out_slices( self, From 3fd3768045422d0aa5c97e90a4e6c659aea9acb9 Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Mon, 20 Apr 2026 02:14:59 +0000 Subject: [PATCH 7/8] feat: support MoE expert LoRA adapter conversion via hook methods Refactor grouped expert adapter handling in peft_bridge to use overridable hook methods instead of hardcoded GPT-OSS-specific logic: - _select_hf_base_param_name: accept HF params that don't end with the expected suffix (e.g. MoE expert names like gate_up_proj) - _resolve_hf_adapter_param_name / _make_lora_param_name: handle param names without .weight suffix gracefully - Add _get_grouped_expert_base_suffixes() hook: default per-expert iteration; GPT-OSS overrides to single .weight0 lookup - Add _prepare_expert_adapter_for_hf() hook: default no-op; GPT-OSS overrides to unsqueeze(0) for 3D shape restoration - Add unit test for grouped expert LoRA adapter export path Made-with: Cursor --- .../bridge/models/conversion/peft_bridge.py | 49 +++++++++++----- .../bridge/models/gpt_oss/gpt_oss_bridge.py | 15 +++++ .../models/test_model_bridge_lora.py | 58 +++++++++++++++++++ 3 files changed, 109 insertions(+), 13 deletions(-) diff --git a/src/megatron/bridge/models/conversion/peft_bridge.py b/src/megatron/bridge/models/conversion/peft_bridge.py index 793b70e911..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.""" @@ -640,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, @@ -690,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, megatron_linear_in_name) - yield HFWeightTuple(linear_out_hf_names[index], current_linear_out_tensor, megatron_linear_out_name) + 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, megatron_linear_in_name) - yield HFWeightTuple(linear_out_hf_names[0], current_linear_out_tensor, megatron_linear_out_name) + 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/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/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", From 65355f07141b37d7a99a9981b80dfcd01385d138 Mon Sep 17 00:00:00 2001 From: Artem Yatsenko Date: Mon, 11 May 2026 11:38:31 -0700 Subject: [PATCH 8/8] add fix for qwen35 mtp layer name --- .../bridge/models/qwen_vl/qwen35_vl_bridge.py | 135 ++++++++++-------- 1 file changed, 76 insertions(+), 59 deletions(-) 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)