- 参考链接：
整体介绍
- DeepSpeed 是由微软开发的开源的深度学习优化库，旨在提高大规模模型训练的效率和可扩展性
- DeepSpeed 框架的核心技术有：
- ZeRO 冗余优化技术 ：通过分布式内存管理，将模型参数、梯度和优化器状态进行分区，大幅降低显存占用，是首次支持千亿级参数模型训练的框架
- 3D 并行策略 ：支持数据并行、流水线并行和张量切片模型并行，并可灵活组合
- 混合精度训练 ：自动混合精度训练（AMP）将单精度和半精度浮点数结合使用，降低内存需求的同时提升计算效率
- 智能推理优化器 ：支持张量并行与异构内存技术，提供低延迟高吞吐的分布式推理服务，可将成本降低 70%
- 全链路内存管理 ：集成 CPU 卸载与显存碎片整理技术，单卡即可训练百亿级模型，资源利用率提升 6 倍
- DeepSpeed 框架的组件构成有：
- Apis ：提供易用的 API 接口
- 运行时组件 ：管理、执行和性能优化，基于 Python 语言实现，负责部署训练任务到分布式设备、数据分区、模型分区等
- 底层内核 ：用 C++ 和 CUDA 实现，优化计算和通信
- DeepSpeed 生态兼容性极好 ：原生兼容 PyTorch 与 Hugging Face 生态，通过简洁 API 可快速迁移项目，开发效率提升 300%
安装 DeepSpeed
仅需一行安装命令即可：
1
pip install deepspeed
暗转完成后，可以使用
ds_report命令验证安装是否成功- 这个命令可以查看环境配置信息
DeepSpeed 的使用
代码修改 ：仅仅需要非常少的代码修改即可将原始训练代码切换到 DeepSpeed 框架上（DeepSpeed 与 PyTorch 无缝集成，只需少量修改即可启用加速）
- 第一步：通过
deepspeed.initialize将模型包装为 DeepSpeed 引擎，自动应用优化 - 第二步：使用
model_engine.backward和model_engine.step替换PyTorch原本的loss.backward()和optimizer.step()
- 第一步：通过
配置 DeepSpeed ：DeepSpeed 的优化行为通过 JSON 配置文件（
ds_config.json）指定- 一般的项目都会自带一些配置好的
.json文件示例，可直接修改使用
- 一般的项目都会自带一些配置好的
运行训练 ：使用 DeepSpeed 的命令行工具启动训练
单节点训练命令为：
1
deepspeed train.py --deepspeed_config ds_config.json
对于多节点集群，使用下面的命令：
1
deepspeed --num_gpus 8 --num_nodes 2 train.py --deepspeed_config ds_config.json
- 其中
--num_gpus指定每节点使用的 GPU 数量，--num_nodes指定集群中的节点数
- 其中
其他常用参数：
- 使用
--log_dir参数启用日志记录，监控内存使用、训练速度等
- 使用
特别说明：使用了 DeepSpeed 框架的代码需要使用
deepspeed命令来启动- 补充：直接使用
python命令启动时会出现deepspeed.initialize调用的错误 - 原因：DeepSpeed 作为一个分布式训练库，需要特殊的启动器来管理多个进程和 GPU 之间的通信和资源分配
- 补充：直接使用
DeepSpeed 使用示例（基于 PyTorch）
- DeepSpeed 使用示例：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
import deepspeed
class DiyModel(nn.Module):
def __init__(self, input_dim, output_dim):
super(DiyModel, self).__init__()
self.fc = nn.Sequential(
nn.Linear(input_dim, 128),
nn.ReLU(),
nn.Linear(128, output_dim)
)
def forward(self, x):
return self.fc(x)
class RandomDataset(Dataset):
def __len__(self):
return 1000
def __getitem__(self, idx):
x = torch.randn(32)
y = torch.randint(0, 10, (1,)).item()
return x, y
# 初始化模型和数据（无需为 DeepSpeed 特别处理）
model = DiyModel(32, 10)
train_dataset = RandomDataset()
train_loader = DataLoader(train_dataset, batch_size=32)
# 初始化 DeepSpeed 引擎（仅需使用 deepspeed.initialize 初始化模型得到 model_engine 即可）
# 注：这一步会自动将模型参数移动到 GPU 上，下面使用的数据也需要将数据移动到对应的 GPU 上才能运行，否则会出现 设备不一致的错误（CPU vs GPU）
# zero_optimization 字段的 stage 键值对应如下效果
## stage: 0：不使用 ZeRO 优化（默认值是 0）
## stage: 1：优化器状态分片
## stage: 2：优化器状态和梯度分片
## stage: 3：优化器状态、梯度和参数分片（最高内存效率）
model_engine, optimizer, _, _ = deepspeed.initialize(
args=None,
model=model,
model_parameters=model.parameters(),
config={
"train_batch_size": 32,
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.001,
"betas": [0.9, 0.999]
}
},
"fp16": {
"enabled": True,
"loss_scale": 0,
"loss_scale_window": 1000,
"initial_scale_power": 16
},
"zero_optimization": {
"stage": 2, # 这里指定Zero层级（0、1、2、3）
"offload_optimizer": {
"device": "cpu" # 可选：指定优化器卸载设备
}
}
}
)
# 训练过程（训练时不再使用原来的模型，使用 model_engine）
for epoch in range(10):
for batch_idx, (data, target) in enumerate(train_loader):
# 将数据挪到和模型相同的 GPU 上
device = model_engine.device
model_dtype = next(model_engine.parameters()).dtype # 通过模型的第一个参数获取dtype
data = data.to(device, dtype=model_dtype) # 将数据移至模型所在设备并转换为与模型相同的dtype
target = target.to(device)
outputs = model_engine(data) # 使用（deepspeed.initialize 初始化得到的）model_engine 来进行前向过程
loss = nn.CrossEntropyLoss()(outputs, target)
model_engine.backward(loss) # 使用 model_engine 来进行后向过程
model_engine.step() # 使用 model_engine 来更新模型参数（注：此时不再需要显示调用 optimizer)
if batch_idx % 100 == 0:
print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item()}")
附录：如何指定目标 GPU
启动时指定 CUDA_VISIBLE_DEVICES 环境变量
在运行命令前设置环境变量来限制 DeepSpeed 可见的 GPU：
1
CUDA_VISIBLE_DEVICES=0,1,2,3 deepspeed your_script.py --args ...
该方案是最常用的方法 ，且适用于常见的很多框架
在 Python 脚本中设置环境变量
也可以在Python脚本中通过
os.environ设置这个环境变量：1
2
3
4
5
6
7import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" # 指定使用GPU 0和1
import deepspeed
import torch
# 训练逻辑注：需要在导入DeepSpeed或PyTorch之前设置
该方案同样适用于常见的很多框架，但因为需要修改代码，不常用
使用 deepspeed 命令的 --include 参数
- 如果使用的是 DeepSpeed 的 launcher，也可以通过
--include参数指定使用的 GPU：1
deepspeed --include localhost:0,1 your_script.py --args ...
多机多卡如何指定 GPU
- 如果你使用的是单机多卡，以上方法都能很好地工作
- 对于多机多卡训练，通常需要结合
deepspeed命令的其他参数如--hostfile等一起使用
附录：数据加载位置管理
- 由于 deepspeed 会自动将模型参数加载到指定 GPU 上，所以数据也要加载到指定 GPU，否则会出现设备不一致的错误
- 加载命令如下（亲测解决方案）：
1
2device = model_engine.device
data = data.to(device)
附录：混合精度训练数据格式管理
- 由于 deepspeed 在启动混合精度训练时，可能会按照指定格式来指定参数形式，此时数据也需要转换为指定类型
- 解决方式如下（亲测遇到错误时的解决方案）：
1
2model_dtype = next(model_engine.parameters()).dtype # 通过模型的第一个参数获取类型，注：写这么复杂的原因是，当前还不支持直接调用 model_engine.dtype()
data = data.to(device, dtype=model_dtype)
