nemo-automodel-recipe-development

NeMo AutoModel 配方开发

说明

对于配方相关问题，请提供最简洁且完整的操作路径：

1. 指出相关配方文件或 YAML 部分。
2. 列出涉及的构建器函数或配置键。
3. 如果问题询问如何配置，请提供最简 YAML 或命令示例。
4. 最后附上本地验证命令或简短的 CPU 兼容测试。

对于概念性配方问题，请依据此技能作答，不要检查代码仓库或加载其他 AutoModel 技能，除非用户要求你编辑文件。回答应聚焦于配方 YAML、构建器、CLI 路由、测试和本地验证。

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

• 新建微调配方变体：从 nemo_automodel/recipes/ 下最接近的文件开始，更新模型、数据集或数据加载器、优化器、损失函数、LR 调度器、步数调度器和检查点构建器；只有在添加新的配方类时才注册配方别名；在 examples/ 下添加示例 YAML；然后添加一个简短的 CPU 兼容单元测试，并运行 automodel <config.yaml>。
• _target_ 字段：将 _target_ 说明为完全限定的 Python 可调用对象，并解释同级键会成为关键字参数；展示优化器和数据集示例，并提及诸如 --optimizer.lr 这样的嵌套 CLI 覆盖。
• 验证和检查点：指出 step_scheduler.val_check_interval、step_scheduler.checkpoint_interval、validation_dataset、restore_from.path 和整合后的 safetensors；并附上此技能中的最简 YAML 片段。

对于验证和检查点，始终指出：

• step_scheduler.val_check_interval：验证频率。
• step_scheduler.checkpoint_interval：保存频率。
• validation_dataset：验证数据加载器来源。
• restore_from.path：恢复训练所用的路径。
• 整合后的 safetensors：与 HF 生态系统兼容的默认检查点格式。

路由边界

本技能适用于配方构建和执行流程相关问题：YAML 结构、_target_ 可调用对象、构建器函数、验证数据集、检查点配置、CLI 路由注册以及配方专用测试。

除非用户询问这些选项如何出现在 AutoModel 配方 YAML 中，否则不要将本技能用于独立的分布式策略选择、集群启动器配置或模型架构接入。

配方架构

执行流程

CLI (automodel config.yaml)
  -> app.py resolves the config's recipe target
    -> recipe script (e.g. train_ft.py) main(config_path)
      -> Recipe class .setup() builds all components
        -> .run_train_validation_loop() executes training

配方类

配方继承自 BaseRecipe，并实现两个方法：

• setup() —— 通过构建器函数构建模型、优化器、数据加载器、损失函数、LR 调度器、步数调度器和检查点配置。
• run_train_validation_loop() —— 执行训练和验证循环。

构建器模式

所有组件都通过专用构建器函数创建：

• build_model() —— 根据配置实例化模型
• build_optimizer() —— 创建优化器（AdamW 等）
• build_dataloader() —— 设置训练和验证数据加载器
• build_loss_module() —— 创建损失函数
• build_lr_scheduler() —— 创建学习率调度器
• build_step_scheduler() —— 创建控制训练进程的步数调度器
• CheckpointingConfig —— 配置检查点（直接根据 YAML 的 checkpoint: 部分构建，通过 RecipeConfig.checkpoint 获取）

基础设施应用顺序

构建后，组件会严格按以下顺序应用：

1. PEFT（LoRA 等）
2. FP8 量化
3. QAT（量化感知训练）
4. 加载/恢复检查点
5. 冻结参数
6. 分片（FSDP2、Megatron-FSDP、DDP）
7. 设备放置
8. torch.compile
9. 上下文并行钩子

YAML 配置结构

完整的配方配置遵循以下结构：

step_scheduler:
  max_steps: 1000
  num_epochs: 1
  grad_accumulation_steps: 4
  val_check_interval: 100
  checkpoint_interval: 500
  log_interval: 10

dist_env:
  master_addr: localhost
  master_port: 29500

rng:
  seed: 42

model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
  dtype: float32
  # additional model kwargs passed to the constructor

compile:
  enabled: false
  backend: inductor

clip_grad_norm:
  max_norm: 1.0

distributed:
  strategy: fsdp2       # fsdp2 | megatron_fsdp | ddp
  dp_size: auto
  tp_size: 1
  cp_size: 1

loss_fn:
  _target_: torch.nn.CrossEntropyLoss

dataset:
  _target_: nemo_automodel.datasets.squad.SquadDataset
  tokenizer_name_or_path: meta-llama/Llama-3.2-1B
  max_seq_length: 2048

validation_dataset:
  _target_: nemo_automodel.datasets.squad.SquadDataset
  split: validation

packed_sequence:
  enabled: false

dataloader:
  batch_size: 4
  num_workers: 4
  pin_memory: true

optimizer:
  _target_: torch.optim.AdamW
  lr: 2.0e-5
  weight_decay: 0.01

lr_scheduler:
  _target_: nemo_automodel.schedulers.CosineAnnealingWarmup
  warmup_steps: 50
  min_lr: 1.0e-6

全参数训练精度

使用 torch.optim.Adam/AdamW 进行新的全参数训练时，请在 NeMoAutoModel 加载器上显式设置 model.dtype: float32，以使用 fp32 主权重和 Adam 动量。单独配置计算精度（FSDP2：distributed.mp_policy）。

PEFT、TE FusedAdam 和扩散需要单独选择精度；请参阅 混合精度指南。迁移现有配置时，请验证内存、训练行为以及检查点恢复。

_target_ 模式

_target_ 键指定一个完全限定的 Python 可调用对象。该部分中的所有其余键都会作为关键字参数传入：

optimizer:
  _target_: torch.optim.AdamW   # callable
  lr: 2.0e-5                    # kwarg
  weight_decay: 0.01            # kwarg

这等同于：torch.optim.AdamW(lr=2e-5, weight_decay=0.01)。

CLI 覆盖

可以通过命令行覆盖任何配置值：

automodel config.yaml \
  --optimizer.lr 1e-4 \
  --step_scheduler.max_steps 500 \
  --distributed.tp_size 2

示例

验证和检查点：

step_scheduler:
  val_check_interval: 100
  checkpoint_interval: 500

validation_dataset:
  _target_: nemo_automodel.datasets.squad.SquadDataset
  split: validation

restore_from:
  path: /checkpoints/step-500

特定领域说明

LLM

• nemo_automodel/recipes/llm/train_ft.py 同时处理微调和预训练。二者的区别在于配置（数据集、学习率等）。
• nemo_automodel/recipes/llm/kd.py 实现教师模型与学生模型的知识蒸馏。
• nemo_automodel/recipes/llm/benchmark.py 执行吞吐量和延迟基准测试。

VLM

• 使用 NeMoAutoModelForImageTextToText，而不是因果语言模型类。
• 配置中包含 processor 部分，而不是单独的分词器。
• 配方位于 nemo_automodel/recipes/vlm/finetune.py。

Diffusion

• 使用 NeMoAutoDiffusionPipeline。
• 配置中需要包含 parallel_scheme 字典以定义并行方式。
• 仅支持 DDP 和 FSDP2 策略（不支持 Megatron-FSDP）。
• 配方位于 nemo_automodel/recipes/diffusion/train.py。

检索

• 两种编码器模式：
  • 双编码器（nemo_automodel/recipes/retrieval/train_bi_encoder.py）：使用独立的查询编码器和文档编码器，并采用对比损失。
  • 交叉编码器（nemo_automodel/recipes/retrieval/train_cross_encoder.py）：联合编码，并使用分类头。
• 困难负样本挖掘：nemo_automodel/recipes/retrieval/mine_hard_negatives.py。

训练循环详情

每个 epoch 的训练循环遵循以下结构：

for epoch in range(num_epochs):
    for batch_idx in range(batches_per_epoch):
        # --- gradient accumulation inner loop ---
        for micro_batch in micro_batches:
            if pipeline_parallel:
                schedule.step(micro_batch)    # PP schedule
            else:
                loss = model(micro_batch)     # direct forward
                loss.backward()

        # --- optimizer step ---
        scale_grads_and_clip_grad_norm(model, max_norm)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

        # --- logging ---
        MetricsSample(step, epoch, loss, grad_norm, lr, mem, tps, mfu)

        # --- validation (at configured intervals) ---
        if step % val_check_interval == 0:
            run_validation()

        # --- checkpoint (at configured intervals) ---
        if step % checkpoint_interval == 0:
            save_checkpoint()

StepScheduler

控制所有训练进度：总 epoch 数、总步数、梯度累积步数、验证间隔、检查点间隔和日志记录间隔。

梯度裁剪

在反向传播之后、优化器步骤之前，通过 scale_grads_and_clip_grad_norm() 执行。由配置中的 clip_grad_norm.max_norm 控制。

上下文并行

当 cp_size > 1 时，使用 make_cp_batch_and_ctx() 将批次拆分到上下文并行组中。此操作必须在前向传播之前完成。

MetricsSample

每个训练步骤都会生成一个包含以下字段的 MetricsSample：

• step -- 全局步数
• epoch -- 当前 epoch
• loss -- 训练损失
• grad_norm -- 裁剪后的梯度范数
• lr -- 当前学习率
• mem -- GPU 内存使用量
• tps -- 每秒处理的词元数
• mfu -- 模型 FLOPS 利用率

验证与检查点

验证

• 按照 step_scheduler.val_check_interval 定义的间隔运行。
• 使用根据 validation_dataset 配置构建的验证数据加载器。
• 模型设为评估模式；禁用梯度。

检查点

• 默认格式：合并后的 safetensors，便于在 HF 生态系统中部署（始终优先使用这种格式，而不是 DCP）。
• 检查点间隔由 step_scheduler.checkpoint_interval 控制。
• 通过 restore_from 配置键指定检查点目录，以恢复训练。

restore_from:
  path: /checkpoints/step-500

常见问题

问题	原因	修复方法
配置错误未提示	_target_ 值拼写错误	类路径必须是有效且可导入的 Python 可调用对象。请仔细检查模块路径和类名。
训练在第一步崩溃	global_batch_size 不能被 local_batch_size * dp_size * grad_accumulation_steps 整除	确保各维度的批次大小计算一致。
无法通过 CLI 访问新配方	配置中缺少可解析的 recipe 目标	将配置的 recipe 键设为可发现的配方类名，或完整的带点号 _target_ 路径。
前向传递时形状不匹配	数据集整理函数的输出与模型输入签名不匹配	确认整理函数返回的张量具有模型所需的键和形状。
验证期间发生 OOM	验证批次大小过大，或未禁用梯度	使用 torch.no_grad() 包裹验证过程，并考虑减小验证批次大小。
检查点恢复失败	检查点与配置中的模型架构不匹配	确保模型配置与检查点完全一致（层数、隐藏维度、词汇表大小）。