拓冰网站定制 
文档教程人工智能【免费下载链接】AISystemAISystem 主要是指AI系统包括AI芯片、AI编译器、AI推理和训练框架等AI全栈底层技术项目地址https://gitcode.com/GitHub_Trending/ai/AISystem点击查看免费下载本篇文章聚焦 AISystem 开源仓库中分布式并行章节的完全分片数据并行FSDP与零冗余优化器ZeRO技术系统讲解精度格式、混合精度训练、损失缩放、显存估算等前置知识以及 ZeRO-DP、ZeRO-R、ZeRO-Infinity 的原理、通信开销分析与 PyTorch FSDP 落地实现。读者学完后将掌握大模型训练中以通信换显存的核心理念能够独立配置 FSDP 自动包装策略、混合精度策略并正确估算大模型训练的内存占用与训练效率。本文内容以 03ZeRODP.md 为主体上一节 02DataParallel.md 讲解了通用数据并行DP与分布式数据并行DDP的实现差异本文在其基础上继续深入权重、优化器、梯度三类模型状态的分布式并行方案。前置知识理解 FSDP 所需的三大基础在介绍针对权重数据、优化器数据和梯度数据进行分布式数据并行的算法 FSDP 前需要先掌握三项前置知识如何执行混合精度训练、如何对显存消耗进行估算以及常用的精度格式。这三者直接决定了 FSDP 中内存分片与通信调度的设计动机。精度格式FP32、FP16、BF16 与 TF32深度学习中常用的精度格式包括FP32、FP16、BF16和TF32四种它们各有不同的位布局、数值范围与适用场景。FP16IEEE 754 标准下的半精度浮点格式。随着深度学习的发展FP16 逐渐取代了 FP32 成为训练的主力类型。因为相较于 FP32更低的精度并不会对神经网络性能产生重大影响额外的精度不会带来任何好处反而会更慢、占用更多内存并降低通信速度。FP16 通常用于混合精度训练MindSpore/PyTorch也用于训练后量化以加快推理速度。其他用于量化的格式还有整数 INT88 位整数、INT44 位甚至 INT1二进制值。BF16谷歌最初开发的另一种 16 位格式全称为 Brain Floating Point Format简称 bfloat16名字源于构思出该格式的谷歌 AI 研究小组谷歌大脑。它最初被使用在谷歌芯片 TPU 中后被广泛使用在 GPU、NPU 等 AI 芯片中。由于具有更多的指数位BF16 常用于处理 FP16 的溢出问题。FP32IEEE 754 标准下的单精度浮点数在很长一段时间内都是深度学习的主力。长期以来神经网络的权重、激活和其他值都默认用 FP32 表示。TF32一种十分特殊的格式无需显示设置而是自动执行。它将 FP32 数据截断为 TF32 进行计算然后再转换回 FP32。这一创新的最大优势在于编译器只需在 AI 编译层提供支持其他代码部分可以继续使用动态范围相同但精度较高的 FP32无需修改。TF32 的快速插入性使得利用 AI 计算加速核心的速度成为可能而无需过多修改现有代码。TF32 采用与 FP16 相同的 10 位尾数这满足了 AI 工作负载的精度要求同时使用与 FP32 相同的 8 位指数因此具有相同的数值范围。从技术上讲它可以视为一种 19 位格式也可以视为扩展精度的 BF16。TF32 的优势在于其格式与 FP32 相同。当使用 TF32 进行计算时输入 FP32 操作数的尾数从 23 位舍入到 10 位然后进行精确乘法运算最后以正常的 FP32 格式进行累加。相比之下FP16 和 BF16 等格式需要更多工作因为它们涉及不同的位布局。尽管如此这些不同的精度格式仍然值得使用因为它们可以减少内存带宽、存储和计算位数从而提升执行速度。在 PyTorch 中TF32 通过以下开关启用torch.backends.npu.matmul.allow_tf32 True torch.backends.npu.matmul.allow_16 True混合精度训练半精度计算 权重备份在当今大模型训练的背景下混合精度训练已然成为一种备受推崇的普遍做法。通过采用混合精度训练能够将训练速度显著提升数倍而又不会对模型的整体性能产生重大影响。在传统科学计算领域人们通常追求较高的精度如 FP128 或 FP64但深度学习中面临的实际上是一个高维函数拟合或近似的优化问题并不需要过于精确的数值表示且使用低精度会带来显著的计算速度提升。使用 FP16 训练有时会出现下溢问题FP16 的有效动态范围约为 $5.96e^{-8} \sim 65504$在训练后期例如激活函数的梯度会非常小甚至在梯度乘以学习率后值更小。由于 FP16 的精度范围有限过小的梯度可能导致更新无效这时就需要使用混合精度训练。混合精度训练可以分为两个部分半精度和权重备份如下图所示这里以 FP16 和 FP32 举例。训练开始时准备两套模型状态一套为 FP32 类型优化器状态和模型参数另一套为 FP16 类型模型参数。前向传播、反向传播时都使用 FP16 类型的模型参数进行计算而在参数更新时将梯度与学习率 $\eta$ 相乘更新到 FP32 类型的模型状态上新一轮训练中再次将 FP32 类型的模型拷贝为 FP16 类型的模型。这个过程就是混合精度训练。由于在计算密集的前向传播、反向传播中使用了 FP16 进行计算与单精度相比训练速度会大幅度提升。另外激活值在训练过程中占用内存的很大部分使用 FP16 储存激活值在大批量训练时也会节省内存在分布式环境下使用 FP16 进行梯度通信也会降低通信量。为了获得最佳性能混合精度中需要额外选择合适的批量大小。通常建议使用 2 的幂次方作为批量大小并与输入/输出神经元的数量相匹配通常为 8 的倍数但也可能更高具体取决于所使用的硬件和模型的数据类型。而在 FSDP 中可以通过在 torch 中指定MixedPrecision配置类进行混合精度的自动配置分别指定参数精度param_dtype、梯度通信精度reduce_dtype和 Buffer 精度buffer_dtypefpSixteen MixedPrecision( param_dtypetorch.float16, # Gradient communication precision. reduce_dtypetorch.float16, # Buffer precision. buffer_dtypetorch.float16, ) bfSixteen MixedPrecision( param_dtypetorch.bfloat16, # Gradient communication precision. reduce_dtypetorch.bfloat16, # Buffer precision. buffer_dtypetorch.bfloat16, ) model FSDP(model, auto_wrap_policyt5_auto_wrap_policy, mixed_precisionbfSixteen)损失缩放解决 FP16 下溢问题的另一手段解决 FP16 下溢问题的另一个方法是损失缩放Loss Scale。训练到后期梯度特别是激活函数平滑段的梯度会特别小FP16 表示容易产生下溢现象。为了解决梯度过小的问题需要对损失进行缩放——由于链式法则的存在损失的缩放也会作用在梯度上。缩放后的梯度就会平移到 FP16 有效的展示范围内。不过缩放并非对所有网络而言都是必须的缩放的取值也会特别大一般在 8 - 32k 之间。在 PyTorch 中可以通过以下方式实现自动损失缩放from torch.cuda.amp import GradScaler, autocast scaler GradScaler() with autocast(): output model(input) loss loss_fn(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这种损失缩放的方式是动态的每当梯度溢出时减少损失缩放规模并且间歇性地尝试增加损失规模从而实现在不引起溢出的情况下使用最高损失缩放因子更好地恢复精度。动态损失缩放的算法从比较高的缩放因子开始如 $2^{24}$然后开始训练并在迭代中检查数是否会溢出Infs/Nans如果没有梯度溢出则不调整缩放因子继续进行迭代如果检测到梯度溢出则缩放因子减半重新确认梯度更新情况直到参数不出现在溢出的范围内在训练后期loss 已经趋近收敛稳定梯度更新的幅度往往变小这时可以允许更高的损失缩放因子来再次防止数据下溢。内存消耗估算模型状态与剩余状态在神经网络模型的训练中合理估算和管理内存消耗非常重要。内存存储主要分为两大块模型状态Model States和剩余状态Residual States。模型状态指和模型本身相关的、必须存储的内容具体包括优化器状态Optimizer StatesAdam 优化算法中的 Momentum 和 Variance梯度Gradients模型梯度 $G$参数Parameters模型参数 $W$。剩余状态是并非模型必须的但在训练过程中会额外产生的内容具体包括激活值Activation在反向传播过程中使用链式法则计算梯度时会用到。有了它算梯度会更快但它不是必须存储的因为可以通过重新前向传播来计算临时存储Temporary Buffers例如把梯度发送到某个 NPU 上进行 All-Reduce 时产生的存储碎片化的存储空间Unusable Fragment Memory虽然总存储空间是够的但是如果取不到连续的存储空间相关的请求也会失败。对这类空间浪费可以通过内存整理来解决。以 FP32 与 FP16 的混合精度训练为例假设模型的参数量是 $\Phi$那么模型状态所消耗的空间为| Model States | Size (Byte) | |-- |--- | | FP32 Parameters | 4$\Phi$ | | FP32 Adam Optimizer Momentum | 4$\Phi$ | | FP32 Adam Optimizer Variance | 4$\Phi$ | | FP16 Gradients | 2$\Phi$ | | FP16 Parameters | 2$\Phi$ | | Total | 16$\Phi$ |而剩余状态与具体模型架构有关需要具体分析。接下来基于 Transformer 的架构进行分析因为所有参数超过 10 亿的 SOTA 模型都遵循这一架构。分析假设使用 Adam 优化器进行混合精度训练因为该优化器是训练基于 Transformer 的大模型较为常用的选择。模型状态基于 Transformer 的模型中的参数总数主要取决于隐藏维度$hd$和 Transformer 层数$nl$。Transformer 块中的几乎所有参数都来自每个块内的四个线性层其大小分别为$hd$,$3hd$、$hd$,$hd$、$hd$,$4hd$和$4hd$,$hd$。因此基于 Transformer 的模型中的总参数可以近似为$$ 12 × nl × hd^2 $$剩余状态剩余状态主要是指激活内存它取决于模型结构、批量大小$bsz$和序列长度$seq$而且可能相当大。不过激活所需的内存可以通过激活检查点activation checkpointing大大减少。假设 $ci$ 是两个激活检查点之间的 Transformer 块数$bsz × seq × hd$ 是每个 Transformer 块的输入大小激活检查点所需的内存估计为$$ 2 × bsz × seq × hd × nl / ci $$激活工作内存激活工作内存是反向传播过程中所需的内存用于在执行实际反向传播之前重新计算激活即两个连续激活检查点之间的激活量。例如如果为每个 Transformer 块创建一个激活检查点那么内存就是每个 Transformer 块的总激活量其字节数约为$$ bsz × seq × ci × (16 × hd 2 × attn_heads × seq) $$模型状态工作内存模型状态工作内存是指在将所有模型状态卸载到 CPU 或 NVMe 之后对模型中最大的单个算子执行前向或后向传播所需的 NPU 内存最小量。这大约由该算子的参数和梯度的大小决定因为必须至少有足够的内存来保存反向传播的参数及其梯度。Transformer 最大的算子是将隐藏状态从 $h$ 转换为 $4h$ 的线性层该线性层的参数和梯度的大小为$$ 4 × hd × 4hd $$完全分片数据并行 FSDP 总览完全分片数据并行Fully Sharded Data ParallelFSDP在分布式 AI 系统中具有重要地位不仅能提高并行效率还能减少显存内存消耗这两方面的优势为模型的大规模训练带来了显著的好处。值得注意的是并行效率和内存消耗之间存在密切关联降低内存占用可以使使用更大的并行度进而提升整体的并行加速效果。完全分片数据并行是基于零冗余优化器ZeRO在 AI 框架中的具体实现主要的实现有微软的DeepSpeed和 Meta 的Fairscale其中 Fairscale 被集成到 PyTorch 中并作为 FSDP 的实现基础。本节将从零冗余优化器的常用技术入手深入剖析如何降低内存开销并提高训练效率。在数据并行中每个 NPU 都需要保存一份完整的参数模型状态和剩余状态。但并不是所有参数在训练的整个过程中都会被使用到而是在特定阶段某个层的前向或反向传播才被使用因此可以在不需要使用的时候将它转移到其他地方以节省内存空间。ZeRO 有两套优化方案ZeRO-DP旨在减少模型状态的内存占用ZeRO-R旨在减少剩余状态内存消耗。这两套方案使 ZeRO 能够在保持高效的同时减少内存占用。此外在混合并行场景中ZeRO 也可以与流水线并行PP、张量并行TP组合使用相关内容可参考 06HybridParallel.md 中的混合并行分析。ZeRO-DP对模型状态进行分区ZeRO-DP 对模型状态进行切分具体来说每个 NPU 只存储 $\frac{1}{N_d}$ 的模型状态其中 $N_d$ 为并行度在需要时通过集合通信 All-Gather 获取参数。ZeRO-DP 保留了数据并行训练DP的高效率同时实现了模型并行MP的内存效率优势。由于数据并行的模型状态在所有数据并行进程中冗余存储因此内存效率低下但数据并行具有更高的计算粒度和更低的通信量从而具有更高的训练效率。模型并行的通信开销很大因此可扩展性比数据并行低但模型并行对模型状态进行分区获得了较高的内存效率。ZeRO-DP 对模型状态进行分区而不是复制它们并使用动态通信调度最小化通信量。通过这样做ZeRO-DP 随着数据并行程度的增加线性减少模型在每块 NPU 的内存占用同时保持通信量接近默认数据并行的通信量从而保持效率。ZeRO-DP 有三个主要优化阶段分别对应于优化器状态、梯度和参数的划分在累积启用时优化状态分区Partition optimizer states$P_{os}$又称为 ZeRO-1将优化器状态按并行度均匀分区每个进程只需存储 $\frac{1}{N_d}$ 的优化器状态。这可将内存消耗减少到 1/4且无额外通信开销。添加梯度分区Partition gradients$P_{osg}$又称为 ZeRO-2在优化器状态分区的基础上对梯度也进行分区。每个进程只需存储用于更新自身参数分区所需的梯度。这可减少 8 倍的内存消耗且无额外通信开销。添加参数分区Partition parameters$P_{osgp}$又称为 ZeRO-3在优化器状态和梯度分区的基础上对参数也进行分区。每个进程只存储自身的参数分区在前向反向传播时需要从其他进程收集所需的参数分区。这会使通信量增加约 50%但可以实现与并行度 $N_d$ 成正比的内存减少。通过这三个阶段的优化ZeRO-DP 最终能够在保持数据并行高效的同时将每个 NPU 的内存消耗降低至 $\frac{1}{N_d}$ 的水平使得利用少量硬件资源训练万亿参数等超大模型成为可能。下面进行每个阶段的详细介绍这里假设模型使用混合精度训练模型参数量为 $4\Psi$。ZeRO-1 内存分析优化器状态是训练过程中 NPU 内存中的主要保存内容但只有在参数更新的时候才会被用到。ZeRO-1 的核心思想是将优化器状态分布到多个 NPU 上减少每个 NPU 所需的显存在需要参数更新时进行聚合。其执行步骤如下数据分片a从优化器状态分片开始将优化器状态分成 N 份每个 NPU 保存一份分片并将训练批次数据batch data分成 N 份每个 NPU 处理一份数据。前向与后向计算b每个 NPU 执行一步前向forward和后向backward计算得到局部梯度 $G_i$。梯度聚合b对各个 NPU 上的局部梯度 $G_i$ 执行 All-Reduce 操作得到完整梯度 $G$。这一步单个 NPU 的通信量为 $2\Phi$。权重更新c使用完整梯度 $G$ 和优化器状态更新权重 $W$。每个 NPU 保存一部分权重 $W$并通过 All-Gather 操作从其他 NPU 获取更新后的部分权重完成权重更新。此时单个 NPU 的通信量为 $\Phi$。在 $P_{os}$ 阶段将 Adam 优化器状态根据数据并行维度 $N_d$ 分成等份。每个 NPU 只需存储和更新总优化器状态的 $1/N_d$并更新对应参数。通过分片和聚合操作显存占用从 $4\Psi K\Psi$ 降低到 $4\Psi K\Psi / N_d$。当 $N_d$ 很大时显存占用接近于 $4\Psi$带来约 4 倍显存节约。ZeRO-2 内存分析ZeRO-2 在 ZeRO-1 的基础上进一步优化通过对梯度Grad也进行切分减少显存占用并提高通信效率。其执行步骤如下数据分片从优化器状态和梯度分片开始将优化器状态和梯度分成 N 份每个 NPU 保存一份分片并将训练批次数据batch data分成 N 份每个 NPU 处理一份数据。前向与后向计算每个 NPU 执行一步前向forward和后向backward计算得到局部梯度 $G_i$。梯度分片与聚合对各块 NPU 上的局部梯度 $G_i$ 执行 Reduce-Scatter 操作确保每个 NPU 只维护自己负责的梯度部分。比如NPU 1 负责维护梯度 $G_1$其他 NPU 只需要将 $G_1$ 对应位置的梯度发送给 NPU 1。聚合完毕后NPU 1 释放无用的显存部分单卡通信量为 $\Phi$。权重更新每个 NPU 使用自身维护的梯度 $G_i$ 和优化器状态 $O_i$ 更新相应的权重 $W_i$。权重聚合对权重 $W_i$ 执行 All-Gather 操作将其他 NPU 的权重 $W_i$ 同步至完整权重 $W$单卡通信量为 $\Phi$。在 $P_{osg}$ 阶段梯度与优化器强相关因此优化器可以更新其独立的梯度参数。更新梯度参数时使用 Reduce-Scatter 操作梯度参数更新后立即释放。具体实现中使用分桶Bucket技术将梯度分到不同的桶中并在桶上进行 Reduce-Scatter 操作。通过移除梯度和优化器状态冗余显存占用从 $4\Psi K\Psi$ 降低到 $2\Psi K\Psi / N_d$。当 $N_d$ 很大时显存占用接近于 $2\Psi$带来约 8 倍显存节约。ZeRO-3 内存分析ZeRO-3 在 ZeRO-1 和 ZeRO-2 的基础上进一步优化通过对优化器状态、梯度和权重进行全面切分最大化显存节约这种优化使得训练超大模型成为可能。其执行步骤如下数据分片对优化器状态、梯度和权重进行全面切分每个 NPU 保存一份分片将训练批次数据batch data分成 N 份每块 NPU 处理一份数据。前向计算在前向计算过程中对权重 $W$ 执行 All-Gather 操作从各 NPU 获取分布在不同 NPU 上的权重得到完整的权重 $W$。不属于自身的权重 $W_{others}$ 被抛弃单卡通信量为 $\Phi$。后向计算在后向计算过程中再次对权重 $W$ 执行 All-Gather 操作取回完整权重并抛弃不属于自身的部分 $W_{others}$单卡通信量为 $\Phi$。梯度聚合后向计算得到各自的梯度 $G_i$ 后对梯度 $G_i$ 执行 Reduce-Scatter 操作从其他 NPU 聚合自身维护的梯度 $G_i$。聚合操作结束后立刻抛弃不是自己维护的梯度单卡通信量为 $\Phi$。权重更新每块 NPU 只保存其权重参数 $W_i$由于只维护部分参数 $W_i$因此无需对 $W_i$ 执行 All-Reduce 操作。在 $P_{osgp}$ 阶段优化器状态、梯度和权重均进行划分。在前向和后向计算过程中通过广播从其他 NPU 中获取参数减少每块 NPU 中的显存占用。通过移除梯度、优化器状态和权重的冗余将显存占用从 $4\Psi K\Psi$ 降低到 $(4\Psi K\Psi) / N_d$。这种方法通过增加通信开销以通信换显存使得显存占用与 $N_d$ 成正比。显存占用的优化带来了 1.5 倍单卡通信量的增加。ZeRO-R优化剩余状态除了优化模型状态优化器状态、梯度和参数的内存利用率ZeRO 还专门针对剩余状态如激活数据、临时缓冲区和内存碎片等进行了优化以进一步减少内存开销。ZeRO-R 对剩余状态进行了切分和优化主要包括以下几个策略分区激活检查点Partitioned Activation Checkpointing$P_{a}$解决了模型并行时激活内存冗余的问题。在模型并行中每个 NPU 需要保存完整的输入激活数据才能计算自己分到的模型部分。ZeRO-R 将激活检查点按模型并行度 $N_m$ 进行分区每个 NPU 只需存储 $\frac{1}{N_m}$ 的激活检查点在需要时通过 All-Gather 操作重构出完整激活数据从而按 $N_m$ 的比例减少激活内存。在极端情况下当模型规模很大时ZeRO-R 甚至可以将分区后的激活检查点卸载到 CPU 内存$P_{acpu}$再次降低 NPU 内存占用代价是额外的 Host-Device 通信开销。该策略在大模型训练时会自动开启以保证足够的 NPU 内存用于计算。恒定大小的缓冲区Constant Size Buffer$C_{b}$一些操作如 All-reduce 需要将张量拼成连续的临时缓冲区。使用恒定大小的缓冲区来避免临时缓冲区随着模型大小的增加而爆炸同时使它们足够大以保持效率。内存碎片化整理Memory Defragmentation$M_{d}$在训练过程中由于激活检查点、梯度等张量生命周期的差异会产生大量内存碎片。ZeRO-R 通过预分配和动态管理这些张量的内存减少了内存碎片和内存分配器的开销提高了内存利用率。通过以上优化策略ZeRO-R 很好地补充和完善了 ZeRO-DP 优化模型状态内存的功能。两者相结合ZeRO 优化器能最大限度减少大模型训练的内存占用为未来万亿参数级别的神经网络模型铺平了道路。ZeRO-Infinity突破 NPU 内存壁垒ZeRO-Infinity 是 ZeRO 的扩展可以将深度学习训练扩展到前所未有的规模。具体来说它突破了 NPU 内存壁垒的限制使得训练具有数万亿参数的模型成为可能这是迄今为止最先进系统所无法企及的量级。此外它为训练具有一千万亿参数的模型铺平了道路——充分利用系统的全部内存容量利用 NPU、CPU 和 Non-Volatile Memory ExpressNVMe等所有异构内存组件的能力。在 ZeRO-Infinity 中参数从较慢的内存源如 CPU 和 NVMe无缝迁移到 NPU在 NPU 上被合并为完整的层。梯度计算完成后这些参数被聚合、重新分区然后重新卸载回较慢的内存组件。内存资源的编排确保了最佳利用和最小的开销。这种创新的方法不仅克服了 NPU 内存的常规限制而且提升了分布式框架的可扩展性。FSDP 简单实现PyTorch 中的落地实践和 DDP 一样可以通过简单的嵌套使用 ZeRO 优化器来实现 FSDP。将模型的参数、梯度和优化器状态分片从而显著减少单个 NPU 的内存占用实现更大模型的训练。以下是一个简单的代码示例展示了如何使用torch.distributed.fsdp中的FullyShardedDataParallelFSDP类来实现完全分片数据并行import torch from torch.distributed.fsdp import FullyShardedDataParallel as FSDP torch.cuda.set_device(device_id) sharded_module FSDP(my_module) optim torch.optim.Adam(sharded_module.parameters(), lr0.0001) x sharded_module(x, y3, ztorch.Tensor([1])) loss x.sum() loss.backward() optim.step()首先设置当前的 NPU然后将my_module包装成一个 FSDP 模块。这会将模型的参数、梯度和优化器状态在多个 NPU 之间进行分片从而减少每个 NPU 上的内存占用。再来看一个更详细的例子使用fsdp_main函数用于分布式训练 T5 模型。函数通过setup_model函数加载模型和分词器并设置分布式训练的相关环境变量包括local_rank、rank和world_size。在数据集和数据加载器设置之后通过functools.partial函数部分应用了transformer_auto_wrap_policy并指定T5Block为要自动包装的 Transformer 层。这一步的目的是定义一个自动包装策略用于后续的模型分片和并行处理。接下来定义了sharding_strategy变量并将其设为ShardingStrategy.SHARD_GRAD_OP。这表示使用 ZeRO-2 分片策略如果要使用 ZeRO-3 策略可以将其设为FULL_SHARD。分片策略决定了在分布式训练中如何管理和分配模型参数以优化内存使用和计算性能。为了支持混合精度训练bf16如果当前 CUDA 版本和 NCCL 版本支持 bf16并且 CUDA 版本大于或等于 11.0那么bf16_ready变量为True并将mp_policy设为bfSixteen。否则mp_policy设为None默认使用 fp32单精度训练。在模型仍在 CPU 上时代码将模型传入FSDP模块配置自动包装策略、混合精度策略mixed_precision以及当前 NPU IDdevice_id。这一步是将模型转换为分布式训练的模式以便在多个 NPU 之间分片和并行计算。def fsdp_main(args): model, tokenizer setup_model(t5-base) local_rank int(os.environ[LOCAL_RANK]) rank int(os.environ[RANK]) world_size int(os.environ[WORLD_SIZE]) # Set dataset and dataloader here t5_auto_wrap_policy functools.partial( transformer_auto_wrap_policy, transformer_layer_cls{ T5Block, }, ) sharding_strategy: ShardingStrategy ShardingStrategy.SHARD_GRAD_OP #for Zero2 and FULL_SHARD for Zero3 torch.cuda.set_device(local_rank) bf16_ready ( torch.version.cuda and torch.cuda.is_bf16_supported() and LooseVersion(torch.version.cuda) 11.0 and dist.is_nccl_available() and nccl.version() (2, 10) ) if bf16_ready: mp_policy bfSixteen else: mp_policy None # defaults to fp32 # model is on CPU before input to FSDP model FSDP(model, auto_wrap_policyt5_auto_wrap_policy, mixed_precisionmp_policy, #sharding_strategysharding_strategy, device_idtorch.cuda.current_device()) optimizer optim.AdamW(model.parameters(), lrargs.lr) scheduler StepLR(optimizer, step_size1, gammaargs.gamma) for epoch in range(1, args.epochs 1): train_accuracy train(args, model, rank, world_size, train_loader, optimizer, epoch, samplersampler1) scheduler.step()当然如果需要更简单而全面的配置可以通过deepspeed库进行便捷的 ZeRO 配置。例如在 DeepSpeed 配置文件中声明zero_optimization的 stage 级别1/2/3即可自动完成对应阶段的优化器状态、梯度与参数分区。ZeRO 通信分析以通信换内存的效率账本无论是零冗余优化还是卸载到 CPU 和 NVMe 内存一个关键问题是它们有限的带宽是否会影响训练效率。很自然地会问是否在用通信量来换取内存效率。换句话说与标准 DP 方法相比ZeRO 驱动的 DP 方法的通信量是多少ZeRO-DP 的通信量最先进的 All-Reduce 实现采用两步法第一步是 Reduce-Scatter 操作第二步是 All-Gather 操作每个流程的总数据移动量为 $\Psi$ 个元素对于 $\Psi$ 个元素的数据。因此标准 DP 在每个训练步骤中会产生 $2\Psi$ 次数据移动。通过梯度分区$P_{osg}$每个进程只存储更新相应参数分区所需的梯度部分。因此ZeRO 只需要对梯度先进行 Reduce-Scatter 操作产生的通信量为 $\Psi$。在每个进程更新完自己负责的参数分区后会执行一次 All-Gather从所有数据并行进程中收集所有更新的参数这也会产生 $\Psi$ 的通信量。因此每个训练步骤的总通信量为 $\Psi \Psi 2\Psi$与标准 DP 相同。在参数分区$P_{osgp}$后每个数据并行进程只存储其更新的参数。因此在前向传播过程中它需要接收所有其他分区的参数。不过这可以通过流水线操作来避免内存开销——在对模型中与特定分区对应的部分进行前向传播计算之前负责该分区的数据并行进程可以向所有数据并行进程广播权重。一旦该分区的前向传播计算完成参数就可以被丢弃。因此总通信量为 $\Psi × N_d / N_d \Psi$。通过在整个前向传播中通过 All-gather 传播参数重新获取参数并在使用完参数后将其丢弃而在后向传播时需要以相反的顺序再次进行参数获取。参数的通信为 $2\Psi$在参数更新时只需要执行一个 Reduce-Scatter 操作通信量为 $\Psi$因此总通信量是 $3\Psi$是标准 DP 的 1.5 倍。ZeRO-R 的通信开销ZeRO-R 的通信开销取决于模型大小、检查点策略和模型并行MP策略。与标准模型并行相比其中没有对激活进行分区ZeRO-R $P_{a}$ 的通信开销通常不到标准模型并行的十分之一。在使用激活检查点的 Megatron-LM 中每个 Transformer 块在前向传播中执行两次大小为 $batch × seq × length × hidden_dim$ 的 All-Reduce 操作然后在反向传播中再执行两次。在使用激活检查点的 ZeRO-R 中每个前向重计算激活之前需要执行一个额外的 All-Gather 操作。通常情况下对每个 Transformer 块的输入激活进行检查点因此每个 Transformer 块需要一个 All-Gather 操作。因此ZeRO-R $P_{a}$ 的通信开销为 $seq_length × hidden_dim$仅增加不到 10%。当 MP 与 DP 一起使用时ZeRO-R $P_{a}$ 可以将数据并行通信量减少一个数量级而模型并行通信量只增加 10%并且当数据并行通信是性能瓶颈时可以显著提高效率。通过模型并行可以减少数据并行的内存消耗从而可以成比例地增加批处理大小。对于大模型MP 可以增加到单节点最大 NPU 数量从而可以将批处理大小增加多达 NPU 数据量的倍数。数据并行训练的通信量与批处理大小成反比由于 $P_{a}$ 导致批处理大小增加一个数量级可能会导致数据并行通信量减少一个数量级。如果应用 $P_{acpu}$则分区激活检查点会被卸载到 CPU将激活内存需求减少到接近零但与 $P_{a}$ 相比往返 CPU 内存的数据移动增加了 2 倍。如果 DP 通信量是主要瓶颈由于批处理大小较小$P_{acpu}$ 也可以通过增加批处理大小来提高效率只要 CPU 数据传输开销小于 DP 通信量开销。ZeRO-Infinity 的效率估算模型可以使用峰值计算吞吐量$peak_{tp}$、数据移动带宽$bw$及其算术强度$ait$来估算 ZeRO-Infinity 的训练效率因为它还涉及到了 NPU 之间的数据移动。工作负载的算术强度AIT是总计算量与计算所需数据量之间的比率。它描述了每次数据移动所需的计算量。AIT 越高意味着对数据移动带宽的要求越低因为每加载一个数据加速器就能完成更多计算。$$ ait \frac{total_computation}{total_data_movement} $$因此 ZeRO-Infinity 的效率可以大致估算为$$ \begin{aligned} compute_time \frac{total_computation}{peak_{tp}} \ communication_time \frac{total_data_movement}{bw} \ \frac{total_computation}{ait × bw} \ efficiency \frac{compute_time}{compute_timecommunication_time} \ \frac{ait × bw}{ait × bw peak_{tp}} \end{aligned} $$同样以 Transformer 为例每次迭代的总计算量可以由参数数量、序列长度和批量大小估算即对于前向传播为 $2 × bsz × seq × params$反向传播的成本大约是前向传播的两倍。因此可以估算计算量$$ computation_per_iter 2 × 4 × bsz × seq × parameters 2 × 4 × 12 × bsz × seq × nl × hd^2 $$在前向和反向传播期间模型参数必须从源位置加载到 NPU 寄存器至少两次前向传播期间和实际后向传播期间导致 2 次的数据移动。在存在激活检查点的情况下可以在后向传播过程中额外加载一次参数以进行重新计算。此外梯度必须至少从 NPU 寄存器存储到其最终位置一次。因此假设参数和梯度存储在相同的最终位置则前向和后向传播期间的总数据移动将为 $4 × parameters$即 $2 × 4 × parameters$以字节为单位。因此参数和梯度的 ait 为$$ seq × bsz $$在优化器迭代期间必须至少读取一次优化器状态并且必须至少写入一次优化器状态。因此总数据移动量为 $2 × optimizer_states$大约为 $2 × 16 × parameters$ 字节。因此在完整的训练迭代期间优化器状态的 ait 为$$ seq × bsz / 4 $$在前向传播期间激活检查点必须保存到其最终位置并且必须在后向传播期间获取。因此激活检查点的总数据移动量以字节为单位为 $2 × total_activation_checkpoints_in_bytes$带入之前计算的激活检查点大小可以得到总数据移动量 $4 × nl/ci × hd × seq × bsz$。所以激活检查点的 ait 为$$ 24 × hd × ci $$模型状态和激活检查点对带宽的要求大不相同前者只取决于批量大小和序列长度而后者只取决于激活检查点的频率和模型的隐藏维度大小。在实际中参数和梯度的带宽超过 70 GB/s即使是最小的批处理量也能实现超过 50% 的效率。在这种带宽下数据移动理论上可以与计算完全重叠从而实现 100% 的效率。与参数和梯度相比优化器状态需要高出近 4 倍的带宽才能达到 50% 的效率。此外优化器状态在前向和后向传播结束时更新不能与计算重叠。因此它们需要更大的带宽来保持整个 DL 工作负载的效率。例如在每个 NPU 的批处理量为 2 的情况下要达到 90% 的效率需要近 1.5 TB/s 的有效带宽甚至超过了 NPU 内存带宽。启用激活检查点后即使隐藏大小为 2K2 GB/s 的微薄带宽也能维持 50% 以上的效率当隐藏大小超过 8K 时带宽需求降至 1 GB/s 以下。小结与思考FSDP 通过在多个 NPU 间分片模型的权重、梯度和优化器状态显著降低了单个 NPU 上的内存占用使得训练更大模型成为可能。ZeRO 技术通过优化器状态、梯度和参数的分区以及将激活检查点卸载到 CPU 或 NVMe 内存进一步提升了内存效率支持了大模型的训练。ZeRO-Infinity 作为 ZeRO 的扩展智能地在 NPU、CPU 和 NVMe 等异构内存组件之间迁移参数消除了 NPU 内存限制为训练高达万亿参数的模型提供了可能。尽管 ZeRO 技术增加了通信开销但通过精心设计的通信策略和利用现代硬件的高带宽可以保持训练效率同时实现内存效率和计算性能的平衡。想要进一步了解与本主题相关的分布式并行全景可继续阅读本目录下的 02DataParallel.mdDP 与 DDP 基础、04TensorParallel.md张量并行、05PipelineParallel.md流水线并行以及 06HybridParallel.mdZeRO 与 PP/TP 的混合组合。赞分享文档教程人工智能【免费下载链接】AISystemAISystem 主要是指AI系统包括AI芯片、AI编译器、AI推理和训练框架等AI全栈底层技术项目地址https://gitcode.com/GitHub_Trending/ai/AISystem点击查看免费下载相关推荐GoJieba高性能中文分词Golang实现完全指南GoJieba高性能中文分词Golang实现完全指南 想要在Go项目中实现高效的中文分词吗 GoJieba正是你需要的终极解决方案作为结巴中文分词混合精度训练nn-zero-to-heroFP16训练的内存与速度优化混合精度训练nn zero to heroFP16训练的内存与速度优化 引言深度学习训练的内存瓶颈与解决方案 在现代深度学习实践中模型规模的爆炸式增长使得示例工程深度学习人工智能DeepSpeed 训练优化完全指南分布式训练、混合精度、ZeRO 与多维并行DeepSpeed 训练优化完全指南分布式训练、混合精度、ZeRO 与多维并行 DeepSpeed 面向“训练先进深度学习模型”这一复杂系统工程问题通过轻量人工智能大模型深度学习分布式训练预训练强化学习模型优化上一篇SymPy 稀疏矩阵工具详解sparsetools 模块的 banded、_doktocsr 与 _csrtodok下一篇SurfSense Google Maps 子代理深度解析基于实时地图数据的本地商家与评论结构化采集创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 AISystem 完全分片数据并行 FSDP 与 ZeRO 优化器深度解析：从混精度训练到万亿参数级模型的内存优化实战
文档教程人工智能【免费下载链接】AISystemAISystem 主要是指AI系统包括AI芯片、AI编译器、AI推理和训练框架等AI全栈底层技术项目地址https://gitcode.com/GitHub_Trending/ai/AISystem点击查看免费下载本篇文章聚焦 AISystem 开源仓库中分布式并行章节的完全分片数据并行FSDP与零冗余优化器ZeRO技术系统讲解精度格式、混合精度训练、损失缩放、显存估算等前置知识以及 ZeRO-DP、ZeRO-R、ZeRO-Infinity 的原理、通信开销分析与 PyTorch FSDP 落地实现。读者学完后将掌握大模型训练中以通信换显存的核心理念能够独立配置 FSDP 自动包装策略、混合精度策略并正确估算大模型训练的内存占用与训练效率。本文内容以 03ZeRODP.md 为主体上一节 02DataParallel.md 讲解了通用数据并行DP与分布式数据并行DDP的实现差异本文在其基础上继续深入权重、优化器、梯度三类模型状态的分布式并行方案。前置知识理解 FSDP 所需的三大基础在介绍针对权重数据、优化器数据和梯度数据进行分布式数据并行的算法 FSDP 前需要先掌握三项前置知识如何执行混合精度训练、如何对显存消耗进行估算以及常用的精度格式。这三者直接决定了 FSDP 中内存分片与通信调度的设计动机。精度格式FP32、FP16、BF16 与 TF32深度学习中常用的精度格式包括FP32、FP16、BF16和TF32四种它们各有不同的位布局、数值范围与适用场景。FP16IEEE 754 标准下的半精度浮点格式。随着深度学习的发展FP16 逐渐取代了 FP32 成为训练的主力类型。因为相较于 FP32更低的精度并不会对神经网络性能产生重大影响额外的精度不会带来任何好处反而会更慢、占用更多内存并降低通信速度。FP16 通常用于混合精度训练MindSpore/PyTorch也用于训练后量化以加快推理速度。其他用于量化的格式还有整数 INT88 位整数、INT44 位甚至 INT1二进制值。BF16谷歌最初开发的另一种 16 位格式全称为 Brain Floating Point Format简称 bfloat16名字源于构思出该格式的谷歌 AI 研究小组谷歌大脑。它最初被使用在谷歌芯片 TPU 中后被广泛使用在 GPU、NPU 等 AI 芯片中。由于具有更多的指数位BF16 常用于处理 FP16 的溢出问题。FP32IEEE 754 标准下的单精度浮点数在很长一段时间内都是深度学习的主力。长期以来神经网络的权重、激活和其他值都默认用 FP32 表示。TF32一种十分特殊的格式无需显示设置而是自动执行。它将 FP32 数据截断为 TF32 进行计算然后再转换回 FP32。这一创新的最大优势在于编译器只需在 AI 编译层提供支持其他代码部分可以继续使用动态范围相同但精度较高的 FP32无需修改。TF32 的快速插入性使得利用 AI 计算加速核心的速度成为可能而无需过多修改现有代码。TF32 采用与 FP16 相同的 10 位尾数这满足了 AI 工作负载的精度要求同时使用与 FP32 相同的 8 位指数因此具有相同的数值范围。从技术上讲它可以视为一种 19 位格式也可以视为扩展精度的 BF16。TF32 的优势在于其格式与 FP32 相同。当使用 TF32 进行计算时输入 FP32 操作数的尾数从 23 位舍入到 10 位然后进行精确乘法运算最后以正常的 FP32 格式进行累加。相比之下FP16 和 BF16 等格式需要更多工作因为它们涉及不同的位布局。尽管如此这些不同的精度格式仍然值得使用因为它们可以减少内存带宽、存储和计算位数从而提升执行速度。在 PyTorch 中TF32 通过以下开关启用torch.backends.npu.matmul.allow_tf32 True torch.backends.npu.matmul.allow_16 True混合精度训练半精度计算 权重备份在当今大模型训练的背景下混合精度训练已然成为一种备受推崇的普遍做法。通过采用混合精度训练能够将训练速度显著提升数倍而又不会对模型的整体性能产生重大影响。在传统科学计算领域人们通常追求较高的精度如 FP128 或 FP64但深度学习中面临的实际上是一个高维函数拟合或近似的优化问题并不需要过于精确的数值表示且使用低精度会带来显著的计算速度提升。使用 FP16 训练有时会出现下溢问题FP16 的有效动态范围约为 $5.96e^{-8} \sim 65504$在训练后期例如激活函数的梯度会非常小甚至在梯度乘以学习率后值更小。由于 FP16 的精度范围有限过小的梯度可能导致更新无效这时就需要使用混合精度训练。混合精度训练可以分为两个部分半精度和权重备份如下图所示这里以 FP16 和 FP32 举例。训练开始时准备两套模型状态一套为 FP32 类型优化器状态和模型参数另一套为 FP16 类型模型参数。前向传播、反向传播时都使用 FP16 类型的模型参数进行计算而在参数更新时将梯度与学习率 $\eta$ 相乘更新到 FP32 类型的模型状态上新一轮训练中再次将 FP32 类型的模型拷贝为 FP16 类型的模型。这个过程就是混合精度训练。由于在计算密集的前向传播、反向传播中使用了 FP16 进行计算与单精度相比训练速度会大幅度提升。另外激活值在训练过程中占用内存的很大部分使用 FP16 储存激活值在大批量训练时也会节省内存在分布式环境下使用 FP16 进行梯度通信也会降低通信量。为了获得最佳性能混合精度中需要额外选择合适的批量大小。通常建议使用 2 的幂次方作为批量大小并与输入/输出神经元的数量相匹配通常为 8 的倍数但也可能更高具体取决于所使用的硬件和模型的数据类型。而在 FSDP 中可以通过在 torch 中指定MixedPrecision配置类进行混合精度的自动配置分别指定参数精度param_dtype、梯度通信精度reduce_dtype和 Buffer 精度buffer_dtypefpSixteen MixedPrecision( param_dtypetorch.float16, # Gradient communication precision. reduce_dtypetorch.float16, # Buffer precision. buffer_dtypetorch.float16, ) bfSixteen MixedPrecision( param_dtypetorch.bfloat16, # Gradient communication precision. reduce_dtypetorch.bfloat16, # Buffer precision. buffer_dtypetorch.bfloat16, ) model FSDP(model, auto_wrap_policyt5_auto_wrap_policy, mixed_precisionbfSixteen)损失缩放解决 FP16 下溢问题的另一手段解决 FP16 下溢问题的另一个方法是损失缩放Loss Scale。训练到后期梯度特别是激活函数平滑段的梯度会特别小FP16 表示容易产生下溢现象。为了解决梯度过小的问题需要对损失进行缩放——由于链式法则的存在损失的缩放也会作用在梯度上。缩放后的梯度就会平移到 FP16 有效的展示范围内。不过缩放并非对所有网络而言都是必须的缩放的取值也会特别大一般在 8 - 32k 之间。在 PyTorch 中可以通过以下方式实现自动损失缩放from torch.cuda.amp import GradScaler, autocast scaler GradScaler() with autocast(): output model(input) loss loss_fn(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这种损失缩放的方式是动态的每当梯度溢出时减少损失缩放规模并且间歇性地尝试增加损失规模从而实现在不引起溢出的情况下使用最高损失缩放因子更好地恢复精度。动态损失缩放的算法从比较高的缩放因子开始如 $2^{24}$然后开始训练并在迭代中检查数是否会溢出Infs/Nans如果没有梯度溢出则不调整缩放因子继续进行迭代如果检测到梯度溢出则缩放因子减半重新确认梯度更新情况直到参数不出现在溢出的范围内在训练后期loss 已经趋近收敛稳定梯度更新的幅度往往变小这时可以允许更高的损失缩放因子来再次防止数据下溢。内存消耗估算模型状态与剩余状态在神经网络模型的训练中合理估算和管理内存消耗非常重要。内存存储主要分为两大块模型状态Model States和剩余状态Residual States。模型状态指和模型本身相关的、必须存储的内容具体包括优化器状态Optimizer StatesAdam 优化算法中的 Momentum 和 Variance梯度Gradients模型梯度 $G$参数Parameters模型参数 $W$。剩余状态是并非模型必须的但在训练过程中会额外产生的内容具体包括激活值Activation在反向传播过程中使用链式法则计算梯度时会用到。有了它算梯度会更快但它不是必须存储的因为可以通过重新前向传播来计算临时存储Temporary Buffers例如把梯度发送到某个 NPU 上进行 All-Reduce 时产生的存储碎片化的存储空间Unusable Fragment Memory虽然总存储空间是够的但是如果取不到连续的存储空间相关的请求也会失败。对这类空间浪费可以通过内存整理来解决。以 FP32 与 FP16 的混合精度训练为例假设模型的参数量是 $\Phi$那么模型状态所消耗的空间为| Model States | Size (Byte) | |-- |--- | | FP32 Parameters | 4$\Phi$ | | FP32 Adam Optimizer Momentum | 4$\Phi$ | | FP32 Adam Optimizer Variance | 4$\Phi$ | | FP16 Gradients | 2$\Phi$ | | FP16 Parameters | 2$\Phi$ | | Total | 16$\Phi$ |而剩余状态与具体模型架构有关需要具体分析。接下来基于 Transformer 的架构进行分析因为所有参数超过 10 亿的 SOTA 模型都遵循这一架构。分析假设使用 Adam 优化器进行混合精度训练因为该优化器是训练基于 Transformer 的大模型较为常用的选择。模型状态基于 Transformer 的模型中的参数总数主要取决于隐藏维度$hd$和 Transformer 层数$nl$。Transformer 块中的几乎所有参数都来自每个块内的四个线性层其大小分别为$hd$,$3hd$、$hd$,$hd$、$hd$,$4hd$和$4hd$,$hd$。因此基于 Transformer 的模型中的总参数可以近似为$$ 12 × nl × hd^2 $$剩余状态剩余状态主要是指激活内存它取决于模型结构、批量大小$bsz$和序列长度$seq$而且可能相当大。不过激活所需的内存可以通过激活检查点activation checkpointing大大减少。假设 $ci$ 是两个激活检查点之间的 Transformer 块数$bsz × seq × hd$ 是每个 Transformer 块的输入大小激活检查点所需的内存估计为$$ 2 × bsz × seq × hd × nl / ci $$激活工作内存激活工作内存是反向传播过程中所需的内存用于在执行实际反向传播之前重新计算激活即两个连续激活检查点之间的激活量。例如如果为每个 Transformer 块创建一个激活检查点那么内存就是每个 Transformer 块的总激活量其字节数约为$$ bsz × seq × ci × (16 × hd 2 × attn_heads × seq) $$模型状态工作内存模型状态工作内存是指在将所有模型状态卸载到 CPU 或 NVMe 之后对模型中最大的单个算子执行前向或后向传播所需的 NPU 内存最小量。这大约由该算子的参数和梯度的大小决定因为必须至少有足够的内存来保存反向传播的参数及其梯度。Transformer 最大的算子是将隐藏状态从 $h$ 转换为 $4h$ 的线性层该线性层的参数和梯度的大小为$$ 4 × hd × 4hd $$完全分片数据并行 FSDP 总览完全分片数据并行Fully Sharded Data ParallelFSDP在分布式 AI 系统中具有重要地位不仅能提高并行效率还能减少显存内存消耗这两方面的优势为模型的大规模训练带来了显著的好处。值得注意的是并行效率和内存消耗之间存在密切关联降低内存占用可以使使用更大的并行度进而提升整体的并行加速效果。完全分片数据并行是基于零冗余优化器ZeRO在 AI 框架中的具体实现主要的实现有微软的DeepSpeed和 Meta 的Fairscale其中 Fairscale 被集成到 PyTorch 中并作为 FSDP 的实现基础。本节将从零冗余优化器的常用技术入手深入剖析如何降低内存开销并提高训练效率。在数据并行中每个 NPU 都需要保存一份完整的参数模型状态和剩余状态。但并不是所有参数在训练的整个过程中都会被使用到而是在特定阶段某个层的前向或反向传播才被使用因此可以在不需要使用的时候将它转移到其他地方以节省内存空间。ZeRO 有两套优化方案ZeRO-DP旨在减少模型状态的内存占用ZeRO-R旨在减少剩余状态内存消耗。这两套方案使 ZeRO 能够在保持高效的同时减少内存占用。此外在混合并行场景中ZeRO 也可以与流水线并行PP、张量并行TP组合使用相关内容可参考 06HybridParallel.md 中的混合并行分析。ZeRO-DP对模型状态进行分区ZeRO-DP 对模型状态进行切分具体来说每个 NPU 只存储 $\frac{1}{N_d}$ 的模型状态其中 $N_d$ 为并行度在需要时通过集合通信 All-Gather 获取参数。ZeRO-DP 保留了数据并行训练DP的高效率同时实现了模型并行MP的内存效率优势。由于数据并行的模型状态在所有数据并行进程中冗余存储因此内存效率低下但数据并行具有更高的计算粒度和更低的通信量从而具有更高的训练效率。模型并行的通信开销很大因此可扩展性比数据并行低但模型并行对模型状态进行分区获得了较高的内存效率。ZeRO-DP 对模型状态进行分区而不是复制它们并使用动态通信调度最小化通信量。通过这样做ZeRO-DP 随着数据并行程度的增加线性减少模型在每块 NPU 的内存占用同时保持通信量接近默认数据并行的通信量从而保持效率。ZeRO-DP 有三个主要优化阶段分别对应于优化器状态、梯度和参数的划分在累积启用时优化状态分区Partition optimizer states$P_{os}$又称为 ZeRO-1将优化器状态按并行度均匀分区每个进程只需存储 $\frac{1}{N_d}$ 的优化器状态。这可将内存消耗减少到 1/4且无额外通信开销。添加梯度分区Partition gradients$P_{osg}$又称为 ZeRO-2在优化器状态分区的基础上对梯度也进行分区。每个进程只需存储用于更新自身参数分区所需的梯度。这可减少 8 倍的内存消耗且无额外通信开销。添加参数分区Partition parameters$P_{osgp}$又称为 ZeRO-3在优化器状态和梯度分区的基础上对参数也进行分区。每个进程只存储自身的参数分区在前向反向传播时需要从其他进程收集所需的参数分区。这会使通信量增加约 50%但可以实现与并行度 $N_d$ 成正比的内存减少。通过这三个阶段的优化ZeRO-DP 最终能够在保持数据并行高效的同时将每个 NPU 的内存消耗降低至 $\frac{1}{N_d}$ 的水平使得利用少量硬件资源训练万亿参数等超大模型成为可能。下面进行每个阶段的详细介绍这里假设模型使用混合精度训练模型参数量为 $4\Psi$。ZeRO-1 内存分析优化器状态是训练过程中 NPU 内存中的主要保存内容但只有在参数更新的时候才会被用到。ZeRO-1 的核心思想是将优化器状态分布到多个 NPU 上减少每个 NPU 所需的显存在需要参数更新时进行聚合。其执行步骤如下数据分片a从优化器状态分片开始将优化器状态分成 N 份每个 NPU 保存一份分片并将训练批次数据batch data分成 N 份每个 NPU 处理一份数据。前向与后向计算b每个 NPU 执行一步前向forward和后向backward计算得到局部梯度 $G_i$。梯度聚合b对各个 NPU 上的局部梯度 $G_i$ 执行 All-Reduce 操作得到完整梯度 $G$。这一步单个 NPU 的通信量为 $2\Phi$。权重更新c使用完整梯度 $G$ 和优化器状态更新权重 $W$。每个 NPU 保存一部分权重 $W$并通过 All-Gather 操作从其他 NPU 获取更新后的部分权重完成权重更新。此时单个 NPU 的通信量为 $\Phi$。在 $P_{os}$ 阶段将 Adam 优化器状态根据数据并行维度 $N_d$ 分成等份。每个 NPU 只需存储和更新总优化器状态的 $1/N_d$并更新对应参数。通过分片和聚合操作显存占用从 $4\Psi K\Psi$ 降低到 $4\Psi K\Psi / N_d$。当 $N_d$ 很大时显存占用接近于 $4\Psi$带来约 4 倍显存节约。ZeRO-2 内存分析ZeRO-2 在 ZeRO-1 的基础上进一步优化通过对梯度Grad也进行切分减少显存占用并提高通信效率。其执行步骤如下数据分片从优化器状态和梯度分片开始将优化器状态和梯度分成 N 份每个 NPU 保存一份分片并将训练批次数据batch data分成 N 份每个 NPU 处理一份数据。前向与后向计算每个 NPU 执行一步前向forward和后向backward计算得到局部梯度 $G_i$。梯度分片与聚合对各块 NPU 上的局部梯度 $G_i$ 执行 Reduce-Scatter 操作确保每个 NPU 只维护自己负责的梯度部分。比如NPU 1 负责维护梯度 $G_1$其他 NPU 只需要将 $G_1$ 对应位置的梯度发送给 NPU 1。聚合完毕后NPU 1 释放无用的显存部分单卡通信量为 $\Phi$。权重更新每个 NPU 使用自身维护的梯度 $G_i$ 和优化器状态 $O_i$ 更新相应的权重 $W_i$。权重聚合对权重 $W_i$ 执行 All-Gather 操作将其他 NPU 的权重 $W_i$ 同步至完整权重 $W$单卡通信量为 $\Phi$。在 $P_{osg}$ 阶段梯度与优化器强相关因此优化器可以更新其独立的梯度参数。更新梯度参数时使用 Reduce-Scatter 操作梯度参数更新后立即释放。具体实现中使用分桶Bucket技术将梯度分到不同的桶中并在桶上进行 Reduce-Scatter 操作。通过移除梯度和优化器状态冗余显存占用从 $4\Psi K\Psi$ 降低到 $2\Psi K\Psi / N_d$。当 $N_d$ 很大时显存占用接近于 $2\Psi$带来约 8 倍显存节约。ZeRO-3 内存分析ZeRO-3 在 ZeRO-1 和 ZeRO-2 的基础上进一步优化通过对优化器状态、梯度和权重进行全面切分最大化显存节约这种优化使得训练超大模型成为可能。其执行步骤如下数据分片对优化器状态、梯度和权重进行全面切分每个 NPU 保存一份分片将训练批次数据batch data分成 N 份每块 NPU 处理一份数据。前向计算在前向计算过程中对权重 $W$ 执行 All-Gather 操作从各 NPU 获取分布在不同 NPU 上的权重得到完整的权重 $W$。不属于自身的权重 $W_{others}$ 被抛弃单卡通信量为 $\Phi$。后向计算在后向计算过程中再次对权重 $W$ 执行 All-Gather 操作取回完整权重并抛弃不属于自身的部分 $W_{others}$单卡通信量为 $\Phi$。梯度聚合后向计算得到各自的梯度 $G_i$ 后对梯度 $G_i$ 执行 Reduce-Scatter 操作从其他 NPU 聚合自身维护的梯度 $G_i$。聚合操作结束后立刻抛弃不是自己维护的梯度单卡通信量为 $\Phi$。权重更新每块 NPU 只保存其权重参数 $W_i$由于只维护部分参数 $W_i$因此无需对 $W_i$ 执行 All-Reduce 操作。在 $P_{osgp}$ 阶段优化器状态、梯度和权重均进行划分。在前向和后向计算过程中通过广播从其他 NPU 中获取参数减少每块 NPU 中的显存占用。通过移除梯度、优化器状态和权重的冗余将显存占用从 $4\Psi K\Psi$ 降低到 $(4\Psi K\Psi) / N_d$。这种方法通过增加通信开销以通信换显存使得显存占用与 $N_d$ 成正比。显存占用的优化带来了 1.5 倍单卡通信量的增加。ZeRO-R优化剩余状态除了优化模型状态优化器状态、梯度和参数的内存利用率ZeRO 还专门针对剩余状态如激活数据、临时缓冲区和内存碎片等进行了优化以进一步减少内存开销。ZeRO-R 对剩余状态进行了切分和优化主要包括以下几个策略分区激活检查点Partitioned Activation Checkpointing$P_{a}$解决了模型并行时激活内存冗余的问题。在模型并行中每个 NPU 需要保存完整的输入激活数据才能计算自己分到的模型部分。ZeRO-R 将激活检查点按模型并行度 $N_m$ 进行分区每个 NPU 只需存储 $\frac{1}{N_m}$ 的激活检查点在需要时通过 All-Gather 操作重构出完整激活数据从而按 $N_m$ 的比例减少激活内存。在极端情况下当模型规模很大时ZeRO-R 甚至可以将分区后的激活检查点卸载到 CPU 内存$P_{acpu}$再次降低 NPU 内存占用代价是额外的 Host-Device 通信开销。该策略在大模型训练时会自动开启以保证足够的 NPU 内存用于计算。恒定大小的缓冲区Constant Size Buffer$C_{b}$一些操作如 All-reduce 需要将张量拼成连续的临时缓冲区。使用恒定大小的缓冲区来避免临时缓冲区随着模型大小的增加而爆炸同时使它们足够大以保持效率。内存碎片化整理Memory Defragmentation$M_{d}$在训练过程中由于激活检查点、梯度等张量生命周期的差异会产生大量内存碎片。ZeRO-R 通过预分配和动态管理这些张量的内存减少了内存碎片和内存分配器的开销提高了内存利用率。通过以上优化策略ZeRO-R 很好地补充和完善了 ZeRO-DP 优化模型状态内存的功能。两者相结合ZeRO 优化器能最大限度减少大模型训练的内存占用为未来万亿参数级别的神经网络模型铺平了道路。ZeRO-Infinity突破 NPU 内存壁垒ZeRO-Infinity 是 ZeRO 的扩展可以将深度学习训练扩展到前所未有的规模。具体来说它突破了 NPU 内存壁垒的限制使得训练具有数万亿参数的模型成为可能这是迄今为止最先进系统所无法企及的量级。此外它为训练具有一千万亿参数的模型铺平了道路——充分利用系统的全部内存容量利用 NPU、CPU 和 Non-Volatile Memory ExpressNVMe等所有异构内存组件的能力。在 ZeRO-Infinity 中参数从较慢的内存源如 CPU 和 NVMe无缝迁移到 NPU在 NPU 上被合并为完整的层。梯度计算完成后这些参数被聚合、重新分区然后重新卸载回较慢的内存组件。内存资源的编排确保了最佳利用和最小的开销。这种创新的方法不仅克服了 NPU 内存的常规限制而且提升了分布式框架的可扩展性。FSDP 简单实现PyTorch 中的落地实践和 DDP 一样可以通过简单的嵌套使用 ZeRO 优化器来实现 FSDP。将模型的参数、梯度和优化器状态分片从而显著减少单个 NPU 的内存占用实现更大模型的训练。以下是一个简单的代码示例展示了如何使用torch.distributed.fsdp中的FullyShardedDataParallelFSDP类来实现完全分片数据并行import torch from torch.distributed.fsdp import FullyShardedDataParallel as FSDP torch.cuda.set_device(device_id) sharded_module FSDP(my_module) optim torch.optim.Adam(sharded_module.parameters(), lr0.0001) x sharded_module(x, y3, ztorch.Tensor([1])) loss x.sum() loss.backward() optim.step()首先设置当前的 NPU然后将my_module包装成一个 FSDP 模块。这会将模型的参数、梯度和优化器状态在多个 NPU 之间进行分片从而减少每个 NPU 上的内存占用。再来看一个更详细的例子使用fsdp_main函数用于分布式训练 T5 模型。函数通过setup_model函数加载模型和分词器并设置分布式训练的相关环境变量包括local_rank、rank和world_size。在数据集和数据加载器设置之后通过functools.partial函数部分应用了transformer_auto_wrap_policy并指定T5Block为要自动包装的 Transformer 层。这一步的目的是定义一个自动包装策略用于后续的模型分片和并行处理。接下来定义了sharding_strategy变量并将其设为ShardingStrategy.SHARD_GRAD_OP。这表示使用 ZeRO-2 分片策略如果要使用 ZeRO-3 策略可以将其设为FULL_SHARD。分片策略决定了在分布式训练中如何管理和分配模型参数以优化内存使用和计算性能。为了支持混合精度训练bf16如果当前 CUDA 版本和 NCCL 版本支持 bf16并且 CUDA 版本大于或等于 11.0那么bf16_ready变量为True并将mp_policy设为bfSixteen。否则mp_policy设为None默认使用 fp32单精度训练。在模型仍在 CPU 上时代码将模型传入FSDP模块配置自动包装策略、混合精度策略mixed_precision以及当前 NPU IDdevice_id。这一步是将模型转换为分布式训练的模式以便在多个 NPU 之间分片和并行计算。def fsdp_main(args): model, tokenizer setup_model(t5-base) local_rank int(os.environ[LOCAL_RANK]) rank int(os.environ[RANK]) world_size int(os.environ[WORLD_SIZE]) # Set dataset and dataloader here t5_auto_wrap_policy functools.partial( transformer_auto_wrap_policy, transformer_layer_cls{ T5Block, }, ) sharding_strategy: ShardingStrategy ShardingStrategy.SHARD_GRAD_OP #for Zero2 and FULL_SHARD for Zero3 torch.cuda.set_device(local_rank) bf16_ready ( torch.version.cuda and torch.cuda.is_bf16_supported() and LooseVersion(torch.version.cuda) 11.0 and dist.is_nccl_available() and nccl.version() (2, 10) ) if bf16_ready: mp_policy bfSixteen else: mp_policy None # defaults to fp32 # model is on CPU before input to FSDP model FSDP(model, auto_wrap_policyt5_auto_wrap_policy, mixed_precisionmp_policy, #sharding_strategysharding_strategy, device_idtorch.cuda.current_device()) optimizer optim.AdamW(model.parameters(), lrargs.lr) scheduler StepLR(optimizer, step_size1, gammaargs.gamma) for epoch in range(1, args.epochs 1): train_accuracy train(args, model, rank, world_size, train_loader, optimizer, epoch, samplersampler1) scheduler.step()当然如果需要更简单而全面的配置可以通过deepspeed库进行便捷的 ZeRO 配置。例如在 DeepSpeed 配置文件中声明zero_optimization的 stage 级别1/2/3即可自动完成对应阶段的优化器状态、梯度与参数分区。ZeRO 通信分析以通信换内存的效率账本无论是零冗余优化还是卸载到 CPU 和 NVMe 内存一个关键问题是它们有限的带宽是否会影响训练效率。很自然地会问是否在用通信量来换取内存效率。换句话说与标准 DP 方法相比ZeRO 驱动的 DP 方法的通信量是多少ZeRO-DP 的通信量最先进的 All-Reduce 实现采用两步法第一步是 Reduce-Scatter 操作第二步是 All-Gather 操作每个流程的总数据移动量为 $\Psi$ 个元素对于 $\Psi$ 个元素的数据。因此标准 DP 在每个训练步骤中会产生 $2\Psi$ 次数据移动。通过梯度分区$P_{osg}$每个进程只存储更新相应参数分区所需的梯度部分。因此ZeRO 只需要对梯度先进行 Reduce-Scatter 操作产生的通信量为 $\Psi$。在每个进程更新完自己负责的参数分区后会执行一次 All-Gather从所有数据并行进程中收集所有更新的参数这也会产生 $\Psi$ 的通信量。因此每个训练步骤的总通信量为 $\Psi \Psi 2\Psi$与标准 DP 相同。在参数分区$P_{osgp}$后每个数据并行进程只存储其更新的参数。因此在前向传播过程中它需要接收所有其他分区的参数。不过这可以通过流水线操作来避免内存开销——在对模型中与特定分区对应的部分进行前向传播计算之前负责该分区的数据并行进程可以向所有数据并行进程广播权重。一旦该分区的前向传播计算完成参数就可以被丢弃。因此总通信量为 $\Psi × N_d / N_d \Psi$。通过在整个前向传播中通过 All-gather 传播参数重新获取参数并在使用完参数后将其丢弃而在后向传播时需要以相反的顺序再次进行参数获取。参数的通信为 $2\Psi$在参数更新时只需要执行一个 Reduce-Scatter 操作通信量为 $\Psi$因此总通信量是 $3\Psi$是标准 DP 的 1.5 倍。ZeRO-R 的通信开销ZeRO-R 的通信开销取决于模型大小、检查点策略和模型并行MP策略。与标准模型并行相比其中没有对激活进行分区ZeRO-R $P_{a}$ 的通信开销通常不到标准模型并行的十分之一。在使用激活检查点的 Megatron-LM 中每个 Transformer 块在前向传播中执行两次大小为 $batch × seq × length × hidden_dim$ 的 All-Reduce 操作然后在反向传播中再执行两次。在使用激活检查点的 ZeRO-R 中每个前向重计算激活之前需要执行一个额外的 All-Gather 操作。通常情况下对每个 Transformer 块的输入激活进行检查点因此每个 Transformer 块需要一个 All-Gather 操作。因此ZeRO-R $P_{a}$ 的通信开销为 $seq_length × hidden_dim$仅增加不到 10%。当 MP 与 DP 一起使用时ZeRO-R $P_{a}$ 可以将数据并行通信量减少一个数量级而模型并行通信量只增加 10%并且当数据并行通信是性能瓶颈时可以显著提高效率。通过模型并行可以减少数据并行的内存消耗从而可以成比例地增加批处理大小。对于大模型MP 可以增加到单节点最大 NPU 数量从而可以将批处理大小增加多达 NPU 数据量的倍数。数据并行训练的通信量与批处理大小成反比由于 $P_{a}$ 导致批处理大小增加一个数量级可能会导致数据并行通信量减少一个数量级。如果应用 $P_{acpu}$则分区激活检查点会被卸载到 CPU将激活内存需求减少到接近零但与 $P_{a}$ 相比往返 CPU 内存的数据移动增加了 2 倍。如果 DP 通信量是主要瓶颈由于批处理大小较小$P_{acpu}$ 也可以通过增加批处理大小来提高效率只要 CPU 数据传输开销小于 DP 通信量开销。ZeRO-Infinity 的效率估算模型可以使用峰值计算吞吐量$peak_{tp}$、数据移动带宽$bw$及其算术强度$ait$来估算 ZeRO-Infinity 的训练效率因为它还涉及到了 NPU 之间的数据移动。工作负载的算术强度AIT是总计算量与计算所需数据量之间的比率。它描述了每次数据移动所需的计算量。AIT 越高意味着对数据移动带宽的要求越低因为每加载一个数据加速器就能完成更多计算。$$ ait \frac{total_computation}{total_data_movement} $$因此 ZeRO-Infinity 的效率可以大致估算为$$ \begin{aligned} compute_time \frac{total_computation}{peak_{tp}} \ communication_time \frac{total_data_movement}{bw} \ \frac{total_computation}{ait × bw} \ efficiency \frac{compute_time}{compute_timecommunication_time} \ \frac{ait × bw}{ait × bw peak_{tp}} \end{aligned} $$同样以 Transformer 为例每次迭代的总计算量可以由参数数量、序列长度和批量大小估算即对于前向传播为 $2 × bsz × seq × params$反向传播的成本大约是前向传播的两倍。因此可以估算计算量$$ computation_per_iter 2 × 4 × bsz × seq × parameters 2 × 4 × 12 × bsz × seq × nl × hd^2 $$在前向和反向传播期间模型参数必须从源位置加载到 NPU 寄存器至少两次前向传播期间和实际后向传播期间导致 2 次的数据移动。在存在激活检查点的情况下可以在后向传播过程中额外加载一次参数以进行重新计算。此外梯度必须至少从 NPU 寄存器存储到其最终位置一次。因此假设参数和梯度存储在相同的最终位置则前向和后向传播期间的总数据移动将为 $4 × parameters$即 $2 × 4 × parameters$以字节为单位。因此参数和梯度的 ait 为$$ seq × bsz $$在优化器迭代期间必须至少读取一次优化器状态并且必须至少写入一次优化器状态。因此总数据移动量为 $2 × optimizer_states$大约为 $2 × 16 × parameters$ 字节。因此在完整的训练迭代期间优化器状态的 ait 为$$ seq × bsz / 4 $$在前向传播期间激活检查点必须保存到其最终位置并且必须在后向传播期间获取。因此激活检查点的总数据移动量以字节为单位为 $2 × total_activation_checkpoints_in_bytes$带入之前计算的激活检查点大小可以得到总数据移动量 $4 × nl/ci × hd × seq × bsz$。所以激活检查点的 ait 为$$ 24 × hd × ci $$模型状态和激活检查点对带宽的要求大不相同前者只取决于批量大小和序列长度而后者只取决于激活检查点的频率和模型的隐藏维度大小。在实际中参数和梯度的带宽超过 70 GB/s即使是最小的批处理量也能实现超过 50% 的效率。在这种带宽下数据移动理论上可以与计算完全重叠从而实现 100% 的效率。与参数和梯度相比优化器状态需要高出近 4 倍的带宽才能达到 50% 的效率。此外优化器状态在前向和后向传播结束时更新不能与计算重叠。因此它们需要更大的带宽来保持整个 DL 工作负载的效率。例如在每个 NPU 的批处理量为 2 的情况下要达到 90% 的效率需要近 1.5 TB/s 的有效带宽甚至超过了 NPU 内存带宽。启用激活检查点后即使隐藏大小为 2K2 GB/s 的微薄带宽也能维持 50% 以上的效率当隐藏大小超过 8K 时带宽需求降至 1 GB/s 以下。小结与思考FSDP 通过在多个 NPU 间分片模型的权重、梯度和优化器状态显著降低了单个 NPU 上的内存占用使得训练更大模型成为可能。ZeRO 技术通过优化器状态、梯度和参数的分区以及将激活检查点卸载到 CPU 或 NVMe 内存进一步提升了内存效率支持了大模型的训练。ZeRO-Infinity 作为 ZeRO 的扩展智能地在 NPU、CPU 和 NVMe 等异构内存组件之间迁移参数消除了 NPU 内存限制为训练高达万亿参数的模型提供了可能。尽管 ZeRO 技术增加了通信开销但通过精心设计的通信策略和利用现代硬件的高带宽可以保持训练效率同时实现内存效率和计算性能的平衡。想要进一步了解与本主题相关的分布式并行全景可继续阅读本目录下的 02DataParallel.mdDP 与 DDP 基础、04TensorParallel.md张量并行、05PipelineParallel.md流水线并行以及 06HybridParallel.mdZeRO 与 PP/TP 的混合组合。赞分享文档教程人工智能【免费下载链接】AISystemAISystem 主要是指AI系统包括AI芯片、AI编译器、AI推理和训练框架等AI全栈底层技术项目地址https://gitcode.com/GitHub_Trending/ai/AISystem点击查看免费下载相关推荐GoJieba高性能中文分词Golang实现完全指南GoJieba高性能中文分词Golang实现完全指南 想要在Go项目中实现高效的中文分词吗 GoJieba正是你需要的终极解决方案作为结巴中文分词混合精度训练nn-zero-to-heroFP16训练的内存与速度优化混合精度训练nn zero to heroFP16训练的内存与速度优化 引言深度学习训练的内存瓶颈与解决方案 在现代深度学习实践中模型规模的爆炸式增长使得示例工程深度学习人工智能DeepSpeed 训练优化完全指南分布式训练、混合精度、ZeRO 与多维并行DeepSpeed 训练优化完全指南分布式训练、混合精度、ZeRO 与多维并行 DeepSpeed 面向“训练先进深度学习模型”这一复杂系统工程问题通过轻量人工智能大模型深度学习分布式训练预训练强化学习模型优化上一篇SymPy 稀疏矩阵工具详解sparsetools 模块的 banded、_doktocsr 与 _csrtodok下一篇SurfSense Google Maps 子代理深度解析基于实时地图数据的本地商家与评论结构化采集创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考