概述
它可以让您通过用额外的计算换取更低的内存使用来训练更深、更大的网络。
深入探讨
Training normally retains intermediate activations needed for backpropagation. Activation checkpointing saves selected boundary tensors and recomputes internal activations during backward. This trades additional computation for a lower activation-memory footprint; the actual savings and runtime overhead depend on checkpoint placement, operations, model graph, and hardware. In a simplified chain of N similarly sized layers, a textbook schedule with checkpoint boundaries spaced about sqrt(N) layers apart can reduce activation storage to order sqrt(N) checkpoints, with recomputation cost governed by that schedule. This asymptotic example is not a universal prediction for framework implementations. PyTorch exposes activation-checkpointing utilities, and its current guidance also describes selective rematerialization and compiler-assisted choices. It can combine with FSDP and mixed precision, but benchmark the full training workload.
现实世界的实施
通过丢弃和重新计算层激活，在单个 GPU 上训练具有更大批量大小的深度转换器。
在高分辨率图像上微调视觉模型，否则激活图会溢出 GPU 内存。
拥抱 Face Transformers 启用gradient_checkpointing=True，以在微调期间适应十亿参数模型。
将检查点与 FSDP 相结合，使参数和激活都保持较小，从而能够训练非常大的语言模型。
梯度检查点的未来
梯度检查点现在是大型模型训练的标准，并且越来越自动化，库会为您选择最佳检查点位置。它与 FSDP、混合精度和卸载自然配合，以提高模型尺寸。期望“选择性”检查点仅重新计算廉价的操作，同时缓存昂贵的操作（如注意力矩阵），加上 PyTorch 的 torch.compile 等工具中的编译器驱动方法，自动决定保存哪些内容与重新计算以获得最佳速度内存平衡。
付诸实践
为什么这很重要
- 成本与预算. 多年来，架构决策决定着性能和运营成本。
- 更清晰的判决. 技术教育帮助团队选择正确的堆栈，而不仅仅是最新的堆栈。
- 质量控制. 更好的工程选择可以减少生产中的可靠性事故。
留意
- 优化一项基准测试可以隐藏更广泛的系统弱点。
- 基础设施和维护成本常常被低估。
- 随着系统变得更加复杂，安全性和可观察性差距可能会扩大。
后续步骤
- 在实施之前定义延迟、质量和成本目标。
- 在实际负载和数据条件下进行基准测试。
- 仪器监控错误、漂移和用户影响。
- 在扩展之前准备回滚和事件响应路径。
不断探索
常见问题
什么是梯度检查点？
梯度检查点（也称为激活检查点）是一种节省内存的技巧，它在前向传播期间丢弃大多数中间激活，并在反向传播期间动态重新计算它们。它可以让您通过用额外的计算换取更低的内存使用来训练更深、更大的网络。
为了节省内存，梯度检查点主要交换什么？
梯度检查点在向后传递过程中重新计算丢弃的激活，花费额外的计算来换取显着减少的内存。
为什么激活通常在前向传递过程中存储？
反向传播使用前向传递的中间激活来计算梯度，因此它们必须可用，除非重新计算。
In a simplified chain of N similarly sized layers, what can checkpoint boundaries spaced about sqrt(N) layers apart achieve in the textbook schedule?
The square-root schedule is a theoretical result for a simplified sequential computation graph, not a universal implementation guarantee.
What most determines gradient-checkpointing runtime overhead in a practical model?
Recomputation overhead varies with checkpoint policy and the cost of operations selected for recomputation.
在 PyTorch 中，哪个实用程序通常用于将梯度检查点应用于模块？
torch.utils.checkpoint 包装一个模块，以便在向后期间重新计算其内部激活而不是存储。