# 数据并行与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 |
| 易用性 | 高（配置驱动） | 中 | 中 |