关于人工智能:使用-DeepSpeed-和-Accelerate-进行超快-BLOOM-模型推理

本文展现了如何应用 1760 亿 (176B) 参数的 BLOOM 模型 生成文本时如何取得超快的词吞吐 (per token throughput)。

因为在应用 bf16 (bfloat16) 权重时该模型内存占用为 352 GB (176*2)，所以最高效的硬件配置是应用 8x80GB 的 A100 GPU。也可应用 2x8x40GB 的 A100 或者 2x8x48GB 的 A6000。应用这些 GPU 的次要起因是截至本文成稿时为止它们是能提供最大显存的 GPU，但你也能够应用其余 GPU。比方，能够应用 24x32GB V100。

一般来讲，应用单节点会带来最快的吞吐，因为大多数时候节点内的 GPU 互联硬件比节点间的快，但未必总是如此。

如果你没有这么高端的硬件或没有这么多硬件，你仍可能通过 CPU 卸载 (CPU offload) 或是 NVMe 卸载 (NVMe offload) 的形式在更小的 GPU 上对 BLOOM 进行推理。当然，生成工夫会慢很多。

咱们打算波及 8 比特量化计划，该计划以稍慢的吞吐为代价将显存需要缩小到一半。咱们还会探讨 BitsAndBytes 和 Deepspeed-Inference 库。

测试基准

事不宜迟，咱们先展现一些数据吧。

为了放弃一致性，除非另有阐明，本文的测试基准都是在雷同的配有 512GB CPU 内存的 8x80GB A100 节点上实现的，该节点来自 法国 Jean Zay 超算核心

所有的测试基准都是应用贪婪搜寻实现最多 100 个词的生成工作:

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

输出提醒词仅蕴含几个词。咱们会缓存先前见到的词，因为每次从新计算它们相当慢。

首先，让咱们疾速看一下从开始到筹备好花了多长时间，即模型加载和筹备花了多长时间：

Deepspeed-Inference 应用了预分片的权重仓库，整个加载工夫大概在 1 分钟。Accelerrate 的加载工夫也很优良，只有大概 2 分钟。其余计划就慢得多。

加载工夫有可能重要也可能并不重要，因为一旦加载胜利你能够一遍遍继续一直地生成词而不再须要额定地加载开销。

接着是最重要的测试基准指标：词生成吞吐 (token generation throughput)。这个吞吐的度量比较简单，即：生成 100 个新词的工夫除以 100 和 batch size (也就是除以生成的总词数)。

上面列出了 8x80GB GPU 的吞吐，单位为毫秒：

这里，当内存耗尽 (Out Of Memory，OOM) 时即表明 batch size 太大 GPU 显寄存不下了。

应用 Deepspeed-Inference 的张量并行 (Tensor Parallelism，TP) 和定制化交融 CUDA 核函数能够失去小于 1 毫秒的吞吐！太棒了！只管应用这个计划去推理那些尚未被验证过的模型时，你可能会须要花一些工夫去开发从而让它工作起来。

Accelerate 也超级快。它应用了非常简单的管线并行 (Pipeline Parallelism，PP)。因为它非常简单，所以它应该对任何模型都是开箱即用的。

因为 Deepspeed-ZeRO 能够并行处理多路生成流，其吞吐能够再除以 8 或者 16，具体数值取决于在调用 generate 时用了 8 个 GPU 还是 16 个 GPU。当然，这也意味着在 8x80GB A100 的状况下 (见上表)，能够解决的 batch size 为 64 且吞吐可至大概 4 毫秒。因而，这 3 种计划的性能是靠近的。

让咱们再从新看一下这些数字是怎么计算出来的。举个例子，应用 Deepspeed-Inference fp16 模式实时生成 batch size 为 128、长度为 100 个新词的文本花了 8832 毫秒，因而咱们能够这样计算吞吐：钟面工夫 / (batch size * 新词数) 或 8821/(128*100) = 0.69。

当初咱们一起看看 Deepspeed-Inference 和 BitsAndBytes 提供的 int8 量化模型的威力，它仅需占用 bfloat16 或 float16 推理所需显存的一半。

以下为 4x80GB GPU 的吞吐，单位为毫秒：

你只需在下述 3 个脚本里增加 --benchmark 即可重现这些测试基准的后果。

计划

首先获取最新的演示代码仓库：

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

cd transformers-bloom-inference

本文咱们筹备应用 bloom-inference-scripts/ 文件夹下的 3 个脚本。

上面咱们按框架的字母序逐个展现相干计划。

HuggingFace Accelerate

Accelerate 按如下步骤进行大模型推理：

1. 用空的权重实例化模型。
2. 剖析每层的大小以及每个设施 (CPU, CPU) 的可用空间，并决定每层应该在哪个设施上推理。
3. 逐比特加载模型 checkpoint 并把权重加载到相应的设施。

而后，它会应用钩子代码 (hook) 来确保模型正确运行，钩子代码被用于在正确的设施间传输输出和输入，并在前向轮运行前加载那些卸载到 CPU (甚至硬盘) 上的权重到 GPU，而后在前向轮完结后再次把权重卸载。

在有多个 GPU 且有足够空间放下整个模型的情景下，该计划在 GPU 间一一切换直至所有层运行结束。每个给定的工夫只有一个 GPU 工作，这听起来很没效率。但只管该计划 GPU 存在闲暇，它的吞吐却相当不错。

因为雷同的代码能够运行在任意给定的设置中，所以本计划非常灵活。Accelerate 首先应用所有可用的 GPU，当显存已满时会卸载到 CPU 内存直至卸载到硬盘。卸载到 CPU 或硬盘会让推理变慢。举个例子，与 8×80 A100 上的 10 毫秒相比，已有用户报告，不作任何代码改变，在 2 个 A100 上运行 BLOOM 吞吐是每词 15 秒。

你能够你从 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 个 GPU：

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

在这个例子中，不 OOM 的最大 batch size 是 40。如果你深入研究脚本，你会看到咱们须要调整显存调配映射从而把第一个 GPU 解放出来去仅解决激活和先前词的缓存。

DeepSpeed-Inference

DeepSpeed-Inference 应用张量并行 (Tensor Parallelism) 以及高效的交融 CUDA 核函数在 128 这个大 batch size 下达到了每词 1 毫秒的超快推理性能。

设置

pip install deepspeed>=0.7.3

运行

1. 最快的办法是应用 TP 预分片 (TP = Tensor Parallel) 的 checkpoint，与非预分片的 bloom checkpoint 相比，它仅需大概 1 分钟即可加载：

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

1a. 如果你想要运行原始 bloom checkpoint，这个 checkpoint 一旦加载就会跟之前的计划跑到雷同的吞吐，但加载须要花 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 checkpoint 并通知脚本跑在 int8 模式。

当然，当初仅需 4x80GB A100 GPU 就够了：

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

这种状况下，不 OOM 的最大 batch size 是 128。

能够看到，本计划中有两个因素在取得更好的性能上起到了主导作用。

1. 本计划的吞吐进步次要来自于张量并行 (Tensor Parallelism，TP) 而不是 Acclerate 的管线并行 (Pipeline Parallelism，PP)。因为 Accelerate 旨在成为十分通用的计划，因而也十分可怜地很难最大化 GPU 使用率。它首先在 GPU 0 上实现所有计算，而后是 GPU 1，等等，始终到 GPU 8，这意味着任何时刻都有 7 个 GPU 是闲暇的。而另一方面，DeepSpeed-Inference 应用了 TP，意味着它会向所有 GPU 发送张量，在每个 GPU 上计算局部生成后果，而后在所有的 GPU 间通信计算结果，并持续做下一层。这就是说 TP 所有的 GPU 都同时是沉闷的，但它须要比 PP 多得多的通信。
2. DeepSpeed-Inference 还应用了定制的 CUDA 核函数以防止调配太多内存以及太多进出 GPU 的张量拷贝。这么做会缩小显存需要及核函数启动次数从而进步吞吐，另外还能够反对更大的 batch size 从而进一步减少总吞吐。

如果你对更多的例子感兴趣，能够看看 在 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 同时创立多个不同的流，因而总性能应该是每秒每词的吞吐除以参加计算的 GPU 的数目，因而依据你是应用 16 个 GPU 还是 8 个 GPU，能够取得 8 倍或者 16 倍的更快性能。

你还能够在一个小型 GPU 上试试卸载计划，运行的工夫会很长，然而如果你没有 8 个巨型 GPU 的话这也是一个聊甚于无的计划。

CPU 卸载 (1x GPUs):

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

NVMe 卸载 (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

请确保在你的高速 NVMe 盘上预留约 400GB 的空间，并把 /path/to/nvme_offload 设成它。

更多客户端及服务器计划

你能够从 transformers-bloom-inference 找到更多十分高效的计划，包含服务器计划。

这里咱们提供一些预览。

服务器计划：

• Mayank Mishra 拿着本博文中探讨的所有演示代码并把它们变成了一个网络服务包，你能够从 这儿 下载。
• Nicolas Patry 开发了一个超高效的 基于 Rust 的网络服务计划。

更多的客户端计划：

• Thomas Wang 正在开发一个很快的 定制 CUDA 核函数的 BLOOM 模型。
• HuggingFace 的 JAX 组已开发了一个 基于 JAX 的计划

因为如果你在本博文公布几个月后读到它，很有可能它曾经不能反映最新的状态了，你能够去 transformers-bloom-inference 的 GitHub 仓库 找到最新的计划。

致谢

万分感激如下这些人，他们提出了好的问题并帮忙进步了本文的可读性：Olatunji Ruwase 和 Philipp Schmid。

英文原文: Incredibly Fast BLOOM Inference with DeepSpeed and Accelerate

译者: Matrix Yao (姚伟峰)，英特尔深度学习工程师，工作方向为 transformer-family 模型在各模态数据上的利用及大规模模型的训练推理。