adding-model-support

# 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 1 — Get the HF model link Ask the user for the HuggingFace model link (e.g. `https://huggingface.co/Qwen/Qwen3.5-VL-27B`). If the model is **not public**, ask the user to provide the `config.json` file directly. ### 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: ```python 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: ```python 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: ```python # 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: ```python # 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: ```python # 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): ```python 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: ```python # 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