使用DeepSpeed和Accelerate实现了极快的BLOOM推断速度

本文介绍了如何在使用176B参数的BLOOM模型进行生成时获得非常快速的每个令牌吞吐量。

由于该模型需要352GB的bf16（bfloat16）权重（176*2），最高效的设置是8x80GB的A100 GPU。也可以使用2x8x40GB的A100或2x8x48GB的A6000。使用这些GPU的主要原因是在撰写本文时，它们提供了最大的GPU内存，但也可以使用其他GPU。例如，可以使用24x32GB的V100。

使用单个节点通常会提供最快的吞吐量，因为大多数时候节点内部的GPU链接硬件比节点间的链接硬件更快，但并非总是如此。

如果您没有那么多硬件，仍然可以在较小的GPU上运行BLOOM推理，通过使用CPU或NVMe卸载，但是生成时间当然会慢得多。

我们还将介绍需要半数GPU内存的8位量化解决方案，尽管吞吐量略有降低。我们将在BitsAndBytes和Deepspeed-Inference库中讨论这些内容。

基准测试

不再拖延，让我们展示一些数字。

为了保持一致性，除非另有说明，本文中的基准测试都是在同一台8x80GB的A100节点上进行的，配备512GB的CPU内存，位于Jean Zay HPC上。JeanZay HPC用户可以享受每秒约3GB的读取速度（GPFS），这对于检查点加载时间非常重要。慢速磁盘会导致较慢的加载时间。特别是因为我们同时在多个进程中进行IO。

所有基准测试都是对100个令牌输出进行贪婪生成：

生成参数{'max_length'：100，'do_sample'：False}

输入提示只包含几个令牌。之前的令牌缓存也开启了，因为始终重新计算它们会非常慢。

首先，让我们快速看一下准备生成所需的时间-即加载和准备模型的时间：

Deepspeed-Inference附带了预分片的权重存储库，在那里加载需要约1分钟。Accelerate的加载时间也非常出色，只需约2分钟。其他解决方案在这方面要慢得多。

加载时间可能重要，也可能不重要，因为加载后，您可以持续生成令牌，而无需额外的加载开销。

接下来是生成令牌吞吐量最重要的基准测试。这里的吞吐量指标很简单-生成100个新令牌所需的时间除以100和批量大小（即除以生成的令牌总数）。

这是在8x80GB GPU上的吞吐量：

其中OOM ==内存不足条件，批量大小过大无法适应GPU内存。

Deepspeed-Inference的张量并行（TP）和自定义融合CUDA内核，吞吐量低于1毫秒！这绝对令人惊讶！不过，对于尚未尝试过的其他模型使用此解决方案可能需要一些开发人员的时间来使其正常工作。

Accelerate也非常快速。它使用了非常简单的纯管道并行（PP）方法，因为它非常简单，所以应该可以与任何模型一起使用。

由于Deepspeed-ZeRO可以并行处理多个生成流，因此其吞吐量可以进一步除以8或16，具体取决于在generate调用期间使用了8个还是16个GPU。当然，这意味着它可以处理8×80 A100的批量大小为64，因此吞吐量约为4毫秒-因此所有3个解决方案非常接近。

让我们再次回顾这些数字是如何计算的。在使用Deepspeed-Inference进行fp16模式的情况下，为批量大小为128生成100个新令牌花费了8832毫秒的实时时间。因此，为了计算吞吐量，我们进行了以下计算：walltime /（batch_size * new_tokens）或8832 /（128 * 100）= 0.69。

现在让我们来看一下由Deepspeed-Inference和BitsAndBytes提供的基于量化int8模型的能力，因为它只需要原始bfloat16或float16推理的一半GPU内存。

在msecs 4x80GB A100上的吞吐量：

要复现基准结果，只需在下面讨论的任何这3个脚本中添加--benchmark。

解决方案

首先检出演示库：

git clone https://github.com/huggingface/transformers-bloom-inference
cd transformers-bloom-inference

在本文中，我们将使用位于bloom-inference-scripts/下的3个脚本。

框架特定的解决方案按字母顺序列出：

HuggingFace Accelerate

Accelerate

Accelerate以以下方式处理推理的大型模型：

使用空权重实例化模型。
分析每个层的大小和每个设备（GPU、CPU）上的可用空间，以决定每个层应该放在哪里。
逐位加载模型检查点，并将每个权重放在其设备上。

然后，它通过在正确设备上传输输入和输出，并确保模型权重在CPU（甚至是磁盘）上卸载的钩子，在正向传递之前在GPU上加载，并在正向传递完成后再次卸载。

在有多个具有足够空间容纳整个模型的GPU的情况下，它会在不同GPU之间切换控制，直到所有层都运行完毕。任何时候只有一个GPU在工作，尽管GPU处于空闲状态，但它确实产生了相当不错的吞吐量。

它也非常灵活，因为相同的代码可以在任何给定的设置上运行。Accelerate首先使用所有可用的GPU，然后卸载到CPU，直到RAM满，最后卸载到磁盘。卸载到CPU或磁盘会使速度变慢。例如，用户报告称，在仅有2个A100的情况下，运行BLOOM而无需更改代码的吞吐量为每个标记15秒，而在8×80 A100的情况下为10毫秒。

您可以在Accelerate文档中了解更多关于此解决方案的信息。

设置

pip install transformers>=4.21.3 accelerate>=0.12.0

运行

简单执行如下：

python bloom-inference-scripts/bloom-accelerate-inference.py --name bigscience/bloom --batch_size 1 --benchmark

要激活BitsAndBytes的8位量化解决方案，首先安装bitsandbytes：

pip install bitsandbytes

然后在前面的命令行中添加--dtype int8：

python bloom-inference-scripts/bloom-accelerate-inference.py --name bigscience/bloom --dtype int8 --batch_size 1 --benchmark

如果您有超过4个GPU，可以告诉它只使用4个：

CUDA_VISIBLE_DEVICES=0,1,2,3 python bloom-inference-scripts/bloom-accelerate-inference.py --name bigscience/bloom --dtype int8 --batch_size 1 --benchmark

在这种情况下，我们能够运行的最高批次大小为40。如果您查看脚本内部，我们不得不调整内存分配映射，以释放第一个GPU，仅处理激活和先前标记的缓存。

DeepSpeed-Inference

DeepSpeed-Inference使用Tensor-Parallelism和高效的融合CUDA内核，在128个大批量大小上实现超快的每个标记推理<1msec。

设置

pip install deepspeed>=0.7.3

运行

最快的方法是使用一个TP-pre-sharded（TP = Tensor Parallel）检查点，该检查点仅需要约1分钟即可加载，而非预分片bloom检查点需要10分钟：

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py --name microsoft/bloom-deepspeed-inference-fp16

1a. 如果您想运行原始的bloom检查点，它的吞吐量与之前的解决方案相同，但加载需要10-20分钟：

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py --name bigscience/bloom

2a. 8位量化版本仅需要正常的半精度版本的一半GPU内存：

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py --name microsoft/bloom-deepspeed-inference-int8 --dtype int8

这里我们使用了microsoft/bloom-deepspeed-inference-int8，并告诉脚本以int8方式运行。

当然，现在只需要4块80GB的A100 GPU即可：

deepspeed --num_gpus 4 bloom-inference-scripts/bloom-ds-inference.py --name microsoft/bloom-deepspeed-inference-int8 --dtype int8

在这种情况下，我们能够运行的最高批次大小是128。

您可以看到这里有两个因素导致了更好的性能。

通过使用张量并行（Tensor Parallelism，TP）而不是加速器的流水线并行（Pipeline Parallelism，PP），提高了吞吐量。因为加速器的设计非常通用，所以很难最大限度地利用GPU。所有计算都会先在GPU 0上进行，然后在GPU 1上进行，依此类推直到GPU 8，这意味着7个GPU一直处于空闲状态。另一方面，DeepSpeed-Inference使用TP，它会将张量发送到所有GPU，每个GPU上计算一部分生成，然后所有GPU之间相互通信结果，然后继续到下一层。这意味着所有GPU同时处于活动状态，但它们需要进行更多的通信。
DeepSpeed-Inference还使用自定义的CUDA内核，以避免分配过多的内存并进行与GPU的张量复制。这样做的效果是减少内存需求和减少内核启动次数，从而提高吞吐量并允许更大的批次大小，从而提高整体吞吐量。

如果您对更多示例感兴趣，可以查看使用DeepSpeed-Inference在GPU上加速加速器GPT-J推断或使用DeepSpeed-Inference在GPU上加速BERT推断。

Deepspeed ZeRO-Inference

Deepspeed ZeRO使用一种神奇的分片方法，可以将几乎任何模型扩展到少数或数百个GPU，并对其进行训练或推断。

设置

pip install deepspeed

运行

请注意，该脚本目前在所有GPU上运行相同的输入，但您可以在每个GPU上运行不同的流，从而获得n_gpu倍的吞吐量。使用Deepspeed-Inference无法做到这一点。

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-zero-inference.py --name bigscience/bloom --batch_size 1 --benchmark

请记住，使用ZeRO，用户可以同时生成多个唯一的流 – 因此，整体性能应该是吞吐量（秒/标记）除以参与的GPU数量 – 因此，速度可以提高8倍到16倍，具体取决于使用了8个还是16个GPU！

您还可以尝试使用仅具有一个较小的GPU的卸载解决方案，这需要很长时间运行，但如果您没有8个巨大的GPU，那么这就是最好的选择。

CPU卸载（1个GPU）：

deepspeed --num_gpus 1 bloom-inference-scripts/bloom-ds-zero-inference.py --name bigscience/bloom --batch_size 8 --cpu_offload --benchmark

NVMe卸载（1个GPU）：

deepspeed --num_gpus 1 bloom-inference-scripts/bloom-ds-zero-inference.py --name bigscience/bloom --batch_size 8 --nvme_offload_path=/path/to/nvme_offload --benchmark

请确保将/path/to/nvme_offload调整为你在快速NVMe驱动器上有约400GB可用内存的位置。

其他客户端和服务器解决方案

在transformers-bloom-inference，你将找到更多非常高效的解决方案，包括服务器解决方案。

下面是一些预览。

服务器解决方案：

Mayank Mishra将本博文中讨论的所有演示脚本转化为一个Web服务器包，你可以从这里下载。
Nicolas Patry开发了一个基于Rust的超高效Web服务器解决方案。

更多客户端解决方案：

Thomas Wang正在开发一个非常快速的自定义CUDA内核BLOOM模型。
HuggingFace的JAX团队已开发了一个基于JAX的解决方案。

由于本博文在发布后几个月后可能已过时，请使用transformers-bloom-inference查找最新的解决方案。

博文致谢

非常感谢以下热心的朋友们提出了好问题并帮助提高了文章的可读性：Olatunji Ruwase和Philipp Schmid。