DeepSpeed分布式训练框架深度解析

内容整理自学习笔记，仅供面试备考参考；不构成录用、培训或考试承诺。

1. 为什么需要DeepSpeed？

常见深度学习框架（TensorFlow、PyTorch、Keras）在面向大规模模型时不方便：

PyTorch的DDP只能做数据并行，模型大于显卡显存时就很难继续使用
大模型（LLM）训练需要大量内存存储中间激活、权重等参数，百亿模型无法在单个GPU上训练

当前训练超大规模语言模型技术路线： GPU + PyTorch + Megatron-LM + DeepSpeed

DeepSpeed由Microsoft提供，核心优势：

特性	说明
3D并行	ZeRO数据并行 + 流水线并行 + 张量切片模型并行，支持万亿参数模型
ZeRO-Offload	单卡V100可训练130亿参数模型（10倍扩展）
Sparse Attention	支持10倍长序列，6倍执行速度提升
1-bit Adam	最多5倍通信量减少，3.5倍分布式训练速度提升

2. DeepSpeed基本概念

2.1 分布式计算基础概念

概念	说明
node_rank	节点编号，区分不同计算机的通信
rank	全局进程编号，区分不同进程的通信
local_rank	局部进程编号，区分同一节点内不同进程
world_size	全局总进程数，确定并行工作量和所需资源
master_ip + master_port	主节点负责协调所有节点和进程

2.2 DeepSpeed支持的功能

功能	说明
Optimizer state partitioning	ZeRO stage 1
Gradient partitioning	ZeRO stage 2
Parameter partitioning	ZeRO stage 3
Custom mixed precision training	FP16/BF16混合精度
Fast CUDA-extension-based optimizers	CUDA加速优化器
ZeRO-Offload to CPU and NVMe	参数卸载到CPU/NVMe

2.3 ZeRO技术

ZeRO（Zero Redundancy Optimizer）将模型参数分成三个部分：Optimizer States、Gradient、Model Parameter。

2.4 混合精度训练

同时使用FP16和FP32，FP16减少内存占用训练更大模型，需使用动态精度缩放和混合精度优化器解决梯度消失和不稳定问题。

3. DeepSpeed通信策略

策略	适用场景
mpi	CPU集群分布式训练
gloo	CPU和GPU分布式训练
nccl	GPU分布式训练（NVIDIA专用，最常用）

4. DeepSpeed安装与使用

4.1 安装

bash

pip install deepspeed==0.8.1
sudo apt-get update
sudo apt-get install openmpi-bin libopenmpi-dev
pip install mpi4py

4.2 使用流程

使用DeepSpeed与写PyTorch模型只有部分区别，核心改动在于初始化方式。

关键代码：

python

import deepspeed
import torch.distributed as dist

# 初始化DeepSpeed引擎
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=model,
    model_parameters=model.parameters()
)
args.local_rank = model_engine.local_rank

# 训练循环
for step, batch_data in enumerate(train_loader):
    model_engine.train()
    logits, label = on_step(batch_data)
    loss = criterion(logits, label)
    model_engine.backward(loss)
    model_engine.step()

注意： 训练时使用 model_engine.backward(loss) 和 model_engine.step()，而非PyTorch的 loss.backward() 和 optimizer.step()。

4.3 DeepSpeed配置示例

json

{
  "train_micro_batch_size_per_gpu": 32,
  "gradient_accumulation_steps": 1,
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 3e-5 }
  },
  "fp16": { "enabled": true },
  "zero_optimization": {
    "stage": 3,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8
  },
  "activation_checkpointing": {
    "partition_activations": true,
    "cpu_checkpointing": true,
    "contiguous_memory_optimization": true
  }
}

注意： ZeRO第3阶段下，模型被划分到不同GPU上。

4.4 运行代码

bash

deepspeed test.py --deepspeed_config config.json

5. ZeRO各Stage配置详解

5.1 ZeRO-2配置

json

{
  "fp16": {
    "enabled": "auto",
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  }
}

重要参数：

参数	说明
overlap_comm	通信与计算重叠，减少通信时间
allgather_bucket_size	Allgather操作的分桶大小，越大通信越快但需更多内存
reduce_bucket_size	Allreduce操作的分桶大小，同理
offload_optimizer	将优化器状态offload到CPU，降低显存需求

内存估算： overlap_comm使用 allgather_bucket_size 和 reduce_bucket_size 的4.5倍。如设5e8需9GB显存；8GB以下GPU应设约2e8（需3.6GB显存）。

5.2 ZeRO-3配置

json

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu", "pin_memory": true },
    "offload_param": { "device": "cpu", "pin_memory": true },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  }
}

ZeRO-3重要参数：

参数	说明
stage3_max_live_parameters	GPU上保留的完整参数数量上限（1e9约消耗2GB）
stage3_max_reuse_distance	参数将来何时再使用的指标，决定丢弃还是保留
stage3_gather_16bit_weights_on_model_save	保存模型时启用fp16权重合并（昂贵操作）
sub_group_size	optimizer steps中更新参数的粒度，控制NVMe offload的CPU内存使用

注意： ZeRO-3中 allgather_partitions、allgather_bucket_size、reduce_scatter 配置参数未被使用。

5.3 ZeRO-Stage-0和Stage-1

Stage	说明
Stage 0	禁用所有分片，相当于DDP
Stage 1	只对优化器参数进行分片，略微加速

json

// Stage 0
{ "zero_optimization": { "stage": 0 } }

// Stage 1
{ "zero_optimization": { "stage": 1 } }

6. ZeRO-3速度优化

ZeRO-3比ZeRO-2慢很多，优化策略：

将 stage3_param_persistence_threshold 设置很大（如 6 * hidden_size * hidden_size）
关闭 offload_params（可极大改善性能）

7. ZeRO Stage与Offload选择指南

速度排名（从快到慢）	显存排名（从少到多）
Stage 0 (DDP)	Stage 3 + offloads
Stage 1	Stage 3
Stage 2	Stage 2 + offload
Stage 2 + offload	Stage 2
Stage 3	Stage 1
Stage 3 + offloads	Stage 0 (DDP)

8. DeepSpeed调参步骤

遇到OOM时按以下顺序调参：

将batch_size设为1，通过梯度累积实现任意有效batch_size
OOM → 设置 gradient_checkpointing
OOM → 尝试ZeRO stage 2
OOM → 尝试ZeRO stage 2 + offload_optimizer
OOM → 尝试ZeRO stage 3
OOM → offload_param到CPU
OOM → offload_optimizer到CPU
OOM → 降低默认参数（如减小beam search范围）
OOM → 混合精度训练（Ampere用bf16，旧GPU用fp16）
OOM → ZeRO-Infinity（offload到NVMe）
batch_size=1无OOM后，测量吞吐量，尽可能增大batch_size
关闭offload或降低ZeRO stage，优化参数，可增加66%性能

9. 显存估算工具

bash

python -c 'from transformers import AutoModel; \
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \
model = AutoModel.from_pretrained("bigscience/T0_3B"); \
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=2, num_nodes=1)'

输出示例：

per CPU	per GPU	配置
70.00GB	0.25GB	offload_param=cpu, offload_optimizer=cpu
62.23GB	2.84GB	offload_param=none, offload_optimizer=cpu
0.74GB	23.58GB	offload_param=none, offload_optimizer=none

10. 训练精度配置

10.1 混合精度选项

精度	启用方式	说明
fp16	`--fp16 --fp16_backend amp`	PyTorch-like AMP或apex-like方式
bf16	`--bf16`	对bf16预训练的模型更稳定
tf32	`--tf32`	Ampere GPU默认启用，结果仍为fp32

建议： fp16混合精度减少内存+提升速度，只在fp16训练表现不佳时考虑fp32（如模型未在fp16中预训练）。

10.2 通信数据类型

默认半精度训练使用fp16作为reduction操作默认值，可设置fp32作为累积数据类型：

json

{ "communication_data_type": "fp32" }

10.3 Apex方式

json

{
  "amp": {
    "enabled": "auto",
    "opt_level": "auto"
  }
}

启用方式：--fp16 --fp16_backend apex --fp16_opt_level O1

11. 获取模型参数

ZeRO模式	模型参数存储
ZeRO-2	fp16形式存储在 `pytorch_model.bin`
ZeRO-3	需设置 `stage3_gather_16bit_weights_on_model_save: true`
优化器参数	fp32存储在 `global_step*/optim_states.pt`

离线获取fp32权重：

bash

python zero_to_fp32.py . pytorch_model.bin

12. 优化器和调度器

12.1 优化器

DeepSpeed主要优化器：Adam、AdamW、OneBitAdam、Lamb，已与ZeRO充分测试。

启用offload_optimizer时可使用非DeepSpeed优化器（需同时有CPU和GPU实现）
offload与DeepSpeed CPU Adam配合效果最佳
0.8.3之后版本使用其他优化器需添加 "zero_force_ds_cpu_optimizer": false

12.2 调度器

DeepSpeed调度器	对应HF参数
WarmupLR	`--lr_scheduler_type constant_with_warmup`
WarmupDecayLR	`--lr_scheduler_type linear`

13. 常见问题与填坑笔记

问题	解决方法
`ModuleNotFoundError: No module named 'torch._six'`	注释掉 `from torch._six import string_classes`，加入 `string_classes = str`、`int_classes = int`；若报 `NameError: name 'inf'` 则加 `import math; inf = math.inf`
单卡也可使用deepspeed	ZeRO-offload将部分数据offload到CPU + 提供显存管理减少碎片
进程被杀死无traceback	GPU显存不够 → 加卡
loss是NaN	训练用bf16推理用fp16导致，改用fp32或bf16
多机配置问题	Hostfile使用IP而非hostname

多机配置步骤

SSH配置：

bash

# hosts中加入节点名称
vim /etc/hosts  # 添加: 10.58.253.27 model1, 10.58.253.26 model2

# 生成ssh key
ssh-keygen -t rsa

# 互相拷贝ssh key
ssh-copy-id ccwork@10.58.253.27
ssh-copy-id ccwork@10.58.253.26

# 测试
ssh model2

安装pdsh：

bash

sudo apt-get install pdsh
export PDSH_RCMD_TYPE=ssh

配置hostfile并运行：

bash

# hostfile
model1 slots=8
model2 slots=8

# 运行
deepspeed --hostfile=hostfile --include="model2:1,2@model1:3" train.py --stage sft

确保一致性

代码路径一致（不一致则建软链接）
虚拟环境路径一致
conda各节点安装路径一致（否则报 return code = 127）
库版本高度一致（否则报 return code = -6）

DeepSpeed分布式训练框架深度解析 ​

1. 为什么需要DeepSpeed？ ​

2. DeepSpeed基本概念 ​

2.1 分布式计算基础概念 ​

2.2 DeepSpeed支持的功能 ​

2.3 ZeRO技术 ​

2.4 混合精度训练 ​

3. DeepSpeed通信策略 ​

4. DeepSpeed安装与使用 ​

4.1 安装 ​

4.2 使用流程 ​

4.3 DeepSpeed配置示例 ​

4.4 运行代码 ​

5. ZeRO各Stage配置详解 ​

5.1 ZeRO-2配置 ​

5.2 ZeRO-3配置 ​

5.3 ZeRO-Stage-0和Stage-1 ​

6. ZeRO-3速度优化 ​

7. ZeRO Stage与Offload选择指南 ​

8. DeepSpeed调参步骤 ​

9. 显存估算工具 ​

10. 训练精度配置 ​

10.1 混合精度选项 ​

10.2 通信数据类型 ​

10.3 Apex方式 ​

11. 获取模型参数 ​

12. 优化器和调度器 ​

12.1 优化器 ​

12.2 调度器 ​

13. 常见问题与填坑笔记 ​

多机配置步骤 ​

确保一致性 ​

DeepSpeed分布式训练框架深度解析

1. 为什么需要DeepSpeed？

2. DeepSpeed基本概念

2.1 分布式计算基础概念

2.2 DeepSpeed支持的功能

2.3 ZeRO技术

2.4 混合精度训练

3. DeepSpeed通信策略

4. DeepSpeed安装与使用

4.1 安装

4.2 使用流程

4.3 DeepSpeed配置示例

4.4 运行代码

5. ZeRO各Stage配置详解

5.1 ZeRO-2配置

5.2 ZeRO-3配置

5.3 ZeRO-Stage-0和Stage-1

6. ZeRO-3速度优化

7. ZeRO Stage与Offload选择指南

8. DeepSpeed调参步骤

9. 显存估算工具

10. 训练精度配置

10.1 混合精度选项

10.2 通信数据类型

10.3 Apex方式

11. 获取模型参数

12. 优化器和调度器

12.1 优化器

12.2 调度器

13. 常见问题与填坑笔记

多机配置步骤

确保一致性