add fix for qwen35 mtp layer name

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,

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(

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,21 +1008,21 @@ 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(
                         "Encountered lm_head.weight when embeddings are tied. This indicates a mapping error."
                     )
                 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.

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.
 

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,

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,

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:

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: