深入探索 deeps

推理和训练大模型通常需要巨大的计算资源和时间。微软推出 DeepSpeed 深度学习优化库，旨在解决这一问题。从本篇博客中，我们将深入了解 deepspeed，并理解微软的工程师们是如何通过对内存、并行度、通信的优化，从而极大地加速了大模型的推理和训练过程。

DeepSpeed 简介

DeepSpeed 是一个专门为深度学习模型训练设计的优化库。它实现了 ZeRO 论文中描述的所有内容。它是目前开源社区中广泛使用的训练大模型的框架。一般地，它支持以下三个阶段（stage）：

• ZeRO stage 1：优化器状态分区
• ZeRO stage 2：梯度分区
• ZeRO stage 3：参数分区

上述三个阶段给了用户更多更灵活的训练选择。用户可参考官方文档，选择适配自己硬件条件的训练方式。一般而言，花费的内存越少，训练时间越长；花费的内存越多，训练时间越短。

Fastest	Memory efficient
ZeRO-1	ZeRO-3 + offload
ZeRO-2	ZeRO-3
ZeRO-2 + offload	ZeRO-2 + offload
ZeRO-3	ZeRO-2
ZeRO-3 + offload	ZeRO-1

此外，它还支持了以下功能：

• 自定义混合精度训练处理。使用类似 PyTorch AMP 的方式，也可以选择使用类似 Apex 的方式
• 基于 CUDA 扩展的快速优化器。主要优化器包括 Adam、AdamW、OneBitAdam 和 Lamb
• 将部分训练参数卸载到 CPU 主存或者 SSD 上，适合于显存空间不足的用户。详细可参考ZeRO-Offload 到 CPU 和 NVMe这两篇论文。

deepspeed 之大模型推理

网络上使用 deepspeed 做训练的博客汗牛充栋，但使用它做推理的博客就比较少，因此我先从推理开始探索 deepspeed 的内部机制。但在这之前，我们需要先用 deepspeed 跑一个模型：

使用 deepspeed 完成模型推理

首先，需要进入 NGC docker，再安装 deepspeed transformers 等库：

docker run --init -it --name ${NAME} nvcr.io/nvidia/pytorch:23.03-py3 /bin/bash
pip install deepspeed transformers sentencepiece mpi4py

然后，准备好模型和权重数据。敲入以下推理代码，以模型 t5-v1_1-small 为例：

# create the model
from transformers import pipeline
from transformers.models.t5.modeling_t5 import T5Block
import os
import torch
import deepspeed

local_rank = int(os.environ["LOCAL_RANK"])
world_size = int(os.environ["WORLD_SIZE"])
deepspeed.init_distributed()
pipe = pipeline("text2text-generation", model="google/t5-v1_1-small", device=local_rank)
# Initialize the DeepSpeed-Inference engine
pipe.model = deepspeed.init_inference(
    pipe.model,
    tensor_parallel={"tp_size": world_size},
    dtype=torch.float,
    injection_policy={T5Block: ('SelfAttention.o', 'EncDecAttention.o', 'DenseReluDense.wo')}
)
output = pipe('Input String')

然后就是启动 deepspeed 完成推理：

deepspeed --num_gpus 2 inference.py

简单地解释一下 deepspeed.init_inference 这个重要的 API。该 API 装入并初始化推理模型。第一个参数是传入的推理模型；tensor_parallel 传入推理时用的张量并行度参数，一般就是显卡的数量；第三个参数是运算类型；第四个 injection_policy 参数给那些不支持 deepspeed 内核的模型准备的，用户需要手动指定模型的“注入策略”，即 Transformer 层的两个特定线性层：1）attention output GeMM 和 2）layer output GeMM。deepspeed 内部会根据用户提供的层，增加必要的 all-reduce 通信以便将各个 GPU 上的计算结果合并起来。上面的 t5-v1_1-small 就是例子。当然不是所有的 transformers 库内的模型都可以这样做，官方文档中列出了目前支持的模型

但具体到 injection_policy 是如何确定的，暂没有确切的定论，我猜测是根据 t5 模型源码里的 T5Block 类确定，最终实现类内的 SelfAttention.o, EncDecAttention.o, DenseReluDense.wo 计算并行。

此外，还有一些比较重要的参数，但在这个例子中没有体现，比如 replace_with_kernel_inject=True 可以将模型内的部分 kernel 替换成 deepspeed 内开发的高性能 kernel。

有时还需要完成 DeepSpeed 配置文件（通常为ds_config.json），指定使用的推理优化策略、参数等资源等，详细的配置说明可参考官方文档。

deepspeed 加速推理的原理

DeepSpeed ZeRO-3 也可以用于推理，因为它允许将单个GPU无法加载的大模型加载到多个GPU上。

DeepSpeed的核心就在于，GPU显存不够，CPU内存来凑。

比方说，我们只有一张 10GB 的 GPU，那么我们很可能需要借助 80GB 的 CPU，才能够训练一个大模型。

看一下官网对于这个理念的描述：

Why would you want to use DeepSpeed with just one GPU?
It has a ZeRO-offload feature which can delegate some computations and memory to the host’s CPU and RAM, and thus leave more GPU resources for model’s needs - e.g. larger batch size, or enabling a fitting of a very big model which normally won’t fit.
It provides a smart GPU memory management system, that minimizes memory fragmentation, which again allows you to fit bigger models and data batches.

具体点说，DeepSpeed将当前时刻，训练模型用不到的参数，缓存到CPU中，等到要用到了，再从CPU挪到GPU。这里的“参数”，不仅指的是模型参数，还指optimizer、梯度等。

越多的参数挪到CPU上，GPU的负担就越小；但随之的代价就是，更为频繁的CPU，GPU交互，极大增加了训练推理的时间开销。因此，DeepSpeed使用的一个核心要义是，时间开销和显存占用的权衡。

有关于offload
上述参数中，最重要的一个就是"offload_optimizer"。如上述所示，我们将其”device“设置成了cpu，DeepSpeed就会按照之前提到过的ZeRO操作，在训练过程中，将优化器状态分配到cpu上。从而降低单张GPU的memory占用。

有关于overlap_comm
另外一个需要提到的参数是overlap_comm。简单地理解，它控制着多个memory上进程之间通信的buffer的大小。这个值越大，进程之间通信越快，模型训练速度也会提升，但相应的显存占用也会变大；反之亦然。

因此，overlap_comm也是一个需要进行一定权衡的参数。

有关于auto
我们可以发现，上述大量参数被设置为auto。由于DeepSpeed目前已经被集成到了HuggingFace Transformer框架。而DeepSpeed的很多参数，和Transformer的Trainer参数设置是一模一样的，例如，“optimizer”，“scheduler”。因此，官方推荐将很多常用的模型训练参数，设置为auto，在使用Trainer进行训练的时候，这些值都会自动更新为Trainer中的设置，或者帮你自动计算。

DeepSpeed的加速原理主要基于以下几个方面的优化：

1. 内存优化：DeepSpeed通过自动混合精度训练、梯度累积和内存卸载等技术，有效减少了训练过程中的内存占用。这使得大型模型能够在有限的硬件资源上进行训练，同时提高了训练速度和稳定性。

2. 模型并行化：DeepSpeed支持多种模型并行化策略，如张量切片和管道并行。通过将模型拆分成多个部分并在多个GPU或节点上进行并行计算，DeepSpeed能够充分利用硬件资源，加速训练过程。

3. 通信优化：针对分布式训练中的通信瓶颈问题，DeepSpeed采用了梯度压缩、稀疏通信等技术，减少了通信开销，提高了训练效率。

4. 硬件感知优化：DeepSpeed能够根据硬件的特性进行相应的优化。例如，它可以利用NVIDIA的Ampere架构中的Tensor Cores进行高效的矩阵乘法运算，或者利用多核CPU进行并行数据处理。

DeepSpeed作为一款强大的深度学习优化库，为研究人员和开发者提供了高效的训练解决方案。它通过内存优化、模型并行化、通信优化以及硬件感知优化等技术，显著提升了深度学习模型的训练速度和效率。使用DeepSpeed，您可以更快速地探索新的模型架构和算法，加速科研和工程实践中的深度学习应用。无论是单个GPU的训练任务还是多节点集群的大规模训练任务，DeepSpeed都能为您带来卓越的性能提升和便捷的使用体验。