本文展现了怎样运用 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}
输入提示词仅包含几个词。咱们会缓存从前见到的词，由于每次从头核算它们适当慢。
首要，让咱们快速看一下从开端到预备好花了多长时刻， 即模型加载和预备花了多长时刻：
| 办法 | 秒 |
|---|---|
| accelerate | 121 |
| ds-inference shard-int8 | 61 |
| ds-inference shard-fp16 | 60 |
| ds-inference unsharded | 662 |
| ds-zero | 462 |
Deepspeed-Inference 运用了预分片的权重库房，整个加载时刻大约在 1 分钟。Accelerrate 的加载时刻也很优异，只要大约 2 分钟。其他计划就慢得多。
加载时刻有或许重要也或许并不重要，由于一旦加载成功你能够一遍遍持续不断地生成词而不再需求额外地加载开支。
接着是最重要的测验基准目标：词生成吞吐 (token generation throughput)。这个吞吐的度量比较简略，即：生成 100 个新词的时刻除以 100 和 batch size (也便是除以生成的总词数)。
下面列出了 8x80GB GPU 的吞吐，单位为毫秒：
| 办法 \ bs | 1 | 8 | 16 | 32 | 64 | 128 | 256 | 512 |
|---|---|---|---|---|---|---|---|---|
| accelerate bf16 | 230.38 | 31.78 | 17.84 | 10.89 | oom | |||
| accelerate int8 | 286.56 | 40.92 | 22.65 | 13.27 | oom | |||
| ds-inference fp16 | 44.02 | 5.70 | 3.01 | 1.68 | 1.00 | 0.69 | oom | |
| ds-inference int8 | 89.09 | 11.44 | 5.88 | 3.09 | 1.71 | 1.02 | 0.71 | oom |
| ds-zero bf16 | 283 | 34.88 | oom |
这里， 当内存耗尽 (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 的吞吐，单位为毫秒：
| 办法 bs | 1 | 8 | 16 | 32 | 64 | 128 |
|---|---|---|---|---|---|---|
| accelerate int8 | 284.15 | 40.14 | 21.97 | oom | ||
| ds-inference int8 | 156.51 | 20.11 | 10.38 | 5.50 | 2.96 | oom |
你只需在下述 3 个脚本里添加 --benchmark 即可重现这些测验基准的成果。
计划
首要获取最新的演示代码库房：
git clone https://github.com/huggingface/transformers-bloom-inference
cd transformers-bloom-inference
本文咱们预备运用 bloom-inference-scripts/ 文件夹下的 3 个脚本。
下面咱们按结构的字母序逐一展现相关计划。
HuggingFace Accelerate
Accelerate 按如下步骤进行大模型推理：
- 用空的权重实例化模型。
- 分析每层的大小以及每个设备 (CPU, CPU) 的可用空间，并决议每层应该在哪个设备上推理。
- 逐比特加载模型 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。
能够看到，本计划中有两个因素在取得更好的功能上起到了主导作用。
-
本计划的吞吐提高首要来自于张量并行 (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 多得多的通信。
-
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 模型在各模态数据上的应用及大规模模型的练习推理。
