通过高效操作融合加速LLM推理
LLM Inference Acceleration via Efficient Operation Fusion
AI总结:
本文提出一种高效操作融合技术,通过将Softmax和层归一化的归一化延迟到线性层之后,与矩阵乘法并行计算集合缩放因子,从而完全隐藏集合操作开销,加速LLM推理,提升硬件利用率并降低延迟。
摘要:
近年来,基于Transformer的大型语言模型(LLMs)的快速发展与其不断增长且已经庞大的规模密切相关。许多LLM包含数千亿参数,并需要专用的硬件资源进行训练和推理。Transformer架构固有的关键挑战之一是支持涉及归一化的众多非线性变换。例如,每个解码器块通常包含至少一个Softmax操作和两个层归一化(Layernorm)操作。相应归一化缩放因子的计算成为主要瓶颈,因为它需要空间集合操作。换句话说,在计算Softmax和层归一化的分母时,所有向量元素必须聚合到单个位置,这需要大量通信。这些集合操作使Transformer上的推理速度降低约20%,破坏了分布式内存计算的全部目的。在这项工作中,我们提出了一种极其高效的技术,可以完全隐藏这类集合操作带来的开销。注意,每个Softmax和层归一化操作通常后跟一个线性层。由于非线性和线性操作在不同的硬件引擎上执行,一旦代数允许这种交换,它们就可以轻松并行化。利用线性操作的内在属性,我们可以将前面的Softmax和层归一化的归一化延迟到线性层计算之后。现在,我们可以与矩阵乘法同时计算集合缩放因子,并将前者的延迟完全隐藏在后者之后。这种并行化保持了数值精度,同时显著提高了硬件利用率并降低了整体延迟。
英文摘要:
The rapid development of the Transformer-based Large Language Models (LLMs) in recent years has been closely linked to their ever-growing and already enormous sizes. Many LLMs contain hundreds of billions of parameters and require dedicated hardware resources for training and inference. One of the key challenges inherent to the Transformer architecture is the requirement to support numerous non-linear transformations that involves normalization. For instance, each decoder block typically contains at least one Softmax operation and two Layernorms. The computation of the corresponding normalization scaling factors becomes a major bottleneck as it requires spatial collective operations. In other words, when it comes to the computation of denominators for Softmax and Layernorm, all vector elements must be aggregated into a single location, requiring significant communication. These collective operations slow down inference on Transformers by approximately 20%, defeating the whole purpose of distributed in-memory compute. In this work, we propose an extremely efficient technique that can completely hide the overhead caused by such collective operations. Note that each Softmax and Layernorm operation is typically followed by a linear layer. Since non-linear and linear operations are performed on different hardware engines, they can be easily parallelized once the algebra allows such commutation. By leveraging the inherent properties of linear operations, we can defer the normalization of the preceding Softmax and Layernorm until after the linear layer is computed. Now we can compute the collective scaling factors concurrently with the matrix multiplication and completely hide the latency of the former behind the latter. Such parallelization preserves the numerical accuracy while significantly improving the hardware utilization and reducing the overall latency.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。