通过操作融合提升LLM推理速度,隐藏归一化计算延迟。
LLM Inference Acceleration via Efficient Operation Fusion
- 将Softmax和层归一化延后至线性层之后计算,实现并行处理。
- 在矩阵乘法期间并发计算归一化缩放因子,降低20%延迟。
- 适合追求低延迟推理的部署场景,尤其适用于分布式系统。
近年来基于Transformer的大语言模型(LLMs)迅速发展,其参数规模已达数百亿,需专用硬件支持训练与推理。Transformer架构中存在大量非线性变换,如每个解码器块通常包含至少一个Softmax和两个Layernorm。这些操作的归一化缩放因子计算依赖空间聚合操作,需将向量所有元素汇聚到单一位置,导致显著通信开销,使推理速度下降约20%,抵消了分布式内存计算的优势。本文提出一种高效操作融合技术,可完全隐藏此类聚合操作的开销。由于Softmax和Layernorm后接线性层,且二者运行于不同硬件引擎,可通过代数性质实现运算顺序交换:将归一化推迟至线性层完成之后。此时,聚合缩放因子的计算可与矩阵乘法并行执行,从而将前者延迟完全掩盖在后者之中。该方法在保持数值精度的同时,显著提升硬件利用率,降低整体推理延迟。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。