← 全部论文 / 2019

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

2019 超大规模稠密 LLM 训练
下载 PDF

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models - 中文验证版

英文原文卡片：zero_2019.md

状态：已翻译。

元数据

  • Slug： zero_2019
  • 年份： 2019
  • 会议： arXiv
  • 作者： Samyam Rajbhandari et al.
  • 阅读状态： read complete
  • 计算范式： 超大规模密集 LLM 训练
  • 主要来源： PDF、抽取文本

计算设置

论文明确给出一个 400-GPU V100 系统：

  • 25 个 DGX-2 节点上共 400 块 V100 GPUs。
  • 800 Gbps 节点间通信带宽。
  • 它分析了 32GB 当代 GPUs 和 1K V100 GPU 设置。

论文中评估的实现是 ZeRO-100B，一个结合了优化器状态和梯度分片以及残差状态内存优化的 PyTorch 系统。对于模型并行实验，ZeRO 与 Megatron-LM 结合。论文还给出了与 Megatron 硬件的具体对比：引用的最新 Megatron 结果使用了 32 个 DGX-2 节点上的 512 块 32GB V100 GPUs，而 ZeRO 的主要实验使用可用的 400-GPU 集群。

瓶颈

瓶颈是数据并行训练中的内存冗余。标准数据并行在每块 GPU 上复制优化器状态、梯度和参数，因此在当前 32GB GPUs 上对于超过约 1.4B 参数的模型会耗尽内存。尽管一个 1.5B 参数的类 GPT-2 模型的 fp16 权重仅约 3GB，但混合精度 Adam 还保留了 fp32 参数、动量、方差、梯度、激活、临时缓冲区和碎片化分配。

论文给出了具体的内存算术。带混合精度的 Adam 需要 fp16 参数和梯度加上 fp32 参数、动量和方差。对于一个 1.5B 参数模型，模型状态本身在激活之前就需要至少 24GB。在序列长度 1K 和 batch size 32 时，同一模型的激活需要约 60GB，或在激活 checkpointing 后约 8GB。对于一个 100B 参数的类 GPT 模型，即使有 checkpointing，激活仍可达到约 60GB。因此限制不仅是参数数量，而是整个训练状态。

模型并行避免了一些内存压力，但可能有较差的通信粒度和低硬件利用率。论文报告 Megatron-LM 在两个 DGX-2 节点上训练 40B 模型时，每块 V100 约 5 TFLOP/s，不到峰值的 5%。ZeRO 的前提是数据并行的可用性和模型并行的内存效率不应互斥。

方法适配

ZeRO 通过跨数据并行秩分片模型状态来适配密集 LLM 训练：

  • Stage 1 分片优化器状态。
  • Stage 2 分片梯度。
  • Stage 3 分片参数。
  • ZeRO-R 处理激活内存、临时缓冲区和碎片化。

该方法在去除冗余内存的同时，保留了数据并行的大部分可用性。Stage 1 优化器状态分片带来 4 倍模型状态内存减少，且通信量与数据并行相同。Stage 2 增加梯度分片达到 8 倍减少，通信量仍相同。Stage 3 增加参数分片，产生与数据并行度成比例的内存减少，通信适度增加，论文描述为约 50%。

ZeRO-R 在模型状态被分片后攻击残差状态内存。它在模型并行设置中分片激活 checkpoint，在适当时可将激活 checkpoint 卸载到 CPU，调整临时缓冲区大小以平衡内存和吞吐量，并主动管理碎片化使训练不会因明显的空闲内存而失败。激活分片尤其与硬件拓扑绑定：可以将激活内存减少模型并行度的倍数，允许更大的每 GPU batch 和更好的算术强度，但 CPU 卸载可能有害，除非更大的可行 batch 带来的好处超过额外的 PCIe 移动成本。

这是一个带宽和 batch 适配。ZeRO 避免将所有模型状态移动到 CPU，因为论文将 CPU 内存带宽和 PCIe 传输视为严重约束。相反，它使用数据并行组内的 GPU 聚合内存。更低的每 GPU 内存允许相同的 GPUs 运行更大的 micro-batch，这提高了吞吐量并可能减少每个 token 的数据并行通信。

证据

标题性证据是 400-GPU V100 评估。ZeRO-100B 在 400 块 GPUs 上运行最多 170B 参数的模型，而基线 Megatron 式模型并行系统在论文的比较中无法有效扩展到 40B 以上。对于 8B 到 100B 模型，ZeRO-100B 平均维持 15 PFLOP/s，超过峰值的 30%。论文将基线退化归因于通信量和离开 DGX-2 节点：带宽从通过 NVSwitch 的每链路 300GB/sec 降到通过 InfiniBand EDR 的每链路 12.5GB/sec。

模型配置表使 batch 效应具体化。一个 100B ZeRO 模型在 400 块 GPUs 上使用 16 路模型并行、125 层、hidden size 8192、64 个头、batch size 32、总 batch size 800。100B 基线使用 384 块 GPUs、128 路模型并行、128 个头、batch size 2、总 batch size 6。对于 170B，ZeRO 运行 400 块 GPUs、16 路模型并行、总 batch size 300，而列出的基线适配使用 256 路模型并行、总 batch size 2。

ZeRO 还报告了仅数据并行的可用性结果：在 128 块 GPUs 上最多 13B 参数，无模型或流水线并行，每 GPU 平均超过 40 TFLOP/s。没有 ZeRO 时，仅数据并行在 1.4B 参数处封顶，每 GPU 不到 20 TFLOP/s。最后，论文报告 Turing-NLG 超过 17B 参数，使用 ZeRO-100B 端到端训练，每 GPU 维持 41.4 TFLOP/s，达到 WebText-103 perplexity 10.21。

万亿参数分析明确是内存可行性论证，而非完整的训练成本结果。论文说明一个 1T 参数模型在 16 位精度下使用 Adam 约需 16TB 用于优化器状态、梯度和参数。分摊到 1024 块 GPUs 上为每 GPU 16GB，在 32GB V100 内存预算内。论文还警告，在当时计算条件下训练这样的模型需要超过一年，且需要 exaflop 系统才能在合理时间内完成。

历史影响

ZeRO 将 LLM 缩放重新表述为内存冗余问题。它使数据并行式训练对于之前需要更困难手工模型并行的模型变得可行，并成为 DeepSpeed 式大模型系统的基础。其历史地位与 Megatron-LM 互补：Megatron 在层内分片张量计算，而 ZeRO 分片数据并行通常复制的训练状态。

本卡片属于超大规模密集 LLM 范式，因为它表明仅加速器数量是不够的。有效资源是聚合 GPU 内存、可用 batch size、互联带宽和优化器状态放置。ZeRO 使这些资源可被访问，而无需将每个模型重写为定制的模型并行程序。

局限

  • 能容纳万亿参数模型不等于能便宜训练它们；计算时间仍然高得惊人。
  • 一些非常大的配置需要 CPU 卸载。
  • 通信和 batch-size 约束仍然存在。
  • 本文中 Stage 3 万亿规模的主张部分是分析性（analytical）的；评估的 ZeRO-100B 系统使用了完整 ZeRO 设计的一个子集。
  • CPU 卸载仅在狭窄情况下有用，因为 GPU-CPU-GPU 传输可能主导开销。
  • 该方法减少了冗余，但不减少每个 token 的密集 Transformer FLOPs。

链接