优化器状态卸载到 CPU 和 NVMe
一种节省内存的技巧,可将繁重的训练记录(优化器状态、梯度,有时是权重)存放在 CPU RAM 或 NVMe SSD 上,而不是稀缺的 GPU 内存中。
概述
它让人们能够训练远超GPU内存容量的模型。
深入探讨
当您使用 Adam 这样的优化器训练神经网络时,每个参数都会带来额外的负担:两个运行统计数据(动量和方差),加上权重的全精度副本及其梯度。在混合精度训练中,每个参数总计大约 16 个字节,使权重本身的 2 个字节相形见绌。卸载将这些负担从 GPU 上移走。 CPU 卸载通过 PCIe 总线将优化器状态流式传输到普通系统 RAM,而 NVMe 卸载将它们一直推送到快速固态磁盘。该技术因 DeepSpeed 的 ZeRO-Infinity 和 ZeRO-Offload 而流行,以原始速度换取容量,让单个 GPU 或小型集群能够微调具有数十亿个参数的模型。
技术洞察
关键是将数据移动与计算重叠。优化器状态位于 CPU/NVMe 中;在向后传递期间,分区在需要之前通过 PCIe 预取,并且优化器步骤本身通常在 CPU 上运行。 ZeRO-Offload 将 float32 主权重和 Adam 矩保留在 CPU 上,因此只有前向和后向数学保留在 GPU 上。 NVMe 添加了分层缓存,因此 TB 级状态会溢出到磁盘,而热分区则保留在 RAM 中。
战略影响
成本与预算
多年来,架构决策决定着性能和运营成本。
更清晰的判决
技术教育帮助团队选择正确的堆栈,而不仅仅是最新的堆栈。
质量控制
更好的工程选择可以减少生产中的可靠性事故。
优化器状态卸载到 CPU 和 NVMe 的未来
随着模型的 GPU 内存不断增长,分层卸载正在成为标准而不是奇特的。预计与更快的互连(例如模糊 CPU-GPU 边界的 NVLink-C2C 和 CXL 内存池)以及更智能的调度程序(可以预测要预取的状态)的更紧密集成。 Grace Hopper 等统一内存架构减少了 PCIe 损失,并且框架正在推动多层卸载几乎透明,以便爱好者可以在适度的硬件上微调大型模型。
现实世界的实施
使用 DeepSpeed ZeRO-Offload 在单个 24 GB 消费级 GPU 上微调 130 亿参数 LLM,将 Adam 状态推送到 CPU RAM。
一个小型研究实验室通过 ZeRO-Infinity 将优化器状态溢出到 NVMe 驱动器,在几个 GPU 上训练数十亿参数模型。
Hugging Face Accelerate 配置可启用 CPU 卸载,以便用户可以运行完整的微调作业,否则会引发内存不足错误。
注重成本的初创公司租用更便宜、内存更低的云 GPU 并卸载到附加的 NVMe,而不是购买顶级 80 GB 卡。
风险与防护栏
优化一项基准测试可以隐藏更广泛的系统弱点。
基础设施和维护成本常常被低估。
随着系统变得更加复杂,安全性和可观察性差距可能会扩大。
实施路线图
在实施之前定义延迟、质量和成本目标。
在实际负载和数据条件下进行基准测试。
仪器监控错误、漂移和用户影响。
在扩展之前准备回滚和事件响应路径。
不断探索
Free newsletter
Get the daily AI briefing
Three verified AI stories every weekday morning, written in plain English. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the Optimizer State Offloading to CPU and NVMe quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
常见问题
什么是向CPU和NVMe进行优化器状态卸载?
一种节省内存的技巧,可将繁重的训练记录(优化器状态、梯度,有时是权重)存放在 CPU RAM 或 NVMe SSD 上,而不是稀缺的 GPU 内存中。它可以让人们训练比 GPU 内存允许的更大的模型。
将优化器状态卸载到 CPU 或 NVMe 的主要目标是什么?
卸载将繁重的优化器状态从 GPU 移至 CPU RAM 或 NVMe,从而释放 GPU 内存,以便可以在同一硬件上训练更大的模型。
对于混合精度的 Adam 优化器,哪些数据通常每个参数消耗最多的内存?
Adam 存储动量、方差和全精度主权重,每个参数总共大约 16 个字节,远远超过 2 字节半精度权重。
哪个框架功能普及了大规模优化器卸载?
DeepSpeed 的 ZeRO-Offload 和 ZeRO-Infinity 大规模引入并推广了 CPU 和 NVMe 的卸载优化器状态。
卸载到 CPU 或 NVMe 的主要性能成本是什么?
将状态移出 GPU 意味着通过 PCIe 或 NVMe 传输数据,这比 GPU 内存上的速度慢,因此除非与计算重叠,否则吞吐量会下降。
卸载系统如何隐藏大部分传输延迟?
当 GPU 仍在计算时,调度程序通过 PCIe 预取所需的分区,将通信与计算重叠以掩盖延迟。