Adding New Model Support in Megatron-Bridge#

Implementation boundaries#

  • Keep model-support changes in Megatron-Bridge. Do not modify Megatron-Core or its submodule to implement a model. Use existing module specs, builders, conversion hooks, and narrow Bridge subclasses or compositions.

  • Before adding a block, inspect both MCore and existing Bridge implementations. Reuse an existing block when it represents the same computation. Differences in accumulation precision, rounding, fused kernels, or mathematically equivalent LayerNorm/RMSNorm operation ordering alone do not justify a replacement implementation.

  • Add model-specific code under the family’s modeling_<model>/ directory only for genuinely missing computation or composition. For example, normalizing each expert before weighted summation differs mathematically from normalizing the combined result. Compose existing normalization and linear operations wherever possible instead of copying their kernels.

  • Invoke Hugging Face vision/video and audio encoders directly; do not port their internals into Bridge. Bridge owns their construction, checkpoint mapping, modality insertion, placement, freezing, and gradient synchronization. Existing native encoder ports are references for previously supported models, not a requirement for new support. Follow an explicit request for a different encoder implementation when the user provides one.

  • Verify support through the requested model API, not only a shared config. For HybridModel, trace the actual layer specs and expand per-layer schedules to Hybrid positions: *E repeated N times has 2N positions. Configure RoPE explicitly, since HybridModel defaults to no positional embedding.

  • Separate weight-independent checks from checkpoint validation. A checkpoint download need not block implementation or small config/block tests; record full conversion, logits, and training validation as pending until they run.

Phase 1: Discovery#

Step 2 — Fetch and analyze config.json#

Read the model’s config.json from HuggingFace (or from the user-provided file). Key fields to extract:

  • model_type — used for @register_bridge(model_type=...)

  • architectures — the HF model class name (used for source=... in registration)

  • tie_word_embeddings — critical for weight tying

  • Architecture fields: num_hidden_layers, hidden_size, intermediate_size, num_attention_heads, num_key_value_heads, vocab_size, max_position_embeddings, rope_theta, etc.

  • MoE fields (if present): num_local_experts, num_experts_per_tok, moe_intermediate_size

  • MLA fields (if present): q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim

For unfamiliar config fields, inspect the supplied Hugging Face modeling_*.py implementation and check existing MCore and Bridge components before concluding that a new block is needed. Ask for missing source only when it is not already available through the user’s supplied artifacts or authorized locations.

Step 3 — Identify supported modalities#

Inspect the config, HF model, and processor for vision/video and audio inputs. Perception configs commonly appear as vision_config or audio_config alongside text_config, but names and nesting vary. Route models with either type of perception encoder through the multimodal wrapper path, including audio-only language models. A text-only model has no perception integration; its language config may be flat or nested.

This distinction determines where to read language config fields, whether a wrapper must combine perception with the language model, and which modality inputs the functional checks must exercise.

Step 4 — Check for quantized weights (FP8 / FP4)#

Inspect the HF checkpoint’s model.safetensors (or model.safetensors.index.json) for quantized weight dtypes such as float8_e4m3fn (FP8) or uint8/uint4 with accompanying *_scale_inv or *_scale tensors. Common signs:

  • config.json mentions quantization_config or dtype fields like "torch_dtype": "float8_e4m3fn"

  • Safetensors contain weight_scale_inv keys alongside the main weight keys

  • The model card mentions FP8/FP4/INT4 weights

Why this matters: The bridge’s import_ckpt path does not automatically dequantize — it loads raw quantized values as-is. This produces a silently broken model (random-level loss, huge grad norms) instead of raising an error.

Fix: Dequantize before or during conversion. The current in-repo pattern is to use a bridge hook plus the shared helpers in src/megatron/bridge/models/conversion/quantization_utils.py. Existing examples include src/megatron/bridge/models/ministral3/ministral3_bridge.py, src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py, and src/megatron/bridge/models/minimax_m2/minimax_m2_bridge.py.

Override maybe_modify_loaded_hf_weight() in the bridge class to dequantize on the fly during import:

def maybe_modify_loaded_hf_weight(self, hf_param, hf_state_dict):
    weight = hf_state_dict[hf_param]
    scale_key = hf_param + "_scale_inv"
    if weight.dtype == torch.float8_e4m3fn and scale_key in hf_state_dict:
        return weight.to(torch.bfloat16) * hf_state_dict[scale_key].to(torch.bfloat16)
    return weight

Always add a sanity check in the verification workflow (e.g., print std of a weight tensor — quantized models typically have std ≈ 13 before dequantization vs std ≈ 0.006 after). Also add or update focused tests when touching export/import quantization paths; see tests/unit_tests/models/test_fp8_param_export.py for current FP8 export coverage.

Phase 2: Bridge Support#

File structure#

LLM — Reference: Qwen3 (src/megatron/bridge/models/qwen/qwen3_bridge.py)

src/megatron/bridge/models/<model>/
├── __init__.py
├── <model>_bridge.py      # Config + weight mappings (no provider or config file needed)
└── modeling_<model>/      # (optional) Custom nn.Module implementations if needed
    └── ...

VLM / audio-language model — Use HF perception encoders and a Bridge wrapper. Reference: Gemma3-VL (src/megatron/bridge/models/gemma_vl/).

src/megatron/bridge/models/<model>/
├── __init__.py
├── <model>_bridge.py         # Config + weight mappings
├── <model>_provider.py       # Only for VLMs that need custom provide()
└── modeling_<model>/         # Missing language blocks and modality integration
    ├── __init__.py
    └── model.py              # Combines vision + language

For a small wrapper, an existing single-file layout can also be followed:

src/megatron/bridge/models/<model>/
├── __init__.py
├── <model>_bridge.py
├── <model>_provider.py       # Only for VLMs that need custom provide()
└── modeling_<model>.py       # HF vision + Megatron language wrapper

Model-specific modeling code: If the model requires custom nn.Module implementations (e.g. a custom RoPE variant, non-standard transformer config, custom thinker/talker architecture), place them in a modeling_<model>/ directory or a single modeling_<model>.py file inside the model family folder. Use a directory when there are multiple files (model, transformer config, custom ops); use a single file when one module suffices. Never put model-specific modeling code in shared directories or as loose files in the bridge family directory — keep them namespaced under the modeling_<model> prefix.

Implementation order#

LLM:

  1. Bridge only — Register bridge, implement hf_config_to_model_config_kwargs() and mapping_registry(). The bridge calls super().hf_config_to_model_config_kwargs(hf_config) to get the fields derived from CONFIG_MAPPING, then update()s the model-specific fields. AutoBridge.get_model_config() turns these kwargs into a builder-backed ModelConfig (BridgeGPTModelConfig by default) that Megatron Core’s ModelBuilder constructs. Set USE_MODEL_CONFIG_FOR_CONVERSION = True so checkpoint conversion builds the model through the same config, and have hf_config_to_provider_kwargs() return the same kwargs so the deprecated provider API stays consistent. Do not override provider_bridge() or create a provider or config file — BridgeGPTModelConfig is usually sufficient for LLMs. Do not add size-specific provider or config classes whose names combine ModelProvider or ModelConfig 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, build the config from a Hugging Face config inside the recipe function (AutoBridge.from_hf_config(hf_config).get_model_config()) instead of exporting a subclass.

VLM:

  1. Bridge — Register bridge, implement config and weight mappings.

  2. Model config and builder — A combined vision+language model cannot be represented by BridgeGPTModelConfig. Define a ModelConfig subclass and a ModelBuilder in the family directory, select the config with MODEL_CONFIG_CLASS, and override hf_config_to_model_config(). Reference: Muse Glimmer (src/megatron/bridge/models/muse_glimmer/). Older VLMs use a provider subclass with a custom provide() instead; see @skills/adding-model-support/vlm-patterns.md.

  3. Model class — Combine vision encoder + language decoder.

For detailed patterns, see:

  • VLM: @skills/adding-model-support/vlm-patterns.md

  • LLM: @skills/adding-model-support/llm-patterns.md

Critical: tie_word_embeddings for VLMs#

For VLMs, tie_word_embeddings lives on the top-level HF config, NOT on text_config. Always read from the parent config:

share_embeddings_and_output_weights = getattr(hf_config, "tie_word_embeddings", False)

Critical: Config field location for VLMs#

When reading HF config for VLMs, check whether each field is in:

  • hf_config (top-level) — e.g. tie_word_embeddings, image_token_id, video_token_id

  • hf_config.text_config — e.g. num_hidden_layers, hidden_size, etc.

  • hf_config.vision_config — e.g. vision encoder dimensions

Encapsulating model-specific layers#

When a new model introduces custom or non-standard layers (novel attention variants, custom normalization, fused expert layouts, MTP heads, etc.), keep all model-specific logic inside the model family directory. Do not modify shared files in src/megatron/bridge/models/conversion/ (e.g. param_mapping.py, model_bridge.py, quant_mapping.py) unless the change is genuinely reusable across multiple model families.

Principle: The bridge and provider files for a model family are your primary extension surface. Shared conversion infrastructure provides hooks and base classes — subclass them locally rather than adding conditionals to shared code.

Prefer Transformer Engine normalization#

Before implementing a model-specific normalization with raw PyTorch operations, inspect the exact Transformer Engine revision in pyproject.toml/uv.lock, its upstream module API, and Megatron-Core’s TENorm adapter. Production Bridge model code should assume Transformer Engine is installed. Use TENorm or the corresponding TE module directly for LayerNorm and RMSNorm; do not add model-local norm wrappers, copy an existing norm implementation, gate the norm on HAVE_TE, or provide a torch.nn fallback. Preserve eps, parameter dtype/device initialization, sequence-parallel metadata, and zero_centered_gamma; standard-gamma and zero-centered-gamma norms are not interchangeable. The attention/linear backend and normalization choice are independent, so a local attention spec should still use TE normalization.

Treat a custom norm implementation as an architectural exception, not a compatibility pattern. It is allowed only when the pinned TE API cannot represent the model’s exact mathematical and parameter contract, such as a genuinely parameter-free RMSNorm while TE requires an affine weight. Prefer a framework-native module such as torch.nn.RMSNorm(elementwise_affine=False) or a direct functional call over a model-local wrapper or copied formula. Document the missing TE capability and test that the exception introduces no checkpoint parameter. Do not create a wrapper or fallback for ordinary affine normalization.

Do not introduce a dummy affine parameter merely to reach a fused kernel. For example, current TE RMSNorm always owns a learnable weight, so a genuinely weightless/scaleless RMSNorm must retain a parameter-free implementation until TE exposes matching semantics. Adding an all-ones frozen weight would change state-dict keys, optimizer state, distributed-checkpoint schema, and conversion coverage.

TE and reference framework kernels can differ numerically even when their architecture and weights match. Require exact HF↔Megatron weight round-trip for unchanged parameter representations and pure layout transformations. Investigate arithmetic parameter transformations and classify discrepancies using the controlled module-swap procedure below. Assess forward behavior with the correlation criteria in the parity-testing skill; do not replace reusable norms solely to reproduce one framework’s operation ordering.

Strategy 1: Create a local mapping subclass#

If the model has a layer whose weight layout doesn’t match any existing mapping class, create a private mapping class in the bridge file or a <model>_mappings.py file in the family directory.

Example — GLM’s fused expert down-projection disables grouped-export transpose:

# src/megatron/bridge/models/glm/glm_moe_mappings.py
class GLMExpertDownProjMapping(FusedExpertMapping):
    def __init__(self, megatron_param, hf_param, permute_dims=None):
        super().__init__(megatron_param, hf_param, permute_dims, transpose_on_export=False)

Example — Nemotron-H’s MTP layers flatten indices during resolve:

# Inside nemotron_h_bridge.py (private to the module)
class _MTPFlatteningMapping(MegatronParamMapping):
    def resolve(self, captures):
        return AutoMapping(self._flatten(captures), ...)

Example — MiniMax-M2’s non-standard QK norm layout:

# Inside minimax_m2_bridge.py (private to the module)
class _FullDimQKNormMapping(MegatronParamMapping):
    def hf_to_megatron(self, hf_weights):
        # Custom scatter logic for full-dim QK norm
        ...
    def megatron_to_hf(self, megatron_weights):
        # Custom gather logic
        ...

Strategy 2: Override bridge hooks#

MegatronModelBridge provides several override hooks — use them instead of modifying the base class:

Hook

When to use

mapping_registry()

Define all weight name mappings (abstract, always overridden)

hf_config_to_model_config_kwargs()

Add model-specific config fields (call super(), then update())

hf_config_to_provider_kwargs()

Return self.hf_config_to_model_config_kwargs(hf_config) so the deprecated provider path matches

hf_config_to_model_config()

Build a custom MODEL_CONFIG_CLASS that flat kwargs cannot express (Muse Glimmer)

maybe_modify_loaded_hf_weight()

Dequantize, rename, or reshape HF weights before conversion

maybe_modify_converted_hf_weight()

Synthesize extra HF keys on export (e.g. inv_freq)

megatron_to_hf_config()

Build HF config.json for export; receives the model config

provider_bridge()

Legacy provider path; do not override in new or migrated bridges

Accessing HF config in mapping_registry(): The bridge instance has self.hf_config available during conversion — it is set automatically by the dispatch system before mapping_registry() is called. Use it when your mapping registry needs config-dependent logic (e.g. dynamic MTP layer count, number of experts):

def mapping_registry(self) -> MegatronMappingRegistry:
    hf_config = getattr(self, "hf_config", None)
    num_mtp_layers = getattr(hf_config, "num_nextn_predict_layers", 0) if hf_config else 0
    ...

Do not override build_conversion_tasks() to stash self._hf_config — that pattern is deprecated.

Strategy 3: Custom model config and builder#

Most models do not need a config file — BridgeGPTModelConfig with its nested Bridge TransformerConfig is usually sufficient for LLMs. Create a ModelConfig subclass and a ModelBuilder only when the model needs fields or construction logic that the GPT config and builder cannot express, such as a combined vision+language model or extra decoder fields:

# src/megatron/bridge/models/<model>/<model>_config.py
@dataclass
class MyTransformerConfig(TransformerConfig):  # Bridge TransformerConfig
    my_decoder_field: float = 1.0


@dataclass(kw_only=True)
class MyModelConfig(ModelConfigOverrideMixin, HybridModelConfig):
    builder: ClassVar[str] = "megatron.bridge.models.<model>.MyModelBuilder"
    transformer_config_class: ClassVar[type[TransformerConfig]] = MyTransformerConfig
    image_token_id: int = 0

The builder subclasses Megatron Core’s GPTModelBuilder or HybridModelBuilder and overrides build_model(). The bridge sets MODEL_CONFIG_CLASS = MyModelConfig and USE_MODEL_CONFIG_FOR_CONVERSION = True, and overrides hf_config_to_model_config() and megatron_to_hf_config(). Follow Muse Glimmer (src/megatron/bridge/models/muse_glimmer/).

Provider subclasses with a custom provide() (PROVIDER_CLASS = MyVLModelProvider) are the legacy equivalent; extend them only in families that have not migrated yet.

When shared file changes ARE justified#

Modify param_mapping.py or model_bridge.py only when the pattern is reusable by 2+ model families. Examples of justified shared changes:

  • FusedExpertMapping / FusedGatedExpertMapping — used by GLM, DeepSeek, OLMoE, etc.

  • RMSNorm2ZeroCenteredRMSNormMapping — used by Gemma, Nemotron, etc.

  • New CONFIG_MAPPING entries — when a standard HF config key maps to a standard Megatron config field

If you’re tempted to add a model-specific if model_type == "..." branch in shared code, or pattern-matching on specific weight names in shared conversion logic, that’s a signal to use a local subclass or hook override instead.

Update FLOPs calculator for new architectural blocks#

If the model introduces a new computational block that differs from standard attention or MLP (e.g., Gated DeltaNet / GDN linear attention, Multi-Token Prediction / MTP heads, Mamba SSM layers), update the FLOPs calculator in src/megatron/bridge/training/utils/flop_utils.py so that training throughput metrics (TFLOPs/GPU) are accurate.

When to update: Any time the new block has different FLOPs-per-token than standard self-attention or standard MLP. Common cases:

  • Linear attention variants (GDN, RetNet, RWKV) — replace the O(s²) attention term with the block’s actual operation count

  • MTP / speculative decoding heads — add FLOPs for the extra projection and norm layers

  • SSM layers (Mamba) — different recurrence FLOPs than attention

  • Novel MoE routing — may change the effective expert count

How to update:

  1. Read the existing transformer_flops() function in flop_utils.py to understand the structure.

  2. Add a conditional block gated on a config attribute (e.g., experimental_attention_variant, mtp_num_layers). Follow the existing MoE pattern for config validation — raise on invalid types, assert list lengths, and use direct attribute access instead of getattr with fallback defaults so that misconfigurations fail explicitly.

  3. Compute the per-layer FLOPs for the new block and blend it with the standard attention term based on the layer pattern.

  4. Add unit tests in tests/unit_tests/training/utils/test_flop_utils.py that verify:

    • New-block FLOPs differ from pure-attention baseline

    • Exact formula matches hand-computed expected values

    • Varying the block ratio (e.g., linear_attention_freq) changes FLOPs

Reference PR: #2925 — GDN FLOPs calculator adds GDN support with both the calculator code and comprehensive tests.

Phase 3: Recipe Support#

Recipes provide pre-configured training settings for each model size.

LLM recipes: src/megatron/bridge/recipes/<family>/<model>.py VLM recipes: src/megatron/bridge/recipes/<family>/<model>.py

Each recipe file defines functions for each model size + training mode:

  • <model>_<size>_sft_config() — Full supervised fine-tuning

  • <model>_<size>_peft_config() — LoRA/DoRA parameter-efficient fine-tuning

  • <model>_<size>_pretrain_config() — Pretraining (LLM only, usually)

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 or config subclasses for recipes; call AutoBridge.from_hf_pretrained(...).get_model_config() to derive the model config from the HF checkpoint, or AutoBridge.from_hf_config(hf_config).get_model_config() with an explicit HF config inside the recipe function. Families that still use the provider path call .to_megatron_provider(load_weights=False) instead.

Export checklist#

  1. Family __init__.py — import and add to __all__

  2. Top-level src/megatron/bridge/recipes/__init__.py — wildcard import

  3. train_any_basic.py — add to config_map, docstring, and --model choices

Phase 4: Tests#

Unit tests (no GPU)#

tests/unit_tests/models/<model>/
├── __init__.py
├── test_<model>_bridge.py    # HF config → verify model config mapping
└── test_<model>_config.py    # (optional) Only if a custom config or builder exists

Functional tests (GPU)#

tests/functional_tests/test_groups/models/<model>/
├── __init__.py
├── test_<model>_conversion.py  # Toy model HF↔Megatron roundtrip
└── test_<model>_provider.py    # compare_provider_configs (optional)

For detailed test patterns, see @skills/adding-model-support/tests-and-examples.md.

Phase 5: Docs and Examples#

Examples#

Model examples: examples/models/<family>/<model>/

examples/models/<family>/<model>/
├── README.md
├── inference.sh         # Generation commands (real model, reasonable output)
├── slurm_sft.sh         # SFT training on SLURM
└── slurm_peft.sh        # PEFT training on SLURM

Key deliverable requirement: The README must include working import and export commands for a real published model (e.g. Qwen/Qwen3-8B, not a toy) using scripts/conversion/convert.sh. Add a model-specific conversion wrapper only when the model requires preparation or verification that the shared CLI cannot express. The inference script must produce reasonable output — for LLMs a coherent text continuation, for VLMs a plausible image description. This is the acceptance bar: conversion runs cleanly and generation makes sense.

Documentation#

Add a model page at docs/models/<type>/<model>.md covering:

  • Supported variants and sizes

  • Conversion commands

  • Training examples (SFT, PEFT)

  • Known limitations

Verification Workflow#

Run the authorized checks in the cluster development container. Do not hand execution back to the user when cluster access and authorization are available.

Discrepancy investigation#

  • Check import, persisted checkpoint reload, export, and roundtrip tensor coverage before interpreting forward differences. Report missing/extra keys, shapes, dtypes, and exact equality separately from forward metrics.

  • Reload in a fresh process as well as within a conversion process. Pass the intended model-parallel overrides explicitly and verify effective TP/PP/EP process-group sizes; existing groups can conceal loader defaults.

  • Exercise every requested modality path using identical processed inputs on both implementations: text-only, image-text, and an appropriate audio-containing case for models that support audio. Preserve input hashes, token IDs, logits, greedy continuations, code revisions, and runtime settings.

  • Isolate each discrepancy to the earliest differing module with the same inputs and weights. Check mapping, layout, masks, positions, routing, and precision rather than assuming a small difference is numerical.

  • Classify a discrepancy as numerical only after a controlled module-swap experiment: replace the suspected module in one execution path with the other implementation, keeping inputs, weights, and other settings fixed. Record baseline and swapped outputs and show that the discrepancy follows the implementation. Use a reverse swap or higher-precision comparison when needed to distinguish incorrect mathematics from floating-point effects. A swap that closes a gap alone does not prove the original computation was correct.

  • Fix confirmed implementation bugs in Bridge and rerun the affected checks. Do not fix confirmed numerical differences by replacing reusable production blocks. Keep diagnostic swaps in validation code, report their measured effects, and leave inconclusive discrepancies unresolved.

1. Smoke test (single GPU)#

uv run python -c "
from megatron.bridge import AutoBridge
bridge = AutoBridge.from_hf_pretrained('<org>/<model>')
model_config = bridge.get_model_config()
model_config.tensor_model_parallel_size = 1
model_config.pipeline_model_parallel_size = 1
model = bridge.get_model(model_config, wrap_with_ddp=False)  # loads the HF weights
for i, (name, tensor) in enumerate(bridge.export_hf_weights(model, cpu=True)):
    print(name, tuple(tensor.shape))
    if i > 10: break
"

2. Conversion roundtrip (multi-GPU)#

./scripts/conversion/convert.sh import \
    --hf-model <org>/<model> \
    --megatron-path /workspace/<model> \
    --torch-dtype bfloat16

./scripts/conversion/convert.sh export \
    --hf-model <org>/<model> \
    --megatron-path /workspace/<model>/iter_0000000 \
    --hf-path /workspace/<model>-hf-export

3. Generation test#

For LLMs:

uv run python examples/conversion/hf_to_megatron_generate_text.py \
    --hf_model_path <org>/<model> --prompt "Hello"

For VLMs:

uv run python examples/conversion/hf_to_megatron_generate_vlm.py \
    --hf_model_path <org>/<model> \
    --image_path "https://example.com/image.jpeg" \
    --prompt "Describe this image."

4. Run tests#

uv run python -m pytest tests/unit_tests/models/<model>/ -v
uv run python -m pytest tests/functional_tests/test_groups/models/<model>/ -v --run-gpu

Quick Decision Tree#

User wants to add a model
│
├─ Has HF link? ─── No ──→ Ask for link (or config.json if private)
│
├─ Has vision/video or audio perception? ─── Yes ──→ Multimodal wrapper path
│   └─ Invoke supplied HF perception encoders directly (Gemma3 wrapper pattern)
│
└─ Text-only ──→ Language model path
    ├─ Standard GPT-style? ──→ Bridge with stock mappings
    └─ Custom layers? ──→ Bridge + local mapping subclasses / hook overrides
        ├─ Custom weight layout? ──→ Local mapping subclass in family dir
        └─ Custom import/export? ──→ Override bridge hooks (maybe_modify_*)