[model] refactor: formalize hf_config on MegatronModelBridge (#3329)

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/models/conversion/model_bridge.py

@@ -287,6 +287,10 @@ def mapping_registry(self) -> MegatronMappingRegistry:
     # Set this in bridge subclasses to include model-specific files beyond standard artifacts
     ADDITIONAL_FILE_PATTERNS = None
 
+    # HuggingFace PretrainedConfig, set by register_bridge_implementation dispatch.
+    # Available in mapping_registry(), stream_weights_*(), and build_conversion_tasks().
+    hf_config = None
+
     # Common bidirectional config field name mapping: (hf_name, megatron_name)
     # Some mappings may not be used by all models - that's fine, unused fields are skipped
     CONFIG_MAPPING = [

src/megatron/bridge/models/deepseek/deepseek_v2_bridge.py

@@ -111,16 +111,7 @@ def megatron_to_hf_config(cls, provider) -> dict:
 
         return hf_cfg
 
-    def build_conversion_tasks(self, hf_pretrained, megatron_model):
-        """Override to store config before mapping_registry is called."""
-        # Store config on instance for use in mapping_registry
-        from transformers import PretrainedConfig
-
-        self._hf_config = hf_pretrained if isinstance(hf_pretrained, PretrainedConfig) else hf_pretrained.config
-        return super().build_conversion_tasks(hf_pretrained, megatron_model)
-
     def mapping_registry(self) -> MegatronMappingRegistry:
-        # Get hf_config if available (set by build_conversion_tasks)
-        hf_config = getattr(self, "_hf_config", None)
+        hf_config = self.hf_config
         mapping_list = get_common_mapping_list(hf_config=hf_config)
         return MegatronMappingRegistry(*mapping_list)

src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py

@@ -120,16 +120,8 @@ def megatron_to_hf_config(cls, provider: MLAModelProvider) -> dict:
 
         return hf_cfg
 
-    def build_conversion_tasks(self, hf_pretrained, megatron_model):
-        """Override to store config before mapping_registry is called."""
-        # Store config on instance for use in mapping_registry
-        from transformers import PretrainedConfig
-
-        self._hf_config = hf_pretrained if isinstance(hf_pretrained, PretrainedConfig) else hf_pretrained.config
-        return super().build_conversion_tasks(hf_pretrained, megatron_model)
-
     def mapping_registry(self) -> MegatronMappingRegistry:
-        hf_config = getattr(self, "_hf_config", None)
+        hf_config = self.hf_config
         mapping_list = get_common_mapping_list(hf_config=hf_config)
         mapping_list.append(
             AutoMapping(

src/megatron/bridge/models/glm/glm45_bridge.py

@@ -100,11 +100,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> GPTModelProvider
         return provider
 
     def build_conversion_tasks(self, hf_pretrained, megatron_model):
-        """Override to store config before mapping_registry is called."""
-        from transformers import PretrainedConfig
-
-        # Store config on instance for use in mapping_registry
-        self._hf_config = hf_pretrained if isinstance(hf_pretrained, PretrainedConfig) else hf_pretrained.config
+        """Override to store HF state source before mapping_registry is called."""
         has_state = hasattr(hf_pretrained, "state") and hasattr(hf_pretrained.state, "source")
         self._hf_state_source = hf_pretrained.state.source if has_state else None
         self._hf_keys = list(self._hf_state_source.get_all_keys()) if self._hf_state_source else None
@@ -208,10 +204,10 @@ def mapping_registry(self) -> MegatronMappingRegistry:
                 ]
             )
         # optionally add MTP mappings
-        if not hasattr(self, "_hf_config"):
+        if self.hf_config is None:
             logger.warning("No HF config found, skipping MTP mappings.")
             return MegatronMappingRegistry(*mapping_list)
-        hf_config = self._hf_config
+        hf_config = self.hf_config
         num_mtp_layers = getattr(hf_config, "num_nextn_predict_layers", 0)
         num_transformer_layers = hf_config.num_hidden_layers
         for mtp_layer in range(num_mtp_layers):

src/megatron/bridge/models/glm_vl/glm_45v_bridge.py

@@ -89,8 +89,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedVLM) -> GLM45VModelProvider:
         return provider
 
     def build_conversion_tasks(self, hf_pretrained, megatron_model):
-        """Override to store config before mapping_registry is called."""
-        self._hf_config = hf_pretrained.config
+        """Override to store HF state source before mapping_registry is called."""
         self._hf_state_source = hf_pretrained.state.source
         self._hf_keys = list(self._hf_state_source.get_all_keys())
         return super().build_conversion_tasks(hf_pretrained, megatron_model)

src/megatron/bridge/models/qwen_audio/modeling_qwen2_audio.py

@@ -149,7 +149,7 @@ def __init__(
 
         # Store audio token id from config
         self.audio_token_id = getattr(config, "audio_token_id", 151646)
-        self.pad_token_id = getattr(config.hf_config, "pad_token_id", -1)
+        self.pad_token_id = getattr(config, "pad_token_id", -1)
 
     def set_input_tensor(self, input_tensor) -> None:
         """Set model chunk input tensor."""