[model, refactor] refactor: Centralize provider_bridge config mapping in base class for VLM models

src/megatron/bridge/models/gemma_vl/gemma3_vl_bridge.py

@@ -30,48 +30,59 @@
 from megatron.bridge.models.hf_pretrained.vlm import PreTrainedVLM
 
 
-@MegatronModelBridge.register_bridge(source=Gemma3ForConditionalGeneration, target=Gemma3VLModel)
+@MegatronModelBridge.register_bridge(
+    source=Gemma3ForConditionalGeneration,
+    target=Gemma3VLModel,
+    provider=Gemma3VLModelProvider,
+    model_type="gemma3_vl",
+)
 class Gemma3VLBridge(MegatronModelBridge):
     """
     Megatron Bridge for Gemma3 VL.
+
+    This bridge handles the conversion between HuggingFace Gemma3ForConditionalGeneration
+    and Megatron-Core Gemma3VLModel formats, including weight mappings and
+    configuration translation for vision-language models.
+
+    Example:
+        >>> from megatron.bridge import AutoBridge
+        >>> bridge = AutoBridge.from_hf_pretrained("google/gemma-3-4b-it")
+        >>> provider = bridge.to_megatron_provider()
     """
 
     def provider_bridge(self, hf_pretrained: PreTrainedVLM) -> Gemma3VLModelProvider:
         hf_config = hf_pretrained.config
         text_config = hf_config.text_config
         vision_config = hf_config.vision_config
 
-        provider = Gemma3VLModelProvider(
-            # Text configuration
-            init_method_std=text_config.initializer_range,
-            hidden_size=text_config.hidden_size,
-            ffn_hidden_size=text_config.intermediate_size,
-            kv_channels=text_config.head_dim,
-            seq_length=text_config.max_position_embeddings,
-            num_attention_heads=text_config.num_attention_heads,
-            num_layers=text_config.num_hidden_layers,
-            num_query_groups=text_config.num_key_value_heads,
-            window_size=text_config.sliding_window,
-            rotary_base=(text_config.rope_local_base_freq, text_config.rope_theta),
-            layernorm_epsilon=text_config.rms_norm_eps,
-            vocab_size=text_config.vocab_size,
-            softmax_scale=1.0 / math.sqrt(text_config.query_pre_attn_scalar),
-            rope_scaling_factor=text_config.rope_scaling["factor"] if text_config.rope_scaling else 1.0,
-            # Vision configuration
-            vision_config=vision_config,
-            mm_tokens_per_image=hf_config.mm_tokens_per_image,
-            # VL-specific token IDs
-            bos_token_id=getattr(hf_config, "bos_token_id", 0),
-            eos_token_id=getattr(hf_config, "eos_token_id", 1),
-            vision_start_token_id=getattr(hf_config, "vision_start_token_id", 255999),
-            vision_end_token_id=getattr(hf_config, "vision_end_token_id", 256000),
-            image_token_id=getattr(hf_config, "image_token_id", 151655),
-            # Precision configuration
-            fp16=(self.dtype_from_hf(hf_config, default=torch.float32) == torch.float16),
-            bf16=(self.dtype_from_hf(hf_config, default=torch.float32) == torch.bfloat16),
-            params_dtype=self.dtype_from_hf(hf_config, default=torch.float32),
-        )
+        # Use base class helper for common config conversion
+        provider_kwargs = self.hf_config_to_provider_kwargs(text_config)
+        provider = Gemma3VLModelProvider(**provider_kwargs)
+
+        # Gemma3-specific features not in CONFIG_MAPPING
+        provider.window_size = text_config.sliding_window
+        provider.rotary_base = (text_config.rope_local_base_freq, text_config.rope_theta)
+        provider.softmax_scale = 1.0 / math.sqrt(text_config.query_pre_attn_scalar)
+        provider.rope_scaling_factor = text_config.rope_scaling["factor"] if text_config.rope_scaling else 1.0
+
+        # Override dtype and vocab settings to match baseline
+        provider.bf16 = True
+        provider.params_dtype = torch.bfloat16
+        provider.autocast_dtype = torch.bfloat16
+        provider.make_vocab_size_divisible_by = 128
+
+        # Vision configuration
+        provider.vision_config = vision_config
+        provider.mm_tokens_per_image = hf_config.mm_tokens_per_image
+
+        # VL-specific token IDs
+        provider.bos_token_id = getattr(hf_config, "bos_token_id", 0)
+        provider.eos_token_id = getattr(hf_config, "eos_token_id", 1)
+        provider.vision_start_token_id = getattr(hf_config, "vision_start_token_id", 255999)
+        provider.vision_end_token_id = getattr(hf_config, "vision_end_token_id", 256000)
+        provider.image_token_id = getattr(hf_config, "image_token_id", 262144)
 
+        # Vision projector configuration
         provider.vision_projector_config.input_size = vision_config.hidden_size
         provider.vision_projector_config.hidden_size = text_config.hidden_size
 

src/megatron/bridge/models/nemotron_vl/nemotron_vl_bridge.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import torch
+from megatron.core.activations import squared_relu
 
 from megatron.bridge.models import ColumnParallelMapping, RowParallelMapping
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
@@ -29,7 +29,12 @@
 from megatron.bridge.models.nemotron_vl.nemotron_vl_provider import NemotronNano12Bv2VLModelProvider
 
 
-@MegatronModelBridge.register_bridge(source="NemotronH_Nano_VL_V2", target=NemotronVLModel)
+@MegatronModelBridge.register_bridge(
+    source="NemotronH_Nano_VL_V2",
+    target=NemotronVLModel,
+    provider=NemotronNano12Bv2VLModelProvider,
+    model_type="nemotron_vl",
+)
 class NemotronVLBridge(MegatronModelBridge):
     """Conversion utilities between HF Nemotron-VL and Megatron-Core format."""
 
@@ -39,25 +44,26 @@ class NemotronVLBridge(MegatronModelBridge):
 
     def provider_bridge(self, hf_pretrained: PreTrainedVLM) -> NemotronNano12Bv2VLModelProvider:  # type: ignore[override]
         hf_config = hf_pretrained.config
+        llm_config = hf_config.llm_config
+
+        # Use base class helper for common config mapping
+        provider_kwargs = self.hf_config_to_provider_kwargs(llm_config)
+
+        # Handle vocab size divisibility
+        provider_kwargs["make_vocab_size_divisible_by"] = self.make_vocab_size_divisible_by(llm_config.vocab_size)
+
+        provider = NemotronNano12Bv2VLModelProvider(**provider_kwargs)
+
+        # Nemotron VL-specific settings
+        # Note: Most defaults come from the provider class hierarchy (NemotronNano12Bv2Provider)
+        provider.scatter_embedding_sequence_parallel = False
+        provider.attention_softmax_in_fp32 = True
+
+        # Override fields that should use NemotronH provider's specialized defaults
+        # instead of HF config values
+        provider.activation_func = squared_relu  # Nemotron uses squared_relu, not HF's hidden_act
+        provider.autocast_dtype = None  # Not set in original code
 
-        provider = NemotronNano12Bv2VLModelProvider(
-            num_layers=hf_config.llm_config.num_hidden_layers,
-            hidden_size=hf_config.llm_config.hidden_size,
-            ffn_hidden_size=hf_config.llm_config.intermediate_size,
-            num_attention_heads=hf_config.llm_config.num_attention_heads,
-            num_query_groups=getattr(
-                hf_config.llm_config, "num_key_value_heads", hf_config.llm_config.num_attention_heads // 2
-            ),
-            init_method_std=hf_config.llm_config.initializer_range,
-            layernorm_epsilon=getattr(hf_config.llm_config, "layer_norm_epsilon", 1e-5),
-            make_vocab_size_divisible_by=self.make_vocab_size_divisible_by(hf_config.llm_config.vocab_size),
-            share_embeddings_and_output_weights=getattr(hf_config.llm_config, "tie_word_embeddings", False),
-            vocab_size=hf_config.llm_config.vocab_size,
-            seq_length=hf_config.llm_config.max_position_embeddings,
-            fp16=(self.dtype_from_hf(hf_config, default=torch.float32) == torch.float16),
-            bf16=(self.dtype_from_hf(hf_config, default=torch.float32) == torch.bfloat16),
-            params_dtype=self.dtype_from_hf(hf_config, default=torch.float32),
-        )
         return provider
 
     # ------------------------------------------------------------------

src/megatron/bridge/models/qwen_vl/__init__.py

@@ -20,7 +20,7 @@
     Qwen3VLMoEModelProvider,
 )
 from megatron.bridge.models.qwen_vl.qwen25_vl_bridge import Qwen25VLBridge
-from megatron.bridge.models.qwen_vl.qwen_vl_provider import (
+from megatron.bridge.models.qwen_vl.qwen25_vl_provider import (
     Qwen25VLModelProvider,
 )
 

src/megatron/bridge/models/qwen_vl/qwen25_vl_bridge.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import torch
 from transformers import Qwen2_5_VLForConditionalGeneration
 
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
@@ -25,10 +24,15 @@
 )
 from megatron.bridge.models.hf_pretrained.vlm import PreTrainedVLM
 from megatron.bridge.models.qwen_vl.modeling_qwen25_vl import Qwen25VLModel
-from megatron.bridge.models.qwen_vl.qwen_vl_provider import Qwen25VLModelProvider
+from megatron.bridge.models.qwen_vl.qwen25_vl_provider import Qwen25VLModelProvider
 
 
-@MegatronModelBridge.register_bridge(source=Qwen2_5_VLForConditionalGeneration, target=Qwen25VLModel)
+@MegatronModelBridge.register_bridge(
+    source=Qwen2_5_VLForConditionalGeneration,
+    target=Qwen25VLModel,
+    provider=Qwen25VLModelProvider,
+    model_type="qwen2_5_vl",
+)
 class Qwen25VLBridge(MegatronModelBridge):
     """
     Megatron Bridge for Qwen2.5-VL Conditional Generation.
@@ -45,35 +49,28 @@ class Qwen25VLBridge(MegatronModelBridge):
 
     def provider_bridge(self, hf_pretrained: PreTrainedVLM) -> Qwen25VLModelProvider:
         hf_config = hf_pretrained.config
+        text_config = hf_config  # Qwen2.5-VL has text config fields directly on main config
 
-        provider = Qwen25VLModelProvider(
-            num_layers=hf_config.num_hidden_layers,
-            hidden_size=hf_config.hidden_size,
-            ffn_hidden_size=hf_config.intermediate_size,
-            num_attention_heads=hf_config.num_attention_heads,
-            num_query_groups=hf_config.num_key_value_heads,
-            init_method_std=hf_config.initializer_range,
-            layernorm_epsilon=hf_config.rms_norm_eps,
-            gated_linear_unit=True,
-            make_vocab_size_divisible_by=self.make_vocab_size_divisible_by(hf_config.vocab_size),
-            rotary_base=hf_config.rope_theta,
-            share_embeddings_and_output_weights=getattr(hf_config, "tie_word_embeddings", False),
-            vocab_size=hf_config.vocab_size,
-            seq_length=hf_config.max_position_embeddings,
-            fp16=(self.dtype_from_hf(hf_config, default=torch.float32) == torch.float16),
-            bf16=(self.dtype_from_hf(hf_config, default=torch.float32) == torch.bfloat16),
-            params_dtype=self.dtype_from_hf(hf_config, default=torch.float32),
-            add_qkv_bias=True,  # Qwen2 has bias in QKV projections
-            vision_config=hf_config.vision_config,
-            # VL-specific token IDs
-            bos_token_id=getattr(hf_config, "bos_token_id", 151643),
-            eos_token_id=getattr(hf_config, "eos_token_id", 151645),
-            vision_start_token_id=getattr(hf_config, "vision_start_token_id", 151652),
-            vision_end_token_id=getattr(hf_config, "vision_end_token_id", 151653),
-            vision_token_id=getattr(hf_config, "vision_token_id", 151654),
-            image_token_id=getattr(hf_config, "image_token_id", 151655),
-            video_token_id=getattr(hf_config, "video_token_id", 151656),
-        )
+        provider_kwargs = self.hf_config_to_provider_kwargs(text_config)
+        provider = Qwen25VLModelProvider(**provider_kwargs)
+
+        # Qwen2-specific settings
+        provider.normalization = "RMSNorm"
+        provider.gated_linear_unit = True
+        provider.add_qkv_bias = True
+        provider.add_bias_linear = False
+        provider.hidden_dropout = 0.0
+
+        # VL-specific overrides
+        provider.position_embedding_type = "mrope"
+        provider.vision_config = hf_config.vision_config
+        provider.bos_token_id = getattr(hf_config, "bos_token_id", 151643)
+        provider.eos_token_id = getattr(hf_config, "eos_token_id", 151645)
+        provider.vision_start_token_id = getattr(hf_config, "vision_start_token_id", 151652)
+        provider.vision_end_token_id = getattr(hf_config, "vision_end_token_id", 151653)
+        provider.vision_token_id = getattr(hf_config, "vision_token_id", 151654)
+        provider.image_token_id = getattr(hf_config, "image_token_id", 151655)
+        provider.video_token_id = getattr(hf_config, "video_token_id", 151656)
 
         return provider
 

src/megatron/bridge/models/qwen_vl/qwen_vl_provider.py → src/megatron/bridge/models/qwen_vl/qwen25_vl_provider.py

@@ -18,25 +18,16 @@
 from megatron.core.models.gpt import GPTModel as MCoreGPTModel
 from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLVisionConfig
 
-from megatron.bridge.models import (
-    Qwen2ModelProvider,
-)
-
-from .modeling_qwen25_vl import Qwen25VLModel
-
-
-# =============================================================================
-# Qwen 2.5 VL Model Providers
-# =============================================================================
+from megatron.bridge.models.gpt_provider import GPTModelProvider
+from megatron.bridge.models.qwen_vl.modeling_qwen25_vl import Qwen25VLModel
 
 
 @dataclass
-class Qwen25VLModelProvider(Qwen2ModelProvider):
+class Qwen25VLModelProvider(GPTModelProvider):
     """
     Base model provider for Qwen 2.5 VL Models.
     """
 
-    # Language configuration inherited from Qwen25ModelProvider3B
     # VL models shouldn't scatter embeddings across sequence parallel regions because
     # the vision embeddings are going to be inserted into the language embeddings.
     scatter_embedding_sequence_parallel: bool = False
@@ -74,4 +65,4 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Qwen25V
         return model
 
     def provide_language_model(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreGPTModel:
-        return super().provide(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage)
+        return GPTModelProvider.provide(self, pre_process=pre_process, post_process=post_process, vp_stage=vp_stage)