快速上手：跑一个 NeMo Megatron 预训练

NeMo 2.x 引入了 nemo.collections.llm 子包和 nemo run CLI，预训练比以前简洁很多。本章在 8×A100 上预训一个 GPT-100M 作为最小可复现例子。

3.1NeMo 2.x 的"recipe" 范式

NeMo 2.x 把"全套训练配置"抽象成 recipe：一个 Python 函数返回 LightningTrainer + Model + Data + Optim 的组装。你不再写 YAML，直接调函数。

# 找到一个内置 recipe
from nemo.collections.llm.recipes import gpt3_5b
recipe = gpt3_5b.pretrain_recipe(
    dir="./pretrain_out",
    name="my_gpt_run",
    num_nodes=1, num_gpus_per_node=8,
)
print(recipe)            # 返回一个完整的训练配方

看 nemo/collections/llm/recipes/，已有 GPT-3 / Llama-2/3 / Mixtral / Mistral / Nemotron / Qwen2 等几十个 recipe 文件。

3.230 行最小预训练脚本

import nemo_run as run
import nemo.collections.llm as nl
from nemo.collections.llm.gpt.model import GPTConfig, GPTModel
from nemo.lightning import Trainer, MegatronStrategy

# 1) 模型配置（小号 GPT）
cfg = GPTConfig(
    num_layers=12, hidden_size=768, ffn_hidden_size=3072,
    num_attention_heads=12, seq_length=2048,
    init_method_std=0.02,
)
model = GPTModel(config=cfg)

# 2) 数据
data = nl.MockDataModule(
    seq_length=2048, global_batch_size=64, micro_batch_size=2,
)

# 3) Trainer
trainer = Trainer(
    devices=8, accelerator="gpu", num_nodes=1,
    strategy=MegatronStrategy(
        tensor_model_parallel_size=2,
        pipeline_model_parallel_size=1,
    ),
    max_steps=1000, val_check_interval=200, log_every_n_steps=10,
)

# 4) 训
nl.train(model=model, data=data, trainer=trainer,
         tokenizer="data", optim=nl.adam.distributed_fused_adam_with_cosine_annealing())

3.3用 nemo run 启动

NeMo 2.x 的 nemo_run 工具支持"声明 recipe + executor"两段式启动：

import nemo_run as run

recipe = gpt3_5b.pretrain_recipe(num_nodes=1, num_gpus_per_node=8, dir="./out")

# Local executor（单机）
executor = run.LocalExecutor(ntasks_per_node=8, launcher="torchrun")
run.run(recipe, executor=executor)

# 或 Slurm executor（多机）
slurm_exec = run.SlurmExecutor(
    account="my_account",
    partition="gpu",
    nodes=4, ntasks_per_node=8,
    time="06:00:00",
)
run.run(recipe, executor=slurm_exec)

3.4训练过程中看什么

# tensorboard 默认开
tensorboard --logdir ./pretrain_out/tb_logs

# 或 wandb（recipe 里默认配了）
export WANDB_API_KEY=...

观察指标：

• train_loss：100M 模型应在前 500 step 降到 ~5；
• tflops_per_gpu：A100 上应 > 100；
• tokens_per_sec_per_gpu：100M 模型应 > 50k tok/s/GPU。

3.5跟旧版 NeMo（1.x）的差别

2.x 的 Python recipe 比 YAML 灵活很多（可以编程 override）；但本书也会在 ch05 讲 Hydra YAML 路径，因为旧项目还有大量用。

3.6常见错误

3.7这章你需要带走的

• NeMo 2.x 用 Python recipe + nemo run 启动训练；
• 核心抽象：Model + Data + Trainer + MegatronStrategy；
• recipes/ 目录有数十个开箱即用配方（GPT-3 / Llama / Mixtral）；
• 本地用 LocalExecutor、集群用 SlurmExecutor。