在 Habana Gaudi2 加速器上进行 BLOOMZ 推理
TL;DR
Hugging Face 已经证明，Habana Gaudi2 加速器相较于当前可用的 GPU，能够为诸如 BLOOMZ 之类的大型语言模型提供更快的推理速度。对于拥有 1760 亿参数的 BLOOMZ 模型，Gaudi2 的速度比 Nvidia A100 80GB 快 1.42 倍。
BLOOMZ 模型能力
BLOOMZ 是 BLOOM 的微调版本，BLOOM 是一个拥有 1760 亿参数的自回归模型，旨在完成 46 种语言和 13 种编程语言的文本序列。BLOOMZ 通过提升泛化能力和零样本能力（即在没有先前训练示例的情况下，对未见输入数据完成任务的能力），对原始 BLOOM 架构进行了专门改进。
部署 BLOOMZ 对计算资源要求极高；在 16 位精度下，单个实例需要 352 GB 内存，这就需要高性能硬件和优化库才能实现低延迟。
Habana Gaudi2 硬件与软件栈
Gaudi2 是 Habana Labs 推出的第二代 AI 硬件加速器。其架构旨在并行执行通用矩阵乘法（GeMM）及其他操作，从而优化深度学习工作流的训练和推理。
硬件规格
单台 Gaudi2 服务器包含八个 Habana Processing Units（HPU），每个 HPU 配备 96 GB 内存，提供足够的容量来容纳超大模型。
软件集成
- SynapseAI™ SDK：支持 PyTorch 和 DeepSpeed。它包含一个图编译器，可通过算子融合、数据布局管理、并行化、流水线以及内存管理来优化执行。
- HPU Graphs：最近在 SynapseAI 中引入，用于支持对延迟敏感的应用。
- Optimum Habana：Hugging Face 的一个库，充当 Gaudi2 硬件与 Transformers 库之间的桥梁，简化部署流程。
推理基准：Gaudi2 vs. A100
为满足 BLOOMZ 的内存需求，Hugging Face 使用 DeepSpeed‑inference（通过 Habana 的 DeepSpeed 分支）在八个设备上实现模型和流水线并行。
延迟结果
基准测试使用 16 位精度、贪婪生成 100 个 token，以及键值缓存进行。结果显示 Gaudi2 在延迟方面始终优于 Nvidia A100 80GB：
| 模型 | 设备数量 | Gaudi2 延迟（秒） | A100-80GB 延迟（秒） | 第一代 Gaudi 延迟（秒） |
|---|---|---|---|---|
| BLOOMZ (176B) | 8 | 3.103 | 4.402 | / |
| BLOOMZ-7B | 8 | 0.734 | 2.417 | 3.321 |
| BLOOMZ-7B | 1 | 0.772 | 2.119 | 2.387 |
关键发现：
- 176B 模型：Gaudi2 的速度比 A100 80GB 快 1.42 倍。
- 7B 模型：Gaudi2 的速度比 A100 80GB 快 2.89 倍。
- 模型并行：Gaudi2 从模型并行中受益显著，而对于较小的 7B 模型，A100 在单设备上更快。
第一代 Gaudi 的性价比
对于 BLOOMZ-7B 模型，第一代 Gaudi 相较于 A100 提供了更优的性价比。A100 的费用超过每小时 $30，而第一代 Gaudi（AWS 上的 DL1 实例）约为每小时 $13，同时保持竞争性的延迟（2.387 秒）。
实现与复现
可以使用 optimum-habana 仓库中提供的脚本在完整数据集上进行推理。基准环境使用了 Transformers v4.28.1、SynapseAI v1.9.0 和 Optimum Habana v1.5.0。
要复现这些结果，用户需安装最新的 SynapseAI 和 Gaudi 驱动程序，然后安装 optimum-habana 库以及 Habana 专用的 DeepSpeed 分支（v1.9.0）。执行命令使用 gaudi_spawn.py，并带有 DeepSpeed、HPU 图和 KV 缓存的标志，以优化生成。