nemo-automodel-model-onboarding

Name: nemo-automodel-model-onboarding
Author: nvidia

github.com/nvidia開発3.8k1.9k 回インストール今日更新

nemo-automodel-model-onboarding について

nemo-automodel-model-onboardingは、NeMo AutoModelにモデルアーキテクチャの対応を追加するためのスキルです。アーキテクチャの分類から実装ファイル、レジストリ更新、重みマッピング、検証テストまでを案内します。

主な機能

01configに基づくアーキテクチャ分類
02実装ファイルとコンポーネントの特定
03チェックポイント重みのマッピング案内
04レジストリと設定の更新
05検証テスト要件の提示

活用シーン

01Dense LLMアーキテクチャの追加
02MoEモデルやチェックポイントの対応
03NeMo AutoModelへのVLM対応追加

インストール

コマンドをコピーしてターミナルで実行します。

$npx skills add https://github.com/nvidia/skills --skill nemo-automodel-model-onboarding

インストール後、Claude Code で/nemo-automodel-model-onboardingを入力して読み込まれたか確認します。

同梱ファイル

SKILL.md と一緒にインストールされる参考資料やスクリプト。

BENCHMARK.md

capabilities-and-precision.md

llm-patterns.md

moe-patterns.md

skill-card.md

skill.oms.sig

vlm-patterns.md

evals/

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 と __init__.py を作成します。チェックポイントの重み変換に必要な場合のみ state_dict_adapter.py を、必要な場合のみ config.py を追加します。_transformers/registry.py で MODEL_ARCH_MAPPING を登録し、サンプル YAML と、小規模設定のユニットテスト、およびレイヤーを書き換えた場合はレイヤー等価性テストを追加します。
MoE state dict：config.json でエキスパートフィールドを特定し、moe-patterns.md を参照します。ルーターのテンソルを別個に対応付け、ルーティング対象エキスパートのインデックス順を維持し、ルーティング対象エキスパート、共有エキスパート、および gate/up/down の射影を対応付けます。アダプターのキー対応付けテストと小規模設定の数値等価性テストを追加し、from_pretrained() のみへの依存や、エラーを通知しないテンソルの形状変更は避けてください。
VLM のオンボーディング：vision_config、text_config、および ForConditionalGeneration アーキテクチャがすべて存在する場合に限り、VLM と分類します。vlm-patterns.md と、mistral4、kimivl、kimi_k25_vl などの既存の VLM 実装を参照します。テキストバックボーン、ビジョンタワー、プロジェクター、プロセッサーの前提条件、テキストおよびビジョンのチェックポイント互換性（必要に応じてアダプターによる対応付け）、レジストリへの登録、チェックポイント全体を使用する前の小規模な画像・テキストテストを確認します。VLM のオンボーディングを因果 LM のみの経路として扱ったり、プロセッサーや画像のテストを省略したりしないでください。

MoE の state dict と VLM に関する質問には、第 2.4 節と第 2.5 節のチェックリストを適用してください。

適用範囲の境界

モデルファイル、カスタムレイヤー、state dict アダプター、Hugging Face の設定マッピング、レジストリエントリー、モデル機能フラグなど、モデルアーキテクチャのサポートを追加または変更する場合にのみ、このスキルを使用してください。

オプティマイザー、データセット、スケジューラー、検証データセット、トレーナーの配線に関する単独のトレーニングレシピ YAML の質問には、新しいモデルアーキテクチャのオンボーディングの一部として明示されている場合を除き、このスキルを使用しないでください。そうしたレシピに関する質問は、nemo-automodel-recipe-development スキルの対象です。

対象となる例：

「新しい Hugging Face causal LM アーキテクチャのサポートを追加する。」
「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 に組み込みの config クラスがない場合、_CUSTOM_CONFIG_REGISTRATIONS でカスタム config を登録するために使用します
hidden_size、intermediate_size、num_hidden_layers、num_attention_heads、num_key_value_heads -- サイズ設定
vocab_size -- 小規模なテスト用 config に必要です
tie_word_embeddings -- サポート対象の各チェックポイントに保存されている設定です。単独の config コンストラクターから推測しないでください
hidden_act -- 活性化関数です（例：SwiGLU の場合は "silu"）

1.2 モデルタイプを特定する

タイプ	判定条件	パターンファイル
Dense LLM	`architectures` に `ForCausalLM` があり、エキスパート関連のフィールドがない	llm-patterns.md
MoE LLM	config に `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 with separate HF-compatible projections
  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 二乗、MoE ルーティング
カスタム config クラス：HF の AutoConfig でモデルの config.json を解析できない場合にのみ必要です（auto_map フィールドを確認してください）

1.5 テスト用 config の次元を記録する

ユニットテスト用に、小規模な config を作成します。目標は約 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 # Only if HF weight names or tensor layouts need conversion
  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 重み変換。チェックポイントの I/O 性能も実装の一部として扱い、セクション 2.6 の説明に従って低メモリ DCP 機能を評価してください。
init.py -- メインのモデルクラスを再エクスポートする

詳しい実装ガイダンスについては、パターンファイルを参照してください。

Dense LLM：llm-patterns.md
MoE：moe-patterns.md
VLM：vlm-patterns.md
機能と fp32 精度：capabilities-and-precision.md

ほとんどのカスタムモデルでは、HF の重み変換に state_dict_adapter.py が必要です。 Llama、Qwen2、Qwen3 のように、サポート対象のバックエンドや config のバリエーション間で HF の名前とテンソルレイアウトがすでに一致している場合に限り、このファイルを省略してください。重みのタイイングは引き続きモデル側の責任です（セクション 2.3）。

2.3 因果 LM の重み共有

因果 lm_head を持つ登録済みモデルクラスは、すべて次の条件を満たす必要があります。

tie_word_embeddings_support: TieSupport を BOTH、TIED_ONLY、または UNTIED_ONLY として宣言します。
__init__ の先頭で、text_config または thinker_config をアンラップする前の元の config を使って、 reject_unsupported_tie_word_embeddings(type(self), config) を呼び出します。

因果 LM ヘッドを持たないクラスのみ、レジストリテストから明示的に除外できます。

ポリシーは、単純な config コンストラクターではなく、実装と実際にサポートされているチェックポイント config に基づいて選択してください。

BOTH：重み共有ありと重み共有なしの両方の構成がサポートされています。
TIED_ONLY：重み共有ありの構成のみがサポートされています。
UNTIED_ONLY：重み共有なしの構成のみがサポートされています。

ランタイムヘルパーでは、TIED_ONLY と UNTIED_ONLY を権威ある設定として扱い、BOTH の場合にのみチェックポイントごとの config フラグを解決してください。現在の BOTH VLM はすべて外側の tie_word_embeddings フラグに従うため、サポート対象の BOTH モデルで別の config パスが実際に必要になるまでは、モデル固有のリゾルバーを追加しないでください。

BOTH と TIED_ONLY では、必ず _tied_weights_keys を宣言し、実際の lm_head と入力埋め込みの FQN を使って tie_weights() を実装してください。継承した Hugging Face の重み共有に依存せず、言語モデルを入れ替えた後は再度重みを共有してください。

ポリシーごとにテストを追加してください。

BOTH：重み共有ありではエイリアスされ、重み共有なしではエイリアスされないこと。
TIED_ONLY：重み共有ありではエイリアスされ、重み共有なしは拒否されること。
UNTIED_ONLY：重みが別々に保たれ、重み共有ありは拒否されること。

意図的にヘッドを分離しているアーキテクチャ、語彙サイズが非対称なアーキテクチャ、または両方のテンソルを所有しないステージでは、重みを共有しないでください。

from_pretrained では、BOTH の場合も含め、チェックポイントに保存された tie_word_embeddings の値を優先します。NeMoAuto* ブリッジは、どちらの方向への切り替えも拒否します。このブリッジを経由しないモデル所有の from_pretrained では、reject_tie_word_embeddings_flip(checkpoint_config, requested_config, model_class_name) を呼び出す必要があります。

2.4 MoE state-dict アダプターのチェックリスト

MoE モデルでは、以下のすべての重みを検証してください。HF とネイティブのレイアウトが異なる場合、アダプターで明示的にマッピングする必要があります。

ルーターの重み。Hugging Face モデルにゲートバイアスまたは補正バイアスのテンソルがある場合は、それらも含めます。
専門家の重み。ローカル専門家とルーティングされた専門家の間で、専門家インデックスの順序を維持します。
ゲート、アップ、ダウンの射影。結合または分割された射影レイアウトも含めます。
アーキテクチャに共有専門家とルーティングされた専門家の両方がある場合は、それぞれを個別に扱います。

期待されるキーのマッピングを検証するテストを追加し、完全なチェックポイントを試す前に、小さな config で数値的な同等性を確認してください。

MoE PEFT では、学習済みのルーター補正バイアスを保持し、ルーターバイアスのチェックを使ってアダプターの保存／再読み込み／再開を検証してください。

次のような近道は使わないでください。

from_pretrained() を呼び出すだけでアダプターを検証しないでください。
明示的なマッピング理由なしに、専門家のキーの欠落や余分なキーを許容しないでください。
HF と NeMo のレイアウト上必要であり、変換が可逆であることをテストで証明できない限り、dtype の変更、次元の転置、テンソルの形状変更を行わないでください。
密結合層のテストに合格したからといって、ルーターや共有専門家のテストを省略しないでください。

2.5 VLM 導入チェックリスト

VLM では、Hugging Face config に vision_config と text_config があり、 architectures が条件付き生成クラスを指していることを確認してください。通常は vlm-patterns.md をはじめとする最も近い VLM パターンファイルを参照し、mistral4、kimivl、kimi_k25_vl などの既存実装と比較してください。

実装では、以下を明示的にカバーする必要があります。

テキストバックボーン、ビジョンタワー、プロジェクター、プロセッサーまたは画像の前処理に関する前提。
必要に応じてアダプターのマッピングを含めた、テキストモジュールとビジョンモジュールの両方に対するチェックポイント互換性。
_transformers/registry.py への ForConditionalGeneration クラスの登録。
画像とテキストの入力を処理し、チェックポイントの読み込み／エクスポートを検証する小規模なテスト。アダプターがある場合は、その往復テストも追加します。

2.6 チェックポイント I/O パフォーマンス

チェックポイントのパフォーマンスは、後から行う最適化ではなく、実装要件として扱ってください。新規または大幅に変更した state-dict アダプターごとに、ロード時と、その変更が影響する保存またはエクスポートの各経路で、レイテンシーとホスト／デバイスメモリのピーク使用量を評価してください。特に、次の点に注意してください。

テンソルを最終的なモデルストレージに直接ロードするか、上限を定めた部分単位で変換できる場合、チェックポイント全体またはモデルサイズの一時コピーを避けてください。
ランクが自身のシャードだけを必要とする場合、分散読み込みと変換はランク単位で行い、すべてのランクでグローバルテンソルを不必要に実体化しないでください。
モデルサイズのループ内で、テンソルのマージ、コピー、ヒープ全体のガベージコレクション、シャードの走査、ファイルのオープンを繰り返さないでください。
最適化した経路について、モデル、dtype、バックエンド、トポロジーを含む、代表的な最適化前後のレイテンシーとピークメモリの根拠を記録してください。

すべてのアダプターで supports_low_memory_dcp_load を評価する必要があります。ほとんどのチェックポイントテンソルが最終的なモデルストレージに直接書き込まれ、残りのすべてのアロケーションを伴う変換で、一時メモリ使用量が小さく、すべての実行時バリアントで上限が定められている場合に限り、_supports_low_memory_dcp_load = True を設定してください。バックエンド、トポロジー、dtype、量子化モード、またはモデルオプションによってモデルサイズの再構築が必要になる場合は、false のままにしてください。false の値は安全なフォールバックを選択します。チェックポイントのロードがサポートされていないという意味ではありません。

オプトインする場合は、直接の書き込み先を通じてセンチネル値を書き込み、最終的なモデルストレージが変化することを証明し、アロケーションを伴う変換に上限があることを確認し、安全でない実行時バリアントがこの機能を false として報告することを検証する、目的を絞ったテストが必要です。このストレージテストは正確性の要件でもあります。誤って true と判定すると、アダプターが一時テンソルをインプレースでロード済みとみなし、実際のパラメーターの再構築をスキップするおそれがあります。

2.7 レジストリへの登録

モデルを _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"),
}

2.8 機能と精度に影響するパラメーターの宣言

MODEL_ARCH_MAPPING に登録されているすべてのクラスは、静的な入れ子の ModelCapabilities データクラス、またはバリアントに対応した get_capabilities(cls, config) メソッドのいずれかで、並列化機能を宣言する必要があります。パターンは必ずどちらか一方を選んでください。機能の宣言には、エンドツーエンドで検証済みのレシピ YAML を反映してください。

Mamba の A_log／dt_bias、MoE のシグモイドゲートバイアス、アテンションシンクのバイアス、ヘッドごとの scale など、精度に影響するパラメーターがある場合は、シャーディング後もこれらのパラメーターの計算精度が fp32 に保たれるよう、_keep_in_fp32_modules_strict を宣言してください。例、バリアントのディスパッチ規則、凍結サブモジュールの dtype に関するガイダンスについては、capabilities-and-precision.md を参照してください。

フェーズ 3：オンボーディング用の設定例

このフェーズは、新たにオンボーディングしたアーキテクチャをロードして実行できることを示す最小限の設定例を追加する場合のみ対象です。一般的なレシピの作成や既存レシピの変更には、nemo-automodel-recipe-development を使用してください。

3.1 設定例の YAML を作成する

examples/llm_finetune/<name>/（または examples/vlm_finetune/<name>/）に設定例を作成します。

新しい全パラメーター Adam／AdamW の設定例では、model.dtype: float32 を設定してください。計算精度やその他の学習モードについては、training precision を参照してください。

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

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 によって、マッピングされた名前、シェイプ、データ型、値が保持されること。
アダプターの有無にかかわらず、HF の読み込み／エクスポートとネイティブ形式での保存／再読み込みによって、重みと共有関係が保持されること。tests/unit_tests/checkpoint/test_native_hf_state_dict.py を参照してください。
書き換えたすべてのアテンション、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 のレイアウトに変換してからエクスポートし、マッピングされたすべてのテンソルについて、参照と名前、シェイプ、データ型が一致し、値が想定の許容誤差内であることを確認します。
コンポーネント単位のパリティ：書き換えたアテンション、MLP、正規化、RoPE、MoE の各コンポーネントを、固定シードと同一のデータ型でHuggingFace実装と比較します。
エンドツーエンドのフォワードパス：同じトークン化済み入力を使って、NeMo AutoModel とHuggingFaceモデルを実行し、ロジット、隠れ状態、損失を比較します。

このフェーズを省略しないでください。単体テストに合格したモデルでも、重み変換の微妙なバグ、バックエンドの違い、RoPE の不一致が原因で、完全なパリティ比較で初めて差異が生じることがあります。

主なファイルの参照先

ファイル	目的
`_transformers/registry.py`	`MODEL_ARCH_MAPPING` と `_CUSTOM_CONFIG_REGISTRATIONS`
`components/models/common/__init__.py`	`BackendConfig`、`HFCheckpointingMixin`、バックエンド構築ユーティリティをエクスポートします
`components/models/llama/model.py`	HF と互換性のある重みを使用した、アテンションと MLP の個別の射影
`components/checkpoint/state_dict_adapter.py`	オプションの `StateDictAdapter` 変換契約
`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`（高密度）
`components/moe/experts.py`	`GroupedExperts`、`GroupedExpertsDeepEP`、`GroupedExpertsTE`

チェックリスト

Markdown 対応

ログインしてコメント

基本情報

作者: nvidia
リポジトリ: nvidia/skills
カテゴリ: 開発
インストール数: 1,935
GitHub スター: 3,778

問題を報告