[models] refactor: Remove size-specific provider classes (#3854)

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

Author: yaoyu-33

docs/nemo2-migration-guide.md

@@ -12,28 +12,29 @@ This guide helps you migrate from NeMo 2.0 training and recipes to Megatron Brid
 
 ## Model Configuration Mapping
 
-Megatron Bridge offers model providers that directly map to NeMo 2.0 model configs.
+Megatron Bridge derives model provider settings from Hugging Face configs through
+`AutoBridge`, or from recipe functions that configure a base provider directly.
 
 ### Examples
 
 | NeMo 2.0 | Megatron Bridge |
 |----------|-----------------|
-| `llm.Llama3Config8B` | {py:class}`bridge.models.Llama3ModelProvider8B` |
-| `llm.Llama31Config70B` | {py:class}`bridge.models.Llama31ModelProvider70B` |
-| `llm.Qwen2Config7B` | {py:class}`bridge.models.Qwen2ModelProvider7B` |
+| `llm.Llama3Config8B` | `AutoBridge.from_hf_pretrained("meta-llama/Meta-Llama-3-8B").to_megatron_provider()` |
+| `llm.Llama31Config70B` | `AutoBridge.from_hf_pretrained("meta-llama/Llama-3.1-70B").to_megatron_provider()` |
+| `llm.Qwen2Config7B` | `AutoBridge.from_hf_pretrained("Qwen/Qwen2-7B").to_megatron_provider()` |
 | `llm.DeepseekV2Config` | {py:class}`bridge.models.DeepseekV2ModelProvider` |
 
 ### Supported Model Families
 
-Megatron Bridge supports the following model families with preset providers:
+Megatron Bridge supports the following model families through bridge mappings and provider classes:
 - **Base Models**: `GPTModelProvider`, `T5ModelProvider`, `MambaModelProvider`
 - **Llama**: Llama2, Llama3, Llama3.1, Llama3.2, CodeLlama, Llama4
 - **Qwen**: Qwen2, Qwen2.5, Qwen3, Qwen3MoE, Qwen2.5VL
 - **DeepSeek**: DeepSeek, DeepSeekV2, DeepSeekV2Lite, DeepSeekV3, Moonlight
 - **Nemotron**: Nemotron3, Nemotron4, NemotronH, NemotronNano
 - **NVIDIA Mamba**: Mamba variants and hybrid models
 
-For a complete list of all model providers and their parameters, see {py:mod}`bridge.models`.
+For a complete list of bridge mappings, providers, and their parameters, see {py:mod}`bridge.models`.
 
 <!-- TODO: Create a dedicated model support table with tested HF checkpoint mappings -->
 
@@ -114,7 +115,7 @@ from megatron.bridge.training.config import (
     LoggerConfig,
     TrainingConfig,
 )
-from megatron.bridge.models import Llama3ModelProvider8B  # Direct equivalent to Llama3Config8B
+from megatron.bridge.models import GPTModelProvider  # Configure directly or derive with AutoBridge
 from megatron.core.optimizer import OptimizerConfig
 from megatron.bridge.training.config import SchedulerConfig
 from megatron.bridge.training.pretrain import pretrain
@@ -123,8 +124,8 @@ from megatron.bridge.training.gpt_step import forward_step
 
 def create_config():
     return ConfigContainer(
-        # Model with parallelism built-in - using preset 8B config
-        model=Llama3ModelProvider8B(
+        # Model with parallelism built in
+        model=GPTModelProvider(
             # Parallelism settings (moved from MegatronStrategy)
             tensor_model_parallel_size=2,
             pipeline_model_parallel_size=2,
@@ -224,13 +225,13 @@ llm.finetune(
 #### Now: Megatron Bridge
 ```python  
 # Megatron Bridge fine-tuning configuration (with optional PEFT)
-from megatron.bridge.models import Llama3ModelProvider8B
+from megatron.bridge.models import GPTModelProvider
 from megatron.bridge.peft import LoRA
 
 def create_finetune_config():
     return ConfigContainer(
-        model=Llama3ModelProvider8B(
-            # Preset config matching Llama3Config8B
+        model=GPTModelProvider(
+            # Configure architecture fields directly or derive with AutoBridge
         ),
         train=TrainingConfig(
             micro_batch_size=1,
@@ -382,7 +383,7 @@ from megatron.bridge.training.config import (
     SchedulerConfig,
 )
 from megatron.core.optimizer import OptimizerConfig
-from megatron.bridge.models import Llama3ModelProvider8B
+from megatron.bridge.models import GPTModelProvider
 from megatron.bridge.training import pretrain
 
 def llama3_8b_config(
@@ -404,9 +405,9 @@ def llama3_8b_config(
 ) -> ConfigContainer:
     """Create a Llama3 8B pretraining configuration."""
     return ConfigContainer(
-        model=Llama3ModelProvider8B(
-            # Preset architecture from Llama3Config8B (num_layers=32, hidden_size=4096, etc.)
-            # Only need to specify parallelism and overrides
+        model=GPTModelProvider(
+            # Configure architecture fields directly or derive this provider with AutoBridge
+            # Then specify parallelism and overrides
             tensor_model_parallel_size=tensor_parallelism,
             pipeline_model_parallel_size=pipeline_parallelism,
         ),
@@ -697,7 +698,7 @@ In Megatron Bridge, vocabulary size can be specified in either the model provide
 ```python
 # Option 1: Let tokenizer determine vocab size
 config = ConfigContainer(
-    model=Llama3ModelProvider8B(
+    model=GPTModelProvider(
         # vocab_size not set - will use tokenizer's vocab size
         vocab_size=None,
     ),
@@ -709,7 +710,7 @@ config = ConfigContainer(
 
 # Option 2: Explicitly set vocab size in model
 config = ConfigContainer(
-    model=Llama3ModelProvider8B(
+    model=GPTModelProvider(
         vocab_size=128256,  # Explicitly set (must be >= tokenizer vocab size)
     ),
     tokenizer=TokenizerConfig(...),
@@ -1154,7 +1155,7 @@ from megatron.bridge.training.config import ConfigContainer, CheckpointConfig
 
 # Include PEFT in ConfigContainer
 config = ConfigContainer(
-    model=Llama3ModelProvider8B(),
+    model=GPTModelProvider(),
     # ... other configs
     checkpoint=CheckpointConfig(
         pretrained_checkpoint="/path/to/megatron/checkpoint",  # Required for PEFT
@@ -1288,7 +1289,7 @@ from megatron.bridge.training import pretrain
 from megatron.bridge.training.gpt_step import forward_step
 
 config = ConfigContainer(
-    model=Llama3ModelProvider8B(
+    model=GPTModelProvider(
         tensor_model_parallel_size=2,
         pipeline_model_parallel_size=2,
     ),
@@ -1324,7 +1325,7 @@ from megatron.bridge.training import finetune
 from megatron.bridge.training.gpt_step import forward_step
 
 config = ConfigContainer(
-    model=Llama3ModelProvider8B(),
+    model=GPTModelProvider(),
     train=TrainingConfig(
         train_iters=1000,
         eval_interval=100,
@@ -1355,7 +1356,7 @@ Add a `peft` configuration to enable parameter-efficient training:
 from megatron.bridge.peft import LoRA
 
 config = ConfigContainer(
-    model=Llama3ModelProvider8B(),
+    model=GPTModelProvider(),
     train=TrainingConfig(
         train_iters=1000,
         eval_interval=100,

skills/adding-model-support/SKILL.md

@@ -130,6 +130,13 @@ directory — keep them namespaced under the `modeling_<model>` prefix.
    then sets model-specific attributes on it. **Do not create a provider file** — the stock
    provider returned by `super().provider_bridge()` is usually sufficient for LLMs
    (e.g., `GPTModelProvider`, or another base provider selected via `PROVIDER_CLASS`).
+   **Do not add size-specific provider classes** whose names combine
+   `ModelProvider` with a model-size suffix. Examples of forbidden suffixes
+   include `7B`, `200M`, and `A3B`. Model size and architecture fields should
+   come from the Hugging Face config through `AutoBridge` /
+   `MegatronModelBridge` config mapping. If a recipe needs a fixed
+   architecture, configure the base provider inside the recipe function instead
+   of exporting a provider subclass.
 
 **VLM:**
 1. **Bridge** — Register bridge, implement config and weight mappings.
@@ -311,6 +318,12 @@ Each recipe file defines functions for each model size + training mode:
 
 For detailed recipe patterns, see @skills/adding-model-support/recipe-patterns.md.
 
+Recipes are the right API surface for model-size presets. Do not create or
+export size-specific provider subclasses for recipes; either call
+`AutoBridge.from_hf_pretrained(...).to_megatron_provider(load_weights=False)` to
+derive the provider from HF config, or instantiate the base provider class with
+explicit architecture fields inside the recipe function.
+
 ### Export checklist
 
 1. Family `__init__.py` — import and add to `__all__`

skills/adding-model-support/llm-patterns.md

@@ -38,8 +38,6 @@ Create a `GPTModelProvider` subclass only when:
    RoPE params, custom MoE fields) that need to serialize into `run_config.yaml`.
 2. **Custom `provide()` logic** — The model needs special instantiation (e.g., TE version
    checks, sink attention, custom layer specs that require runtime logic).
-3. **Predefined size variants for recipes** — Hardcoded configs like `LlamaModelProvider8B`
-   used by recipe functions (not by the bridge itself).
 
 ```python
 @dataclass
@@ -62,19 +60,35 @@ def provider_bridge(self, hf_pretrained) -> MyModelProvider:
     return provider
 ```
 
-### Predefined size variants (for recipes only)
+### No Size-Specific Provider Classes
 
-Size-specific subclasses are used by recipes, not by the bridge:
+Do not create provider subclasses whose names combine `ModelProvider` with a
+model-size suffix. Examples of forbidden suffixes include `7B`, `200M`, and
+`A3B`. The bridge should derive architecture fields from the Hugging Face config
+via `AutoBridge` / `MegatronModelBridge` config mapping.
+
+For recipe presets, keep the size in the recipe function name and configure a
+base provider inside the function:
 
 ```python
-@dataclass
-class MyModelProvider7B(MyModelProvider):
-    num_layers: int = 32
-    hidden_size: int = 4096
-    num_attention_heads: int = 32
-    num_query_groups: int = 8
-    ffn_hidden_size: int = 14336
-    vocab_size: int = 128256
+def my_model_7b_pretrain_config() -> ConfigContainer:
+    cfg = _pretrain_common()
+    cfg.model = MyModelProvider(
+        num_layers=32,
+        hidden_size=4096,
+        num_attention_heads=32,
+        num_query_groups=8,
+        ffn_hidden_size=14336,
+        vocab_size=128256,
+    )
+    return cfg
+```
+
+When the recipe targets an existing HF checkpoint, prefer deriving the provider
+from HF config:
+
+```python
+cfg.model = AutoBridge.from_hf_pretrained("org/my-model-7b").to_megatron_provider(load_weights=False)
 ```
 
 ## Bridge Pattern

src/megatron/bridge/diffusion/models/wan/wan_provider.py

@@ -89,35 +89,3 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> WanMode
             fp16_lm_cross_entropy=self.fp16_lm_cross_entropy,
             parallel_output=self.parallel_output,
         )
-
-
-@dataclass
-class WanModelProvider1_3B(WanModelProvider):
-    """WAN 1.3B model configuration.
-
-    Architecture: 30 layers, hidden_size=1536, 12 attention heads,
-    ffn_hidden_size=8960. Default seq_length=1024.
-    """
-
-    num_layers: int = 30
-    hidden_size: int = 1536
-    ffn_hidden_size: int = 8960
-    num_attention_heads: int = 12
-    crossattn_emb_size: int = 1536
-    seq_length: int = 1024
-
-
-@dataclass
-class WanModelProvider14B(WanModelProvider):
-    """WAN 14B model configuration.
-
-    Architecture: 40 layers, hidden_size=5120, 40 attention heads,
-    ffn_hidden_size=13824. Default seq_length=1024.
-    """
-
-    num_layers: int = 40
-    hidden_size: int = 5120
-    ffn_hidden_size: int = 13824
-    num_attention_heads: int = 40
-    crossattn_emb_size: int = 5120
-    seq_length: int = 1024

src/megatron/bridge/diffusion/recipes/wan/wan.py

@@ -18,7 +18,7 @@
 from megatron.core.distributed import DistributedDataParallelConfig
 
 from megatron.bridge.diffusion.data.wan.wan_energon_datamodule import WanDatasetConfig
-from megatron.bridge.diffusion.models.wan.wan_provider import WanModelProvider1_3B, WanModelProvider14B
+from megatron.bridge.diffusion.models.wan.wan_provider import WanModelProvider
 from megatron.bridge.training.config import (
     CheckpointConfig,
     ConfigContainer,
@@ -49,7 +49,14 @@ def wan_1_3b_pretrain_config() -> ConfigContainer:
     tensorboard_dir = os.path.join(run_output_dir, "tb_logs")
 
     # Model configuration
-    model_cfg = WanModelProvider1_3B()
+    model_cfg = WanModelProvider(
+        num_layers=30,
+        hidden_size=1536,
+        ffn_hidden_size=8960,
+        num_attention_heads=12,
+        crossattn_emb_size=1536,
+        seq_length=1024,
+    )
     model_cfg.tensor_model_parallel_size = 1
     model_cfg.pipeline_model_parallel_size = 1
     model_cfg.pipeline_dtype = torch.bfloat16
@@ -148,7 +155,14 @@ def wan_14b_pretrain_config() -> ConfigContainer:
     tensorboard_dir = os.path.join(run_output_dir, "tb_logs")
 
     # Model configuration
-    model_cfg = WanModelProvider14B()
+    model_cfg = WanModelProvider(
+        num_layers=40,
+        hidden_size=5120,
+        ffn_hidden_size=13824,
+        num_attention_heads=40,
+        crossattn_emb_size=5120,
+        seq_length=1024,
+    )
     model_cfg.tensor_model_parallel_size = 2
     model_cfg.pipeline_model_parallel_size = 1
     model_cfg.pipeline_dtype = torch.bfloat16

src/megatron/bridge/models/__init__.py

@@ -37,20 +37,9 @@
     FalconH1ModelProvider,
 )
 from megatron.bridge.models.gemma import (
-    CodeGemmaModelProvider2B,
-    CodeGemmaModelProvider7B,
     Gemma2ModelProvider,
-    Gemma2ModelProvider2B,
-    Gemma2ModelProvider9B,
-    Gemma2ModelProvider27B,
     Gemma3ModelProvider,
-    Gemma3ModelProvider1B,
-    Gemma3ModelProvider4B,
-    Gemma3ModelProvider12B,
-    Gemma3ModelProvider27B,
     GemmaModelProvider,
-    GemmaModelProvider2B,
-    GemmaModelProvider7B,
 )
 from megatron.bridge.models.gemma_vl import (
     Gemma3VLBridge,
@@ -99,13 +88,9 @@
     Ministral3Bridge,
     Ministral3Model,
     Ministral3ModelProvider,
-    Ministral3ModelProvider3B,
-    Ministral3ModelProvider8B,
-    Ministral3ModelProvider14B,
 )
 from megatron.bridge.models.mistral import (
     MistralModelProvider,
-    MistralSmall3ModelProvider24B,
 )
 from megatron.bridge.models.nemotron import (
     NemotronBridge,
@@ -115,9 +100,9 @@
     NemotronOmniModel,
 )
 from megatron.bridge.models.nemotron_vl import (
-    NemotronNano12Bv2VLModelProvider,
     NemotronVLBridge,
     NemotronVLModel,
+    NemotronVLModelProvider,
 )
 from megatron.bridge.models.nemotronh import (
     NemotronHBridge,
@@ -185,19 +170,8 @@
     "FalconH1Bridge",
     "FalconH1ModelProvider",
     "Gemma3ModelProvider",
-    "Gemma3ModelProvider1B",
-    "Gemma3ModelProvider4B",
-    "Gemma3ModelProvider12B",
-    "Gemma3ModelProvider27B",
-    "CodeGemmaModelProvider2B",
-    "CodeGemmaModelProvider7B",
     "GemmaModelProvider",
-    "GemmaModelProvider2B",
-    "GemmaModelProvider7B",
     "Gemma2ModelProvider",
-    "Gemma2ModelProvider2B",
-    "Gemma2ModelProvider9B",
-    "Gemma2ModelProvider27B",
     "GLM45Bridge",
     "GLM47FlashBridge",
     "GLM5Bridge",
@@ -214,14 +188,10 @@
     "LlamaNemotronHeterogeneousProvider",
     "LlamaNemotronBridge",
     "MistralModelProvider",
-    "MistralSmall3ModelProvider24B",
     # Ministral 3 Models
     "Ministral3Bridge",
     "Ministral3Model",
     "Ministral3ModelProvider",
-    "Ministral3ModelProvider3B",
-    "Ministral3ModelProvider8B",
-    "Ministral3ModelProvider14B",
     "MiniMaxM2Bridge",
     "OlMoEBridge",
     "OlMoEModelProvider",
@@ -255,7 +225,7 @@
     "Gemma4VLModelProvider",
     "NemotronVLModel",
     "NemotronVLBridge",
-    "NemotronNano12Bv2VLModelProvider",
+    "NemotronVLModelProvider",
     "NemotronOmniBridge",
     "NemotronOmniModel",
     # ASR Models