惊人快速的 BLOOM 推理：使用 DeepSpeed 与 Accelerate

TL;DR

Hugging Face 表明，在 8×80 GB A100 GPU 单节点上，使用 DeepSpeed‑Inference 张量并行（或使用 Accelerate 流水线并行）时，1760 亿参数的 BLOOM 模型可以在不到 1 ms 的时间内生成标记，并且还提供了量化至 8‑bit 的方案，可将显存占用减半。

硬件需求与配置

最佳配置：8 × 80 GB A100（共 352 GB bf16 权重）。可选方案包括 2 × 8 × 40 GB A100、2 × 8 × 48 GB A6000，或 24 × 32 GB V100。
单节点优势：节点内部 GPU 互连速度快于跨节点链路，通常能获得更高吞吐量。
低端选项：CPU 或 NVMe 卸载可以在较小的 GPU 上运行 BLOOM，但生成延迟会显著增加。
量化推理：8‑bit 模型（通过 BitsAndBytes）大约只需一半的 GPU 显存，吞吐量略有下降。

基准环境

节点：Jean Zay HPC，8 × 80 GB A100，512 GB CPU RAM，GPFS 存储（约 3 GB/s 读取速度）。
任务：使用贪婪解码生成 100 个标记（max_length=100, do_sample=False），使用短提示并开启 KV‑cache。
指标：模型加载时间（秒）以及每标记吞吐量（ms per token = 实际时间 ÷ (batch × tokens)）。

模型加载时间

方案	加载时间 (s)
Accelerate	121
DeepSpeed‑Inference (shard‑int8)	61
DeepSpeed‑Inference (shard‑fp16)	60
DeepSpeed‑Inference (unsharded)	662
DeepSpeed‑ZeRO	462

预分片的 DeepSpeed 检查点约在 1 分钟内加载完成，而未分片的检查点可能需要超过 10 分钟。

标记生成吞吐量（8 × 80 GB A100）

方案 (dtype)	Batch 1	Batch 8	Batch 16	Batch 32	Batch 64	Batch 128	Batch 256	Batch 512
Accelerate bf16	230.38 ms	31.78 ms	17.84 ms	10.89 ms	OOM	–	–	–
Accelerate int8	286.56 ms	40.92 ms	22.65 ms	13.27 ms	OOM	–	–	–
DeepSpeed‑Inference fp16	44.02 ms	5.70 ms	3.01 ms	1.68 ms	1.00 ms	0.69 ms	OOM	–
DeepSpeed‑Inference int8	89.09 ms	11.44 ms	5.88 ms	3.09 ms	1.71 ms	1.02 ms	0.71 ms	OOM
DeepSpeed‑ZeRO bf16	283 ms	34.88 ms	OOM	–	–	–	–	–

关键观察

DeepSpeed‑Inference 在 batch 128 时实现了亚毫秒级每标记性能，得益于张量并行（TP）和自定义融合 CUDA 核心。
Accelerate 使用朴素的流水线并行（PP），在 batch 32 时约为 10 ms/标记，但在 batch 64 以上因显存限制无法继续扩展。
量化至 int8 可将显存占用减半；DeepSpeed‑Inference int8 在 batch 128 仍能达到约 1 ms/标记，而 Accelerate int8 更早出现 OOM。

量化 8‑bit 吞吐量（4 × 80 GB A100）

方案	Batch 1	Batch 8	Batch 16	Batch 32	Batch 64	Batch 128
Accelerate int8	284.15 ms	40.14 ms	21.97 ms	OOM	–	–
DeepSpeed‑Inference int8	156.51 ms	20.11 ms	10.38 ms	5.50 ms	2.96 ms	OOM

方案细节

HuggingFace Accelerate

方法：根据层大小和可用显存，惰性地将模型权重加载到设备上；使用简单的流水线并行，每层仅在一块 GPU 上激活。
优点：开箱即用，适配任何硬件配置；当显存不足时可卸载到 CPU 或磁盘。
缺点：前向传播大部分时间 GPU 处于空闲，导致峰值吞吐受限；较大 batch 会快速触发 OOM。

使用方式：

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
# 8‑bit 量化运行
pip install bitsandbytes
python bloom-inference-scripts/bloom-accelerate-inference.py \
  --name bigscience/bloom --dtype int8 --batch_size 1 --benchmark

DeepSpeed‑Inference

方法：张量并行将每层切分到多块 GPU；自定义融合核降低内存拷贝和 kernel 启动次数。
性能驱动因素：
1. TP vs PP – 所有 GPU 同时计算，提高利用率。
2. 融合核 – 降低内存开销并减少 kernel 调用。
预分片检查点（microsoft/bloom-deepspeed-inference-fp16）约在 1 分钟内加载完成；未分片检查点加载需 10‑20 分钟，但加载后运行速度相同。
量化 int8（microsoft/bloom-deepspeed-inference-int8）显存需求减半，在 8 × 80 GB A100 上仍能在 batch 128 达到亚毫秒级每标记，或在 4 × 80 GB A100 上运行。

使用方式：

pip install deepspeed>=0.7.3
# 快速 TP‑预分片 fp16
deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py \
  --name microsoft/bloom-deepspeed-inference-fp16
# 原始检查点（加载慢）
deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py \
  --name bigscience/bloom
# 8‑bit 版本（显存减半）
deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py \
  --name microsoft/bloom-deepspeed-inference-int8 --dtype int8

DeepSpeed‑ZeRO Inference

方法：将模型状态在 GPU 之间分片（ZeRO‑3 方式），并可并行运行多个独立的生成流，理论上可获得与 GPU 数量成比例的加速。
局限：提供的脚本在所有 GPU 上运行相同输入；要实现 8× 或 16× 加速需自行实现每块 GPU 的独立流。
卸载选项：CPU 卸载或 NVMe 卸载可以在单卡上进行推理，但会带来较大延迟。

使用方式：

pip install deepspeed
# 多 GPU ZeRO 推理
deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-zero-inference.py \
  --name bigscience/bloom --batch_size 1 --benchmark
# CPU 卸载（单 GPU）
deepspeed --num_gpus 1 bloom-inference-scripts/bloom-ds-zero-inference.py \
  --name bigscience/bloom --batch_size 8 --cpu_offload --benchmark
# NVMe 卸载（单 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

社区服务器与客户端扩展

服务器实现：
- Mayank Mishra 将演示脚本打包成可直接运行的 Web 服务器。
- Nicolas Patry 开发了高性能的 Rust 版服务器。
客户端项目：
- Thomas Wang 正在开发自定义 CUDA 核心加速的 BLOOM 模型。
- HuggingFace JAX 团队发布了针对 BLOOM 的 JAX 推理后端。
保持更新：transformers-bloom-inference 仓库聚合了最新的脚本和服务器实现。

实践要点

追求最高原始吞吐 时，单节点使用 DeepSpeed‑Inference 加预分片 fp16 检查点；在 batch 128 可实现亚毫秒级每标记。
显存受限 时，使用 BitsAndBytes 的 8‑bit 量化（DeepSpeed 或 Accelerate），显存减半，速度略有下降。
Accelerate 是最灵活的方案，适配异构硬件，但其流水线并行导致的延迟高于 DeepSpeed 的张量并行。
ZeRO 推理 为跨多 GPU 扩展或在单 GPU 上通过卸载提供了路径，适用于没有大规模 GPU 集群的情况。

“我们在此情况下能够运行的最大 batch 大小而不出现 OOM 为 128。” – Hugging Face 博客，描述在 8 × 80 GB A100 上的 DeepSpeed‑Inference int8。

本文反映了截至 2022 年 9 月的 BLOOM 推理性能状态。欲了解更新的优化，请查阅 transformers-bloom-inference GitHub 仓库。

惊人快速的 BLOOM 推理：使用 DeepSpeed 与 Accelerate

TL;DR

硬件需求与配置

基准环境

模型加载时间

标记生成吞吐量（8 × 80 GB A100）

量化 8‑bit 吞吐量（4 × 80 GB A100）

方案细节

HuggingFace Accelerate

DeepSpeed‑Inference

DeepSpeed‑ZeRO Inference

社区服务器与客户端扩展

实践要点

Sources

标记生成吞吐量（8 × 80 GB A100）

量化 8‑bit 吞吐量（4 × 80 GB A100）