⚡

NeMoAutoModel模型接入Skill nemo-automodel-model-onboarding

该技能为NeMo AutoModel添加新模型架构提供分阶段指南，包括架构识别、实现模式、注册与验证，并涵盖MoE和VLM的专用适配清单，适用于模型文件、自定义层、state dict adapter和registry配置等。关键词：NeMo AutoModel、模型接入、模型架构、MoE、VLM、state_dict适配、注册。

大模型训练框架 0 次安装 0 次浏览 更新于 9/7/2026
名称 nemo-automodel-model-onboarding
描述 在 NeMo AutoModel 中接入新模型架构的指南，包括架构发现、实现模式、注册和验证。 when_to_use: 在 NeMo AutoModel 中添加或修改模型架构支持时，如 LLM/VLM/MoE 模型文件、自定义层、state-dict 适配器、注册条目、Hugging Face 配置映射或能力标志。
开源协议 Apache-2.0 metadata:
作者 NVIDIA tags: - nemo-automodel - model-onboarding

向 NeMo AutoModel 添加模型支持

目的

本技能指导在 NeMo AutoModel 中实现新的模型架构。请按顺序遵循以下五个阶段。

说明

回答接入问题时，请按以下顺序组织回复：

  1. 根据 config.json 对架构进行分类。
  2. 说明 components/models/<name>/ 下的具体实现文件。
  3. 指出注册表及可选的自定义配置更新。
  4. 说明在使用完整检查点之前必须添加的验证测试。

对于概念性接入问题，直接基于本技能回答，除非用户要求修改代码，否则不要打开模式文件。可引用模式文件名作为参考，然后给出直接清单。

使用直接动作动词：对模型进行分类、命名文件、映射权重、注册类、添加测试。除非用户明确将其与新架构接入关联，否则不要讨论分布式策略、启动器配置或通用 recipe 编写。

示例

使用以下紧凑回答模式回答常见问题：

  • 稠密因果 LM：仅当 architectures 包含 ForCausalLM 类且没有专家字段（如 num_local_experts、n_routed_experts 或 num_experts_per_tok）时，才归类为稠密。创建 components/models/<name>/model.py、state_dict_adapter.py、init.py，以及可选的 config.py，在 _transformers/registry.py 中注册 MODEL_ARCH_MAPPING，添加示例 YAML，并添加微型配置单元测试以及重写层的等效性测试。
  • MoE 状态字典：识别 config.json 中的专家字段，参考 moe-patterns.md，分别映射 router 张量，保留 routed-expert 索引顺序，映射 routed experts、shared experts 以及 gate/up/down 投影，添加适配器键映射测试和微型配置数值等效测试；不要仅依赖 from_pretrained() 或静默张量 reshape。
  • VLM 接入：仅当 vision_config、text_config 和 ForConditionalGeneration 架构都存在时，才归类为 VLM。参考 vlm-patterns.md 和现有 VLM 实现，如 mistral4、kimivl 或 kimi_k25_vl；检查文本骨干、视觉塔、projector、processor 假设、文本和视觉 state_dict_adapter.py 映射、注册表注册，以及微型图像-文本测试后再使用完整检查点。不要将 VLM 接入视为纯因果 LM 路径，也不要跳过 processor/图像测试。

对于 MoE state-dict 和 VLM 问题，应用第 2.4 和 2.5 节中的清单。

路由边界

仅当用户添加或修改模型架构支持时使用本技能：模型文件、自定义层、状态字典适配器、Hugging Face 配置映射、注册表条目或模型能力标志。

不要将本技能用于独立的训练 recipe YAML 问题，例如优化器、数据集、调度器、验证数据集或 trainer 接线，除非它们明确属于新模型架构接入的一部分。这些问题属于 nemo-automodel-recipe-development 技能。

范围内的示例：

  • 为新的 Hugging Face 因果 LM 架构添加支持。
  • 映射 Hugging Face checkpoint 中的 MoE router 和 expert 权重。
  • 在 NeMo AutoModel 中注册新的模型类。

范围外的示例：

  • 编写包含优化器和数据集部分的微调 recipe YAML。
  • 选择 FSDP2、DDP、张量并行或上下文并行设置。
  • 配置 Slurm、SkyPilot、容器、挂载或启动调度。

Phase 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 —— 每个支持的 checkpoint 中保存的设置；不要从裸配置构造函数推断
  • hidden_act —— 激活函数（例如 silu 对应 SwiGLU）

1.2 确定模型类型

类型 指示 模式文件
稠密 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/           # 标准 GQA + SwiGLU (CombinedQKV + CombinedGateUpMLP)
  qwen2/           # 与 Llama 相同，但有注意力偏置 + QKV 偏置
  baichuan/        # ALiBi 注意力变体
  deepseek_v3/     # MLA 注意力 + MoE（DeepSeek 风格分组专家）
  mistral4/        # MLA + MoE + VLM（Pixtral 视觉）
  kimivl/          # DeepSeek-V3 骨干 + MoonVit 视觉
  kimi_k25_vl/     # 更新的 KimiVL，使用不同的 projector
  qwen3_moe/       # Qwen3 使用 MoE 层
  nemotron_v3/     # 混合 mamba-attention

1.4 识别自定义组件

检查模型是否需要：

  • 自定义注意力：GQA（标准）、MLA（DeepSeek/Mistral4）、滑动窗口、双向
  • 自定义 RoPE：标准（Llama）、YaRN 缩放、NTK-aware、复数（DeepSeek）
  • 自定义归一化：RMSNorm（标准）、LayerNorm、不同 eps 值
  • 自定义 MLP：SwiGLU（标准）、GeGLU、ReLU-squared、MoE routing
  • 自定义配置类：仅当 HF AutoConfig 无法解析模型的 config.json 时需要（检查 auto_map 字段）

1.5 记录测试配置的维度

对于单元测试，创建微配置。目标：约 1M 参数或更少。

# 类似 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: 实现

2.1 创建目录结构

components/models/<name>/
  __init__.py
  model.py
  state_dict_adapter.py
  config.py            # 仅当 HF 配置不足时
  layers.py            # 仅用于 MoE / MLA / 其他非标准层
  rope_utils.py        # 仅用于自定义 RoPE

2.2 实现顺序

按依赖顺序实现文件：

  1. config.py（如果需要）——自定义 PretrainedConfig 子类
  2. rope_utils.py（如果需要）——RoPE 实现
  3. layers.py（如果需要）——注意力、MLP、decoder block 类
  4. model.py——主要的 ForCausalLM（或 ForConditionalGeneration）类
  5. state_dict_adapter.py——HF 权重转换
  6. init.py——重新导出主要模型类

有关详细实现指导，请参阅模式文件：

2.3 因果 LM 权重绑定

每个注册的具有因果 lm_head 的模型类必须：

  • 将 tie_word_embeddings_support 声明为 TieSupport，值为 BOTH、TIED_ONLY 或 UNTIED_ONLY。
  • 在 init 顶部调用 reject_unsupported_tie_word_embeddings(type(self), config)，使用原始 config，在解包 text_config 或 thinker_config 之前。

只有没有因果 LM 头的类才能被注册表测试显式豁免。

根据实现和实际支持的 checkpoint 配置选择策略，而不是根据裸配置构造函数：

  • BOTH：支持绑定和未绑定配置。
  • TIED_ONLY：仅支持绑定配置。
  • UNTIED_ONLY：仅支持未绑定配置。

运行时辅助程序必须将 TIED_ONLY 和 UNTIED_ONLY 视为权威，仅对 BOTH 解析每 checkpoint 的配置标志。所有当前 BOTH VLM 都遵循外部 tie_word_embeddings 标志，因此在支持的 BOTH 模型实际需要另一个配置路径之前，不要添加模型特定的解析器。

对于 BOTH 和 TIED_ONLY，始终声明 _tied_weights_keys 并使用实际的 lm_head 和输入嵌入 FQN 实现 tie_weights()。不要依赖继承的 Hugging Face 绑定，并在任何语言模型交换后重新绑定。

添加策略特定测试：

  • BOTH：绑定别名；未绑定不别名。
  • TIED_ONLY：绑定别名；未绑定被拒绝。
  • UNTIED_ONLY：权重保持分离；绑定被拒绝。

不要绑定具有有意分离的头部、不对称 vocab 大小或不拥有两个张量的阶段的架构。

对于 from_pretrained，checkpoint 中保存的 tie_word_embeddings 值是权威的，即使对于 BOTH 也是如此。NeMoAuto* 桥接器在任一方向都拒绝翻转。模型拥有的绕过该桥接器的 from_pretrained 必须调用 reject_tie_word_embeddings_flip(checkpoint_config, requested_config, model_class_name)。

2.4 MoE state-dict 适配器检查清单

对于 MoE 模型，不要停留在通用加载。适配器必须显式映射：

  • Router 权重，当 Hugging Face 模型有 gate bias 或 correction-bias 张量时。
  • 专家权重，保留本地和路由专家的专家索引顺序。
  • Gate/up/down 投影，包括组合或拆分投影布局。
  • Shared experts 与 routed experts 分开，当架构两者都有时。

添加断言预期键映射的测试，并使用微配置在完整 checkpoint 前运行数值等效测试。

不要使用这些捷径：

  • 不要仅通过调用 from_pretrained() 验证适配器。
  • 不要在没有明确映射原因的情况下接受缺失或多余的专家键。
  • 不要更改 dtype、转置维度或 reshape 张量，除非 HF 和 NeMo 布局需要且测试证明转换可逆。
  • 不要因为稠密层测试通过而跳过 router 或 shared-expert 测试。

2.5 VLM 接入检查清单

对于 VLM，确认 Hugging Face config 具有 vision_config 和 text_config，并且 architectures 指向条件生成类。从最接近的 VLM 模式文件开始，通常是 vlm-patterns.md，并比较现有实现，例如 mistral4、kimivl 或 kimi_k25_vl。

实现应显式涵盖：

  • 文本骨干、视觉塔、projector 和 processor 或图像预处理假设。
  • state_dict_adapter.py 中文本和视觉模块的权重映射。
  • 在 _transformers/registry.py 中注册 ForConditionalGeneration 类。
  • 使用图像-文本输入并验证适配器往返的微测试。

2.6 在注册表中注册

在 _transformers/registry.py 的 MODEL_ARCH_MAPPING 中添加模型：

# 在 _transformers/registry.py 中
MODEL_ARCH_MAPPING = OrderedDict([
    # ... 现有条目 ...
    (
        "NewModelForCausalLM",
        ("nemo_automodel.components.models.new_model.model", "NewModelForCausalLM"),
    ),
])

如果模型具有自定义配置类且在其 config.json 中有 auto_map，则也在 _CUSTOM_CONFIG_REGISTRATIONS 中注册：

_CUSTOM_CONFIG_REGISTRATIONS: Dict[str, Tuple[str, str]] = {
    # ... 现有条目 ...
    "new_model": ("nemo_automodel.components.models.new_model.configuration", "NewModelConfig"),
}

2.7 声明能力和精度敏感参数

MODEL_ARCH_MAPPING 中注册的每个类必须声明并行能力，可以使用静态嵌套 ModelCapabilities dataclass 或变体感知的 get_capabilities(cls, config) 方法。选择恰好一种模式。能力应反映已经端到端验证的 recipe YAML。

如果模型有精度敏感参数，如 Mamba A_log / dt_bias、MoE sigmoid gate bias、attention-sink bias 或每头 scale，则声明 _keep_in_fp32_modules_strict，以便分片将这些参数保持在 fp32 计算中。参见 capabilities-and-precision.md 获取示例、变体分派规则和冻结子模块 dtype 指导。

Phase 3: 接入示例配置

本阶段仅为添加一个最小示例配置，证明新接入的架构可以加载和运行。对于通用 recipe 编写或现有 recipe 修改，使用 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

# ... 数据、优化器配置 ...

3.2 验证模型加载

测试模型能从 HuggingFace checkpoint 加载：

from nemo_automodel import NeMoAutoModelForCausalLM

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

3.3 先使用微配置测试

在使用完整模型之前，使用微配置（1-2 层、小隐藏维度）验证，以尽早发现形状不匹配。

Phase 4: 测试

在加载完整 checkpoint 之前，创建 tests/unit_tests/models/<name>/ 并覆盖以下检查：

  • 使用微配置的前向形状冒烟测试。
  • state-dict 适配器往返：from_hf -> to_hf 保留映射名称、形状、dtype 和值。
  • 对所有重写的注意力、MLP、归一化、RoPE 或 MoE 层的层等效测试。使用配置中的模型 dtype、相同的播种权重、相同的输入和 dtype 适当的 torch.allclose 容差。
  • 验证训练几步后损失下降的简短功能测试。

Phase 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）以及任何限制。

Phase 6: 奇偶校验测试

在实现和单元测试完成后，运行完整的奇偶校验工作流，以验证新模型产生的数值结果与参考 HuggingFace 实现一致。

运行三个级别的比较：

  1. state-dict 往返：加载参考 HuggingFace checkpoint，转换为 NeMo AutoModel 布局，导出回去，并验证所有映射张量在预期容差内匹配参考名称、形状、dtype 和值。
  2. 组件级奇偶校验：在固定种子和相同 dtype 下，将重写的注意力、MLP、归一化、RoPE 和 MoE 组件与 HuggingFace 实现进行比较。
  3. 端到端前向传递：在相同的 tokenized 输入上运行完整的 NeMo AutoModel 和 HuggingFace 模型，比较 logits、隐藏状态和 loss。

不要跳过此阶段。通过单元测试的模型仍可能因微妙的权重转换错误、后端差异或仅在完整奇偶校验比较中出现的 RoPE 不匹配而与 HF 发散。

关键文件参考

文件 用途
_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 CombinedQKVAttentionMixin，包含 setup_qkv_projection() 和 compute_qkv()
components/models/common/combined_projection/combined_mlp.py CombinedGateUpMLP，具有交错 gate/up 布局
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 dataclass
components/moe/fsdp_mixin.py MoEFSDPSyncMixin，用于分布式专家处理
components/moe/layers.py MoE 层，MoE 块的 MLP（稠密）
components/moe/experts.py GroupedExperts、GroupedExpertsDeepEP、GroupedExpertsTE

检查清单

  • [ ] 获取并分析来自 HuggingFace 的 config.json
  • [ ] 确定模型类型（稠密 LLM / MoE / VLM）
  • [ ] 识别自定义组件（注意力、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 中注册自定义配置（如适用）
  • [ ] 声明 ModelCapabilities 嵌套 dataclass（静态）或 get_capabilities(cls, config) classmethod（变体分派，例如 ERNIE-4.5 MoE 与稠密）——两者不能同时，也不能都无
  • [ ] 为每个带有因果 lm_head 的类声明 TieSupport 并调用构造函数守卫（或添加显式无头豁免）——参见 §2.3
  • [ ] 为 BOTH / TIED_ONLY 添加显式 _tied_weights_keys 和 tie_weights()，并添加策略特定的别名和拒绝测试——参见 §2.3
  • [ ] 如果任何绕过 NeMoAuto* 桥接器的模型自有 from_pretrained 防止 checkpoint 翻转，则添加守卫——参见 §2.3
  • [ ] 创建示例 YAML 配置
  • [ ] 通过 NeMoAutoModelForCausalLM.from_pretrained() 验证模型加载
  • [ ] 创建单元测试（前向形状、state_dict 往返）
  • [ ] 为每个内部 fp32 参数声明 _keep_in_fp32_modules_strict（SSM A_log/dt_bias、Mamba D（当参考为 fp32 时）、MoE gate bias、attention-sink bias、scale 等）——参见 §2.7
  • [ ] 为每个重写的层创建层等效测试（匹配模型 dtype）
  • [ ] 创建功能测试（训练 loss 下降）
  • [ ] 更新 docs/model-coverage 页面
  • [ ] 运行 state-dict 往返、组件奇偶校验和 E2E 前向传递奇偶校验
  • [ ] 在模块底部设置 ModelClass = <Name>ForCausalLM