大批量训练的梯度累积技术 | AI生成和翻译
Andrej Karpathy 对梯度累积步骤的解析
在其 YouTube 视频《Let’s build GPT: from scratch, in code, spelled out》（关于从零开始实现和训练 GPT 类模型的详细教程）中，Andrej Karpathy 在 2:51:40 左右讨论了梯度累积。他主要将其作为在 GPU 内存有限的硬件上训练大语言模型的实用解决方案，同时仍能模拟原始 GPT-2 和 GPT-3 论文中描述的大规模有效批次大小（例如 50 万个 token）。以下是其核心观点的梳理：
为何使用梯度累积？
-
硬件限制：由于 Transformer 模型的计算需求，现代 GPU（如消费级显卡）无法在内存中容纳巨大的批次大小。例如，Karpathy 的目标是约 50 万个 token 的批次大小以匹配研究设置，但单个 GPU 可能仅能处理 4-8 个样本的微批次。若不使用累积，您将受限于极小的批次，导致梯度噪声大且收敛缓慢。
-
更好的训练稳定性和性能：通过累积实现的有效批次大小增大可减少梯度方差，使更新更可靠并加速训练。Karpathy 强调这是为了可复现性——紧贴 GPT 论文的超参数（例如与特定学习率调度搭配）比天真的小批次训练效果更好。
-
Karpathy 的引述（约 2:51:40）：“我们想使用约 50 万的批次大小，但问题是我们如何在小 GPU 上实现这一点？为此，我们需要使用所谓的梯度累积……接下来我们将转向这个方法，它允许我们以串行方式模拟任意设定的批次大小。”
本质上，这是一种并行性的“串行模拟”：您以小数据块处理数据，但像处理一个大批次那样累积梯度，从而无需数据中心即可实现可扩展的训练。
其工作原理（在他的实现中）
Karpathy 在 PyTorch 中显式实现了它：
- 设置一个小的微批次大小（例如 4-8 个样本），使其适应 GPU 内存。
- 运行
accumulation_steps次迭代（例如 32 步）的前向和反向传播，每次元素级求和梯度而不更新权重。 - 所有步骤完成后，使用累积的梯度执行一次优化器步骤。
- 将损失按
1/accumulation_steps缩放以进行归一化。
这确保了有效批次大小为 微批次大小 × 累积步数 × 序列长度。他结合了混合精度训练、FlashAttention 和 torch.compile 等技巧以提高效率。
额外内容：通过权重绑定的隐式累积
Karpathy 还将此与权重绑定（在 GPT-2 风格中共享 token 嵌入矩阵 wte 与输出头）联系起来。在反向传播期间，来自输入嵌入和输出 logits 的梯度“自然累积”到同一张量中——节省约 30% 的参数，并添加了输入/输出表示应对齐的归纳偏置。引述：“在反向传播中，我们将从两个分支获得梯度贡献……这些梯度将在 wte 张量上累加。”
总之，这些技术让您能在单 GPU 上训练 GPT 规模的模型，同时忠实复现论文结果。完整上下文请观看视频中约 2:50:00 之后的部分。
Let’s build GPT: from scratch, in code, spelled out