Ajouter une prise en charge de modèle à NeMo AutoModel

Objectif

Cette skill guide l'implémentation de nouvelles architectures de modèles dans NeMo AutoModel. Suivez les cinq phases dans l'ordre.

Instructions

Quand vous répondez à une question d'onboarding, gardez la réponse dans cet ordre :

1. Classifiez l'architecture à partir de config.json.
2. Nommez les fichiers d'implémentation exacts sous components/models/<name>/.
3. Identifiez le registry et les mises à jour de config personnalisée optionnelles.
4. Énoncez les tests de validation qui doivent être ajoutés avant l'utilisation complète du checkpoint.

Pour les questions d'onboarding conceptuel, répondez à partir de cette skill sans ouvrir les fichiers de pattern sauf si l'utilisateur vous demande d'éditer du code. Mentionnez les noms de fichiers pattern comme références, puis donnez la checklist directe.

Utilisez des verbes d'action directs : classifiez le modèle, nommez les fichiers, mappez les poids, enregistrez la classe, ajoutez des tests. Ne discutez pas de stratégie distribuée, configuration du launcher, ou création générale de recettes sauf si l'utilisateur le connecte explicitement à l'onboarding d'une nouvelle architecture.

Exemples

Utilisez ces patterns de réponse compacts pour les questions courantes :

• Dense causal LM : classifiez comme dense uniquement quand architectures contient une classe ForCausalLM et que les champs d'experts tels que num_local_experts, n_routed_experts, ou num_experts_per_tok sont absents. Créez components/models/<name>/model.py, state_dict_adapter.py, __init__.py, et optionnel config.py, enregistrez MODEL_ARCH_MAPPING dans _transformers/registry.py, ajoutez un exemple YAML, et ajoutez des tests unitaires tiny-config plus des tests d'équivalence de couche pour les couches réécrites.
• MoE state dict : identifiez les champs d'expert dans config.json, référencez moe-patterns.md, mappez les tenseurs router séparément, préservez l'ordre d'index expert routé, mappez les experts routés, experts partagés, et les projections gate/up/down, ajoutez des tests key-map adapter et des tests d'équivalence numérique tiny-config, et ne comptez pas uniquement sur from_pretrained() ou des reshapes de tenseur silencieux.
• VLM onboarding : classifiez comme VLM uniquement quand vision_config, text_config, et une architecture ForConditionalGeneration sont présentes. Référencez vlm-patterns.md et les implémentations VLM existantes telles que mistral4, kimivl, ou kimi_k25_vl ; vérifiez le backbone de texte, la vision tower, le projector, les hypothèses du processor, les mappages state_dict_adapter.py de texte et vision, l'enregistrement du registry, et les petits tests image-texte avant les checkpoints complets. Ne traitez pas l'onboarding VLM comme un chemin causal-LM pur ou ne sautez pas les tests de processor/image.

Pour les questions MoE state-dict, toujours inclure la checklist de sécurité :

• Mappez les tenseurs router séparément des tenseurs d'expert.
• Préservez l'ordre d'index expert routé ; ne jamais trier, supprimer, fusionner, ou reshaper silencieusement les poids d'expert pour faire passer le chargement.
• Mappez explicitement les projections gate, up, et down, incluant les layouts de projection combinés et les experts partagés quand présents.
• Ajoutez des tests key-map adapter et des tests d'équivalence numérique tiny-config avant de vous fier au chargement de checkpoint complet.

Pour les questions VLM, vérifiez explicitement vision_config, text_config, l'architecture conditional-generation, le backbone de texte, la vision tower, le projector, les hypothèses du processor, l'entrée du registry, et les petits tests image-texte.

Limite de routage

Utilisez cette skill uniquement quand l'utilisateur ajoute ou modifie la prise en charge de l'architecture du modèle : fichiers de modèle, couches personnalisées, adaptateurs state-dict, mappage de config Hugging Face, entrées du registry, ou flags de capacité du modèle.

N'utilisez pas cette skill pour les questions YAML de recettes d'entraînement autonomes sur les optimiseurs, datasets, schedulers, datasets de validation, ou wiring du trainer sauf s'ils font explicitement partie de l'onboarding d'une nouvelle architecture de modèle. Ces questions de recette appartiennent à la skill nemo-automodel-recipe-development.

Exemples dans le scope :

• « Ajouter le support d'une nouvelle architecture causal LM Hugging Face. »
• « Mapper les poids router et expert MoE depuis un checkpoint Hugging Face. »
• « Enregistrer une nouvelle classe de modèle dans NeMo AutoModel. »

Exemples hors du scope :

• « Écrire une recette YAML de finetuning avec sections optimiseur et dataset. »
• « Choisir FSDP2, DDP, tensor parallel, ou context parallel settings. »
• « Configurer Slurm, SkyPilot, containers, mounts, ou dispatch de lancement. »

Phase 1 : Découverte

Avant d'écrire du code, rassemblez des informations sur le modèle cible.

1.1 Récupérer config.json HuggingFace

Téléchargez le config.json du modèle depuis le HuggingFace Hub (ou utilisez AutoConfig.from_pretrained). Champs clés à extraire :

• architectures -- détermine le nom de la classe et la clé d'enregistrement (ex. "LlamaForCausalLM", "Qwen3MoeForCausalLM", "Mistral3ForConditionalGeneration")
• model_type -- utilisé pour l'enregistrement de config personnalisée dans _CUSTOM_CONFIG_REGISTRATIONS si HF n'a pas de classe config intégrée
• hidden_size, intermediate_size, num_hidden_layers, num_attention_heads, num_key_value_heads -- dimensionnement
• vocab_size -- nécessaire pour les configs de test tiny
• tie_word_embeddings -- si lm_head partage les poids avec embed_tokens
• hidden_act -- fonction d'activation (ex. "silu" pour SwiGLU)

1.2 Déterminer le type de modèle

1.3 Vérifier les architectures similaires existantes

Cherchez dans components/models/ des architectures avec des patterns d'attention ou MLP similaires :

components/models/
  llama/           # GQA standard + SwiGLU (CombinedQKV + CombinedGateUpMLP)
  qwen2/           # Identique à Llama mais avec attention bias + QKV bias
  baichuan/        # Variante ALiBi attention
  deepseek_v3/     # MLA attention + MoE (experts groupés style DeepSeek)
  mistral4/        # MLA + MoE + VLM (vision Pixtral)
  kimivl/          # Backbone DeepSeek-V3 + vision MoonVit
  kimi_k25_vl/     # KimiVL mis à jour avec projector différent
  qwen3_moe/       # Qwen3 avec couches MoE
  nemotron_v3/     # Mamba-attention hybride

1.4 Identifier les composants personnalisés

Vérifiez si le modèle a besoin de :

• Attention personnalisée : GQA (standard), MLA (DeepSeek/Mistral4), sliding window, bidirectionnelle
• RoPE personnalisé : Standard (Llama), scaling YaRN, NTK-aware, nombre complexe (DeepSeek)
• Normalisation personnalisée : RMSNorm (standard), LayerNorm, valeurs eps différentes
• MLP personnalisé : SwiGLU (standard), GeGLU, ReLU-squared, routage MoE
• Classe config personnalisée : Nécessaire uniquement si AutoConfig HF ne peut pas parser le config.json du modèle (vérifiez le champ auto_map)

1.5 Noter les dimensions pour le config de test

Pour les tests unitaires, créez un config tiny. Objectif : ~1M paramètres ou moins.

# Exemple config tiny pour un modèle de type Llama :
tiny_config = LlamaConfig(
    hidden_size=64,
    intermediate_size=128,
    num_hidden_layers=2,
    num_attention_heads=4,
    num_key_value_heads=2,
    vocab_size=256,
    max_position_embeddings=128,
)

Phase 2 : Implémentation

2.1 Créer la structure de répertoire

components/models/<name>/
  __init__.py
  model.py
  state_dict_adapter.py
  config.py            # Uniquement si le config HF est insuffisant
  layers.py            # Uniquement pour MoE / MLA / autres couches non-standard
  rope_utils.py        # Uniquement pour RoPE personnalisé

2.2 Ordre d'implémentation

Implémentez les fichiers dans l'ordre des dépendances :

1. config.py (si nécessaire) -- Sous-classe PretrainedConfig personnalisée
2. rope_utils.py (si nécessaire) -- Implémentation RoPE
3. layers.py (si nécessaire) -- Classes attention, MLP, decoder block
4. model.py -- La classe principale ForCausalLM (ou ForConditionalGeneration)
5. state_dict_adapter.py -- Conversion des poids HF
6. init.py -- Ré-exporter la classe modèle principale

Voir les fichiers pattern pour une guidance détaillée :

• Dense LLM : llm-patterns.md
• MoE : moe-patterns.md
• VLM : vlm-patterns.md

2.3 Checklist adaptateur MoE state-dict

Pour les modèles MoE, n'arrêtez pas au chargement générique. L'adaptateur doit mapper explicitement :

• Les poids du router, incluant le gate bias ou les tenseurs correction-bias quand le modèle Hugging Face les a.
• Les poids d'expert, préservant l'ordre d'index expert sur les experts locaux et routés.
• Les projections gate/up/down, incluant les layouts de projection combinés ou divisés.
• Les experts partagés séparément des experts routés quand l'architecture les a tous les deux.

Ajoutez des tests qui affirment les mappages clés attendus et exécutent l'équivalence numérique avec des configs tiny avant d'essayer des checkpoints complets.

Ne pas utiliser ces raccourcis :

• Ne validez pas l'adaptateur uniquement en appelant from_pretrained().
• N'acceptez pas les clés d'expert manquantes ou supplémentaires sans une raison de mapping explicite.
• Ne changez pas dtype, transposez dimensions, ou reshapez les tenseurs à moins que les layouts HF et NeMo ne l'exigent et qu'un test prove que la conversion est réversible.
• Ne sautez pas les tests du router ou expert partagé parce que les tests de couche dense passent.

2.4 Checklist onboarding VLM

Pour les VLMs, confirmez que le config Hugging Face a vision_config et text_config et que architectures pointe vers une classe conditional-generation. Commencez par le fichier pattern VLM le plus proche, généralement vlm-patterns.md, et comparez les implémentations existantes comme mistral4, kimivl, ou kimi_k25_vl.

L'implémentation doit explicitement couvrir :

• Le backbone de texte, vision tower, projector, et hypothèses du processor ou prétraitement d'image.
• Le mappage de poids pour les modules texte et vision dans state_dict_adapter.py.
• L'enregistrement de la classe ForConditionalGeneration dans _transformers/registry.py.
• Les petits tests qui exercent les entrées image-texte et vérifient le round-trip adapter.

2.5 Enregistrer dans le registry

Ajoutez le modèle à MODEL_ARCH_MAPPING dans _transformers/registry.py :

# Dans _transformers/registry.py
MODEL_ARCH_MAPPING = OrderedDict([
    # ... entrées existantes ...
    (
        "NewModelForCausalLM",
        ("nemo_automodel.components.models.new_model.model", "NewModelForCausalLM"),
    ),
])

Si le modèle a une classe config personnalisée avec auto_map dans son config.json, enregistrez aussi dans _CUSTOM_CONFIG_REGISTRATIONS :

_CUSTOM_CONFIG_REGISTRATIONS: Dict[str, Tuple[str, str]] = {
    # ... entrées existantes ...
    "new_model": ("nemo_automodel.components.models.new_model.configuration", "NewModelConfig"),
}

Phase 3 : Exemple Config Onboarding

Cette phase est uniquement pour ajouter un config d'exemple minimal qui prouve que l'architecture nouvellement embarquée peut charger et tourner. Utilisez nemo-automodel-recipe-development pour la création générale de recette ou les modifications de recette existantes.

3.1 Créer un config YAML d'exemple

Créez un config d'exemple sous examples/llm_finetune/<name>/ (ou examples/vlm_finetune/<name>/) :

model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: <org>/<model-name>

trainer:
  max_steps: 100
  gradient_clip_val: 1.0
  accumulate_grad_batches: 1

# ... config data, optimizer ...

3.2 Vérifier le chargement du modèle

Testez que le modèle charge depuis un checkpoint HuggingFace :

from nemo_automodel import NeMoAutoModelForCausalLM

model = NeMoAutoModelForCausalLM.from_pretrained("<org>/<model-name>")

3.3 Tester d'abord avec un config tiny

Avant d'utiliser des modèles de taille complète, vérifiez avec un config tiny (1-2 couches, petite hidden dim) pour détecter les erreurs de shape tôt.

Phase 4 : Tests

Créez tests/unit_tests/models/<name>/ et couvrez les vérifications ci-dessous avant de charger les checkpoints complets :

• Test smoke forward-shape avec un config tiny.
• Adapter state-dict round-trip : from_hf -> to_hf préserve les noms mappés, shapes, dtypes, et valeurs.
• Tests d'équivalence de couche pour chaque attention réécrite, MLP, normalisation, RoPE, ou couche MoE. Utilisez le dtype du modèle depuis le config, les poids identiques avec seed, les entrées identiques, et les tolérances torch.allclose appropriées au dtype.
• Test fonctionnel court qui vérifie que la perte décroît sur quelques étapes d'entraînement.

Phase 5 : Documentation

5.1 Mettre à jour la page de couverture des modèles

Éditez le fichier approprié dans docs/model-coverage/ :

• LLM/MoE : docs/model-coverage/llm/index.md
• VLM : docs/model-coverage/vlm/index.md

Ajoutez une ligne avec le nom du modèle, les features supportées (TP, PP, FSDP, LoRA, QLoRA), et toute limitation.

Phase 6 : Tests de parité

Après que l'implémentation et les tests unitaires soient complets, exécutez le workflow complet de tests de parité pour vérifier que le nouveau modèle produit des résultats numériquement équivalents à l'implémentation de référence HuggingFace.

Exécutez trois niveaux de comparaison :

1. Round-trip state-dict : chargez un checkpoint HuggingFace de référence, convertissez-le dans la disposition NeMo AutoModel, exportez-le en arrière, et vérifiez que tous les tenseurs mappés correspondent aux noms, shapes, dtypes, et valeurs de référence dans la tolérance attendue.
2. Parité au niveau composant : comparez les composants attention, MLP, normalisation, RoPE, et MoE réécrits contre l'implémentation HuggingFace avec seeds fixes et dtype identique.
3. Forward pass end-to-end : exécutez le modèle NeMo AutoModel complet et le modèle HuggingFace sur la même entrée tokenisée et comparez les logits, hidden states, et loss.

Ne sautez pas cette phase. Un modèle qui passe les tests unitaires peut quand même diverger de HF à cause de bugs subtils de conversion de poids, différences de backend, ou erreurs de RoPE qui ne surgissent que dans une comparaison de parité complète.

Référence fichiers clés

Checklist

• [ ] Récupéré et analysé config.json de HuggingFace
• [ ] Déterminé le type de modèle (dense LLM / MoE / VLM)
• [ ] Identifié les composants personnalisés (attention, RoPE, normalisation, MLP)
• [ ] Créé le répertoire components/models/<name>/
• [ ] Implémenté config.py (si config personnalisé nécessaire)
• [ ] Implémenté layers.py (si couches personnalisées nécessaires)
• [ ] Implémenté rope_utils.py (si RoPE personnalisé nécessaire)
• [ ] Implémenté model.py avec HFCheckpointingMixin
• [ ] Implémenté state_dict_adapter.py
• [ ] Implémenté init.py avec ré-export
• [ ] Enregistré dans MODEL_ARCH_MAPPING dans _transformers/registry.py
• [ ] Enregistré config personnalisé dans _CUSTOM_CONFIG_REGISTRATIONS (si applicable)
• [ ] Créé config YAML d'exemple
• [ ] Vérifié que le modèle charge via NeMoAutoModelForCausalLM.from_pretrained()
• [ ] Créé tests unitaires (forward shape, round-trip state_dict)
• [ ] Créé tests d'équivalence de couche pour chaque couche réécrite (correspondant au dtype du modèle)
• [ ] Créé tests fonctionnels (loss d'entraînement décroît)
• [ ] Mis à jour la page docs/model-coverage
• [ ] Exécuté les vérifications de round-trip state-dict, parité de composant, et forward-pass E2E
• [ ] Fixé ModelClass = <Name>ForCausalLM au bas du module