使用 DeepSpeed 和 Accelerate 实现快得惊人的 BLOOM 推理

本文展示了在使用拥有 176B 参数的 BLOOM 模型进行文本生成时，如何获得快得惊人的单 token 吞吐量。

由于该模型在 bf16 (bfloat16) 精度下需要 352GB 的权重 (176*2)，最高效的配置是 8x80GB 的 A100 GPU。此外，2x8x40GB 的 A100s 或 2x8x48GB 的 A6000 也可以使用。使用这些 GPU 的主要原因是，在撰写本文时，它们提供了最大的 GPU 显存，但其他 GPU 也可以使用。例如，24x32GB 的 V100s 也是可以的。

使用单个节点通常能提供最快的吞吐量，因为大多数情况下，节点内 GPU 之间的连接硬件比节点间连接更快，但这并非总是如此。

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

我们还将介绍 8 比特量化解决方案，这些方案需要的 GPU 显存减半，但代价是吞吐量略有降低。我们将讨论 BitsAndBytes 和 Deepspeed-Inference 这两个库。

基准测试

话不多说，我们先来看一些数据。

为了保持一致性，除非另有说明，本文中的所有基准测试都是在 Jean Zay HPC 上的同一个 8x80GB A100 节点上完成的，该节点配有 512GB 的 CPU 内存。JeanZay HPC 用户享有约 3GB/s 的快速 IO 读取速度（GPFS）。这对于检查点加载时间非常重要。慢速磁盘会导致加载时间变慢。尤其是在我们并发地在多个进程中进行 IO 操作时。

所有基准测试都在进行 贪心生成 100 个 token 的输出。

Generate args {'max_length': 100, 'do_sample': False}

输入提示仅包含几个 token。同时也开启了前一个 token 的缓存，因为每次都重新计算会相当慢。

首先，我们快速看一下准备生成需要多长时间——即加载和准备模型需要多长时间

DeepSpeed-Inference 附带了预分片的权重仓库，加载时间大约为 1 分钟。Accelerate 的加载时间也非常出色，仅需约 2 分钟。其他解决方案在这里则慢得多。

加载时间可能重要，也可能不重要，因为一旦加载完成，你就可以一次又一次地持续生成 token，而无需额外的加载开销。

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

以下是在 8x80GB GPU 上的吞吐量（毫秒）

其中 OOM == Out of Memory（内存不足），表示批处理大小太大而无法装入 GPU 内存。

使用 Deepspeed-Inference 的张量并行（Tensor Parallelism，TP）和定制的融合 CUDA 内核，实现了低于 1 毫秒的吞吐量！这绝对是惊人的！不过，将此解决方案用于其他未经测试的模型可能需要一些开发时间才能使其工作。

Accelerate 也超级快。它使用一种非常简单的朴素流水线并行（Pipeline Parallelism，PP）方法，由于其简单性，它应该可以与任何模型开箱即用。

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

让我们再回顾一下这些数字是如何计算的。在使用 Deepspeed-Inference 的 fp16 模式时，为批大小为 128 的批次生成 100 个新 token 的实际时间为 8832 毫秒。因此，为了计算吞吐量，我们做了：walltime/(batch_size*new_tokens) 或 `8832/(128*100) = 0.69`。

现在让我们看看 Deepspeed-Inference 和 BitsAndBytes 提供的基于 int8 量化模型的强大功能，因为它只需要 bfloat16 或 float16 推理所需 GPU 显存的一半。

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 以下列方式处理大模型的推理

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

然后，它通过钩子确保模型正常运行，这些钩子将输入和输出传输到正确的设备上，并将卸载到 CPU（甚至磁盘）上的模型权重在正向传播之前加载到 GPU 上，然后在正向传播结束后再次卸载。

在有多个 GPU 且空间足以容纳整个模型的情况下，它会从一个 GPU 切换到下一个 GPU，直到所有层都运行完毕。任何给定时间只有一个 GPU 在工作，这听起来效率很低，但尽管 GPU 处于空闲状态，它仍能产生不错的吞吐量。

它也非常灵活，因为相同的代码可以在任何给定的设置上运行。Accelerate 会首先使用所有可用的 GPU，然后在 RAM 满时卸载到 CPU，最后卸载到磁盘。卸载到 CPU 或磁盘会使速度变慢。例如，有用户报告称，在仅有 2 个 A100 的情况下，无需更改代码即可运行 BLOOM，吞吐量为每 token 15 秒，而在 8x80 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，没有出现 OOM。如果你查看脚本内部，我们不得不调整内存分配映射，以释放第一个 GPU 来专门处理激活和前一个 token 的缓存。

DeepSpeed-Inference

DeepSpeed-Inference 使用张量并行和高效的融合 CUDA 内核，在大批处理大小为 128 的情况下，提供了超快的每 token 不到 1 毫秒的推理速度。

设置

pip install deepspeed>=0.7.3

运行

1. 最快的方法是使用一个经过 TP-pre-sharded（TP = 张量并行）的检查点，加载只需约 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` 模式运行。

当然，现在只需要 4x80GB 的 A100 GPU 就足够了

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

在这种情况下，我们能够运行的最高批处理大小是 128，没有出现 OOM。

你可以看到这里有两个因素共同作用，带来了更好的性能。

1. 这里的吞吐量通过使用张量并行（TP）而不是 Accelerate 的流水线并行（PP）得到了提升。因为 Accelerate 的目标是具有很高的通用性，不幸的是，这也使得最大化 GPU 使用率变得困难。所有计算首先在 GPU 0 上完成，然后在 GPU 1 上，依此类推，直到 GPU 8，这意味着 7 个 GPU 一直处于空闲状态。另一方面，DeepSpeed-Inference 使用 TP，这意味着它会将张量发送到所有 GPU，在每个 GPU 上计算生成的一部分，然后所有 GPU 相互通信结果，再进行下一层。这意味着所有 GPU 同时处于活动状态，但它们需要更多的通信。

2. DeepSpeed-Inference 还使用定制的 CUDA 内核来避免分配过多内存以及在 GPU 之间进行张量复制。这样做的效果是减少了内存需求和内核启动次数，从而提高了吞吐量，并允许更大的批处理大小，最终带来更高的整体吞吐量。

如果你对更多示例感兴趣，可以查看 在 GPU 上使用 DeepSpeed-Inference 加速 GPT-J 推理 或 在 GPU 上使用 DeepSpeed-Inference 加速 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，用户可以同时生成多个独立的流，因此整体性能应该是以秒/token 计算的吞吐量除以参与的 GPU 数量——因此，根据使用的是 8 个还是 16 个 GPU，速度会快 8 到 16 倍！

你也可以只用一个较小的 GPU 尝试卸载解决方案，这会花费很长时间，但如果你没有 8 个大型 GPU，这已经是最好的选择了。

CPU-Offload (1x GPUs)

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

NVMe-Offload (1x GPUs)

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` 调整到一个你有约 400GB 可用空间的快速 NVMe 驱动器上。

额外的客户端和服务器解决方案

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

以下是一些预览。

服务器解决方案

• Mayank Mishra 将本博客文章中讨论的所有演示脚本转化为一个网络服务器包，你可以从这里下载

• Nicolas Patry 开发了一个超高效的基于 Rust 的网络服务器解决方案。

更多客户端解决方案

• Thomas Wang 正在开发一个非常快的定制 CUDA 内核 BLOOM 模型。

• HuggingFace 的 JAX 团队开发了一个基于 JAX 的解决方案

由于这篇博客文章在你阅读时可能已经过时，如果是在发布几个月后阅读，请使用 transformers-bloom-inference 寻找最新的解决方案。

博客致谢

非常感谢以下友好人士，他们提出了很好的问题，并帮助提高了文章的可读性：Olatunji Ruwase 和 Philipp Schmid。