完全分片数据并行
完全分片数据并行 (FSDP) 是一种分布式训练技术,可将模型的参数、梯度和优化器状态拆分到多个 GPU 上,因此每个设备仅保留一个切片。
概述
It makes training huge models possible on hardware that could never fit the whole model in one GPU's memory.
深入探讨
传统的数据并行性在每个 GPU 上保留模型的完整副本,这会浪费内存并限制模型大小。 FSDP 由 Meta 的 PyTorch 普及,并受到 Microsoft 的 ZeRO 的启发,而是跨设备分片三件事:参数、梯度和优化器状态。在前向传递过程中,每个 GPU 通过全收集临时收集其正在计算的层的全部权重,运行计算,然后立即释放收集的副本。向后传递的工作原理类似,然后是减少分散,将梯度切片分配回其所属的 GPU。由于每个设备仅永久存储模型的一小部分,因此内存使用量与 GPU 数量大致呈线性下降,从而使团队可以训练具有数百或数千亿参数的模型。
技术洞察
FSDP 用额外的通信来节省内存。每层的权重根据需要在使用前进行全聚集重建,并在使用后立即丢弃,而梯度则通过减少分散进行组合和分割。通过在当前层运行时预取下一层的参数,通信可以与计算重叠,从而隐藏大部分网络延迟。调整分片粒度(包装策略)可以平衡内存占用与通信开销。
战略影响
成本与预算
多年来,架构决策决定着性能和运营成本。
更清晰的判决
技术教育帮助团队选择正确的堆栈,而不仅仅是最新的堆栈。
质量控制
更好的工程选择可以减少生产中的可靠性事故。
完全分片数据并行的未来
FSDP 正在成为开放大型模型训练的默认设置,PyTorch 中的 FSDP2 提高了可用性和每参数分片。预计与万亿参数模型的张量和管道并行性更紧密的集成,对混合精度和 fp8 的更好支持,以及为您选择分片边界的更智能的自动包装。随着 NVLink 和 InfiniBand 等 GPU 间互连变得越来越快,分片的通信成本不断缩小,使其在更大的规模上变得实用。
现实世界的实施
在 8 个单独无法承受全部权重的 GPU 上微调 700 亿参数的 Llama 模型。
通过在数百个加速器上分片优化器状态(Adam 主导内存),在人工智能实验室预训练大型语言模型。
研究人员使用 PyTorch 的 FSDP 包装器在大学集群上训练视觉变压器,而无需购买旗舰 80GB GPU。
将 FSDP 与混合精度 bfloat16 相结合,可将内存大致减半并加快多模式模型的训练吞吐量。
风险与防护栏
优化一项基准测试可以隐藏更广泛的系统弱点。
基础设施和维护成本常常被低估。
随着系统变得更加复杂,安全性和可观察性差距可能会扩大。
实施路线图
在实施之前定义延迟、质量和成本目标。
在实际负载和数据条件下进行基准测试。
仪器监控错误、漂移和用户影响。
在扩展之前准备回滚和事件响应路径。
不断探索
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 Fully Sharded Data Parallel 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
常见问题
What is Fully Sharded Data Parallel?
完全分片数据并行 (FSDP) 是一种分布式训练技术,可将模型的参数、梯度和优化器状态拆分到多个 GPU 上,因此每个设备仅保留一个切片。它使得在硬件上训练大型模型成为可能,而单个 GPU 的内存永远无法容纳整个模型。
FSDP 可以跨 GPU 分片,而标准数据并行性则不能?
FSDP 跨设备对模型的参数、梯度和优化器状态进行分片,而标准数据并行性则在每个 GPU 上复制完整模型。
FSDP 在计算之前使用哪种集体操作来重建层的完整权重?
在层运行之前,FSDP 会执行一次全收集,以临时收集所有分片中的完整参数,然后释放它们。
为什么 FSDP 在一层计算后立即释放收集到的全部权重?
永久仅保留一个分片并暂时收集完整权重可以使内存使用率保持较低水平并大致与模型的一小部分成比例。
FSDP 主要受到哪种早期内存优化方法的启发?
FSDP 的参数、梯度和优化器状态的分片密切遵循 DeepSpeed 库中 Microsoft 的 ZeRO 中引入的思想。
FSDP 如何隐藏大部分网络延迟以防止收集权重?
FSDP 在当前层仍在计算时预取下一层的参数,将全收集通信与有用的工作重叠。