nemo-automodel-model-onboarding について
NeMo AutoModel に新しいモデルアーキテクチャを追加する開発者向けのスキルです。Dense LLM、MoE、VLM ごとに、調査から実装、登録、テストまでの手順を示します。
主な機能
- 01config.json によるアーキテクチャ分類
- 02モデルファイルとカスタム層の計画
- 03Hugging Face 重みのマッピング
- 04モデルとカスタム設定の登録
- 05小規模設定と検証テストの設計
活用シーン
- 01因果言語モデルのアーキテクチャ追加
- 02MoE エキスパート重みの対応確認
- 03VLM サポートの実装
インストール
コマンドをコピーしてターミナルで実行します。
npx skills add https:/ / github.com/ promptingcompany/ nv-skills --skill nemo-automodel-model-onboardingインストール後、Claude Code で/nemo-automodel-model-onboardingを入力して読み込まれたか確認します。
同梱ファイル
SKILL.md と一緒にインストールされる参考資料やスクリプト。
BENCHMARK.mdllm-patterns.mdmoe-patterns.mdskill-card.mdskill.oms.sigvlm-patterns.mdevals/evals.json詳細
NeMo AutoModel へのモデルサポートの追加
目的
このスキルでは、NeMo AutoModel に新しいモデルアーキテクチャを実装する手順を説明します。5 つのフェーズを順番に進めてください。
手順
オンボーディングに関する質問に回答するときは、次の順序を守ってください。
config.jsonからアーキテクチャを分類します。components/models/<name>/内の実装ファイルを正確に示します。- レジストリと、必要に応じてカスタム設定の更新箇所を特定します。
- チェックポイントを完全に利用する前に追加すべき検証テストを示します。
概念的なオンボーディングの質問には、ユーザーからコードの編集を依頼されない限り、パターンファイルを開かずにこのスキルに基づいて回答してください。参考資料としてパターンのファイル名を示したうえで、直接的なチェックリストを提示します。
直接的な動詞を使ってください。モデルを分類し、ファイル名を示し、重みを対応付け、クラスを登録し、テストを追加してください。ユーザーが新しいアーキテクチャのオンボーディングに明示的に関連付けない限り、分散ストラテジー、ランチャーの設定、一般的なレシピ作成については説明しないでください。
例
よくある質問には、次の簡潔な回答パターンを使ってください。
- Dense causal LM:
architecturesにForCausalLMクラスが含まれ、num_local_experts、n_routed_experts、num_experts_per_tokなどのエキスパート関連フィールドが存在しない場合にのみ、Dense として分類します。components/models/<name>/model.py、state_dict_adapter.py、__init__.py、必要に応じてconfig.pyを作成し、_transformers/registry.pyにMODEL_ARCH_MAPPINGを登録して、サンプル YAML を追加します。また、小さな設定でのユニットテストと、書き換えたレイヤーの同等性テストを追加します。 - MoE の state dict:
config.json内のエキスパート関連フィールドを特定し、moe-patterns.mdを参照します。ルーターのテンソルを個別に対応付け、ルーティングされたエキスパートのインデックス順を維持し、ルーティングされたエキスパート、共有エキスパート、gate/up/down の射影を対応付けます。アダプターのキー対応テストと、小さな設定での数値的同等性テストを追加します。from_pretrained()だけに頼ったり、テンソルを通知なしに reshape したりしないでください。 - VLM のオンボーディング:
vision_config、text_config、ForConditionalGenerationアーキテクチャがすべて存在する場合にのみ、VLM として分類します。vlm-patterns.mdとmistral4、kimivl、kimi_k25_vlなどの既存の VLM 実装を参照します。テキストバックボーン、ビジョンタワー、プロジェクター、プロセッサーの前提条件、テキストおよびビジョンのstate_dict_adapter.pyにおける対応付け、レジストリへの登録、チェックポイントを完全に利用する前の小規模な画像・テキストテストを確認します。VLM のオンボーディングを純粋な causal-LM の手順として扱ったり、プロセッサーや画像のテストを省略したりしないでください。
MoE の state dict に関する質問には、必ず次の安全チェックリストを含めてください。
- ルーターのテンソルをエキスパートのテンソルとは別に対応付けます。
- ルーティングされたエキスパートのインデックス順を維持します。ロードを成功させるために、エキスパートの重みを並べ替えたり、削除したり、統合したり、通知なしに reshape したりしないでください。
- gate、up、down の射影を、結合された射影レイアウトや共有エキスパートがある場合も含めて、明示的に対応付けます。
- チェックポイント全体のロードに頼る前に、アダプターのキー対応テストと、小さな設定での数値的同等性テストを追加します。
VLM に関する質問では、vision_config、text_config、条件付き生成アーキテクチャ、テキストバックボーン、ビジョンタワー、プロジェクター、プロセッサーの前提条件、レジストリエントリー、小規模な画像・テキストテストを明示的に確認してください。
ルーティング境界
このスキルは、モデルファイル、カスタムレイヤー、state-dict アダプター、Hugging Face 設定のマッピング、レジストリエントリー、モデルの機能フラグなど、モデルアーキテクチャのサポートを追加または変更する場合にのみ使用してください。
オプティマイザー、データセット、スケジューラー、検証データセット、トレーナーの配線に関する独立したトレーニングレシピ YAML の質問には、新しいモデルアーキテクチャの導入に明示的に関係する場合を除き、このスキルを使用しないでください。こうしたレシピに関する質問は、nemo-automodel-recipe-development スキルの対象です。
対象となる例:
- 「新しい Hugging Face 因果言語モデルアーキテクチャのサポートを追加する。」
- 「Hugging Face チェックポイントから MoE ルーターとエキスパートの重みをマッピングする。」
- 「NeMo AutoModel に新しいモデルクラスを登録する。」
対象外となる例:
- 「オプティマイザーとデータセットのセクションを含むファインチューニングレシピ YAML を作成する。」
- 「FSDP2、DDP、テンソル並列、またはコンテキスト並列の設定を選択する。」
- 「Slurm、SkyPilot、コンテナー、マウント、または起動ディスパッチを設定する。」
フェーズ 1:調査
コードを書く前に、対象モデルに関する情報を集めてください。
1.1 HuggingFace の config.json を取得する
HuggingFace Hub からモデルの config.json をダウンロードするか、AutoConfig.from_pretrained を使用してください。抽出する主なフィールドは次のとおりです。
architectures-- クラス名と登録キーを決定します(例:"LlamaForCausalLM"、"Qwen3MoeForCausalLM"、"Mistral3ForConditionalGeneration")model_type-- HF に組み込みの設定クラスがない場合、_CUSTOM_CONFIG_REGISTRATIONSでカスタム設定を登録する際に使用しますhidden_size、intermediate_size、num_hidden_layers、num_attention_heads、num_key_value_heads-- サイズ設定vocab_size-- 小規模なテスト用設定に必要ですtie_word_embeddings-- lm_head が embed_tokens と重みを共有するかどうかhidden_act-- 活性化関数(例:SwiGLU の場合は"silu")
1.2 モデルの種類を特定する
| 種類 | 指標 | パターンファイル |
|---|---|---|
| Dense LLM | architectures に ForCausalLM があり、エキスパート関連フィールドがない | llm-patterns.md |
| MoE LLM | 設定に n_routed_experts、num_local_experts、num_experts_per_tok がある | moe-patterns.md |
| VLM | architectures に ForConditionalGeneration があり、vision_config と text_config がある | vlm-patterns.md |
1.3 類似する既存アーキテクチャを確認する
components/models/ 内で、類似したアテンションまたは MLP パターンを持つアーキテクチャを探してください。
components/models/
llama/ # Standard GQA + SwiGLU (CombinedQKV + CombinedGateUpMLP)
qwen2/ # Same as Llama but with attention bias + QKV bias
baichuan/ # ALiBi attention variant
deepseek_v3/ # MLA attention + MoE (DeepSeek-style grouped experts)
mistral4/ # MLA + MoE + VLM (Pixtral vision)
kimivl/ # DeepSeek-V3 backbone + MoonVit vision
kimi_k25_vl/ # Updated KimiVL with different projector
qwen3_moe/ # Qwen3 with MoE layers
nemotron_v3/ # Hybrid mamba-attention
1.4 カスタムコンポーネントを特定する
モデルに次のものが必要か確認してください。
- カスタムアテンション:GQA(標準)、MLA(DeepSeek/Mistral4)、スライディングウィンドウ、双方向
- カスタム RoPE:標準(Llama)、YaRN スケーリング、NTK 対応、複素数(DeepSeek)
- カスタム正規化:RMSNorm(標準)、LayerNorm、異なる eps 値
- カスタム MLP:SwiGLU(標準)、GeGLU、ReLU-squared、MoE ルーティング
- カスタム設定クラス:HF の
AutoConfigでモデルのconfig.jsonを解析できない場合にのみ必要です(auto_mapフィールドを確認してください)
1.5 テスト用設定の次元を記録する
単体テスト用に、小規模な設定を作成してください。目標は約 1M パラメーター以下です。
# Example tiny config for a Llama-like model:
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,
)
フェーズ 2:実装
2.1 ディレクトリ構造を作成する
components/models/<name>/
__init__.py
model.py
state_dict_adapter.py
config.py # Only if HF config is insufficient
layers.py # Only for MoE / MLA / other non-standard layers
rope_utils.py # Only for custom RoPE
2.2 実装順序
依存関係の順にファイルを実装してください。
- config.py(必要な場合)-- カスタム
PretrainedConfigサブクラス - rope_utils.py(必要な場合)-- RoPE の実装
- layers.py(必要な場合)-- アテンション、MLP、デコーダーブロックのクラス
- model.py -- メインの
ForCausalLM(またはForConditionalGeneration)クラス - state_dict_adapter.py -- HF の重み変換
- init.py -- メインのモデルクラスを再エクスポート
実装の詳細なガイダンスについては、パターンファイルを参照してください。
- Dense LLM:llm-patterns.md
- MoE:moe-patterns.md
- VLM:vlm-patterns.md
2.3 MoE state-dict アダプタのチェックリスト
MoE モデルでは、一般的な読み込みだけで終わらせないでください。アダプタでは、以下を明示的にマッピングする必要があります。
- Hugging Face モデルに存在する場合は、ゲートのバイアスや補正バイアスのテンソルを含むルーターの重み。
- ローカルエキスパートとルーティングされたエキスパート間でエキスパートのインデックス順を維持した、エキスパートの重み。
- 結合形式または分割形式のどちらの射影レイアウトにも対応する、ゲート/アップ/ダウン射影。
- アーキテクチャに両方ある場合は、共有エキスパートとルーティングされたエキスパートを個別に扱うこと。
期待されるキーのマッピングを検証するテストを追加し、完全なチェックポイントを試す前に、小さな設定で数値的な等価性を確認してください。
次のような近道は使わないでください。
from_pretrained()を呼び出すだけでアダプタを検証しないでください。- 明確なマッピング理由なしに、不足または余分なエキスパートのキーを許容しないでください。
- HF と NeMo のレイアウト上必要であり、かつテストで変換が可逆であると証明されない限り、データ型の変更、次元の転置、テンソルの形状変更を行わないでください。
- 全結合層のテストに合格したからといって、ルーターや共有エキスパートのテストを省略しないでください。
2.4 VLM 導入チェックリスト
VLM では、Hugging Face の設定に vision_config と text_config があり、architectures が条件付き生成クラスを指していることを確認してください。最も近い VLM パターンのファイル(通常は vlm-patterns.md)を起点にして、mistral4、kimivl、kimi_k25_vl などの既存の実装と比較してください。
実装では、以下を明示的に扱う必要があります。
- テキストバックボーン、ビジョンタワー、プロジェクター、プロセッサーまたは画像前処理に関する前提。
state_dict_adapter.pyにおける、テキストモジュールとビジョンモジュールの両方の重みマッピング。_transformers/registry.pyにおけるForConditionalGenerationクラスの登録。- 画像とテキストの入力を処理し、アダプタの往復変換を検証する小規模なテスト。
2.5 レジストリへの登録
モデルを _transformers/registry.py の MODEL_ARCH_MAPPING に追加します。
# In _transformers/registry.py
MODEL_ARCH_MAPPING = OrderedDict([
# ... existing entries ...
(
"NewModelForCausalLM",
("nemo_automodel.components.models.new_model.model", "NewModelForCausalLM"),
),
])
モデルのカスタム設定クラスの config.json に auto_map がある場合は、_CUSTOM_CONFIG_REGISTRATIONS にも登録します。
_CUSTOM_CONFIG_REGISTRATIONS: Dict[str, Tuple[str, str]] = {
# ... existing entries ...
"new_model": ("nemo_automodel.components.models.new_model.configuration", "NewModelConfig"),
}
フェーズ 3:導入例の設定
このフェーズは、新たに導入したアーキテクチャが読み込みと実行が可能であることを示す、最小限の設定例を追加する場合にのみ行います。一般的なレシピの作成や既存レシピの変更には、nemo-automodel-recipe-development を使用してください。
3.1 設定例の YAML を作成する
examples/llm_finetune/<name>/(または 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
# ... data, optimizer config ...
3.2 モデルの読み込みを確認する
HuggingFace のチェックポイントからモデルを読み込めることをテストします。
from nemo_automodel import NeMoAutoModelForCausalLM
model = NeMoAutoModelForCausalLM.from_pretrained("<org>/<model-name>")
3.3 まず小さな設定でテストする
完全なサイズのモデルを使う前に、小さな設定(1~2 層、小さい隠れ次元)で検証し、形状の不一致を早期に検出します。
フェーズ 4:テスト
完全なチェックポイントを読み込む前に、tests/unit_tests/models/<name>/ を作成し、以下の確認項目を網羅します。
- 小さな設定でのフォワード処理の形状を確認するスモークテスト。
- state-dict アダプタの往復変換:
from_hf -> to_hfで、マッピングされた名前、形状、データ型、値が維持されること。 - 書き換えたすべてのアテンション、MLP、正規化、RoPE、MoE レイヤーのレイヤー等価性テスト。設定のモデルデータ型、同一のシード済み重み、同一の入力、およびデータ型に適した
torch.allcloseの許容誤差を使用します。 - 数回の学習ステップで損失が減少することを確認する短い機能テスト。
フェーズ 5:ドキュメント
5.1 モデル対応状況のページを更新する
docs/model-coverage/ 内の該当するファイルを編集します。
- LLM/MoE:
docs/model-coverage/llm/index.md - VLM:
docs/model-coverage/vlm/index.md
モデル名、サポートされる機能(TP、PP、FSDP、LoRA、QLoRA)、および制限事項を記載した行を追加します。
フェーズ 6:同等性テスト
実装とユニットテストが完了したら、完全な同等性テストのワークフローを実行して、新しいモデルが参照先のHuggingFace実装と数値的に同等な結果を生成することを検証します。
次の3段階で比較します。
- State-dictのラウンドトリップ:参照先のHuggingFaceチェックポイントを読み込み、NeMo AutoModelのレイアウトに変換してから再度エクスポートし、マッピングされたすべてのテンソルについて、名前、形状、データ型、値が許容誤差の範囲内で参照先と一致することを検証します。
- コンポーネントレベルの同等性:固定シードと同一のデータ型を使用して、書き換えたAttention、MLP、正規化、RoPE、MoEの各コンポーネントをHuggingFaceの実装と比較します。
- エンドツーエンドのフォワードパス:同じトークン化済み入力を使用してNeMo AutoModelとHuggingFaceモデルを実行し、ロジット、隠れ状態、損失を比較します。
このフェーズを省略しないでください。ユニットテストに合格したモデルでも、微妙な重み変換のバグ、バックエンドの違い、または完全な同等性比較で初めて表面化するRoPEの不一致により、結果が異なる場合があります。
主要ファイルの一覧
| ファイル | 目的 |
|---|---|
_transformers/registry.py | MODEL_ARCH_MAPPINGと_CUSTOM_CONFIG_REGISTRATIONS |
components/models/common/__init__.py | CombinedQKVAttentionMixin、CombinedGateUpMLP、BackendConfig、HFCheckpointingMixinなどをエクスポートします |
components/models/common/combined_projection/combined_qkv.py | setup_qkv_projection()とcompute_qkv()を備えたCombinedQKVAttentionMixin |
components/models/common/combined_projection/combined_mlp.py | インターリーブされたgate/upレイアウトのCombinedGateUpMLP |
components/models/common/combined_projection/state_dict_adapter.py | CombinedProjectionStateDictAdapterの基底クラス |
components/models/common/hf_checkpointing_mixin.py | 保存/読み込み用のHFCheckpointingMixin |
components/models/common/utils.py | BackendConfig、initialize_rms_norm_module、initialize_linear_module、get_rope_config |
components/moe/config.py | MoEConfigデータクラス |
components/moe/fsdp_mixin.py | 分散エキスパート処理用のMoEFSDPSyncMixin |
components/moe/layers.py | MoEブロック用のMoEレイヤーとMLP(Dense) |
components/moe/experts.py | GroupedExperts、GroupedExpertsDeepEP、GroupedExpertsTE |
チェックリスト
- HuggingFaceから
config.jsonを取得して分析した - モデルタイプ(Dense LLM/MoE/VLM)を特定した
- カスタムコンポーネント(Attention、RoPE、正規化、MLP)を特定した
-
components/models/<name>/ディレクトリを作成した - config.pyを実装した(カスタム設定が必要な場合)
- layers.pyを実装した(カスタムレイヤーが必要な場合)
- rope_utils.pyを実装した(カスタムRoPEが必要な場合)
-
HFCheckpointingMixinを使用してmodel.pyを実装した - state_dict_adapter.pyを実装した
- 再エクスポートを含む__init__.pyを実装した
-
_transformers/registry.pyのMODEL_ARCH_MAPPINGに登録した - カスタム設定を
_CUSTOM_CONFIG_REGISTRATIONSに登録した(該当する場合) - サンプル用のYAML設定を作成した
-
NeMoAutoModelForCausalLM.from_pretrained()でモデルを読み込めることを確認した - ユニットテストを作成した(フォワードパスの形状、state_dictのラウンドトリップ)
- 書き換えたすべてのレイヤーについて、レイヤー同等性テストを作成した(モデルのデータ型を一致させる)
- 機能テストを作成した(学習損失が減少する)
- docs/model-coverageページを更新した
- state_dictのラウンドトリップ、コンポーネントの同等性、E2Eフォワードパスの同等性チェックを実行した
- モジュール末尾に
ModelClass = <Name>ForCausalLMを設定した