论文信息

  • 标题: ZeRO： Memory Optimizations Toward Training Trillion Parameter Models
  • 作者: Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He (Microsoft)
  • 发表年份: 2019
  • 链接: arXiv
    .02054

核心贡献

ZeRO (Zero Redundancy Optimizer) 是一种革命性的内存优化技术，通过消除数据并行训练中的内存冗余，使得训练超大规模模型成为可能。

主要创新

  1. ZeRO-DP: 将优化器状态、梯度和参数分片到多个 GPU
  2. ZeRO-R: 优化激活值、临时缓冲区和内存碎片
  3. 三阶段优化: 渐进式减少内存占用

内存优化阶段

Stage 1： 优化器状态分片 → 4x 内存减少
Stage 2： 梯度分片 → 8x 内存减少
Stage 3： 参数分片 → 线性内存扩展

技术原理

传统数据并行的问题

在传统数据并行中，每个 GPU 都保存完整的模型副本：

  • 模型参数
  • 梯度
  • 优化器状态 (如 Adam 的动量和方差)

这导致大量内存冗余。

ZeRO 的解决方案

ZeRO 将这些状态分片到不同的 GPU 上，每个 GPU 只保存一部分，需要时通过通信获取。

影响与应用

ZeRO 技术已被广泛应用于：

  • DeepSpeed: Microsoft 的深度学习优化库
  • GPT-3/4 训练: 支持千亿参数模型
  • 开源社区: Hugging Face、PyTorch 等集成

这项工作使得在有限硬件资源上训练超大模型成为可能，是大模型时代的关键基础设施。

Paper Information

  • Title: ZeRO： Memory Optimizations Toward Training Trillion Parameter Models
  • Authors: Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He (Microsoft)
  • Year: 2019
  • Link: arXiv
    .02054

Core Contributions

ZeRO (Zero Redundancy Optimizer) is a revolutionary memory optimization technique that enables training of extremely large models by eliminating memory redundancy in data parallel training.

Key Innovations

  1. ZeRO-DP: Partitions optimizer states, gradients, and parameters across GPUs
  2. ZeRO-R: Optimizes activations, temporary buffers, and memory fragmentation
  3. Three-stage optimization: Progressive memory reduction

Technical Approach

Traditional data parallelism requires each GPU to hold complete copies of model parameters, gradients, and optimizer states. ZeRO partitions these across GPUs, with each GPU holding only a fraction and communicating when needed.

Impact

ZeRO has become foundational infrastructure for large model training, integrated into DeepSpeed and enabling models like GPT-3/4. It democratized large-scale training by making it possible on limited hardware resources.