技术指南

完全分片数据并行

完全分片数据并行 (FSDP) 是一种分布式训练技术,可将模型的参数、梯度和优化器状态拆分到多个 GPU 上,因此每个设备仅保留一个切片。

概述

完全分片数据并行 (FSDP) 是一种分布式训练技术,可将模型的参数、梯度和优化器状态拆分到多个 GPU 上,因此每个设备仅保留一个切片。它使得在硬件上训练大型模型成为可能,而单个 GPU 的内存永远无法容纳整个模型。

完全分片数据并行是一个技术构建块,会大规模影响模型质量、基础设施成本、延迟和可靠性。

深入探讨

传统的数据并行性在每个 GPU 上保留模型的完整副本,这会浪费内存并限制模型大小。 FSDP 由 Meta 的 PyTorch 普及,并受到 Microsoft 的 ZeRO 的启发,而是跨设备分片三件事:参数、梯度和优化器状态。在前向传递过程中,每个 GPU 通过全收集临时收集其正在计算的层的全部权重,运行计算,然后立即释放收集的副本。向后传递的工作原理类似,然后是减少分散,将梯度切片分配回其所属的 GPU。由于每个设备仅永久存储模型的一小部分,因此内存使用量与 GPU 数量大致呈线性下降,从而使团队可以训练具有数百或数千亿参数的模型。

技术洞察

FSDP 用额外的通信来节省内存。每层的权重根据需要在使用前进行全聚集重建,并在使用后立即丢弃,而梯度则通过减少分散进行组合和分割。通过在当前层运行时预取下一层的参数,通信可以与计算重叠,从而隐藏大部分网络延迟。调整分片粒度(包装策略)可以平衡内存占用与通信开销。

掌握完全分片数据并行

完全分片数据并行 (FSDP) 是一种分布式训练技术,可将模型的参数、梯度和优化器状态拆分到多个 GPU 上,因此每个设备仅保留一个切片。它使得在硬件上训练大型模型成为可能,而单个 GPU 的内存永远无法容纳整个模型。完全分片数据并行是一个技术构建块,会大规模影响模型质量、基础设施成本、延迟和可靠性。为了建立深入的理解,请将完全分片数据并行视为一种操作模型,而不是单个功能:定义所需的结果,澄清假设,并将系统可以可靠地执行的操作与仍需要专家判断的操作分开。

在实践中,使用完全分片数据并行的强大团队根据可靠性和成本优化架构、数据和基础设施选择。他们记录明确的成功标准,根据实际数据和工作流程进行测试,并根据观察到的失败模式而不是一次性基准测试胜利进行迭代。这就是理论理解转变为跨产品、政策和运营的持久能力的地方。

多年来,架构决策决定着性能和运营成本。与此同时,优化一个基准测试可以隐藏更广泛的系统弱点。最具弹性的方法是将实验速度与治理规则结合起来:运行试点、捕获证据、发布决策日志,并随着模型行为、用户期望和监管要求的发展不断更新保障措施。

战略影响

多年来,架构决策决定着性能和运营成本。

多年来,架构决策决定着性能和运营成本。在高质量部署中,这会转化为可衡量的操作规则、所有权边界和定期审查仪式,以便团队可以增强信心,而不是扩大模糊性。

技术教育帮助团队选择正确的堆栈,而不仅仅是最新的堆栈。

技术教育帮助团队选择正确的堆栈,而不仅仅是最新的堆栈。在高质量部署中,这会转化为可衡量的操作规则、所有权边界和定期审查仪式,以便团队可以增强信心,而不是扩大模糊性。

更好的工程选择可以减少生产中的可靠性事故。

更好的工程选择可以减少生产中的可靠性事故。在高质量部署中,这会转化为可衡量的操作规则、所有权边界和定期审查仪式,以便团队可以增强信心,而不是扩大模糊性。

完全分片数据并行的未来

FSDP 正在成为开放大型模型训练的默认设置,PyTorch 中的 FSDP2 提高了可用性和每参数分片。预计与万亿参数模型的张量和管道并行性更紧密的集成,对混合精度和 fp8 的更好支持,以及为您选择分片边界的更智能的自动包装。随着 NVLink 和 InfiniBand 等 GPU 间互连变得越来越快,分片的通信成本不断缩小,使其在更大的规模上变得实用。

现实世界的实施

在 8 个单独无法承受全部权重的 GPU 上微调 700 亿参数的 Llama 模型。

通过在数百个加速器上分片优化器状态(Adam 主导内存),在人工智能实验室预训练大型语言模型。

研究人员使用 PyTorch 的 FSDP 包装器在大学集群上训练视觉变压器,而无需购买旗舰 80GB GPU。

将 FSDP 与混合精度 bfloat16 相结合,可将内存大致减半并加快多模式模型的训练吞吐量。

实施模式

实践中的完全分片数据并行

在 8 个单独无法承受全部权重的 GPU 上微调 700 亿参数的 Llama 模型。

在 8 个 GPU 上微调 700 亿参数的 Llama 模型,这些 GPU 单独无法承受全部权重。当团队预先定义质量阈值、为边缘情况保留人工升级路径并随着时间的推移跟踪生产力增益和错误成本时,通常会获得更好的结果。

实践中的完全分片数据并行

通过在数百个加速器上分片优化器状态(Adam 主导内存),在人工智能实验室预训练大型语言模型。

通过在数百个加速器上分片优化器状态(Adam 主导内存),在人工智能实验室中预训练大型语言模型。当团队预先定义质量阈值、为边缘情况保留人工升级路径并随着时间的推移跟踪生产力增益和错误成本时,通常会获得更好的结果。

实践中的完全分片数据并行

研究人员使用 PyTorch 的 FSDP 包装器在大学集群上训练视觉变压器,而无需购买旗舰 80GB GPU。

研究人员使用 PyTorch 的 FSDP 包装器在大学集群上训练视觉转换器,而无需购买旗舰 80GB GPU。当团队预先定义质量阈值、为边缘情况保留人工升级路径并随着时间的推移跟踪生产力增益和错误成本时,通常会获得更好的结果。

实践中的完全分片数据并行

将 FSDP 与混合精度 bfloat16 相结合,可将内存大致减半并加快多模式模型的训练吞吐量。

将 FSDP 与混合精度 bfloat16 相结合,可将内存大致减半并加快多模式模型的训练吞吐量。当团队预先定义质量阈值、为边缘情况保留人工升级路径并随着时间的推移跟踪生产力增益和错误成本时,通常会获得更好的结果。

风险与防护栏

!

优化一项基准测试可以隐藏更广泛的系统弱点。

!

基础设施和维护成本常常被低估。

!

随着系统变得更加复杂,安全性和可观察性差距可能会扩大。

实施路线图

1

在实施之前定义延迟、质量和成本目标。

在实施之前定义延迟、质量和成本目标。将每个步骤视为证据门:如果不满足标准,则暂停推出,缩小差距,然后再扩大使用。

2

在实际负载和数据条件下进行基准测试。

在实际负载和数据条件下进行基准测试。将每个步骤视为证据门:如果不满足标准,则暂停推出,缩小差距,然后再扩大使用。

3

仪器监控错误、漂移和用户影响。

仪器监控错误、漂移和用户影响。将每个步骤视为证据门:如果不满足标准,则暂停推出,缩小差距,然后再扩大使用。

4

在扩展之前准备回滚和事件响应路径。

在扩展之前准备回滚和事件响应路径。将每个步骤视为证据门:如果不满足标准,则暂停推出,缩小差距,然后再扩大使用。

不断探索