第 9 章 Accelerate 手动训练

学习目标

• 理解 Trainer 之外「手动训练循环」存在的意义
• 掌握 Accelerator 的设备管理与混合精度
• 掌握梯度累积的写法
• 了解 accelerate launch 与分布式训练

9.1 为什么需要手动循环

Trainer 适合标准流程;当你需要控制每一步(自定义损失、条件更新、在循环里做评估)时,就要自己写训练循环。手写循环的痛点是要处理设备、混合精度、梯度累积、多卡同步——accelerate 专门解决这些。

Accelerate 的核心是 Accelerator 对象:调用 prepare 后,你的模型、优化器、数据加载器自动被放到正确的设备、包装成正确的精度与并行方式。

9.2 第一个手动训练循环

import torch
from torch.utils.data import DataLoader
from datasets import load_dataset
from transformers import (BertConfig, BertForSequenceClassification,
                          AutoTokenizer, DataCollatorWithPadding, set_seed)
from accelerate import Accelerator

set_seed(42)
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")

def tokenize(examples):
    return tokenizer(examples["sentence"], truncation=True, max_length=64)

ds = load_dataset("nyu-mll/glue", "sst2", split="train[:200]")
ds = (ds.map(tokenize, batched=True, remove_columns=["sentence", "idx"])
        .rename_column("label", "labels"))

config = BertConfig.from_pretrained("prajjwal1/bert-tiny", num_labels=2)
model = BertForSequenceClassification.from_pretrained(
    "prajjwal1/bert-tiny", config=config)

collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="pt")
train_loader = DataLoader(ds, batch_size=16, shuffle=True, collate_fn=collator)

optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
accelerator = Accelerator(mixed_precision="fp16")
model, optimizer, train_loader = accelerator.prepare(
    model, optimizer, train_loader)

print("设备:", accelerator.device, "| 混合精度:", accelerator.mixed_precision)

model.train()
for step, batch in enumerate(train_loader):
    outputs = model(**batch)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()
    if step % 5 == 0 or step == len(train_loader) - 1:
        print(f"step {step + 1}/{len(train_loader)} loss: {loss.item():.4f}")

输出(真实运行,省略加载与 LOAD REPORT):

设备: cuda | 混合精度: fp16
step 1/13 loss: 0.6935
step 6/13 loss: 0.6646
step 11/13 loss: 0.5968
step 13/13 loss: 0.6050

对比第 7 章的 Trainer 输出:loss 数值一致地下降,但这里每一步都是你控制的。

9.3 三个关键机制

1. 设备放置

accelerator.prepare(...) 自动把模型和数据放到 accelerator.device。代码里不再出现 model.to("cuda")、batch.to(device)——加速器帮你做。

2. 混合精度

Accelerator(mixed_precision="fp16") 让前向用 fp16 计算、优化器状态保持 fp32,在几乎不损失精度的情况下省显存、提速。backward 也要走 accelerator.backward(loss),它会处理梯度缩放。

3. 梯度累积

小 batch 在显存不足时用梯度累积模拟大 batch:连续几步反向但不更新,攒够再更新一次。

gradient_accumulation_steps = 2

for step, batch in enumerate(train_loader):
    outputs = model(**batch)
    loss = outputs.loss / gradient_accumulation_steps  # 平均分摊
    accelerator.backward(loss)
    if (step + 1) % gradient_accumulation_steps == 0 or \
            step == len(train_loader) - 1:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()

要点:

• 每步的 loss 除以累积步数,保证「2 步小 batch」的梯度总和等价于「1 步大 batch」。
• 只在达到累积步数时才 optimizer.step() 与 zero_grad()。
• clip_grad_norm_ 限制梯度范数,防止梯度爆炸。

9.4 训练一个 epoch 的完整循环

标准结构:外层 epoch、内层 step、epoch 结束评估:

for epoch in range(2):
    model.train()
    for step, batch in enumerate(train_loader):
        outputs = model(**batch)
        accelerator.backward(outputs.loss)
        optimizer.step()
        optimizer.zero_grad()

    # 评估
    model.eval()
    total, correct = 0, 0
    for batch in eval_loader:
        with torch.no_grad():
            logits = model(**batch).logits
        preds = logits.argmax(dim=-1)
        correct += (preds == batch["labels"]).sum().item()
        total += batch["labels"].numel()
    print(f"epoch {epoch + 1} 准确率: {correct / total:.3f}")

注意 torch.no_grad() 与 model.eval():推理时不做反向、关闭 dropout。

9.5 从脚本到多卡:accelerate launch

Accelerator 写的代码天然支持多卡:在命令行用 accelerate launch 启动,加速器自动配置进程数:

$ accelerate launch train.py

也可以用 accelerate config 交互式生成配置文件,或用 --num_processes 指定卡数:

$ accelerate launch --num_processes=2 train.py

代码里唯一要注意的:多卡时 batch["labels"] 等张量需要先 accelerator.gather 再统计——第 12 章评估时会遇到。

动手实践

1. 把 9.2 的循环改成 3 个 epoch,并在每个 epoch 结束后打印平均 loss。
2. 实现 9.3 的梯度累积版本,对比相同 batch 数下的 loss 曲线。
3. 尝试 Accelerator(mixed_precision="bf16"),记录输出中的 loss 数值是否与 fp16 一致。

常见错误

错误 1:手动 model.to("cuda") 后又 prepare。

prepare 会覆盖设备状态,提前手动移动可能导致重复包装。让 accelerate 全权处理设备。

错误 2:忘了 accelerator.backward,直接 loss.backward()。

fp16 下梯度缩放由 accelerator.backward 管理,直接用 loss.backward() 会破坏混合精度(或数值不稳)。

错误 3:梯度累积时忘了除累积步数。

不除的话,等效学习率被放大了 N 倍,训练可能发散。

章末练习

基础

1. Accelerator.prepare 接受哪三类对象?分别做什么?
2. 写出梯度累积 4 步的标准代码片段。

提高

3. 在手动循环里加入 clip_grad_norm_(model.parameters(), 1.0),比较加与不加的 grad_norm 与 loss。
4. 用 accelerator.print 替代 print,说明多卡下它的作用。

挑战

5. 用 accelerate launch --num_processes=2 跑通 9.2 的脚本(环境只有一张卡时,可以模拟多进程),记录加速器打印的进程信息。

章末自测

1. Accelerator(mixed_precision="fp16") 中的 fp16 是什么?
2. 反向传播应该调用 loss.backward() 还是 accelerator.backward(loss)?
3. 梯度累积的目的是什么?
4. 多卡训练通常用什么命令启动脚本?
5. 判断:prepare 之后还需要手动 model.to("cuda")。
6. 评估时应该用 torch.no_grad() 吗?为什么?