《AI系统》分布式优化器精讲

《AI系统》分布式优化器精讲

核心目标

分布式优化器的核心是用通信换内存，通过精心设计的数据分布和通信策略，在保证训练效果的前提下，显著降低单设备内存占用，从而支持更大模型训练。其技术演进如下图所示：

数据并行(DP) → ZeRO-1 → ZeRO-2 → ZeRO-3 → ZeRO-Infinity
内存效率: 1/N → 1/N → 1/N → 1/N → 突破设备内存限制
通信开销: 1x → 1x → 1x → 1.5x → 异构内存管理

---

一、ZeRO技术详解（核心思想：分而治之）

1. ZeRO-DP（模型状态优化）

解决的问题：传统数据并行中，每个GPU都完整存储：

• 模型参数（Parameters）
• 梯度（Gradients）
• 优化器状态（Optimizer States）
  造成大量冗余存储。

三级优化策略：

阶段	切分对象	内存降幅	通信开销变化	典型场景
ZeRO-1	优化器状态	4倍	无增加	常规大模型训练
ZeRO-2	优化器状态+梯度	8倍	无增加	千亿参数模型
ZeRO-3	全部模型状态	N倍	增加50%（1.5x）	万亿参数级模型

关键技术：

• 动态获取机制：参数使用时通过All-Gather临时重构完整层
• 梯度同步优化：Reduce-Scatter替代All-Reduce，通信量不变
• 计算通信重叠：在反向传播同时进行梯度聚合

2. ZeRO-R（剩余状态优化）

针对对象：

• 激活值（Activations）
• 临时缓冲区（Temporary Buffers）
• 内存碎片（Memory Fragmentation）

三大杀招：

1. 激活值分区：

   • 原始：每个设备存储完整激活 → 内存占用O(L)
   • 优化：按模型并行度分片存储 → 内存占用O(L/N)
   • 代价：增加约10%的All-Gather通信

2. 恒定缓冲区：

   • 传统：缓冲区随模型增大而膨胀
   • 优化：固定大小的环形缓冲区
   • 效果：抑制内存波动，提升利用率

3. 内存整理：

   • 问题：PyTorch原生分配器产生碎片
   • 方案：预分配+动态内存池管理
   • 收益：碎片减少80%以上

---

二、ZeRO-Infinity（突破物理限制）

创新点：将异构内存统一管理

GPU显存 → CPU内存 → NVMe存储
访问速度: 100% → 10% → 1%
容量比例: 1x → 10x → 100x

工作流程：

1. 前向传播时：按需将参数从NVMe→CPU→GPU
2. 计算完成后：立即将数据回传至NVMe
3. 梯度更新时：仅保持当前层参数在GPU

关键技术：

• 智能预取（Prefetching）：重叠数据传输与计算
• 异步流水线：CPU-GPU-NVMe三级流水
• 带宽优化：参数聚合传输（减少小数据包）

实测效果：

• 可训练参数量提升1000倍
• 在A100上实现45%计算效率（相比纯GPU训练的70%）

---

三、LOMO（极限内存优化）

设计哲学：即时更新代替梯度存储

与传统优化器对比

	普通SGD	LOMO
梯度存储	O(N)	O(1)
内存占用	3×参数大小	≈参数大小
更新时机	反向传播结束后	反向传播过程中

实现魔法：

1. 反向钩子（Hook）机制：

   def gradient_hook(param):
       if param.grad is not None:
           param.data -= lr * param.grad  # 立即更新
           param.grad = None  # 立即释放
   

2. 混合精度策略：

   • 计算时：FP16加速
   • 更新时：FP32稳定

3. 梯度归一化：

   • 在更新前对梯度进行裁剪/缩放
   • 避免即时更新带来的数值不稳定

适用场景：

• 单卡微调70B+大模型
• 边缘设备轻量化训练

---

四、技术选型指南

决策树

graph TD
    A[模型大小] -->|＜10B| B[传统DDP]
    A -->|10B-1T| C[ZeRO-2/3]
    A -->|＞1T| D[ZeRO-Infinity]
    B --> E{显存是否不足}
    E -->|是| F[+激活检查点]
    C --> G{是否需要跨节点}
    G -->|是| H[+模型并行]
    D --> I[+NVMe卸载]

性能权衡表

技术	内存节省	通信开销	实现复杂度	典型加速比
DDP	1x	1x	★★☆	8卡×0.9
ZeRO-2	8x	1x	★★★	64卡×0.85
ZeRO-3	Nx	1.5x	★★★★	512卡×0.7
LOMO	≈1x参数	1x	★★☆	单卡×0.6

---

五、前沿趋势

1. 通信压缩：

   • 梯度量化（1-bit Adam）
   • 稀疏通信（DeepSpeed稀疏注意力）

2. 异构计算：

   • 光计算加速通信
   • CXL内存池化技术

3. 智能调度：

   • 动态参数重要性感知分配
   • 基于RL的卸载策略

这些技术正在推动分布式训练向"万亿参数+千卡规模+90%效率"的目标迈进。