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

DeepSpeed

概述

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

ZeRO实现

核心架构

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

ZeRO-1: 优化器状态分片

# 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: 梯度分片

{
    "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: 参数分片

{
    "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

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

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

3D并行配置

并行策略组合

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

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

设备映射

# 假设有 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权重

通信域配置

# 通信域层级
# 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）通过在前向传播时不保存中间激活值，而在反向传播时重新计算来节省显存。

# 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结合

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

混合精度训练支持

BF16训练

{
    "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

# 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模型单节点多卡

{
    "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模型多节点

{
    "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关键技术细节

参数分片管理

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)

通信重叠

# 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()

混合分片策略

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

DeepSpeed训练流程

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