196 lines
6.0 KiB
Markdown
196 lines
6.0 KiB
Markdown
# 数据并行与ZeRO
|
||
|
||
## 数据并行回顾
|
||
|
||
数据并行（Data Parallelism，DP）是最基础的分布式训练范式：
|
||
|
||
1. **模型复制**：每张GPU持有完整模型副本
|
||
2. **数据分片**：不同GPU处理不同数据batch
|
||
3. **梯度同步**：每轮迭代后对各GPU的梯度做AllReduce同步
|
||
4. **参数更新**：所有GPU使用相同的更新后参数继续下一轮
|
||
|
||
```
|
||
GPU 0: [模型副本] + [数据分片0] → [本地梯度] → AllReduce → [更新后模型]
|
||
GPU 1: [模型副本] + [数据分片1] → [本地梯度] → AllReduce → [更新后模型]
|
||
...
|
||
```
|
||
|
||
### 朴素数据并行的局限
|
||
|
||
- **显存浪费**：每张GPU保存完整模型、激活值、优化器状态，显存占用极大
|
||
- **通信瓶颈**：梯度AllReduce通信量 = 模型参数量，对大模型不现实
|
||
- **扩展性差**：模型参数量超过单卡显存时无法运行
|
||
|
||
## ZeRO（Zero Redundancy Optimizer）
|
||
|
||
ZeRO是微软DeepSpeed团队提出的显存优化技术，通过分片（sharding）消除数据并行中的显存冗余。
|
||
|
||
### Stage 1：优化器状态分片（Optimizer State Sharding）
|
||
|
||
**原理**：将Adam优化器的状态（动量、方差）在数据并行进程间分片，每个GPU只保存 1/N 的优化器状态。
|
||
|
||
```
|
||
优化器状态总量: 2 * 参数总量（动量 + 方差）
|
||
Stage 1后: 每GPU只保存 2 * 参数总量 / N
|
||
```
|
||
|
||
**通信开销**：
|
||
- 通信量：与标准数据并行相同（梯度AllReduce）
|
||
- 额外通信：参数分片 scatter/gather（可选）
|
||
|
||
**收益**：约4倍显存降低（假设N=8，则显存降至1/4）
|
||
|
||
### Stage 2：梯度分片（Gradient Sharding）
|
||
|
||
**原理**：在Stage 1基础上，进一步将梯度分片存储。
|
||
|
||
```
|
||
标准DP: 每GPU保存完整梯度 → AllReduce → 每GPU有完整梯度 → 参数更新
|
||
Stage 2: 每GPU保存 1/N 梯度 → 需要时 AllReduce（分片通信）
|
||
```
|
||
|
||
**实现机制**：
|
||
1. 反向传播时本地计算梯度
|
||
2. 执行Reduce-Scatter，每个GPU获得其他GPU数据的部分和
|
||
3. 本地更新对应参数分片
|
||
4. 下次前向传播前执行All-Gather重建完整梯度（可选）
|
||
|
||
**通信模式**：
|
||
- 通信量：与Stage 1相同，但通信模式不同（从AllReduce变为Reduce-Scatter + AllGather）
|
||
- 可与通信计算重叠进一步优化
|
||
|
||
**收益**：约8倍显存降低（结合Stage 1）
|
||
|
||
### Stage 3：参数分片（Parameter Sharding）
|
||
|
||
**原理**：将模型参数也分片存储，每个GPU只保存 1/N 的参数。
|
||
|
||
```
|
||
完整分片流程:
|
||
前向/反向: 动态 All-Gather 获取所需参数分片 → 计算 → 丢弃非必要分片
|
||
```
|
||
|
||
**实现机制**：
|
||
1. 前向传播时：参数分片本地计算，需要邻 GPU 参数时通过All-Gather获取
|
||
2. 反向传播时：梯度分片本地更新，触发Reduce-Scatter同步
|
||
3. 每轮迭代结束后参数分片自动同步
|
||
|
||
**通信开销**：
|
||
- 每层前向/反向都需要 All-Gather（获取参数）和 Reduce-Scatter（同步梯度）
|
||
- 通信量显著增加，但换取极大显存节省
|
||
|
||
**收益**：理论上显存降低与GPU数成线性关系
|
||
|
||
### ZeRO 通信量分析
|
||
|
||
| Stage | 通信量（相比DP） | 显存收益 |
|
||
|-------|-----------------|---------|
|
||
| Stage 1 | 1x (相同) | 4x |
|
||
| Stage 2 | 1x (相同) | 8x |
|
||
| Stage 3 | ~2x (AllGather+ReduceScatter) | O(N) |
|
||
|
||
### ZeRO-DP 与标准DP对比
|
||
|
||
```
|
||
标准DP:
|
||
每GPU内存占用 = 参数 + 梯度 + 优化器状态 + 激活值
|
||
≈ 模型大小 × (1 + 2 + 4) ≈ 7× 模型大小（16位存储）
|
||
|
||
ZeRO-3:
|
||
每GPU内存占用 ≈ (参数 + 梯度 + 优化器状态) / N + 激活值
|
||
≈ 模型大小 / N + 激活值
|
||
```
|
||
|
||
## DeepSpeed实现
|
||
|
||
### 配置文件示例
|
||
|
||
```python
|
||
{
|
||
"zero_optimization": {
|
||
"stage": 3,
|
||
"offload_optimizer": {
|
||
"device": "cpu",
|
||
"pin_memory": true
|
||
},
|
||
"offload_param": {
|
||
"device": "cpu",
|
||
"pin_memory": true
|
||
},
|
||
"overlap_comm": true,
|
||
"contiguous_gradients": true,
|
||
"reduce_bucket_size": 1e8,
|
||
"stage3_prefetch_bucket_size": 1e9,
|
||
"stage3_param_persistence_threshold": 1e5
|
||
},
|
||
"bf16": {
|
||
"enabled": true
|
||
}
|
||
}
|
||
```
|
||
|
||
### 核心API
|
||
|
||
```python
|
||
from deepspeed import DeepSpeedConfig, initialize
|
||
model, optimizer, _, _ = initialize(
|
||
args=args,
|
||
model=model,
|
||
model_parameters=model.parameters(),
|
||
config=deepspeed_config
|
||
)
|
||
```
|
||
|
||
### ZeRO-3 关键技术
|
||
|
||
1. **分片通信与计算重叠**：利用CUDA stream在计算时通信
|
||
2. **Gradient Accumulation**：梯度累积避免流水线空闲
|
||
3. **Weight Gradients Partitioning**：参数梯度自动分片
|
||
4. **Partition Communication Group**：自定义通信域
|
||
|
||
### ZeRO-Stage 3与PipeDream的结合
|
||
|
||
DeepSpeed支持将ZeRO与流水线并行结合：
|
||
|
||
```python
|
||
{
|
||
"zero_optimization": {"stage": 3},
|
||
"pipeline": {"stages": 4}
|
||
}
|
||
```
|
||
|
||
每个pipeline stage内部使用ZeRO-3分片，进一步降低单卡显存。
|
||
|
||
### Offload技术
|
||
|
||
DeepSpeed支持将部分数据卸载到CPU或NVMe：
|
||
|
||
- **ZeRO-Offload**：优化器状态卸载到CPU，配合Stage 2
|
||
- **ZeRO-Infinity**：参数和优化器状态卸载到NVMe，支持超大规模模型
|
||
|
||
### 性能调优
|
||
|
||
1. **通信瓶颈**：TP/PP机器与DP机器分离，避免跨域通信
|
||
2. **通信域配置**：优化 NCCL 组配置
|
||
3. **Mixed Precision**：BF16训练减少通信量
|
||
4. **Async IO**：参数卸载异步化避免阻塞
|
||
|
||
## 与其他并行策略的关系
|
||
|
||
| 策略 | 维度 | 分片对象 |
|
||
|------|------|---------|
|
||
| 张量并行 | 层内 | 参数/激活切分 |
|
||
| 流水线并行 | 层间 | 模型层分Stage |
|
||
| 数据并行 | batch | 梯度/参数同步 |
|
||
| ZeRO | 数据并行维度 | 参数/梯度/优化器状态分片 |
|
||
|
||
三者可组合：TP=8, PP=4, ZeRO-3 + DP = N/(8*4)
|
||
|
||
## 框架对比
|
||
|
||
| 特性 | DeepSpeed ZeRO | Megatron-FSDP | PyTorch FSDP |
|
||
|------|---------------|---------------|--------------|
|
||
| 分片方式 | 优化器状态+梯度+参数 | 参数+梯度+优化器状态 | 参数+梯度 |
|
||
| Offload | CPU/NVMe | CPU | CPU |
|
||
| 通信优化 | Overlap + Partition | Async | Overlap |
|
||
| 易用性 | 高（配置驱动） | 中 | 中 | |