Ajouter le support d'un nouveau modèle dans Megatron-Bridge

Phase 1 : Découverte

Étape 1 — Obtenir le lien du modèle HF

Demandez à l'utilisateur le lien du modèle HuggingFace (par ex. https://huggingface.co/Qwen/Qwen3.5-VL-27B).

Si le modèle n'est pas public, demandez à l'utilisateur de fournir directement le fichier config.json.

Étape 2 — Récupérer et analyser config.json

Lisez le fichier config.json du modèle depuis HuggingFace (ou depuis le fichier fourni par l'utilisateur). Champs clés à extraire :

model_type — utilisé pour @register_bridge(model_type=...)
architectures — le nom de la classe HF du modèle (utilisé pour source=... dans l'enregistrement)
tie_word_embeddings — critique pour le partage de poids
Champs d'architecture : num_hidden_layers, hidden_size, intermediate_size, num_attention_heads, num_key_value_heads, vocab_size, max_position_embeddings, rope_theta, etc.
Champs MoE (si présents) : num_local_experts, num_experts_per_tok, moe_intermediate_size
Champs MLA (si présents) : q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim

S'il existe des champs de config que vous ne reconnaissez pas par rapport aux modèles précédemment pris en charge (vérifiez CONFIG_MAPPING dans model_bridge.py et les bridges existants), cela indique probablement un nouveau bloc architectural (par ex., une variante d'attention novatrice, une normalisation personnalisée, ou un nouveau type de couche). Demandez à l'utilisateur de fournir l'implémentation HuggingFace modeling_*.py de ce bloc pour que vous puissiez comprendre le calcul et créer le bon mappage côté Megatron ou le module personnalisé.

Étape 3 — Déterminer VLM vs LLM

VLM (Vision-Language Model) si config.json contient :

Des sous-configs text_config ET vision_config
Remarque : les VLM peuvent ou non avoir « VL » dans le nom

LLM (texte uniquement) si :

Pas de text_config / vision_config
Config plate unique pour le modèle de langage

Cette distinction affecte :

Les fichiers à créer (les VLM ont besoin d'un model.py combinant vision + langage)
Où lire les champs de config (text_config vs niveau supérieur pour les VLM)
Les motifs de test (les VLM ont besoin d'entrées de vision dans les tests fonctionnels)

Étape 4 — Vérifier les poids quantifiés (FP8 / FP4)

Inspectez model.safetensors du checkpoint HF (ou model.safetensors.index.json) pour les dtypes de poids quantifiés comme float8_e4m3fn (FP8) ou uint8/uint4 avec les tenseurs *_scale_inv ou *_scale associés. Signes courants :

config.json mentionne quantization_config ou des champs dtype comme "torch_dtype": "float8_e4m3fn"
Safetensors contiennent des clés weight_scale_inv aux côtés des clés de poids principales
La carte du modèle mentionne les poids FP8/FP4/INT4

Pourquoi cela importe : Le chemin import_ckpt du bridge ne dézantifte pas automatiquement — il charge les valeurs quantifiées brutes telles quelles. Cela produit un modèle silencieusement cassé (perte au niveau aléatoire, normes de gradient énormes) au lieu de lever une erreur.

Correctif : Dézantifiez avant la conversion. Deux approches :

Script autonome (recommandé pour les modèles orientés utilisateur) — Écrivez un dequant_fp8_for_bridge.py dans le dossier examples du modèle. Référence : examples/models/vlm/ministral3/dequant_fp8_for_bridge.py. Le motif est : w_bf16 = fp8_weight.to(bfloat16) * weight_scale_inv.

Hook in-bridge — Remplacez maybe_modify_loaded_hf_weight() dans la classe bridge pour dézantifier à la volée lors de l'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

Ajoutez toujours une vérification de cohérence dans le flux de vérification (par ex., imprimez std d'un tenseur de poids — les modèles quantifiés ont généralement std ≈ 13 avant dézantification vs std ≈ 0,006 après).

Phase 2 : Support du Bridge

Structure des fichiers

LLM — Référence : Qwen2 (src/megatron/bridge/models/qwen/qwen2_bridge.py)

src/megatron/bridge/models/<model>/
├── __init__.py
├── <model>_bridge.py      # Config + mappages de poids (fichier provider non nécessaire)
└── modeling_<model>/      # (optionnel) Implémentations personnalisées nn.Module si nécessaire
    └── ...

VLM — Référence : Qwen3.5-VL (src/megatron/bridge/models/qwen_vl/)

src/megatron/bridge/models/<model>/
├── __init__.py
├── <model>_bridge.py         # Config + mappages de poids
├── <model>_provider.py       # Uniquement pour les VLM qui ont besoin d'un provide() personnalisé
└── modeling_<model>/         # Si utilisation d'un encodeur de vision Megatron
    ├── __init__.py
    └── model.py              # Combine vision + langage

OU avec encodeur de vision HF (Référence : Gemma3-VL) :

src/megatron/bridge/models/<model>/
├── __init__.py
├── <model>_bridge.py
├── <model>_provider.py       # Uniquement pour les VLM qui ont besoin d'un provide() personnalisé
└── modeling_<model>.py       # Wrapper vision HF + langage Megatron

Code modeling spécifique au modèle : Si le modèle nécessite des implémentations nn.Module personnalisées (par ex. une variante RoPE personnalisée, config transformer non standard, architecture thinker/talker personnalisée), placez-les dans un répertoire modeling_<model>/ ou un fichier modeling_<model>.py unique à l'intérieur du dossier de la famille de modèles. Utilisez un répertoire s'il y a plusieurs fichiers (modèle, config transformer, ops personnalisées) ; utilisez un fichier unique si un seul module suffit. Ne mettez jamais de code modeling spécifique au modèle dans des répertoires partagés ou comme fichiers libres dans le répertoire de la famille bridge — gardez-les avec un espace de noms sous le préfixe modeling_<model>.

Ordre d'implémentation

LLM :

Bridge uniquement — Enregistrez le bridge, implémentez provider_bridge() et mapping_registry(). Le bridge appelle super().provider_bridge() pour obtenir un GPTModelProvider de CONFIG_MAPPING, puis définit des attributs spécifiques au modèle. Ne créez pas de fichier provider — le provider stock retourné par super().provider_bridge() est généralement suffisant pour les LLM (par ex., GPTModelProvider, ou une autre base sélectionnée via PROVIDER_CLASS).

VLM :

Bridge — Enregistrez le bridge, implémentez config et mappages de poids.
Provider (si nécessaire) — Seuls les VLM qui nécessitent un provide() personnalisé pour instancier un modèle vision+langage combiné ont besoin d'une sous-classe provider. Le bridge appelle manuellement hf_config_to_provider_kwargs(text_config) et instancie le provider personnalisé.
Classe de modèle — Combine encodeur de vision + décodeur de langage.

Pour les motifs détaillés, voir :

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

Critique : `tie_word_embeddings` pour les VLM

Pour les VLM, tie_word_embeddings se trouve au niveau supérieur de la config HF, PAS sur text_config. Lisez toujours depuis la config parent :

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

Critique : Localisation du champ de config pour les VLM

Lors de la lecture de la config HF pour les VLM, vérifiez si chaque champ se trouve dans :

hf_config (niveau supérieur) — par ex. tie_word_embeddings, image_token_id, video_token_id
hf_config.text_config — par ex. num_hidden_layers, hidden_size, etc.
hf_config.vision_config — par ex. dimensions de l'encodeur de vision

Encapsuler les couches spécifiques au modèle

Quand un nouveau modèle introduit des couches personnalisées ou non standard (variantes d'attention novatrice, normalisation personnalisée, mises en page d'experts fusionnées, têtes MTP, etc.), gardez toute la logique spécifique au modèle à l'intérieur du répertoire de la famille de modèles. Ne modifiez pas les fichiers partagés dans src/megatron/bridge/models/conversion/ (par ex. param_mapping.py, model_bridge.py, quant_mapping.py) sauf si la modification est véritablement réutilisable par plusieurs familles de modèles.

Principe : Les fichiers bridge et provider d'une famille de modèles sont votre surface d'extension principale. L'infrastructure de conversion partagée fournit des hooks et des classes de base — créez-en des sous-classes localement plutôt que d'ajouter des conditionnelles au code partagé.

Stratégie 1 : Créer une sous-classe de mappage locale

Si le modèle a une couche dont la mise en page des poids ne correspond à aucune classe de mappage existante, créez une classe de mappage privée dans le fichier bridge ou un fichier <model>_mappings.py dans le répertoire de la famille.

Exemple — La projection down d'expert fusionnée de GLM désactive la transposition export groupée :

# 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)

Exemple — Les couches MTP de Nemotron-H aplati les indices lors de la résolution :

# À l'intérieur de nemotron_h_bridge.py (privé au module)
class _MTPFlatteningMapping(MegatronParamMapping):
    def resolve(self, captures):
        return AutoMapping(self._flatten(captures), ...)

Exemple — Mise en page de norme QK non standard de MiniMax-M2 :

# À l'intérieur de minimax_m2_bridge.py (privé au module)
class _FullDimQKNormMapping(MegatronParamMapping):
    def hf_to_megatron(self, hf_weights):
        # Logique de scatter personnalisée pour la norme QK pleine dimension
        ...
    def megatron_to_hf(self, megatron_weights):
        # Logique de gather personnalisée
        ...

Stratégie 2 : Remplacer les hooks du bridge

MegatronModelBridge fournit plusieurs hooks de remplacement — utilisez-les à la place de modifier la classe de base :

Hook	Quand l'utiliser
`mapping_registry()`	Définissez tous les mappages de noms de poids (abstrait, toujours remplacé)
`provider_bridge()`	Configurez le provider avec des drapeaux spécifiques au modèle (appelez `super()` puis setattr)
`maybe_modify_loaded_hf_weight()`	Dézantifiez, renommez ou remodellez les poids HF avant la conversion
`maybe_modify_converted_hf_weight()`	Synthétisez des clés HF supplémentaires à l'export (par ex. `inv_freq`)
`megatron_to_hf_config()`	Construisez `config.json` HF pour l'export
`hf_config_to_provider_kwargs()`	Remplacez le comportement de CONFIG_MAPPING pour des champs spécifiques

Accès à la config HF dans mapping_registry() : L'instance du bridge a self.hf_config disponible lors de la conversion — elle est définie automatiquement par le système de dispatch avant l'appel de mapping_registry(). Utilisez-la quand votre registre de mappage a besoin de logique dépendante de la config (par ex. nombre de couches MTP dynamiques, nombre d'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
    ...

Ne remplacez pas build_conversion_tasks() pour stocker self._hf_config — ce motif est déprécié.

Stratégie 3 : Sous-classe provider personnalisée (VLM uniquement)

La plupart des modèles n'ont pas besoin de fichier provider — le provider stock (par ex., GPTModelProvider, ou une autre base sélectionnée via PROVIDER_CLASS) est généralement suffisant pour les LLM. Créez une sous-classe provider uniquement quand un VLM a besoin de logique provide() personnalisée pour instancier un modèle vision+langage combiné :

# src/megatron/bridge/models/<model>/<model>_provider.py
class MyVLModelProvider(GPTModelProvider):
    image_token_id: int = 0

    def provide(self, ...):
        # Construction personnalisée du modèle combinant encodeur de vision + décodeur de langage
        ...

Le bridge le référence ensuite via PROVIDER_CLASS = MyVLModelProvider ou l'instancie directement dans provider_bridge().

Quand les modifications de fichiers partagés SONT justifiées

Modifiez param_mapping.py ou model_bridge.py uniquement quand le motif est réutilisable par 2+ familles de modèles. Exemples de modifications partagées justifiées :

FusedExpertMapping / FusedGatedExpertMapping — utilisées par GLM, DeepSeek, OLMoE, etc.
RMSNorm2ZeroCenteredRMSNormMapping — utilisées par Gemma, Nemotron, etc.
Nouvelles entrées CONFIG_MAPPING — quand une clé config HF standard mappe à un attribut provider standard

Si vous êtes tenté d'ajouter une branche if model_type == "..." spécifique au modèle dans le code partagé, ou un pattern matching sur des noms de poids spécifiques dans la logique de conversion partagée, c'est un signal pour utiliser une sous-classe locale ou un hook override à la place.

Mettre à jour la calculatrice FLOPs pour les nouveaux blocs architecturaux

Si le modèle introduit un nouveau bloc computationnel qui diffère de l'attention standard ou du MLP (par ex., Gated DeltaNet / GDN linear attention, Multi-Token Prediction / MTP heads, couches Mamba SSM), mettez à jour la calculatrice FLOPs dans src/megatron/bridge/training/utils/flop_utils.py de sorte que les métriques de débit d'entraînement (TFLOPs/GPU) soient précises.

Quand mettre à jour : Chaque fois que le nouveau bloc a des FLOPs par token différents de l'attention standard ou du MLP standard. Cas courants :

Variantes d'attention linéaire (GDN, RetNet, RWKV) — remplacez le terme d'attention O(s²) par le compte d'opérations réel du bloc
Têtes MTP / speculative decoding — ajoutez les FLOPs pour les couches de projection et norm supplémentaires
Couches SSM (Mamba) — FLOPs de récurrence différents de l'attention
Routage MoE novel — peut changer le nombre effectif d'experts

Comment mettre à jour :

Lisez la fonction transformer_flops() existante dans flop_utils.py pour comprendre la structure.
Ajoutez un bloc conditionnel gated sur un attribut de config (par ex., experimental_attention_variant, mtp_num_layers). Suivez le motif MoE existant pour la validation de config — levez sur les types invalides, affirmez les longueurs de liste, et utilisez l'accès direct aux attributs au lieu de getattr avec des valeurs par défaut de fallback pour que les mésconfigurations échouent explicitement.
Calculez les FLOPs par couche pour le nouveau bloc et mélangez-le avec le terme d'attention standard en fonction du motif de couche.
Ajoutez des tests unitaires dans tests/unit_tests/training/utils/test_flop_utils.py qui vérifient :
- Les FLOPs du nouveau bloc diffèrent de la baseline d'attention pure
- La formule exacte correspond aux valeurs attendues calculées à la main
- Varier le ratio de bloc (par ex., linear_attention_freq) change les FLOPs

Référence PR : #2925 — GDN FLOPs calculator ajoute le support GDN avec à la fois le code calculatrice et les tests complets.

Phase 3 : Support des Recettes

Les recettes fournissent des paramètres d'entraînement préconfigurés pour chaque taille de modèle.

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

Chaque fichier de recette définit des fonctions pour chaque taille de modèle + mode d'entraînement :

<model>_<size>_sft_config() — Fine-tuning supervisé complet
<model>_<size>_peft_config() — LoRA/DoRA fine-tuning efficace en paramètres
<model>_<size>_pretrain_config() — Préentraînement (LLM uniquement, généralement)

Pour les motifs de recette détaillés, voir @skills/adding-model-support/recipe-patterns.md.

Checklist d'export

Family __init__.py — importez et ajoutez à __all__
Top-level src/megatron/bridge/recipes/__init__.py — wildcard import
train_any_basic.py — ajoutez à config_map, docstring, et choix --model

Phase 4 : Tests

Tests unitaires (pas de GPU)

tests/unit_tests/models/<model>/
├── __init__.py
├── test_<model>_bridge.py    # Config HF simulée → vérifiez mappage provider
└── test_<model>_provider.py  # (optionnel) Uniquement si sous-classe provider personnalisée

Tests fonctionnels (GPU)

tests/functional_tests/models/<model>/
├── __init__.py
├── test_<model>_conversion.py  # Roundtrip modèle toy HF↔Megatron
└── test_<model>_provider.py    # compare_provider_configs (optionnel)

Pour les motifs de test détaillés, voir @skills/adding-model-support/tests-and-examples.md.

Phase 5 : Docs et Exemples

Exemples

Exemples LLM : examples/models/<model>/ Exemples VLM : examples/models/vlm/<model>/

examples/models/<model>/          # LLM
examples/models/vlm/<model>/      # VLM
├── README.md
├── conversion.sh        # Commandes de conversion HF↔Megatron (modèle réel)
├── inference.sh         # Commandes de génération (modèle réel, sortie raisonnable)
├── slurm_sft.sh         # Entraînement SFT sur SLURM
└── slurm_peft.sh        # Entraînement PEFT sur SLURM

Exigence livrables clés : conversion.sh et inference.sh doivent cibler un modèle publié réel (par ex. Qwen/Qwen3-8B, pas un toy). Le script d'inférence doit produire une sortie raisonnable — pour les LLM une continuation de texte cohérente, pour les VLM une description d'image plausible. C'est la barre d'acceptation : la conversion s'exécute proprement et la génération a du sens.

Documentation

Ajoutez une page de modèle à docs/models/<type>/<model>.md couvrant :

Variantes et tailles supportées
Commandes de conversion
Exemples d'entraînement (SFT, PEFT)
Limitations connues

Flux de vérification

Après implémentation du support du bridge, invitez l'utilisateur à exécuter ces commandes sur le cluster :

1. Test de fumée (GPU unique)

uv run python -c "
from megatron.bridge import AutoBridge
bridge = AutoBridge.from_hf_pretrained('<org>/<model>')
provider = bridge.to_megatron_provider()
provider.tensor_model_parallel_size = 1
provider.pipeline_model_parallel_size = 1
provider.finalize()
model = provider.provide_distributed_model(wrap_with_ddp=False)
bridge.load_hf_weights(model)
for i, (name, tensor) in enumerate(bridge.export_hf_weights(model, cpu=True)):
    print(name, tuple(tensor.shape))
    if i > 10: break
"

2. Roundtrip de conversion (multi-GPU)

uv run python examples/conversion/convert_checkpoints.py import \
    --hf-model <org>/<model> \
    --megatron-path /workspace/<model> \
    --torch-dtype bfloat16

uv run python examples/conversion/convert_checkpoints.py export \
    --hf-model <org>/<model> \
    --megatron-path /workspace/<model>/iter_0000000 \
    --hf-path /workspace/<model>-hf-export

3. Test de génération

Pour les LLM :

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

Pour les VLM :

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. Exécuter les tests

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

Arbre de décision rapide

L'utilisateur veut ajouter un modèle
│
├─ A un lien HF? ─── Non ──→ Demander le lien (ou config.json si privé)
│
├─ A text_config + vision_config? ─── Oui ──→ Chemin VLM
│   ├─ A un encodeur de vision Megatron? ──→ Encodeur Megatron (motif Qwen3.5)
│   └─ Pas d'encodeur Megatron ──→ Encodeur HF (motif Gemma3)
│
└─ Pas de vision config ──→ Chemin LLM (bridge uniquement, pas de fichier provider)
    ├─ Style GPT standard? ──→ Bridge avec mappages stock
    └─ Couches personnalisées? ──→ Bridge + sous-classes de mappage locales / hook overrides
        ├─ Mise en page des poids personnalisée? ──→ Sous-classe de mappage locale dans répertoire family
        └─ Import/export personnalisé? ──→ Remplacer les hooks du bridge (maybe_modify_*)