Skip to content

Commit e05ec7a

Browse files
docs(model): teach builder-backed model configs in the adding-model-support skill (#6402)
Signed-off-by: Kamran Jafari <kjafarisadeg@nvidia.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
1 parent f46e89e commit e05ec7a

6 files changed

Lines changed: 349 additions & 185 deletions

File tree

‎skills/adding-model-support/SKILL.md‎

Lines changed: 62 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
---
22
name: 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.
44
metadata:
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
```
121121
src/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:**
176180
1. **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.
180187
3. **Model class** — Combine vision encoder + language decoder.
181188

182189
For detailed patterns, see:
@@ -188,7 +195,7 @@ For detailed patterns, see:
188195
For 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`
295304
available during conversion — it is set automatically by the dispatch system before
@@ -306,24 +315,34 @@ def mapping_registry(self) -> MegatronMappingRegistry:
306315
Do **not** override `build_conversion_tasks()` to stash `self._hf_config` — that pattern is
307316
deprecated.
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

337356
If you're tempted to add a model-specific `if model_type == "..."` branch in shared code, or
338357
pattern-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:
385404
For detailed recipe patterns, see @skills/adding-model-support/recipe-patterns.md.
386405

387406
Recipes 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
404424
tests/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.
486506
uv run python -c "
487507
from megatron.bridge import AutoBridge
488508
bridge = 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
495513
for 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

Comments
 (0)