nemo-automodel-model-onboarding

为 NeMo AutoModel 添加模型支持

目的

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

说明

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

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

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

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

示例

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

• 稠密因果语言模型：仅当 architectures 中包含 ForCausalLM 类，且不存在 num_local_experts、n_routed_experts 或 num_experts_per_tok 等专家字段时，才归类为纯稠密模型。创建 components/models/<name>/model.py 和 __init__.py；仅在需要转换检查点权重时添加 state_dict_adapter.py，仅在需要时添加 config.py。在 _transformers/registry.py 中注册 MODEL_ARCH_MAPPING，添加示例 YAML，并添加小型配置单元测试，以及针对重写层的等效性测试。
• MoE 状态字典：识别 config.json 中的专家字段，并参考 moe-patterns.md；单独映射路由器张量，保留路由专家的索引顺序，映射路由专家、共享专家以及门控／上投影／下投影；添加适配器键映射测试和小型配置数值等效性测试；不要只依赖 from_pretrained()，也不要静默重塑张量。
• VLM 入门：仅当存在 vision_config、text_config 和 ForConditionalGeneration 架构时，才归类为 VLM。参考 vlm-patterns.md 和现有 VLM 实现（例如 mistral4、kimivl 或 kimi_k25_vl）；检查文本主干、视觉塔、投影器、处理器假设、文本和视觉检查点兼容性（必要时添加适配器映射）、注册信息，以及使用完整检查点前的小型图文测试。 不要将 VLM 入门当作纯因果语言模型流程，也不要跳过处理器和图像测试。

对于 MoE 状态字典和 VLM 问题，请应用第 2.4 节和第 2.5 节中的检查清单。

适用范围边界

仅当用户要添加或修改模型架构支持时使用本技能：模型文件、自定义层、状态字典适配器、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 —— 每个受支持检查点中保存的设置；不要根据空白配置构造函数推断此值
• 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 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 路由
• 自定义配置类：仅当 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 # 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 实现顺序

按依赖顺序实现文件：

1. config.py（如需要）——自定义 PretrainedConfig 子类
2. rope_utils.py（如需要）——RoPE 实现
3. layers.py（如需要）——注意力、MLP、解码器块类
4. model.py ——主要的 ForCausalLM（或 ForConditionalGeneration）类
5. state_dict_adapter.py（如需要）——HF 权重转换。将检查点 I/O 性能视为实现的一部分，并按照第 2.6 节所述评估低内存 DCP 能力。
6. init.py ——重新导出主模型类

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

• Dense LLM：llm-patterns.md
• MoE：moe-patterns.md
• VLM：vlm-patterns.md
• 能力和 fp32 精度：capabilities-and-precision.md

大多数自定义模型都需要 state_dict_adapter.py 来转换 HF 权重。 只有当受支持的后端／配置变体中的 HF 名称和张量布局已经一致时，才省略此文件和属性，例如 Llama、Qwen2 和 Qwen3。 权重绑定仍由模型负责（第 2.3 节）。

2.3 因果 LM 权重绑定

每个注册的、带因果 lm_head 的模型类都必须：

• 声明 tie_word_embeddings_support: TieSupport，取值为 BOTH、TIED_ONLY 或 UNTIED_ONLY。
• 在 __init__ 开头调用 reject_unsupported_tie_word_embeddings(type(self), config)，使用尚未解包 text_config 或 thinker_config 的原始 config。

只有不带因果 LM head 的类才能明确豁免注册表测试。

应根据实现和实际支持的 checkpoint config 选择策略，而不是依据一个基础 config 构造函数：

• BOTH：同时支持绑定和非绑定配置。
• TIED_ONLY：仅支持绑定配置。
• UNTIED_ONLY：仅支持非绑定配置。

运行时辅助函数必须将 TIED_ONLY 和 UNTIED_ONLY 视为权威策略，仅在 BOTH 时解析每个 checkpoint 的 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：权重保持分离；绑定时拒绝。

对于有意保持 head 分离、词表大小不对称或不同时拥有这两个张量的阶段，不要绑定权重。

对于 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 模型，检查以下所有权重。当其 HF 布局与原生布局不同时，适配器必须显式映射：

• 路由器权重，包括 Hugging Face 模型具有的 gate 偏置或 correction-bias 张量。
• 专家权重，并保留本地专家和路由专家之间的专家索引顺序。
• Gate/up/down 投影，包括组合式或拆分式投影布局。
• 当架构同时包含共享专家和路由专家时，分别处理两者。

添加测试，断言预期的键映射，并在尝试完整 checkpoint 前，使用小型 config 运行数值等价性测试。

对于 MoE PEFT，保留预训练路由器 correction bias，并使用路由器偏置检查验证适配器的保存、重新加载和恢复。

不要使用以下捷径：

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

2.5 VLM 接入检查清单

对于 VLM，确认 Hugging Face config 包含 vision_config 和 text_config，且 architectures 指向一个条件生成类。从最接近的 VLM 模式文件入手，通常是 vlm-patterns.md，并对照 mistral4、kimivl 或 kimi_k25_vl 等现有实现。

实现应明确涵盖：

• 文本主干、视觉塔、投影器，以及处理器或图像预处理假设。
• 文本和视觉模块的 checkpoint 兼容性，并在需要时提供适配器映射。
• 在 _transformers/registry.py 中注册 ForConditionalGeneration 类。
• 使用微型测试处理图文输入并验证 checkpoint 加载/导出；若存在适配器，还要进行往返测试。

2.6 检查点 I/O 性能

将检查点性能视为实现要求，而不是后续优化。对于每个新建或有重大变更的 state-dict 适配器，都要评估加载以及该变更影响的保存或导出路径的延迟和主机／设备峰值内存。具体来说：

• 避免创建完整检查点或模型大小的临时副本；如果张量可以直接加载到最终模型存储中，或分批进行转换，就应采用这些方式。
• 当某个 rank 只需要自己的分片时，分布式读取和转换应仅在该 rank 上进行；不要在每个 rank 上不必要地构建全局张量。
• 避免在模型大小的循环中重复合并张量、复制张量、执行完整堆垃圾回收、扫描分片或打开文件。
• 记录优化路径具有代表性的优化前后延迟和峰值内存数据，并注明模型、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"),
    ),
])

如果模型有带 auto_map 的自定义配置类（位于其 config.json 中），还要在 _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 sigmoid 门控偏置、注意力汇聚偏置或每个注意力头的 scale，则应声明 _keep_in_fp32_modules_strict，确保分片时这些参数保留 fp32 计算精度。有关示例、变体分派规则和冻结子模块的 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。有关计算精度和其他训练模式，请参阅训练精度。

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 层、较小的 hidden dim）进行验证，以便尽早发现形状不匹配问题。

阶段 4：测试

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

• 使用微型配置进行前向传播形状冒烟测试。
• 状态字典适配器往返测试（如有）：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 参考实现等价。

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

1. 状态字典往返：加载 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	导出 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

---

检查清单

• 从 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；依据第 2.6 节评估了检查点延迟和峰值内存
• 对于适配器，评估了 supports_low_memory_dcp_load；任何选择启用该功能的实现都已证明直接目标会到达模型存储、限制了分配转换，并且对不安全的变体报告 False
• 实现了带有重新导出的 init.py
• 在 _transformers/registry.py 的 MODEL_ARCH_MAPPING 中注册了模型
• 在 _CUSTOM_CONFIG_REGISTRATIONS 中注册了自定义配置（如适用）
• 声明了嵌套的 ModelCapabilities 数据类（静态），或 get_capabilities(cls, config) 类方法（变体分发，例如 ERNIE-4.5 MoE 与密集型）——不得同时声明，也不得两者皆不声明
• 为每个带有因果 lm_head 的类声明了 TieSupport 并调用构造函数防护（或明确豁免无头模型）——参见 §2.3
• 为 BOTH / TIED_ONLY 明确添加了 _tied_weights_keys 和 tie_weights()，并添加了针对特定策略的别名与拒绝测试——参见 §2.3
• 对于任何绕过 NeMoAuto* 桥接的模型自有 from_pretrained，已添加防护以避免检查点切换——参见 §2.3
• 创建了示例 YAML 配置
• 验证模型可通过 NeMoAutoModelForCausalLM.from_pretrained() 加载
• 创建了单元测试（前向传播形状、state_dict 往返）
• 根据 MoE 的 LoRA 验证了 MoE PEFT 路由器冻结，以及适配器的保存 / 重新加载 / 恢复
• 为每个本质上为 fp32 的参数声明了 _keep_in_fp32_modules_strict（SSM A_log / dt_bias、参考实现为 fp32 时的 Mamba D、MoE gate bias、attention-sink bias、scale 等）——参见 §2.8
• 为每个重写的层创建了等价性测试（匹配模型 dtype）
• 创建了功能测试（训练损失下降）
• 更新了 docs/model-coverage 页面
• 运行了 state-dict 往返、组件一致性和端到端前向传播一致性检查
• 在模块底部设置了 ModelClass = <Name>ForCausalLM