nemo-automodel-model-onboarding

向 NeMo AutoModel 添加模型支持

目的

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

说明

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

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

对于概念性入门问题，请依据本技能作答；除非用户要求编辑代码，否则不要打开模式文件。提及相关模式文件名作为参考，然后直接给出检查清单。

使用直接的行动动词：对模型进行分类、列出文件、映射权重、注册类并添加测试。除非用户明确将分布式策略、启动器配置或通用配方编写与新架构入门联系起来，否则不要讨论这些内容。

示例

常见问题可使用以下简洁的回答模式：

• Dense 因果 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 状态字典：识别 config.json 中的专家字段，参考 moe-patterns.md，单独映射路由器张量，保留路由专家的索引顺序，映射路由专家、共享专家以及 gate/up/down 投影；添加适配器键映射测试和小型配置数值等效性测试；不要只依赖 from_pretrained() 或静默重塑张量。
• VLM 入门：仅当存在 vision_config、text_config 和 ForConditionalGeneration 架构时，才将其归类为 VLM。参考 vlm-patterns.md 和现有 VLM 实现，例如 mistral4、kimivl 或 kimi_k25_vl；检查文本骨干网络、视觉塔、投影器、处理器假设、文本和视觉 state_dict_adapter.py 映射、注册表注册情况，以及完整检查点使用前的小型图文测试。 不要将 VLM 入门当作纯因果 LM 路径，也不要跳过处理器和图像测试。

对于 MoE 状态字典问题，务必包含以下安全检查清单：

• 将路由器张量与专家张量分开映射。
• 保留路由专家的索引顺序；绝不要对专家权重排序、丢弃或合并，也不要为了让加载通过而静默重塑权重。
• 明确映射 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 确定模型类型

类型	指标	模式文件
稠密 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 平方、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 实现顺序

按依赖顺序实现文件：

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

详细实现指南请参阅模式文件：

• 稠密 LLM：llm-patterns.md
• MoE：moe-patterns.md
• VLM：vlm-patterns.md

2.3 MoE 状态字典适配器检查清单

对于 MoE 模型，不要止步于通用加载。适配器必须明确映射：

• 路由器权重，包括 Hugging Face 模型中存在的门控偏置或校正偏置张量。
• 专家权重，在本地专家和路由专家之间保留专家索引顺序。
• 门控／上投影／下投影，包括合并或拆分的投影布局。
• 如果架构同时包含共享专家和路由专家，需分别处理。

添加测试以断言预期的键映射，并先使用微型配置运行数值等价性测试，再尝试加载完整检查点。

不要使用以下捷径：

• 不要只通过调用 from_pretrained() 来验证适配器。
• 除非有明确的映射理由，否则不要接受缺失或多余的专家键。
• 除非 HF 和 NeMo 的布局要求这样做，且测试证明转换可逆，否则不要更改 dtype、转置维度或重塑张量。
• 不要因为密集层测试通过，就跳过路由器或共享专家测试。

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>/，并在加载完整检查点之前覆盖以下检查：

• 使用微型配置进行前向传播形状冒烟测试。
• 状态字典适配器往返测试：from_hf -> to_hf 保留映射后的名称、形状、dtype 和数值。
• 对每个重写的注意力、MLP、归一化、RoPE 或 MoE 层进行层等价性测试。使用配置中的模型 dtype、相同的种子权重、相同的输入，以及适合该 dtype 的 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 参考实现等效。

运行以下三个层级的比较：

1. State-dict 往返：加载 HuggingFace 参考检查点，将其转换为 NeMo AutoModel 布局，再导出并转换回来，验证所有映射后的张量在名称、形状、数据类型和值上是否与参考实现匹配，且误差在预期容差范围内。
2. 组件级一致性：使用固定随机种子和相同的数据类型，将重写后的注意力、MLP、归一化、RoPE 和 MoE 组件与 HuggingFace 实现进行比较。
3. 端到端前向传播：对相同的分词后输入，分别运行完整的 NeMo AutoModel 和 HuggingFace 模型，并比较 logits、隐藏状态和损失。

切勿跳过此阶段。即使模型通过了单元测试，也仍可能因细微的权重转换错误、后端差异或 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	提供 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（稠密）
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 中注册自定义配置（如适用）
• 创建示例 YAML 配置
• 验证模型可通过 NeMoAutoModelForCausalLM.from_pretrained() 加载
• 创建单元测试（前向传播形状、state_dict 往返）
• 为每个重写的层创建层等效性测试（匹配模型数据类型）
• 创建功能测试（训练损失下降）
• 更新 docs/model-coverage 页面
• 运行 state-dict 往返、组件一致性和 E2E 前向传播一致性检查
• 在模块底部设置 ModelClass = <Name>ForCausalLM