目录

数据并行（DDP）与显存冗余

1. DDP 流程：

2. ZeRO 冗余优化器：

激活值检查点

1. 定义：

2. 实现原理：

3. 子图切分策略：

模型并行策略

张量并行（TP）：参数按维度切分

流水并行（PP）：模型按层切分

序列并行（SP）：沿序列维度切分

Offloading 技术

其他优化技术

总结

---

本文将系统讲解大模型训练中通过 ZeRO 分片、激活值重计算、模型并行等核心技术，以通信和计算为代价，系统性优化显存占用的方法。

数据并行（DDP）与显存冗余

1. DDP 流程：

• • 每张卡上模型相同，数据分片。

• • 反向计算后，梯度通过 AllReduce 汇总。

• • DDP 主要解决算力扩展问题，对显存占用的优化有限。

2. ZeRO 冗余优化器：

• • 目标： 在保持数据并行语义不变的前提下，系统性地消除参数、梯度和优化器状态的冗余复制。

• • 三级优化方案：

ZeRO Stage 1：

• • 优化器状态按 DP size 切分。

• • 前向、反向流程不变，梯度更新时使用各自优化器状态切片。

• • 各卡仅更新其负责的参数分片，随后通过 AllGather 同步更新后的参数。

ZeRO Stage 2：

• • 梯度状态也按 DP size 切分。

• • 反向计算中，梯度 AllReduce 替换为 ReduceScatter。

• • 优化器更新时使用各自梯度切片。

ZeRO Stage 3：

• • 参数状态也按 DP size 切分。

• • 前向、反向计算中涉及完整参数时，临时进行 AllGather。

• • 优化器更新后无需再进行 AllGather。

激活值检查点

1. 定义：

• • 前向过程中只保留少数激活值检查点，反向过程中重新计算未保存的激活值。

2. 实现原理：

• • 将前向图分为若干子图，子图边缘保存检查点。

• • 反向过程中，未找到的激活值从最近检查点重新计算。

3. 子图切分策略：

• • Selective Checkpointing：基于先验知识或人为调优选择性地保存激活值。

• • Structured Checkpointing：根据模型结构进行结构化划分。

模型并行策略

张量并行（TP）：参数按维度切分

• • 对模型内单个参数进行切分，下发到不同 rank。

• • 减少参数、梯度和优化器状态的显存占用。

流水并行（PP）：模型按层切分

• • 对模型不同阶段进行切分，下发到不同 rank。

• • 减少单个 rank 上的模型参数量，进而减少显存占用。

序列并行（SP）：沿序列维度切分

• • 沿序列维度或上下文维度进行切分，减少激活值显存占用。

Offloading 技术

• • 将优化器状态移至 CPU 或硬盘

• • 针对消费级显卡和小显存场景

• • 性能影响较大，需异步拷贝优化

其他优化技术

• • 混合精度训练：FP16/BF16+FP32 组合使用

• • 梯度累积：减小 Micro-batch size

• • 算子融合：将多个算子合并执行，减少中间激活占用

• • 内存布局优化减少碎片

总结

本文系统讲解了大模型训练中通过 ZeRO 分片、激活值重计算、模型并行等核心技术，以通信和计算为代价，系统性优化显存占用的方法。