训练循环的关键实现#

本页内容5 个部分

1. 一次有效更新的约定#

这份实现把前面三篇的规则放进同一训练循环：先按实际样本数归一化整个累积窗口，再完成反向、解除缩放、裁剪与更新，最后推进学习率。代码是单进程 CUDA 的分类示例，接口按 PyTorch 2.9 编写；model、criterion、optimizer、scheduler 与 dataloader 由调用者提供。

criterion 返回当前 batch 的样本平均损失，数据为 CPU 上的 (images, labels)。每个样本权重相同；若类别权重、mask 或有效 token 改变分母，应使用对应的损失和与有效权重总数。窗口只把当前 microbatch 搬到 GPU。

2. AMP、尾部累积与裁剪放在一起#

from itertools import islice
import torch

accum_steps = 4
use_fp16 = True
amp_dtype = torch.float16 if use_fp16 else torch.bfloat16
scaler = torch.amp.GradScaler("cuda", enabled=use_fp16)
model.train()
batches = iter(dataloader)

while window := list(islice(batches, accum_steps)):
    total_samples = sum(len(images) for images, _ in window)
    if total_samples == 0:
        continue
    optimizer.zero_grad(set_to_none=True)

    for images, labels in window:
        batch_size = len(images)
        if batch_size == 0:
            continue
        images, labels = images.cuda(), labels.cuda()
        with torch.autocast("cuda", dtype=amp_dtype):
            mean_loss = criterion(model(images), labels)
            weighted_loss = mean_loss * batch_size / total_samples
        scaler.scale(weighted_loss).backward()

    scaler.unscale_(optimizer)
    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    old_scale = scaler.get_scale()
    scaler.step(optimizer)
    scaler.update()
    # 对标准单优化器 GradScaler，溢出会跳步并减小 scale。
    updated = not use_fp16 or scaler.get_scale() >= old_scale
    if updated:
        scheduler.step()

尾部窗口只有 2 批时，仍执行一次更新，且使用这 2 批的实际样本总数。所有反向完成后只 unscale 一次，grad_norm 是裁剪前范数。BF16 分支关闭 scaler 的缩放与溢出跳步功能，仍需在实际训练中监控 loss 和梯度的有限性。

用 scale 判断跳步依赖这里的单优化器、标准 scaler 更新约定；多优化器、手动指定 scale 或自定义 scaler 要显式跟踪各优化器是否成功更新。这里的 scheduler 按成功优化步计数；若预算按已消费 token 计数，应另行维护 token 日程。

3. 核对权重与等价条件#

手算窗口有两批，样本数为 3 与 1，平均损失为 2 与 6。窗口损失为 $(3\times2+1\times6)/4=3$ ；直接平均两份 mean loss 会得到 4，改变目标。两份梯度也必须乘 $3/4$ 与 $1/4$ 。

梯度累积与大 batch 的等价条件见梯度管理。本实现没有 DDP；加入分布式同步时，还需处理各 rank 的有效分母与 no_sync() 范围。不能仅把模型包装为 DDP 就默认 token 权重已经正确。

4. Warmup 与 cosine 的位置函数#

下列纯函数使用连续位置 $t\in[0,T]$ ，便于核对起点、峰值与终点；它不是另一个训练循环。

import math

def learning_rate(t, total, warmup, peak, minimum=0.0):
    if not 0 <= warmup < total:
        raise ValueError("require 0 <= warmup < total")
    if not 0 <= minimum <= peak:
        raise ValueError("require 0 <= minimum <= peak")
    t = min(max(t, 0), total)
    if warmup and t < warmup:
        return peak * t / warmup
    progress = (t - warmup) / (total - warmup)
    return minimum + (peak - minimum) * (1 + math.cos(math.pi * progress)) / 2

取 $T=100000$ 、 $T_w=2000$ ，函数在 $t=0$ 返回零、 $t=2000$ 返回峰值、 $t=51000$ 返回峰值与最低值的均值、 $t=T$ 返回最低值。训练中的离散 scheduler 还要明确首次参数更新采用哪个位置，并把其状态连同 optimizer 和 scaler 一起保存，恢复后继续核对。

参考资料#

PyTorch 2.9 AMP examples：梯度累积、unscale 与更新边界。
PyTorch 调度器调用约定：optimizer 与 scheduler 的调用顺序。