Notes/机器学习/深度学习模型/07-计算架构/5-分布式训练框架/1-DeepSpeed.md

# DeepSpeed

## 概述

DeepSpeed是微软开发的深度学习训练优化库，专注于大规模模型的分布式训练加速和显存优化。核心特性包括ZeRO（Zero Redundancy Optimizer）优化器、3D并行支持、混合精度训练等。

## ZeRO实现

### 核心架构

DeepSpeed的ZeRO实现分为三个阶段（详见"数据并行与ZeRO"文档），这里重点介绍实现细节。

### ZeRO-1: 优化器状态分片

```python
# DeepSpeed ZeRO-1配置
{
    "zero_optimization": {
        "stage": 1,
        "contiguous_gradients": true,
        "reduce_scatter": true,
        "allgather_partitions": true
    }
}
```

**实现机制**：
- 每个GPU持有优化器状态的1/N分片
- 梯度计算后执行Reduce-Scatter，每个GPU只保留对应分片的梯度
- 参数更新只涉及本地分片
- 需要时通过AllGather重建完整梯度

### ZeRO-2: 梯度分片

```json
{
    "zero_optimization": {
        "stage": 2,
        "stage1_param": false,
        "stage2_param": false,
        "stage2_cpu_offload": false,
        "contiguous_gradients": true,
        "overlap_comm": true
    }
}
```

**额外特性**：
- 梯度分片存储，避免完整梯度复制
- Overlap通信与计算：利用CUDA Stream技术在反向传播时同步梯度
- Gradient accumulation支持

### ZeRO-3: 参数分片

```json
{
    "zero_optimization": {
        "stage": 3,
        "stage3_param_persistence_threshold": 1e5,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e5
    }
}
```

**关键参数**：
- `stage3_max_live_parameters`：GPU上同时保持的最大参数分片数
- `stage3_param_persistence_threshold`：超过此阈值的参数保持在GPU

### ZeRO-Inference

```python
# ZeRO-Inference配置
{
    "zero_optimization": {
        "stage": 3,
        "stage3_offload_param": {
            "device": "cpu"
        }
    },
    "inference": {
        "enabled": true
    }
}
```

推理时将参数卸载到CPU，通过动态加载满足计算需求。

## 3D并行配置

### 并行策略组合

DeepSpeed支持数据并行、流水线并行、张量并行的组合配置：

```json
{
    "parallelism": {
        "tensor_parallel": {
            "enabled": true,
            "size": 8,
            "mpu": "MegatronMPU"
        },
        "pipeline": {
            "enabled": true,
            "stages": 4,
            "pipe_schedule": "GPipe"
        },
        "data_parallel": {
            "degree": 32
        }
    }
}
```

### 设备映射

```python
# 假设有 256 个 GPU
# 配置: TP=8, PP=4, DP=8
# 总GPU需求: 8 × 4 × 8 = 256

# 数据并行度 = 总GPU数 / (TP × PP) = 256 / 32 = 8
data_parallel_size = 8
tensor_parallel_size = 8
pipeline_parallel_size = 4

# 模型切分
model = transformer_layer / pipeline_parallel_size  # 每Stage 1/4层
layer = weight / tensor_parallel_size  # 每GPU 1/8权重
```

### 通信域配置

```python
# 通信域层级
# 1. TP通信域（张量并行）：all-reduce within 8 GPUs
# 2. PP通信域（流水线并行）：send/recv between stages  
# 3. DP通信域（数据并行）：all-reduce across DP replicas

# DeepSpeed自动构建通信域
ds_config = DeepSpeedConfig(config)
```

## Gradient Checkpointing

### 原理

Gradient Checkpointing（也称Activation Checkpointing）通过在前向传播时不保存中间激活值，而在反向传播时重新计算来节省显存。

```python
# DeepSpeed Gradient Checkpointing配置
{
    "gradient_checkpointing": {
        "enabled": true,
        "checkpoint_every_n_steps": 1,
        "contiguous": true,
        "checkpoint_in_cpu": false
    }
}
```

### 内存-计算权衡

```
标准前向：
- 前向计算: O(1)
- 激活存储: O(N)  # N=层数

Checkpointing：
- 前向计算: O(N)  # 每层重新计算
- 激活存储: O(1)  # 只保存输入输出
- 反向计算: O(N)  # 重新计算中间激活

节省显存: ~70%（取决于模型结构）
额外计算开销: ~30%
```

### 与ZeRO结合

```python
{
    "gradient_checkpointing": true,
    "zero_optimization": {
        "stage": 3
    }
}
# ZeRO-3 + Checkpointing: 显存 ≈ 模型大小 / N + 激活 / N
```

## 混合精度训练支持

### BF16训练

```json
{
    "bf16": {
        "enabled": true,
        "loss_scaling": "dynamic"
    },
    "fp16": {
        "enabled": false
    }
}
```

**关键特性**：
- 前向/反向传播使用BF16
- 优化器状态保持FP32精度
- Loss scaling动态调整
- 支持ZeRO-3与BF16结合

### 混合精度流程

```
1. 前向传播: BF16
   ↓
2. Loss计算: BF16
   ↓
3. 反向传播: BF16梯度
   ↓
4. 梯度转FP32: 用于优化器更新
   ↓
5. 参数更新: FP32优化器 → FP16/BF16参数
```

### Loss Scaling

```python
# DeepSpeed自动loss scaling
class BF16Scaler:
    def __init__(self):
        self.loss_scale = None
        self.scale_factor = 1.0
        
    def set_scale(self, loss_scale):
        self.loss_scale = loss_scale
        
    def update_scale(self, found_inf):
        # 发现inf时缩小scale
        if found_inf:
            self.loss_scale /= 2
        else:
            self.loss_scale = self.loss_scale  # 保持或增长
```

## 训练配置示例

### 7B模型单节点多卡

```json
{
    "train_batch_size": 32,
    "train_micro_batch_size_per_gpu": 4,
    "gradient_accumulation_steps": 8,
    "steps_per_print": 10,
    
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        }
    },
    
    "bf16": {
        "enabled": true
    },
    
    "gradient_clipping": 1.0,
    
    "WallClockBreakdown": true,
    "memory_breakdown": false
}
```

### 65B模型多节点

```json
{
    "train_batch_size": 1536,
    "train_micro_batch_size_per_gpu": 1,
    "gradient_accumulation_steps": 1536,
    
    "parallelism": {
        "tensor_parallel": 8,
        "pipeline": 4,
        "data_parallel": 64
    },
    
    "zero_optimization": {
        "stage": 3,
        "stage3_param_persistence_threshold": 1e5,
        "stage3_max_live_parameters": 1e9
    },
    
    "bf16": {
        "enabled": true
    },
    
    "通信优化": {
        "overlap_comm": true,
        "contiguous_gradients": true
    }
}
```

## ZeRO-Stage 3关键技术细节

### 参数分片管理

```python
class ParameterShardingManager:
    """管理参数分片的加载和卸载"""
    
    def __init__(self, model, device):
        self.param_shapes = {}  # 参数形状信息
        self.param_offsets = {}  # 分片偏移信息
        self.param缓存 = {}      # 已加载的分片
        
    def get_param_shard(self, param, shard_id):
        """获取参数分片"""
        shape = param.shape
        num_shards = self.get_num_shards(shape)
        
        # 计算分片范围
        start = shard_id * (shape[0] // num_shards)
        end = (shard_id + 1) * (shape[0] // num_shards)
        
        return param[start:end].clone()
    
    def load_param_async(self, param, shard_id):
        """异步加载参数分片到GPU"""
        shard = self.get_param_shard(param, shard_id)
        self.param_cache[param] = shard.to(device)
        
    def scatter_param(self, param):
        """将参数分片分发到各GPU"""
        # 使用集合通信
        return self.all_gather(param)
```

### 通信重叠

```python
# Overlap策略：反向传播与梯度同步重叠
backward_pass():
    for layer in reversed(layers):
        # 1. 计算本地梯度
        grad = layer.compute_gradient()
        
        # 2. 触发异步梯度同步
        if should_sync(layer):
            dist.all_reduce_async(grad, op=dist.ReduceOp.SUM)
        
        # 3. 继续计算后续层梯度
        next_layer.compute_gradient(grad)
    
    # 等待所有同步完成
    dist.synchronize()
```

### 混合分片策略

```python
# 部分参数使用不同分片策略
{
    "zero_optimization": {
        "stage": 3,
        "param_sharding_strategy": "order_based",
        "stage3_order": [
            "embeddings",
            "attention",
            "mlp"
        ]
    }
}
```

## DeepSpeed训练流程

```python
import deepspeed

# 1. 初始化DeepSpeed
deepspeed.init_distributed()

# 2. 创建模型
model = MyModel()

# 3. 创建优化器（可选，DeepSpeed内部创建）
optimizer = torch.optim.Adam(model.parameters())

# 4. 初始化DeepSpeed
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    optimizer=optimizer,
    config=ds_config
)

# 5. 训练循环
while training:
    # 前向
    loss = model_engine.forward(batch)
    
    # 反向
    model_engine.backward(loss)
    
    # 更新
    model_engine.step()
```

## 与PyTorch FSDP对比

| 特性 | DeepSpeed ZeRO-3 | PyTorch FSDP |
|------|-----------------|--------------|
| 参数分片 | ZeRO-3分片 | Full Sharded |
| 通信模式 | AllGather + ReduceScatter | AllGather |
| CPU Offload | 原生支持 | 需要额外配置 |
| 优化器卸载 | 原生支持 | 部分支持 |
| 易用性 | 配置驱动 | 代码驱动 |
| 生态 | 成熟度高 | 官方维护 |

## 性能调优建议

1. **通信优化**：启用`overlap_comm`和`contiguous_gradients`
2. **内存优化**：适当使用CPU卸载缓解GPU压力
3. **Batch Size**：配合gradient accumulation调整
4. **通信带宽**：确保使用高速互联（NVLink）
5. **混合精度**：使用BF16而非FP16

## 常见问题

| 问题 | 解决方案 |
|------|---------|
| OOM错误 | 降低batch size或启用ZeRO-3 |
| 通信瓶颈 | 检查网络拓扑，启用overlap |
| 精度问题 | 使用BF16，检查loss scaling |
| 速度慢 | 调整num_workers，增加warmup |