Hugging Face BLOOM 推理优化

Hugging Face 通过一系列迭代优化，将 BLOOM 模型的推理延迟降低了 5 倍，吞吐量提升了 50 倍。最终的架构使用 PyTorch 结合张量并行（TP）、用于注意力的自定义 CUDA 内核，以及 torch.jit.script 进行内核融合。

从流水线并行转向张量并行

BLOOM（1760 亿参数，352GB bf16）最初的推理是通过 accelerate 库的 device_map="auto" 使用流水线并行（PP）实现的。在 PP 中，每个 GPU 拥有一组特定的层，顺序处理数据并将其交给下一个 GPU。

为了降低延迟，Hugging Face 转向张量并行（TP），在该模式下每个 GPU 拥有每层权重的一部分，使所有 GPU 能够同时工作。此转变带来了显著的性能提升：

  • Latency（延迟）: 从 300ms/标记 降至 91ms/标记。
  • Throughput（吞吐量）: 提升至每秒 10 请求（RPS）。

虽然 TP 通过 ncclAllReduce 引入了通信开销，但批处理的能力（批大小为 1 与 32 时延迟往往相近）显著提升了整体吞吐量。

基于 PyTorch 的优化

除了并行策略之外，还实现了多项低层次的 PyTorch 优化，以消除通过 TensorBoard 进行性能分析时发现的瓶颈。

使用 torch.jit.script 的内核融合

Gelu 操作符最初会启动多个逐元素内核，导致大量张量拷贝。通过对 bloom_gelu_forward 函数使用 @torch.jit.script，Hugging Face 将其融合为单个内核操作，使延迟从 91ms/标记 降至 81ms/标记。

高效的 PyTorch 实现

  • ALiBi 优化: 位置嵌入之前在太多位置被计算。将其集中计算后，使该特定操作加速了 10 倍。
  • 减少张量拷贝: 性能分析显示注意力路径受到 reshape 和 transpose 操作的严重负担。重新组织权重和 KV 缓存（即 “past”）消除了这些不必要的拷贝。

自定义 CUDA 内核与硬件加速

为了进一步优化 torch.jit.script 无法覆盖的热点路径，Hugging Face 开发了自定义 CUDA 内核，将 masked fill 与 softmax 操作融合。

具体而言，该内核优化了以下序列：

  1. 使用注意力掩码对注意力分数进行 masked_fill_。
  2. 在 float32 上计算 softmax 以保证数值稳定性。

通过仅在内核内部对必要的求和和累加进行上溢（upcasting），延迟进一步从 81ms/标记 降至 71ms/标记。

Web 服务器架构与请求处理

为了服务多样化的用户请求，参数和长度各不相同，Hugging Face 实现了灵活的批处理系统：

  • 进程间通信: 由于 torch.distributed 需要独立的进程，服务器使用 Redis 的发布/订阅机制将原始字符串分发给所有进程。
  • 自定义生成循环: 标准的 generate 函数被自定义实现取代，该实现对批次中的每个成员应用不同的参数（例如采样、top-p）。
  • 动态批次提取: 为防止短请求被同批次的长请求拖慢，服务器在请求达到其 token 限制时立即提取并返回已完成的请求，而不是等待整个批次结束。

已评估但被放弃的方法

在整个优化过程中，探索了多条其他路线：

  • 在 TPU 上使用 JAX/Flax: 虽然并行实现更容易，但团队遇到了 Ray 与 TPU 工作节点通信的严重稳定性问题，并且缺乏对编译的细粒度控制。
  • DeepSpeed: 虽然提供了与最终迭代相似的出色结果，但在高负载下出现了稳定性问题，包括频繁的内核崩溃（CUDA illegal access）。
  • Rust 实现: 使用 tch-rs 编写了 Rust 版本，以获得更好的并发控制。然而，发现所谓的性能提升实际上是因为在 PyTorch 基准测试中留下了活跃的性能分析器。
  • ONNX/TensorRT: 被认为在文本生成循环所需的灵活性以及保持张量在 GPU 上进行 logits 计算的需求上过于僵硬。

Sources