训练资源工程：Batch、梯度累积、混合精度与 Checkpoint

问题与目标

训练 OOM 不是简单地“显卡太小”。权重、梯度、优化器状态、激活和临时缓冲区拥有不同生命周期。本篇建立资源排查顺序，并说明 DeepSpeed、ZeRO 与 Unsloth 应在什么时候引入。

核心概念

训练显存可以粗略拆为：

峰值显存 ≈ 权重 + 梯度 + 优化器状态 + 激活 + 临时缓冲区

全参 AdamW 中，梯度和优化器状态往往比权重更贵；LoRA 显著缩小前三者中的可训练部分，但激活仍随 Batch、序列长度和层数增长。

调整顺序建议是：缩短异常长样本、减小单卡 Batch、梯度累积、混合精度、梯度检查点、LoRA/QLoRA，最后再扩展到多卡与状态切分。

可运行实现

下面计算有效 Batch，并在实验开始前拒绝不一致配置：

from dataclasses import dataclass


@dataclass(frozen=True)
class TrainPlan:
    per_device_batch: int
    accumulation_steps: int
    world_size: int
    max_length: int

    @property
    def effective_batch(self) -> int:
        return self.per_device_batch * self.accumulation_steps * self.world_size


def compare(old: TrainPlan, new: TrainPlan) -> None:
    print({"old_batch": old.effective_batch, "new_batch": new.effective_batch})
    if old.effective_batch != new.effective_batch:
        print("警告：有效 Batch 改变，学习率和实验结论需要重新检查")
    if new.max_length > old.max_length:
        print("警告：序列变长会显著增加激活显存")


compare(
    TrainPlan(4, 1, 1, 1024),
    TrainPlan(1, 4, 1, 1024),
)

工具定位：

• DeepSpeed 是训练运行与分布式优化系统；
• ZeRO 按阶段切分优化器状态、梯度和参数；
• Unsloth 提供面向部分模型与硬件的优化训练路径；
• Accelerate 负责设备和分布式启动抽象。

不要在同一次实验同时更换数据、模型、精度、框架和并行策略。先用单卡小模型验证正确性，再扩大资源。

常见问题与排查

Batch 减小后效果突然变化

检查有效 Batch 是否保持一致，学习率、Warmup 和梯度噪声是否需要调整。

开启梯度检查点后训练变慢

它通过反向时重新计算激活换取显存，速度下降是预期权衡。记录峰值显存与每步时间。

ZeRO 阶段越高越好

切分越多，通信和复杂度通常越高。根据模型状态是否放不下选择阶段，而不是追求最大数字。

Checkpoint 能保存但无法恢复

验证模型、优化器、调度器、随机状态和数据位置是否一同恢复。训练前做一次短程中断—恢复测试。

小结

训练资源优化应从显存组成与可复现实验出发。先控制序列、Batch 和精度，再选择 PEFT、DeepSpeed 或 Unsloth，才能知道性能变化来自哪里。