add specific providers

Signed-off-by: Ananth Subramaniam <[email protected]>

Author: ananthsub

src/megatron/bridge/models/__init__.py

@@ -186,6 +186,11 @@
     # Nemotron Models
     "NemotronBridge",
     "NemotronModelProvider",
+    "Nemotron3ModelProvider4B",
+    "Nemotron3ModelProvider8B",
+    "Nemotron3ModelProvider22B",
+    "Nemotron4ModelProvider15B",
+    "Nemotron4ModelProvider340B",
     # VL Models
     "Qwen25VLModel",
     "Qwen25VLBridge",

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

@@ -13,10 +13,22 @@
 # limitations under the License.
 
 from megatron.bridge.models.nemotron.nemotron_bridge import NemotronBridge
-from megatron.bridge.models.nemotron.nemotron_provider import NemotronModelProvider
+from megatron.bridge.models.nemotron.nemotron_provider import (
+    Nemotron3ModelProvider4B,
+    Nemotron3ModelProvider8B,
+    Nemotron3ModelProvider22B,
+    Nemotron4ModelProvider15B,
+    Nemotron4ModelProvider340B,
+    NemotronModelProvider,
+)
 
 
 __all__ = [
     "NemotronBridge",
     "NemotronModelProvider",
+    "Nemotron3ModelProvider4B",
+    "Nemotron3ModelProvider8B",
+    "Nemotron3ModelProvider22B",
+    "Nemotron4ModelProvider15B",
+    "Nemotron4ModelProvider340B",
 ]

src/megatron/bridge/models/nemotron/nemotron_provider.py

@@ -14,7 +14,7 @@
 
 import logging
 from dataclasses import dataclass, field
-from typing import Callable
+from typing import Callable, Optional
 
 import torch
 
@@ -50,3 +50,101 @@ class NemotronModelProvider(GPTModelProvider):
     layernorm_zero_centered_gamma: bool = True
     cross_entropy_loss_fusion: bool = True
     apply_rope_fusion: bool = field(default_factory=fusions.can_enable_apply_rope_fusion)
+
+    # Nemotron3Config4B as default configs
+    num_layers: int = 32
+    seq_length: int = 4096
+    hidden_size: int = 3072
+    ffn_hidden_size: int = 9216
+    num_attention_heads: int = 24
+    num_query_groups: Optional[int] = 8
+    kv_channels: Optional[int] = 128
+    init_method_std: float = 0.0134
+
+    # Data type settings to match HF models
+    bf16: bool = True
+    fp16: bool = False
+    params_dtype: torch.dtype = torch.bfloat16
+    autocast_dtype: torch.dtype = torch.bfloat16
+
+
+@dataclass
+class Nemotron3ModelProvider4B(NemotronModelProvider):
+    """
+    Configuration class for the Nemotron3 4B model, inheriting from NemotronModelProvider.
+    Maps to: nvidia/Minitron-4B-Base, nvidia/Nemotron-Mini-4B-Instruct
+    """
+
+    num_layers: int = 32
+    seq_length: int = 4096
+    hidden_size: int = 3072
+    ffn_hidden_size: int = 9216
+    num_attention_heads: int = 24
+    num_query_groups: int = 8
+    kv_channels: Optional[int] = 128
+    init_method_std: float = 0.0134
+
+
+@dataclass
+class Nemotron3ModelProvider8B(NemotronModelProvider):
+    """
+    Configuration class for the Nemotron3 8B model, inheriting from NemotronModelProvider.
+    Maps to: nvidia/Minitron-8B-Base
+    """
+
+    num_layers: int = 32
+    seq_length: int = 4096
+    hidden_size: int = 4096
+    ffn_hidden_size: int = 16384
+    num_attention_heads: int = 48  # Updated to match HF model (was 32)
+    num_query_groups: int = 8  # Updated to match HF model (was None)
+    kv_channels: Optional[int] = 128  # Updated to match HF model (was None)
+    init_method_std: float = 0.010
+
+
+@dataclass
+class Nemotron3ModelProvider22B(NemotronModelProvider):
+    """
+    Configuration class for the Nemotron3 22B model, inheriting from NemotronModelProvider.
+    """
+
+    num_layers: int = 40
+    seq_length: int = 4096
+    hidden_size: int = 6144
+    ffn_hidden_size: int = 24576
+    num_attention_heads: int = 48
+    num_query_groups: Optional[int] = None
+    kv_channels: Optional[int] = None
+    init_method_std: float = 0.008
+
+
+@dataclass
+class Nemotron4ModelProvider15B(NemotronModelProvider):
+    """
+    Configuration class for the Nemotron4 15B model, inheriting from NemotronModelProvider.
+    """
+
+    num_layers: int = 32
+    seq_length: int = 4096
+    hidden_size: int = 6144
+    ffn_hidden_size: int = 24576
+    num_attention_heads: int = 48
+    num_query_groups: Optional[int] = 8
+    kv_channels: Optional[int] = None
+    init_method_std: float = 0.0134
+
+
+@dataclass
+class Nemotron4ModelProvider340B(NemotronModelProvider):
+    """
+    Configuration class for the Nemotron4 340B model, inheriting from NemotronModelProvider.
+    """
+
+    num_layers: int = 96
+    seq_length: int = 4096
+    hidden_size: int = 18432
+    ffn_hidden_size: int = 73728
+    num_attention_heads: int = 96
+    num_query_groups: Optional[int] = 8
+    kv_channels: Optional[int] = None
+    init_method_std: float = 0.0063

tests/functional_tests/models/test_nemotron_provider.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from megatron.bridge.models.conversion.auto_bridge import AutoBridge
+from megatron.bridge.models.nemotron import (
+    Nemotron3ModelProvider4B,
+    Nemotron3ModelProvider8B,
+)
+from tests.functional_tests.utils import compare_provider_configs
+
+
+HF_MODEL_ID_TO_BRIDGE_MODEL_PROVIDER = {
+    "nvidia/Minitron-4B-Base": Nemotron3ModelProvider4B,
+    "nvidia/Minitron-8B-Base": Nemotron3ModelProvider8B,
+    "nvidia/Nemotron-Mini-4B-Instruct": Nemotron3ModelProvider4B,
+}
+
+
+class TestNemotronModelProviderMapping:
+    """Test that bridge provider configs are equivalent to predefined provider configs."""
+
+    @pytest.mark.parametrize("hf_model_id,provider_class", list(HF_MODEL_ID_TO_BRIDGE_MODEL_PROVIDER.items()))
+    def test_bridge_vs_predefined_provider_config_equivalence(self, hf_model_id, provider_class):
+        """Test that bridge converted provider config matches predefined provider config."""
+        # Create bridge from HF model
+        bridge = AutoBridge.from_hf_pretrained(hf_model_id)
+        converted_provider = bridge.to_megatron_provider(load_weights=False)
+        converted_provider.finalize()
+
+        # Create predefined provider
+        predefined_provider = provider_class()
+        predefined_provider.finalize()
+
+        # Compare configs
+        compare_provider_configs(converted_provider, predefined_provider, hf_model_id)

tests/unit_tests/models/nemotron/test_nemotron_bridge.py

@@ -21,7 +21,14 @@
 from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
 from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
 from megatron.bridge.models.nemotron.nemotron_bridge import NemotronBridge
-from megatron.bridge.models.nemotron.nemotron_provider import NemotronModelProvider
+from megatron.bridge.models.nemotron.nemotron_provider import (
+    Nemotron3ModelProvider4B,
+    Nemotron3ModelProvider8B,
+    Nemotron3ModelProvider22B,
+    Nemotron4ModelProvider15B,
+    Nemotron4ModelProvider340B,
+    NemotronModelProvider,
+)
 
 
 class TestNemotronBridge:
@@ -132,3 +139,97 @@ def test_dtype_configuration(self, mock_pretrained_nemotron):
         assert provider.bf16 is True
         assert provider.fp16 is False
         assert provider.params_dtype == torch.bfloat16
+
+
+class TestNemotronSpecificProviders:
+    """Test cases for specific Nemotron model provider configurations."""
+
+    def test_nemotron3_4b_config(self):
+        """Test Nemotron3 4B provider configuration matches HF model specs."""
+        provider = Nemotron3ModelProvider4B()
+
+        # Should match nvidia/Minitron-4B-Base and nvidia/Nemotron-Mini-4B-Instruct
+        assert provider.hidden_size == 3072
+        assert provider.num_layers == 32
+        assert provider.num_attention_heads == 24
+        assert provider.num_query_groups == 8
+        assert provider.ffn_hidden_size == 9216
+        assert provider.kv_channels == 128
+        assert provider.seq_length == 4096
+        assert provider.init_method_std == 0.0134
+
+    def test_nemotron3_8b_config(self):
+        """Test Nemotron3 8B provider configuration matches HF model specs."""
+        provider = Nemotron3ModelProvider8B()
+
+        # Should match nvidia/Minitron-8B-Base
+        assert provider.hidden_size == 4096
+        assert provider.num_layers == 32
+        assert provider.num_attention_heads == 48  # Updated to match HF
+        assert provider.num_query_groups == 8  # Updated to match HF
+        assert provider.ffn_hidden_size == 16384
+        assert provider.kv_channels == 128  # Updated to match HF
+        assert provider.seq_length == 4096
+        assert provider.init_method_std == 0.010
+
+    def test_nemotron3_22b_config(self):
+        """Test Nemotron3 22B provider configuration."""
+        provider = Nemotron3ModelProvider22B()
+
+        assert provider.hidden_size == 6144
+        assert provider.num_layers == 40
+        assert provider.num_attention_heads == 48
+        assert provider.num_query_groups is None
+        assert provider.ffn_hidden_size == 24576
+        assert provider.kv_channels is None
+        assert provider.seq_length == 4096
+        assert provider.init_method_std == 0.008
+
+    def test_nemotron4_15b_config(self):
+        """Test Nemotron4 15B provider configuration."""
+        provider = Nemotron4ModelProvider15B()
+
+        assert provider.hidden_size == 6144
+        assert provider.num_layers == 32
+        assert provider.num_attention_heads == 48
+        assert provider.num_query_groups == 8
+        assert provider.ffn_hidden_size == 24576
+        assert provider.kv_channels is None
+        assert provider.seq_length == 4096
+        assert provider.init_method_std == 0.0134
+
+    def test_nemotron4_340b_config(self):
+        """Test Nemotron4 340B provider configuration."""
+        provider = Nemotron4ModelProvider340B()
+
+        # Should match nvidia/Nemotron-4-340B-Base/Instruct (if available)
+        assert provider.hidden_size == 18432
+        assert provider.num_layers == 96
+        assert provider.num_attention_heads == 96
+        assert provider.num_query_groups == 8
+        assert provider.ffn_hidden_size == 73728
+        assert provider.kv_channels is None
+        assert provider.seq_length == 4096
+        assert provider.init_method_std == 0.0063
+
+    def test_all_providers_have_nemotron_defaults(self):
+        """Test that all specific providers inherit Nemotron-specific defaults."""
+        providers = [
+            Nemotron3ModelProvider4B(),
+            Nemotron3ModelProvider8B(),
+            Nemotron3ModelProvider22B(),
+            Nemotron4ModelProvider15B(),
+            Nemotron4ModelProvider340B(),
+        ]
+
+        for provider in providers:
+            # Check Nemotron-specific defaults
+            assert provider.normalization == "LayerNorm"
+            assert provider.position_embedding_type == "rope"
+            assert provider.share_embeddings_and_output_weights is False
+            assert provider.add_bias_linear is False
+            assert provider.hidden_dropout == 0.0
+            assert provider.attention_dropout == 0.0
+            assert provider.rotary_percent == 0.5
+            assert provider.layernorm_zero_centered_gamma is True
+            assert provider.cross_entropy_loss_fusion is True