一文详解解 DeepSpeed

大模型训练热潮中，DeepSpeed 这个框架的名字几乎成了绕不开的话题。无论是想训一个百亿参数的模型，还是想在一张单卡上跑出几十亿参数的推理，DeepSpeed 的出现都让那些看似遥不可及的目标变得触手可及。下面就从它的设计初衷、核心技术，到实际使用和性能表现，把它的底细彻底理清楚。

1. 概览

1.1 背景（Why）

过去两年，ChatGPT 类模型掀起的技术浪潮，让整个行业都陷入了对大模型的狂热追逐。几乎每一家做 AI 的公司，都想用更低的成本、更快的速度训练出自己的大模型。可现实是，模型参数规模越堆越大，训练数据量也在指数级增长，动辄就需要昂贵的多卡多节点 GPU 集群。硬件成本已经让人头疼了，更麻烦的是，现有的开源训练系统在效率上表现并不理想——大多数系统的机器利用率甚至不到硬件极限的 50%。换句话说，砸了更多的钱买更好的硬件，并不意味着就能跑出更高的吞吐量；系统吞吐量上去了，也不代表模型就能训练得更准、收敛更快，更不说明这套软件就好用。

在这种背景下，市场迫切需要一套真正高效、有效且易用的开源系统，帮开发者实实在在提升生产力。微软开源的 DeepSpeed 框架，正是在这种诉求下进入了大家的视野。

模型规模和训练数据规模变化图

模型规模变化及当前模型训练框架面临的挑战

所谓高效性、有效性和易于使用，可以拆成三句话来看：高效，就是能充分利用硬件资源，实现高吞吐和可扩展性；有效，就是高精度、快速收敛、低成本；易用，说白了就是让开发者能专注算法而非落地细节。

1.2 简介（What）

DeepSpeed 是一个由微软开发的开源深度学习优化库，核心目标是提升大规模模型训练的效率和可扩展性。它通过多种技术手段来加速训练，比如模型并行、梯度累积、动态精度缩放、混合精度训练等等。除此之外，DeepSpeed 还提供了分布式训练管理、内存优化和模型压缩等一系列辅助工具，帮助开发者更好地管理和优化大规模深度学习训练任务。需要特别提一句，DeepSpeed 是构建在 PyTorch 之上的，这意味着迁移成本极低，通常只需要做简单修改就能把现有训练代码搬过去。目前，DeepSpeed 已经在语言模型、图像分类、目标检测等多类大规模深度学习项目中得到了应用。

DeepSpeed 在整个深度学习模型软件体系架构中处于什么位置？

图3 DeepSpeed 在深度学习模型软件体系架构中所处的位置

简单来说，DeepSpeed 是一个大模型训练加速库，它位于模型训练框架和模型本身之间，专门用来提升训练和推理的效率。

1.3 软件架构

DeepSpeed 的架构可以分成三个核心部分：

• Apis
  ：提供易用的接口，训练模型和推理模型只需要简单调用几个接口。其中最重要的就是 initialize 接口，用来初始化引擎，并传入训练参数及优化技术的配置。这些配置参数通常保存在 config.json 文件中。

• Runtime
  ：运行时组件，是 DeepSpeed 管理、执行和性能优化的核心。它负责部署训练任务到分布式设备、数据分区、模型分区、系统优化、微调、故障检测、checkpoints 的保存和加载等。这一部分用 Python 实现。

• Ops
  ：用 C++ 和 CUDA 实现的底层内核，用来优化计算和通信，比如超快速的 Transformer kernels、fuse LAN kernels 和一些自定义的算子。

那么，为什么采用这种架构？

• 可以在训练框架上分别对训练和推理两部分进行优化，互不干扰；
• 相比紧耦合的结构，这种松耦合方式能更好地利用整个生态，维护成本也低得多；
• 与基础设施无关，用户可以选择自己喜欢的平台，比如 Azure ML、Azure VMs 等，不受限制。

2. 核心技术

下面来细看 DeepSpeed 用了哪些核心技术，来实现它宣称的高效和有效。

2.1 ZeRO（零冗余优化器）

ZeRO（Zero Redundancy Optimizer，零冗余优化器）是一种专门为大规模分布式深度学习设计的新型内存优化技术。它的目标很明确：在当前一代 GPU 集群上，以当前最佳系统吞吐量的 3 到 5 倍速度，训练拥有 1000 亿参数的深度学习模型。不仅如此，它还为训练数万亿参数的模型提供了一条清晰可行的路径。作为 DeepSpeed 的一部分，ZeRO 同时提升了显存效率和计算效率。

ZeRO 克服了数据并行和模型并行各自的局限性，同时吸收了二者的优点。具体做法是，在数据并行进程中划分模型状态（参数、梯度、优化器状态），来消除数据并行进程中的内存冗余，而不是像传统做法那样复制它们。训练过程中，通过动态通信调度在分布式设备之间共享必要的状态，从而保持数据并行的计算粒度和通信量。

ZeRO 的三个阶段以及对应的功能

ZeRO 有三个主要的优化阶段，分别对应优化器状态、梯度和参数的划分：

1. Optimizer State Partitioning（Pos）
   ：减少 4 倍内存，通信量与数据并行相同。

2. 添加梯度分区（Pos+g）
   ：减少 8 倍内存，通信量与数据并行相同。

3. 添加参数分区（Pos+g+p）
   ：内存减少量与数据并行度 Nd 呈线性关系。比如在 64 个 GPU（Nd=64）之间进行拆分，就能实现 64 倍的内存缩减。通信量有 50% 的适度增长。

ZeRO 消除了内存冗余，让集群的全部聚合内存容量变得可用。在启用所有三个阶段的情况下，ZeRO 能够在 1024 个 NVIDIA GPU 上训练万亿参数模型。如果用 Adam 加 16 位精度，一个万亿参数模型大约需要 16 TB 内存来保存优化器状态、梯度和参数。16 TB 除以 1024 约等于 16 GB，这已经落在了主流 GPU 的合理显存范围内。

优化 1：ZeRO-2

ZeRO-2 扩展了 ZeRO-1，在减少梯度内存占用的同时，还增加对激活内存和碎片内存的优化。相比 ZeRO-1，ZeRO-2 将 DeepSpeed 可训练的模型规模扩大了一倍，并显著提升了训练效率。使用 ZeRO-2，1000 亿参数模型的训练速度，可以比仅使用模型并行性的现有技术快 10 倍。

优化 2：ZeRO-3 offload

ZeRO-3 offload 是 ZeRO Stage 3 和 ZeRO offload 结合的一种高效且易用的实现方式，它的核心目标是推动 AI 训练的民主化——让更多人能够高效地进行大规模深度学习训练。它的主要优势包括：

• 极高的内存效率
  ：可以在有限的 GPU 资源上运行非常大的模型，例如在单个 GPU 上运行超 400 亿参数的模型，或在 512 个 GPU 上运行 2 万亿参数的微调模型。

• 极易使用
  ：扩展到万亿参数级别，不需要复杂地组合多种并行技术。对于现有 DeepSpeed 用户，只需在配置文件中添加几个标志就能开启 ZeRO-3 offload。

• 高性能吞吐量
  ：在 512 个 NVIDIA V100 GPU 上使用 1 万亿参数进行训练，计算性能可维持在 25 PetaFlops，相当于每颗 GPU 达到 49 TFlop。相比纯 ZeRO 方案，单个 GPU 的吞吐量还能增加 2 倍。

2.2 基于 ZeRO 的 3D 并行化实现万亿参数模型训练

DeepSpeed 通过 ZeRO 支持的数据并行、流水线并行和张量切片模型并行，可以灵活组合使用。这三种并行方式各有所长：

• 数据并行
  ：当模型和参数较小时，单张 GPU 就能加载完整模型。但大模型训练需要海量数据，一张 GPU 放不下，就需要对数据分块，放到不同 GPU 上处理。反向传播后，通过通信归约梯度，保证优化器在各机器上做相同的更新。数据并行的优点是计算效率高，但它的 batch size 会随着 worker 数量增大，很多时候无法在不影响收敛的情况下一直增大 batch size。

• 模型并行
  ：数据并行时，每张 GPU 都会加载完整的模型结构。但当模型参数太多，单张 GPU 放不下时，就需要对模型进行分层，每张 GPU 处理其中一层。DeepSpeed 借用了英伟达的 Megatron-LM，基于 Transformer 提供大规模模型并行功能。模型并行的显存效率最高，代价是计算效率最低。

• 流水线并行
  ：将模型的各层划分后并行处理。完成一个 micro-batch 的正向传递时，激活内存会被通信到流水线的下一个阶段。下一阶段完成反向传播后，通过管道反向通信梯度。为了确保流水线各阶段能并行计算，需要同时处理多个 micro-batch。DeepSpeed 通过梯度累积实现并行，保持与传统数据并行和模型并行训练在相同的总 batch size 下收敛。

传统的数据并行和模型并行会保存模型运行时的全部状态，造成大量内存冗余。而 DeepSpeed 采用的基于 ZeRO 的 3D 并行，能极大优化显存利用率和计算效率。

2.3 ZeRO-Offload：单卡训练 10 倍大的模型

ZeRO-Offload 能同时利用 CPU 和 GPU 内存来训练大型模型，核心思路是把优化器状态和梯度卸载到 CPU 内存中。用户在使用单张 NVIDIA V100 GPU 时，可以在不耗尽显存的情况下运行多达 130 亿个参数的模型，模型规模扩展至现有方法的 10 倍，同时保持有竞争力的吞吐量。这个功能让数十亿参数模型的训练变得更接地气。

2.4 DeepSpeed Sparse Attention：6 倍速度处理 10 倍长序列

基于注意力的深度学习模型（如 Transformer）能很好地捕捉输入序列中 token 之间的关系，即使是长距离的。这让它们很适合处理文本、图像和声音这类输入，这些场景中序列长度动辄成千上万个 token。但实践中，长序列输入的应用受到注意力的计算和内存资源限制——注意力的计算需求随序列长度 n 呈二次增长，即 O(n²)。

DeepSpeed 提供的稀疏注意力核（Sparse Attention Kernel）解决了这个问题。它通过块稀疏计算，把注意力计算的计算和内存需求降低几个数量级。不仅缓解了内存瓶颈，还能高效执行稀疏计算。API 设计得很灵活，方便集成到任何基于 Transformer 的模型中。具体来说，稀疏注意力（SA）可以计算相邻 token 之间的局部注意力，再用局部注意力计算的摘要 token 来计算全局注意力。它还支持随机注意力，或者局部、全局和随机注意力的任意组合。这使得内存占用降到 O(wn)，其中 w 是一个远小于 n 的参数。

2.5 1 比特 Adam：减少 5 倍通信量，提升 3.4 倍训练速度

Adam 是目前大规模深度学习模型训练中最有效（也可能是应用最广泛）的优化器。但它通常与通信效率优化算法不太兼容。在跨设备分布式扩展时，通信开销很容易成为瓶颈。1 比特 Adam 新算法高效地解决了这个问题。它把通信量最多减少 5 倍，同时保持与 Adam 相似的收敛率。在通信受限的场景下，分布式训练速度提升了 3.5 倍，让该算法可以轻松扩展到不同类型的 GPU 集群和网络环境。

3. 如何使用

3.1 安装前须知

• 在安装 DeepSpeed 之前，必须安装 PyTorch，推荐使用版本 >= 1.9。
• 需要 CUDA 或 ROCm 编译器（如 nvcc 或 hipcc），用来编译 C++ / CUDA / HIP 扩展。
• DeepSpeed 最完整的测试在以下 GPU 上进行：NVIDIA 的 Pascal、Volta、Ampere 和 Hopper 架构，以及 AMD 的 MI100 和 MI200。

3.2 安装

安装非常简单，一行命令搞定：

pip install deepspeed

安装完成后，可以通过环境报告验证安装并检查计算机与哪些扩展/操作兼容：

ds_report

3.3 模型训练

训练流程可以归纳为三个步骤：

1. 初始化引擎
   ：初始化分布式环境、数据并行、混合精度训练设置等。系统会根据配置文件参数构建和管理训练优化器、数据加载器和学习率调度器。

2. 训练
   ：引擎初始化后，只需通过 3 个 API 就能完成前向传播、后向传播和权重更新。

训练中的模型保存和加载通过 sa ve_checkpoint 和 load_checkpoint 函数完成。这两个函数主要涉及两个参数：ckpt_dir（checkpoints 保存路径）和 ckpt_id（唯一标识符，代码中常用 loss 值作为标识）。DeepSpeed 可以自动保存和恢复模型、优化器及学习率调度器的状态。如果需要保存用户自定义的训练数据，可以通过客户端状态字典 client_sd 保存。

训练参数通过 config.json 配置文件进行管理，用户可以灵活设置 batch size、梯度累积步数、优化器类型、学习率、训练精度、是否启用 ZeRO 优化等。

3. 启动训练程序

对于多机环境，DeepSpeed 使用 hostfile 配置多节点计算资源。hostfile 包含主机名（或 SSH 别名）列表，这些机器支持无密码 SSH 访问，以及指定系统上可用 GPU 数量的槽计数。

启动命令如下，它会自动在 hostfile 中指定的所有可用节点和 GPU 上启动 PyTorch 训练脚本：

其中，client_entry.py 是模型入口脚本，client_args 是 argparse 命令行参数，ds_config.json 是 DeepSpeed 的配置文件。

对于多节点环境，用户可以在 .deepspeed_env 文件中定义需要传播的环境变量（key=value 格式），例如：

对于单节点环境（一个或多个 GPU），不需要添加 hostfile，系统会自动查询本地 GPU 数量和可用槽数。用户只需将 localhost 指定为主机名。注意：不需要用 CUDA_VISIBLE_DEVICES 来控制设备，它有专门的参数。

例如，只使用当前节点的 GPU 1，用以下命令启动：

deepspeed --include localhost:1 client_entry.py

也可以指定 GPU 节点和卡号，比如指定 worker-2 上的 0,1 卡：

3.4 模型推理

推理阶段同样简洁：

1. 加载模型
   ：使用 init_inference API 加载模型并配置推理参数。可以在初始化接口参数中指定模型并行度（MP 度）、提供 checkpoint 描述（通过 json 文件或 checkpoint 路径），也可以通过设置 replace_with_kernel_inject=True 注入高性能内核。需要注意的是，只有那些 DeepSpeed 官方支持的模型（如 Megatron-Turing NLG 530B、Jurassic-1 178B、BLOOM 176B、GLM 130B、YaLM 100B、GPT-NeoX 20B、AlexaTM 20B、Turing NLG 17B、METRO-LM 5.4B 等）才能使用 replace_with_kernel_inject=True。对 DeepSpeed 当前不支持的模型，用户可以通过 PR 提交一个 replace_policy 类来自定义注意力和前馈部分的参数，实现加速推理。

以下是用 DeepSpeed 推理 T5 模型实现文本生成任务的示例：

对于 HuggingFace 训练的模型，可以直接用 from_pretrained API 预加载模型 checkpoint。对于使用模型并行训练的 Megatron-LM 模型，需要在 JSON 配置中传递所有模型并行 checkpoint 的列表。

对于 DeepSpeed 训练的模型，checkpoint 只需提供保存路径：

2. 启动推理
   ：启动命令与训练类似，用 deepspeed 命令运行推理脚本即可。

3.5 DeepSpeed 与 HuggingFace 结合使用

PyTorch 训练的模型同样可以使用 DeepSpeed 加速推理。把 DeepSpeed 推理与 HuggingFace Pipeline 结合，可以非常方便地实现端到端生成。下面是用 GPT-NEO-2.7B 模型生成文本的客户端代码：

运行脚本：

输出结果：

3.6 实例：基于 PyTorch 和 DeepSpeed 实现 BERT 与 GPT 对比

使用纯 PyTorch 训练 BERT 的伪代码：

# 构建分布式模型
model = BertMultiTask()
model = DistributedDataParallel(model)
...
# 构建 FP16 优化器
optimizer = FusedAdam(model_parameters, ...)
optimizer = FP16_Optimizer(optimizer, ...)
# 前向传播
loss = model(batch)
# 后向传播
optimizer.backward(loss)
# 参数更新
optimizer.step()

使用 DeepSpeed 后的伪代码：

# 构建分布式模型
model = BertMultiTask()
...
# 构建 FP16 优化器
model, optimizer, _, _ = deepspeed.initialize(
    args=args, 
    model=model,
    model_parameters=model_parameters)
# 前向传播
loss = model(batch)
# 后向传播
model.backward(loss)
# 参数更新
model.step()

使用纯 PyTorch 训练 GPT-2 的伪代码：

# 构建分布式模型
model = GPT2Model(num_layers=args.num_layers, ...)
model = FP16_Module(model)
model = DistributedDataParallel(model, ...)
...
# 构建 FP16 优化器
optimizer = FusedAdam(model_parameters, ...)
optimizer = FP16_Optimizer(optimizer, ...)
# 前向传播
loss = model(batch)
# 后向传播
optimizer.backward(loss)
# 参数更新
optimizer.step()

使用 DeepSpeed 后的伪代码：

# 构建 GPT 模型
model = GPT2Model(num_layers=args.num_layers, ...)
...
# 构建 Adam 优化器
optimizer = Adam(param_groups, ...)
# 封装模型、优化器和学习率调度器
model, optimizer, lr_scheduler, _ = deepspeed.initialize(
    args=args, model=model, optimizer=optimizer,
    lr_scheduler=lr_scheduler, mpu=mpu)
# 前向传播
output = model(tokens, ...)
# 后向传播
model.backward(loss)
# 参数更新
model.step()

4. 性能分析

（1）模型规模

从模型规模来看，使用 ZeRO-Offload 技术的 DeepSpeed 可以训练 10 倍大的模型。用户在单张 NVIDIA V100 GPU 上，可以运行多达 130 亿个参数的模型，且保持有竞争力的吞吐量。这大大降低了数十亿参数模型训练的门槛。

上图可以直观看出，在 NVIDIA V100 GDX-2 卡节点上，不同节点数量下，与 3D 并行相比，使用 ZeRO-Infinity 能够支持的模型参数量更大，且随着节点数增加差距越明显。在 512 个 NVIDIA V100 GPU 上，ZeRO-Infinity 为数十亿和数万亿参数模型提供了超过 25 PB 的持续性能，训练效率非常出色。而传统 3D 并行由于 GPU 内存限制，遇到超过 1T 的模型就会内存溢出，ZeRO-Infinity 却能一路支持到 20T 参数的模型。

与 PyTorch 相比，DeepSpeed 的推理延迟降低了约 6 倍，吞吐量提升了 7 倍以上。

（2）速度

在速度方面，选择 GPT 家族中 40B-170B 参数的大模型，超过 25B 参数的模型时，DeepSpeed 相比 baseline 加速了 10 倍以上。

（3）扩展性

扩展性表现上，小规模时实现了超线性扩展，大规模时则呈现近线性或次线性扩展。这主要取决于内存可用性和最大全局 batch size 之间的相互作用。

上图显示，DeepSeed-RLHF 在最多 64 个 GPU 的集群上实现了良好的整体扩展。小规模时的超线性扩展，是因为随着 GPU 数量增加，每颗 GPU 的内存消耗减少，支持了更大的 batch 大小。而在大规模时，全局 batch size 的限制使扩展逐渐趋于线性或次线性。

（4）吞吐量

与 baseline 相比，MII-Public 提供了超过 15 倍的吞吐量提升，MII-Azure 则超过 40 倍。

（5）推理延迟

DeepSpeed 在推理延迟上的优势也很突出：

• 在 Stable Diffusion 图像生成任务上，延迟降低约 2 倍。
• 使用 MII 库后，Bloom、OPT 和 NeoX 等大模型的多 GPU 推理延迟降低高达 5.7 倍。
• OPT 和 BLOOM（7B 参数）等较小模型在单 GPU 上运行时，延迟降低高达 3 倍。
• RoBERTa 和 BERT 模型在文本表示任务上（如填充掩码、文本分类、问答、标记分类），延迟降低高达 9 倍。

（6）成本

在 Azure 上使用 MII 生成 1M tokens 的成本降低了 40 倍。MII（Model Implementations for Inference）是 DeepSpeed 的一个新开源 Python 库，它让大模型的低延迟、低成本推断变得可行且易于访问。

（7）模型压缩

DeepSpeed 提供的先进压缩技术（如 XTC）可以将模型尺寸缩小 32 倍，ZeROQuant 可降低压缩成本 5000 倍，单个 GPU 平均提升 2 倍效率。

5. 优势/特点

• 存储效率
  ：ZeRO 可将内存状态进行分区以节省大量内存，当前实现 ZeRO-2 相比现有技术可减少内存多达 8 倍。ZeRO-Offload 利用 CPU 和 GPU 内存，单卡即可运行 >130 亿参数的模型，规模扩大 10 倍。

• 可扩展性
  ：支持高效的数据并行、模型并行、管道并行及三者组合的 3D 并行。

• 通信效率
  ：Pipeline 并行减少了通信量，在网络带宽有限的集群上也能以 2-7 倍速度训练数十亿参数模型。1-bit Adam 和 0/1 Adam 可将通信量减少多达 26 倍。

• 数据效率
  ：通过课程学习提供高效数据采样，通过随机分层令牌丢弃提供高效数据路由。组合方案可在 GPT-3/BERT 预训练和 GPT/ViT 微调期间节省 2 倍数据和时间。

• 支持长序列
  ：稀疏注意力内核支持的输入序列长度是经典稠密 Transformer 的 10 倍，速度快 6 倍，精度相当。

• 快速收敛
  ：支持 LAMB 等优化器，提高训练效果并减少所需样本数量。

• 易用性
  ：仅需修改几行代码即可将 PyTorch 模型迁移到 DeepSpeed 和 ZeRO。

6. 参考文献

[1] microsoft/DeepSpeed
[2] DeepSpeed-MII: instant speedup on 24,000+ open-source DL models with up to 40x cheaper inference
[3] ZeRO & Fastest BERT: Increasing the scale and speed of deep learning training in DeepSpeed
[4] ZeRO-Inference: Democratizing massive model inference
[5] [译] DeepSpeed：所有人都能用的超大规模模型训练工具
[6] ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters