NeMo AutoModel继续预训练

准备阶段

下载并启动AutoModel官方镜像

sh 复制代码

docker pull nvcr.io/nvidia/nemo-automodel:26.04

sudo docker run -it --rm --ipc=host --gpus all --network=host \
    -v /data/hmlp:/hmlp/data \
    --name nemo-automodel-test \
    nvcr.io/nvidia/nemo-automodel:26.04 \
    /bin/bash
    
    
sudo docker run -it --ipc=host --gpus=all --network=host \
    -v /data/hmlp:/hmlp/data \
    -e NCCL_P2P_DISABLE=1 \
    -e NCCL_SHM_DISABLE=1 \
    -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
    -w /hmlp/data/finetune/196d63aa-84f8-42f6-9cff-722758276518/f1a28b50-b8e5-4318-8545-5a8b57b67be8 \
    --name train-text \
    nvcr.io/nvidia/nemo-automodel:26.04 \
    /bin/bash

模型准备

shell 复制代码

mkdir -p /hmlp/data/finetune/nemo-automodel-sft-test
cd /hmlp/data/finetune/nemo-automodel-sft-test

# data: 存储数据集; output: 存储微调后的checkpoint
mkdir -p data output

# 模型使用本地Qwen2.5-0.5B-Instruct
/hmlp/data/storage/Qwen2.5-0.5B-Instruct/V1/Qwen2.5-0.5B-Instruct

文本格式数据集继续预训练

准备数据集

可以通过hfd.sh脚本下载

预处理文本格式数据集为megatron格式

shell 复制代码

uv run /opt/Automodel/tools/preprocess_megatron_dataset.py \
  --input "/hmlp/data/finetune/nemo-automodel-sft-test/data/c4_demo.jsonl" \
  --json-keys text \
  --output-prefix domain_corpus \
  --output-path /hmlp/data/finetune/nemo-automodel-sft-test/data/megatron \
  --workers 8 \
  --pretrained-model-name-or-path /hmlp/data/storage/Qwen2.5-0.5B-Instruct/V1/Qwen2.5-0.5B-Instruct \
  --append-eod

创建数据集缓存索引目录

每次启动预训练，Megatron都需要为庞大的数据集构建一套随机访问的索引（包含文档索引、样本索引和打乱索引），这一过程在TB级数据上可能耗时超过30分钟。

index_mapping_dir 参数的作用正是 "缓存" 这些复杂的索引文件。

避免重复计算：一旦指定了目录，系统会优先尝试从该路径加载已有的.npy格式索引文件。只有首次运行或数据集发生变化时，才会重新构建。
加速下次启动：这种机制尤其适用于多次启动且训练数据不变的场景，能大幅缩短后续的训练启动时间。

如下创建index_mapping_dir参数需要使用的缓存文件

shell 复制代码

mkdir -p /hmlp/data/finetune/nemo-automodel-sft-test/data/megatron/mapping_dir

准备好训练配置文件-使用文本数据集

shell 复制代码

vim train-Qwen-Qwen-2.5-0.5B-Instruct-V1-text.yaml
# 文件内容如下
recipe: TrainFinetuneRecipeForNextTokenPrediction

step_scheduler:
  global_batch_size: 32
  local_batch_size: 1
  ckpt_every_steps: 200
  val_every_steps: 100
  num_epochs: 1
  max_steps: 1000

dist_env:
  backend: nccl
  timeout_minutes: 30

rng:
  _target_: nemo_automodel.components.training.rng.StatefulRNG
  seed: 1111
  ranked: true

model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: /hmlp/data/storage/Qwen2.5-0.5B-Instruct/V1/Qwen2.5-0.5B-Instruct
  torch_dtype: bf16
  trust_remote_code: true

checkpoint:
  enabled: true
  checkpoint_dir: /hmlp/data/finetune/nemo-automodel-sft-test/output
  model_save_format: safetensors
  save_consolidated: true
  # 恢复训练时打开下面这一行
  # restore_from: LATEST

distributed:
  strategy: fsdp2
  dp_size: none
  tp_size: 1
  cp_size: 1
  pp_size: 1
  sequence_parallel: false
  activation_checkpointing: true

loss_fn:
  _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy

dataset:
  _target_: nemo_automodel.components.datasets.llm.megatron_dataset.MegatronPretraining
  paths: /hmlp/data/finetune/aa15502a-09a5-41bf-a37e-42a657dfd4f3/f1a28b50-b8e5-4318-8545-5a8b57b67be8/data/train-megatron/domain_corpus_*_text_document*
  index_mapping_dir: /hmlp/data/finetune/aa15502a-09a5-41bf-a37e-42a657dfd4f3/f1a28b50-b8e5-4318-8545-5a8b57b67be8/data/train-megatron/mapping_dir
  tokenizer:
    _target_: nemo_automodel._transformers.auto_tokenizer.NeMoAutoTokenizer.from_pretrained
    pretrained_model_name_or_path: /hmlp/data/storage/Qwen2.5-0.5B-Instruct/V1/Qwen2.5-0.5B-Instruct
    trust_remote_code: true
  seq_length: 2048
  split: "1, 0, 0"
  splits_to_build: "train"

dataloader:
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
  collate_fn: torch.utils.data.default_collate
  dataloader_type: single

validation_dataset:
  _target_: nemo_automodel.components.datasets.llm.megatron_dataset.MegatronPretraining
  paths: /hmlp/data/finetune/aa15502a-09a5-41bf-a37e-42a657dfd4f3/f1a28b50-b8e5-4318-8545-5a8b57b67be8/data/eval-megatron/domain_corpus_*_text_document*
  index_mapping_dir: /hmlp/data/finetune/aa15502a-09a5-41bf-a37e-42a657dfd4f3/f1a28b50-b8e5-4318-8545-5a8b57b67be8/data/eval-megatron/mapping_dir
  tokenizer:
    _target_: nemo_automodel._transformers.auto_tokenizer.NeMoAutoTokenizer.from_pretrained
    pretrained_model_name_or_path: /hmlp/data/storage/Qwen2.5-0.5B-Instruct/V1/Qwen2.5-0.5B-Instruct
    trust_remote_code: true
  seq_length: 2048
  split: "0, 1, 0"
  splits_to_build: "validation"

validation_dataloader:
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
  collate_fn: torch.utils.data.default_collate
  dataloader_type: single

optimizer:
  _target_: torch.optim.AdamW
  lr: 5.0e-6
  betas: [0.9, 0.95]
  eps: 1.0e-8
  weight_decay: 0.1

lr_scheduler:
  lr_decay_style: cosine
  lr_warmup_steps: 100
  min_lr: 1.0e-6

执行继续预训练

shell 复制代码

# 推荐命令
CUDA_VISIBLE_DEVICES=0,1 \
  automodel \
  --nnodes=1 \
  --nproc-per-node=2 \
  --node_rank=0 \
  --master_addr=10.10.0.1 \
  --master_port=39500 \
  train-Qwen-Qwen-2.5-0.5B-Instruct-V1-text.yaml
  
  
CUDA_VISIBLE_DEVICES=0,1 automodel --nproc-per-node=2 train-Qwen-Qwen-2.5-0.5B-Instruct-V1-text.yaml

# 不推荐只是简单测试
CUDA_VISIBLE_DEVICES=0,1 \
  uv run torchrun \
  --nnodes=1 \
  --nproc-per-node=2 \
  --node_rank=0 \
  --master_addr=10.10.0.1 \
  --master_port=29500 \
  /opt/Automodel/examples/llm_pretrain/pretrain.py \
  --config train-Qwen-Qwen-2.5-0.5B-Instruct-V1-text.yaml

# 也可以用这个
CUDA_VISIBLE_DEVICES=0,1 torchrun \
  --nnodes=1 \
  --nproc-per-node=2 \
  --node_rank=0 \
  --master_addr=10.10.0.1 \
  --master_port=29500 \ 
  -m nemo_automodel._cli.app \
  pretrain llm -c train-Qwen-Qwen-2.5-0.5B-Instruct-V1-text.yaml

nnodes：总节点数
node_rank：当前节点是第几台节点，从0开始
master_addr：主节点地址
master_port：主节点端口

支持的完整微调配置文件

yaml 复制代码

# 配方
# 选择哪个配方类来运行训练循环。
# 使用短名称（自动发现）或完整的 Python 路径：
#   recipe: nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction
recipe: TrainFinetuneRecipeForNextTokenPrediction

# 训练计划
# 控制 epoch 数量、批量大小以及保存检查点/验证的频率。
# 没有 _target_ --- 这些是配方直接读取的普通值。
step_scheduler:
  grad_acc_steps: 4       # 每次优化器步骤之前累积的微批次数。有效批量大小 = grad_acc_steps × batch_size。
  ckpt_every_steps: 10    # 每 N 个梯度步骤保存一个检查点
  val_every_steps: 10     # 每 N 个梯度步骤运行一次验证循环
  num_epochs: 1           # 对训练数据集进行多少次完整遍历

# 进程组
# 初始化 PyTorch 分布式进程组。
# 没有 _target_ --- 由配方直接使用。
# 通常不需要调整此项。
dist_env:
  backend: nccl           # 通信后端："nccl"（GPU，推荐）或 "gloo"（CPU）
  timeout_minutes: 1      # 集合通信操作的超时时间；对于初始化时间较长的大模型可增加此值

# 分布式策略
# 确定模型权重、数据和计算如何在 GPU 之间拆分。
# 没有 _target_ --- 由配方直接使用。
# 详细信息请参阅"高级主题"中的"分布式训练：TP、PP、CP 和 EP"。
distributed:
  strategy: fsdp2         # 并行策略："fsdp2"（推荐）、"megatron_fsdp" 或 "ddp"。
                          # FSDP2 在数据并行组中对参数和优化器状态进行分片。
  dp_size: null           # 数据并行组大小。null = 从 world_size ÷ (tp_size × cp_size × pp_size) 自动检测。
  tp_size: 1              # 张量并行大小：在 GPU 之间拆分权重矩阵。
                          # 如果模型无法放在单个 GPU 上，则设置为 2、4 或 8。
                          # 应能整除注意力头的数量。
  cp_size: 1              # 上下文并行大小：在 GPU 之间拆分输入序列。
                          # 对于非常长的上下文（例如 32k+ token），请增加此值。
  sequence_parallel: false # 当为 true 时，扩展 TP 以同时沿序列维度分片激活值，从而进一步节省内存。

# 随机数生成器
# _target_ → StatefulRNG：一个可保存检查点的 RNG，确保训练重启时的序列相同。
# seed 和 ranked 是传递给 StatefulRNG() 的关键字参数。
rng:
  _target_: nemo_automodel.components.training.rng.StatefulRNG
  seed: 1111              # 用于可重现性的全局随机种子
  ranked: true            # 当为 true 时，每个 GPU 等级获得一个从种子派生的唯一 RNG 流，
                          # 因此每个 GPU 的数据打乱方式不同

# 模型
# _target_ → NeMoAutoModelForCausalLM.from_pretrained：下载（或从缓存加载）预训练的 HuggingFace 模型，
# 并包装以支持 NeMo 分布式训练。接受任何 from_pretrained 关键字参数（cache_dir、torch_dtype 等）。
model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B

# PEFT（如需全参数 SFT，请删除或注释整个部分）
# _target_ → PeftConfig：一个描述哪些层获得 LoRA 适配器的数据类。
# 配方将此配置传递给 build_model()，后者将适配器附加到匹配的层上。
peft:
  _target_: nemo_automodel.components._peft.lora.PeftConfig
  target_modules: "*.proj" # 与全限定层名称匹配的 glob 模式；
                           # "*.proj" 匹配每个以 "proj" 结尾的层
  dim: 8                   # 低秩维度 r --- 控制适配器的容量。
                           # 值越大表达能力越强，但使用更多内存。
  alpha: 32                # LoRA 缩放因子：适配器输出乘以 alpha/dim。
                           # 值越高，适配器在训练期间的影响越大。
  use_triton: True         # 使用优化的 Triton 内核进行 LoRA 前向/反向传播
                           # （需要安装 triton 包）

# 检查点
# 没有 _target_ --- 由配方直接使用的普通键值组。
checkpoint:
  enabled: true            # 设置为 false 以完全跳过保存检查点
  checkpoint_dir: checkpoints/  # 输出目录。Docker 用户：请绑定挂载此路径
                                # （例如 -v $(pwd)/checkpoints:/workspace/checkpoints）
                                # 以在容器重启后保留检查点。
  model_save_format: safetensors  # "safetensors"（推荐，更快更安全）或
                                  # "torch_save"（传统的基于 pickle 的格式）
  save_consolidated: True  # 当为 true 时，将一个与 HuggingFace 兼容的单个检查点写入 model/consolidated/，
                           # 可直接由 Transformers、vLLM 等加载。需要 safetensors 格式。

# 训练数据集
# _target_ → make_squad_dataset：一个工厂函数，用于下载 SQuAD 数据集、进行标记化，并返回一个 torch Dataset。
# 要使用不同的数据集，请将 _target_ 更改为另一个工厂函数（参见数据集指南）。
dataset:
  _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
  dataset_name: rajpurkar/squad  # HuggingFace Hub 数据集 ID
  split: train                   # 使用哪个拆分（train、validation、test）

# 验证数据集
validation_dataset:
  _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
  dataset_name: rajpurkar/squad
  split: validation
  limit_dataset_samples: 64  # 将验证集限制为 64 个样本以加快评估循环；
                             # 删除此行以使用完整验证集

# 训练 DataLoader
# _target_ → StatefulDataLoader：来自 torchdata 的可保存检查点的 DataLoader，
# 在训练重启时保存和恢复迭代状态，因此恢复的运行不会重新处理已见过的批次。
dataloader:
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
  collate_fn: nemo_automodel.components.datasets.utils.default_collater
                               # 将单个样本填充并批处理为张量的函数；可替换为自定义整理函数
  batch_size: 8                # 每个 GPU 每个微批次的样本数
  shuffle: true                # 每个 epoch 是否打乱数据集

# 验证 DataLoader
validation_dataloader:
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
  collate_fn: nemo_automodel.components.datasets.utils.default_collater
  batch_size: 8

# 损失函数
# _target_ → MaskedCrossEntropy：标准的交叉熵损失，自动忽略填充 token，使其不影响梯度。
# 其他可用的损失函数（替换 _target_ 以使用）：
#   - nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy
#       沿序列维度分块计算 CE，以降低峰值内存。对于超长序列很有用。接受 chunk_len（默认 32）。
#   - nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy
#       将最终的线性投影（lm_head）与 CE 计算融合，避免生成完整的 logit 张量。
#       对于大词汇表可显著节省**内存**。
#   - nemo_automodel.components.loss.te_parallel_ce.TEParallelCrossEntropy
#       基于 TransformerEngine 的并行 CE，使用 Triton 内核。专为 logits 在 TP 等级间分片的张量并行设置设计。
loss_fn:
  _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy

# 优化器
# _target_ → torch.optim.Adam：此处可使用任何 torch.optim 类（例如 AdamW、SGD）。
# 其余所有键成为构造函数的参数。
optimizer:
  _target_: torch.optim.Adam
  lr: 1.0e-5               # 学习率 --- 最重要的可调超参数
  betas: [0.9, 0.999]      # Adam 动量系数（β₁ 用于均值，β₂ 用于方差）
  eps: 1e-8                 # 为保证数值稳定性加到分母上的小常数
  weight_decay: 0           # L2 正则化强度（0 = 无正则化）

# 日志记录（可选）
# 取消注释以启用 Weights & Biases 实验跟踪。
# wandb:
#   project: <your_wandb_project>    # W&B 项目名称
#   entity: <your_wandb_entity>      # W&B 团队或用户名
#   name: <your_wandb_exp_name>      # 本次运行的显示名称
#   save_dir: <your_wandb_save_dir>  # W&B 工件的本地目录