11---
22name : adding-model-support
3- description : Guide for adding support for new LLM or VLM models in Megatron-Bridge. Covers bridge, provider , recipe, tests, docs, and examples.
3+ description : Guide for adding support for new LLM or VLM models in Megatron-Bridge. Covers bridge, builder-backed model config , recipe, tests, docs, and examples.
44metadata :
55 when_to_use : User asks to add, onboard, or integrate a new model family; 'add Qwen4 support', 'onboard Llama 5', 'create a bridge for X', 'write a recipe for Y'.
66---
@@ -115,12 +115,12 @@ Also add or update focused tests when touching export/import quantization paths;
115115
116116### File structure
117117
118- ** LLM** — Reference: Qwen2 (` src/megatron/bridge/models/qwen/qwen2_bridge .py ` )
118+ ** LLM** — Reference: Qwen3 (` src/megatron/bridge/models/qwen/qwen3_bridge .py ` )
119119
120120```
121121src/megatron/bridge/models/<model>/
122122├── __init__.py
123- ├── <model>_bridge.py # Config + weight mappings (no provider file needed)
123+ ├── <model>_bridge.py # Config + weight mappings (no provider or config file needed)
124124└── modeling_<model>/ # (optional) Custom nn.Module implementations if needed
125125 └── ...
126126```
@@ -159,24 +159,31 @@ directory — keep them namespaced under the `modeling_<model>` prefix.
159159### Implementation order
160160
161161** LLM:**
162- 1 . ** Bridge only** — Register bridge, implement ` provider_bridge() ` and ` mapping_registry() ` .
163- The bridge calls ` super().provider_bridge() ` to get a ` GPTModelProvider ` from ` CONFIG_MAPPING ` ,
164- then sets model-specific attributes on it. ** Do not create a provider file** — the stock
165- provider returned by ` super().provider_bridge() ` is usually sufficient for LLMs
166- (e.g., ` GPTModelProvider ` , or another base provider selected via ` PROVIDER_CLASS ` ).
167- ** Do not add size-specific provider classes** whose names combine
168- ` ModelProvider ` with a model-size suffix. Examples of forbidden suffixes
162+ 1 . ** Bridge only** — Register bridge, implement ` hf_config_to_model_config_kwargs() ` and
163+ ` mapping_registry() ` . The bridge calls ` super().hf_config_to_model_config_kwargs(hf_config) `
164+ to get the fields derived from ` CONFIG_MAPPING ` , then ` update() ` s the model-specific fields.
165+ ` AutoBridge.get_model_config() ` turns these kwargs into a builder-backed ` ModelConfig `
166+ (` BridgeGPTModelConfig ` by default) that Megatron Core's ` ModelBuilder ` constructs. Set
167+ ` USE_MODEL_CONFIG_FOR_CONVERSION = True ` so checkpoint conversion builds the model through the
168+ same config, and have ` hf_config_to_provider_kwargs() ` return the same kwargs so the deprecated
169+ provider API stays consistent. ** Do not override ` provider_bridge() ` or create a provider or
170+ config file** — ` BridgeGPTModelConfig ` is usually sufficient for LLMs.
171+ ** Do not add size-specific provider or config classes** whose names combine
172+ ` ModelProvider ` or ` ModelConfig ` with a model-size suffix. Examples of forbidden suffixes
169173 include ` 7B ` , ` 200M ` , and ` A3B ` . Model size and architecture fields should
170174 come from the Hugging Face config through ` AutoBridge ` /
171175 ` MegatronModelBridge ` config mapping. If a recipe needs a fixed
172- architecture, configure the base provider inside the recipe function instead
173- of exporting a provider subclass.
176+ architecture, build the config from a Hugging Face config inside the recipe function
177+ ( ` AutoBridge.from_hf_config(hf_config).get_model_config() ` ) instead of exporting a subclass.
174178
175179** VLM:**
1761801 . ** Bridge** — Register bridge, implement config and weight mappings.
177- 2 . ** Provider** (when needed) — Only VLMs that require a custom ` provide() ` to instantiate a
178- combined vision+language model need a provider subclass. The bridge manually calls
179- ` hf_config_to_provider_kwargs(text_config) ` and instantiates the custom provider.
181+ 2 . ** Model config and builder** — A combined vision+language model cannot be represented by
182+ ` BridgeGPTModelConfig ` . Define a ` ModelConfig ` subclass and a ` ModelBuilder ` in the family
183+ directory, select the config with ` MODEL_CONFIG_CLASS ` , and override
184+ ` hf_config_to_model_config() ` . Reference: Muse Glimmer
185+ (` src/megatron/bridge/models/muse_glimmer/ ` ). Older VLMs use a provider subclass with a custom
186+ ` provide() ` instead; see @skills/adding-model-support /vlm-patterns.md.
1801873 . ** Model class** — Combine vision encoder + language decoder.
181188
182189For detailed patterns, see:
@@ -188,7 +195,7 @@ For detailed patterns, see:
188195For VLMs, ` tie_word_embeddings ` lives on the ** top-level** HF config, NOT on ` text_config ` . Always read from the parent config:
189196
190197``` python
191- provider. share_embeddings_and_output_weights = getattr (hf_config, " tie_word_embeddings" , False )
198+ share_embeddings_and_output_weights = getattr (hf_config, " tie_word_embeddings" , False )
192199```
193200
194201### Critical: Config field location for VLMs
@@ -285,11 +292,13 @@ class _FullDimQKNormMapping(MegatronParamMapping):
285292| Hook | When to use |
286293| ------| -------------|
287294| ` mapping_registry() ` | Define all weight name mappings (abstract, always overridden) |
288- | ` provider_bridge() ` | Configure the provider with model-specific flags (call ` super() ` then setattr) |
295+ | ` hf_config_to_model_config_kwargs() ` | Add model-specific config fields (call ` super() ` , then ` update() ` ) |
296+ | ` hf_config_to_provider_kwargs() ` | Return ` self.hf_config_to_model_config_kwargs(hf_config) ` so the deprecated provider path matches |
297+ | ` hf_config_to_model_config() ` | Build a custom ` MODEL_CONFIG_CLASS ` that flat kwargs cannot express (Muse Glimmer) |
289298| ` maybe_modify_loaded_hf_weight() ` | Dequantize, rename, or reshape HF weights before conversion |
290299| ` maybe_modify_converted_hf_weight() ` | Synthesize extra HF keys on export (e.g. ` inv_freq ` ) |
291- | ` megatron_to_hf_config() ` | Build HF ` config.json ` for export |
292- | ` hf_config_to_provider_kwargs ()` | Override CONFIG_MAPPING behavior for specific fields |
300+ | ` megatron_to_hf_config() ` | Build HF ` config.json ` for export; receives the model config |
301+ | ` provider_bridge ()` | Legacy provider path; do not override in new or migrated bridges |
293302
294303** Accessing HF config in ` mapping_registry() ` :** The bridge instance has ` self.hf_config `
295304available during conversion — it is set automatically by the dispatch system before
@@ -306,24 +315,34 @@ def mapping_registry(self) -> MegatronMappingRegistry:
306315Do ** not** override ` build_conversion_tasks() ` to stash ` self._hf_config ` — that pattern is
307316deprecated.
308317
309- #### Strategy 3: Custom provider subclass (VLMs only)
318+ #### Strategy 3: Custom model config and builder
310319
311- Most models do ** not** need a provider file — the stock provider (e.g., ` GPTModelProvider ` , or
312- another base selected via ` PROVIDER_CLASS ` ) is usually sufficient for LLMs. Only create a provider subclass when a VLM needs custom ` provide() ` logic to instantiate
313- a combined vision+language model:
320+ Most models do ** not** need a config file — ` BridgeGPTModelConfig ` with its nested Bridge
321+ ` TransformerConfig ` is usually sufficient for LLMs. Create a ` ModelConfig ` subclass and a
322+ ` ModelBuilder ` only when the model needs fields or construction logic that the GPT config and
323+ builder cannot express, such as a combined vision+language model or extra decoder fields:
314324
315325``` python
316- # src/megatron/bridge/models/<model>/<model>_provider.py
317- class MyVLModelProvider (GPTModelProvider ):
318- image_token_id: int = 0
326+ # src/megatron/bridge/models/<model>/<model>_config.py
327+ @dataclass
328+ class MyTransformerConfig (TransformerConfig ): # Bridge TransformerConfig
329+ my_decoder_field: float = 1.0
319330
320- def provide (self , ...):
321- # Custom model construction combining vision encoder + language decoder
322- ...
331+
332+ @dataclass (kw_only = True )
333+ class MyModelConfig (ModelConfigOverrideMixin , HybridModelConfig ):
334+ builder: ClassVar[str ] = " megatron.bridge.models.<model>.MyModelBuilder"
335+ transformer_config_class: ClassVar[type[TransformerConfig]] = MyTransformerConfig
336+ image_token_id: int = 0
323337```
324338
325- The bridge then references it via ` PROVIDER_CLASS = MyVLModelProvider ` or instantiates it directly
326- in ` provider_bridge() ` .
339+ The builder subclasses Megatron Core's ` GPTModelBuilder ` or ` HybridModelBuilder ` and overrides
340+ ` build_model() ` . The bridge sets ` MODEL_CONFIG_CLASS = MyModelConfig ` and
341+ ` USE_MODEL_CONFIG_FOR_CONVERSION = True ` , and overrides ` hf_config_to_model_config() ` and
342+ ` megatron_to_hf_config() ` . Follow Muse Glimmer (` src/megatron/bridge/models/muse_glimmer/ ` ).
343+
344+ Provider subclasses with a custom ` provide() ` (` PROVIDER_CLASS = MyVLModelProvider ` ) are the
345+ legacy equivalent; extend them only in families that have not migrated yet.
327346
328347#### When shared file changes ARE justified
329348
@@ -332,7 +351,7 @@ families**. Examples of justified shared changes:
332351
333352- ` FusedExpertMapping ` / ` FusedGatedExpertMapping ` — used by GLM, DeepSeek, OLMoE, etc.
334353- ` RMSNorm2ZeroCenteredRMSNormMapping ` — used by Gemma, Nemotron, etc.
335- - New ` CONFIG_MAPPING ` entries — when a standard HF config key maps to a standard provider attribute
354+ - New ` CONFIG_MAPPING ` entries — when a standard HF config key maps to a standard Megatron config field
336355
337356If you're tempted to add a model-specific ` if model_type == "..." ` branch in shared code, or
338357pattern-matching on specific weight names in shared conversion logic, that's a signal to use a
@@ -385,10 +404,11 @@ Each recipe file defines functions for each model size + training mode:
385404For detailed recipe patterns, see @skills/adding-model-support /recipe-patterns.md.
386405
387406Recipes are the right API surface for model-size presets. Do not create or
388- export size-specific provider subclasses for recipes; either call
389- ` AutoBridge.from_hf_pretrained(...).to_megatron_provider(load_weights=False) ` to
390- derive the provider from HF config, or instantiate the base provider class with
391- explicit architecture fields inside the recipe function.
407+ export size-specific provider or config subclasses for recipes; call
408+ ` AutoBridge.from_hf_pretrained(...).get_model_config() ` to derive the model config
409+ from the HF checkpoint, or ` AutoBridge.from_hf_config(hf_config).get_model_config() `
410+ with an explicit HF config inside the recipe function. Families that still use the
411+ provider path call ` .to_megatron_provider(load_weights=False) ` instead.
392412
393413### Export checklist
394414
@@ -403,8 +423,8 @@ explicit architecture fields inside the recipe function.
403423``` text
404424tests/unit_tests/models/<model>/
405425├── __init__.py
406- ├── test_<model>_bridge.py # Mock HF config → verify provider mapping
407- └── test_<model>_provider .py # (optional) Only if custom provider subclass exists
426+ ├── test_<model>_bridge.py # HF config → verify model config mapping
427+ └── test_<model>_config .py # (optional) Only if a custom config or builder exists
408428```
409429
410430### Functional tests (GPU)
@@ -486,12 +506,10 @@ execution back to the user when cluster access and authorization are available.
486506uv run python -c "
487507from megatron.bridge import AutoBridge
488508bridge = AutoBridge.from_hf_pretrained('<org>/<model>')
489- provider = bridge.to_megatron_provider()
490- provider.tensor_model_parallel_size = 1
491- provider.pipeline_model_parallel_size = 1
492- provider.finalize()
493- model = provider.provide_distributed_model(wrap_with_ddp=False)
494- bridge.load_hf_weights(model)
509+ model_config = bridge.get_model_config()
510+ model_config.tensor_model_parallel_size = 1
511+ model_config.pipeline_model_parallel_size = 1
512+ model = bridge.get_model(model_config, wrap_with_ddp=False) # loads the HF weights
495513for i, (name, tensor) in enumerate(bridge.export_hf_weights(model, cpu=True)):
496514 print(name, tuple(tensor.shape))
497515 if i > 10: break
0 commit comments