[论文解读] DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale

Reza Yazdani Aminabadi, Samyam Rajbhandari+9 more|arXiv (Cornell University)|Jun 30, 2022

Advanced Neural Network Applications被引用 8

一句话总结

DeepSpeed Inference 引入了一个统一的系统，用于在各种硬件上高效、可扩展地推理大规模密集型和稀疏型（专家混合模型，MoE）Transformer 模型。通过结合多 GPU 优化、基于 ZeRO-Inference 的 CPU/NVMe 内存卸载，以及专用内核，其推理延迟相比最先进水平降低高达 7.3 倍，吞吐量提升 1.5 倍，使数百张 GPU 上实现万亿参数模型的实时推理成为可能。

ABSTRACT

The past several years have witnessed the success of transformer-based models, and their scale and application scenarios continue to grow aggressively. The current landscape of transformer models is increasingly diverse: the model size varies drastically with the largest being of hundred-billion parameters; the model characteristics differ due to the sparsity introduced by the Mixture-of-Experts; the target application scenarios can be latency-critical or throughput-oriented; the deployment hardware could be single- or multi-GPU systems with different types of memory and storage, etc. With such increasing diversity and the fast-evolving pace of transformer models, designing a highly performant and efficient inference system is extremely challenging. In this paper, we present DeepSpeed Inference, a comprehensive system solution for transformer model inference to address the above-mentioned challenges. DeepSpeed Inference consists of (1) a multi-GPU inference solution to minimize latency while maximizing the throughput of both dense and sparse transformer models when they fit in aggregate GPU memory, and (2) a heterogeneous inference solution that leverages CPU and NVMe memory in addition to the GPU memory and compute to enable high inference throughput with large models which do not fit in aggregate GPU memory. DeepSpeed Inference reduces latency by up to 7.3X over the state-of-the-art for latency-oriented scenarios and increases throughput by over 1.5x for throughput-oriented scenarios. Moreover, it enables trillion parameter scale inference under real-time latency constraints by leveraging hundreds of GPUs, an unprecedented scale for inference. It can inference 25x larger models than with GPU-only solutions, while delivering a high throughput of 84 TFLOPS (over $50\%$ of A6000 peak).

研究动机与目标

• 解决在异构硬件上部署大规模、多样化 Transformer 模型（从密集模型到 MoE 架构）所面临的日益严峻挑战，这些硬件具有不同的内存和延迟约束。
• 克服现有推理系统在小批量、低延迟关键工作负载下表现不佳，且无法在超出 GPU 内存限制后高效扩展的局限性。
• 通过利用 CPU 和 NVMe 存储作为内存层次结构的扩展，实现对超出 GPU 内存容量的模型的高吞吐量、低延迟推理。
• 通过多 GPU 和异构系统中的自适应调度、内核优化和内存管理，同时支持对延迟敏感和吞吐量导向的工作负载。

提出的方法

• 提出 ZeRO-Inference，一种参数卸载策略，将模型权重移至 CPU 或 NVMe 存储，同时将激活值保留在 GPU 内存中，从而实现超出 GPU 内存容量的模型推理。
• 采用 DeepFusion 技术，将多个算子（如线性层、层归一化、Softmax）联合定位并融合为单个内核，减少内核启动和数据移动开销，从而在小批量输入下提升延迟性能。
• 针对推理工作负载优化 GeMM 内核，以最大化内存带宽利用率，尤其在小批量推理中带宽是瓶颈时更为关键。
• 采用混合调度策略，结合流水线并行和张量并行，并引入动态负载均衡，以保持所有设备持续工作，这对具有 token 依赖关系的自回归生成尤为重要。
• 通过 all-to-all 通信优化和专家并行技术，高效地在数百张 GPU 上扩展稀疏 MoE 模型，实现高内存带宽可扩展性。
• 利用多张 GPU 之间的 PCIe 带宽聚合，实现接近线性的吞吐量扩展，即使在处理超大模型时也能实现高吞吐量推理。

实验结果

研究问题

• RQ1 推理系统如何高效支持不同规模和稀疏度的大型 Transformer 模型在各种硬件上的低延迟关键和吞吐量导向工作负载？
• RQ2 在小批量输入下，推理延迟受内存带宽限制时，需要哪些系统级优化才能实现高内存带宽利用率和低内核启动开销？
• RQ3 如何在不牺牲吞吐量或引入不可接受延迟的前提下，有效将模型参数卸载到 CPU 和 NVMe 存储？
• RQ4 在自回归生成中，每个 token 依赖于前一个 token 的情况下，需要哪些并行策略和调度技术来保持高设备利用率？
• RQ5 在使用异构内存（GPU、CPU、NVMe）和混合精度计算时，推理性能在数百张 GPU 上的可扩展性能达到何种程度？

主要发现

• DeepSpeed Inference 在低延迟关键场景下，相比最先进水平，推理延迟最高降低 7.3 倍，主要得益于内核融合和优化的内存访问。
• 通过 ZeRO-Inference 的 CPU/NVMe 卸载机制，即使模型可完全容纳于 GPU 内存，也能显著提升批量大小，使吞吐量提升超过 1.5 倍，适用于吞吐量导向工作负载。
• 该系统在使用数百张 GPU 时，支持万亿参数模型在实时延迟约束下的推理，实现 84 TFLOPS 的吞吐量（超过 A6000 峰值的 50%）。
• 对于可完全容纳于 CPU 内存的模型，ZeRO-Inference 的吞吐量比纯 CPU 推理高出 25 倍以上，证明了混合内存卸载的有效性。
• 在 8× A100 节点上，混合调度相比 FasterTransformer 将提示处理延迟降低 1.18 倍（PP+MP）和 3.06 倍（仅 MP），显著提升了自回归推理速度。
• DeepSpeed Inference 在稀疏 MoE 模型上实现了优越的内存带宽可扩展性，优于 PyTorch 基线，在高达 128 张 GPU 时仍能保持高带宽，得益于优化的 all-to-all 和 MoE 内核。