1.背景

這次指的大型語言模型（LLMs）主要說的是采用Transformers架構的模型，該架構在訓練階段具有高度的并行性。然而，在實際應用中，高效地部署這些模型面臨一些挑戰。這是因為生成推理是逐個token進行的，而每個token的計算都依賴于先前生成的標記。因此，支持大規模訓練的模型需要仔細考慮并行布局和內存優化，以實現所需的可擴展性和高效的低延遲推理。從而更好地支持大規模訓練的模型在實際應用中實現高效的低延遲推理，從而滿足對于實時性和可擴展性的需求。

然而，在實際的工程化環境中，模型性能優化從未是一個簡單的最優化問題。在大型語言模型（LLMs）的生成推理中，我們面臨著諸多挑戰，包括龐大的內存占用、嚴格的延遲目標以及長序列長度，不同指標質檢的優化方法還存在一定的沖突。這時就需要引入了帕累托前沿的概念，用以描述一組解決方案，即在不犧牲其他目標的前提下，無法改善一個特定目標。

• 大模型訓練/推理的主要特點：large deep models（模型大）, with tight latency targets（時間緊） 以及 long sequence lengths（理解長度長）；
• Pareto frontier：核心思想是，在多目標決策問題中，存在一組解決方案，這些解決方案在某些目標上的改善是不可能的，而在其他目標上的改善是可能的。這意味著你不能在一個目標上取得改進而不犧牲其他目標。這些解決方案構成了 Pareto frontier，也稱為 Pareto 邊界。在 Pareto frontier 上的解決方案通常被認為是非劣解（non-dominated solutions）或 Pareto 最優解（Pareto optimal solutions）。舉例來說，考慮一個制造公司需要在成本和質量兩個目標上做決策。如果提高產品質量會增加成本，那么在成本和質量之間存在一個權衡。Pareto frontier 就是描述了在不犧牲一方面目標的情況下，如何最大化或最小化另一目標。

參考《Efficiently Scaling Transformer Inference》提供的一套工程原則，可以引導Transformer結構的推理并行策略優化，以幫助我們在考慮模型規模和特定應用需求的情況下，經驗性地推導出最佳的模型并行策略。對于不同的推理需求，可以通過實驗列表進行按需調整，選擇合適的核心數和批處理大小，如下圖，可以根據在不同場景中的延遲要求，選擇最優的核心數（加速卡數量）以及批處理大小：

▲ 圖 1：《Efficiently Scaling Transformer Inference》

2.推理成本權衡

一次端到端的推理延遲需要考慮到兩個主要方面：數據處理（prefill，大部分在CPU上進行處理）和逐詞生成（decode，大部分在GPU上進行）。

衡量指標：為了衡量吞吐量，我們可以使用MFU（The Model FLOPS Utilization）作為指標，即實際吞吐量與理論最大吞吐量的比值（也可以用于衡量計算效率）。

內存成本：在prefill和decoder兩個階段都會存在，主要體現在將數據從HBM（high-bandwidth memory）傳輸到計算核心的成本上。

計算成本：一般涉及到對于一個包含N個參數的decoder-only模型來說，生成一個token需要進行2N次矩陣乘法運算。計算成本即指這些乘法運算的代價。

3.分布式推理

1）符號表示

（1）  ：表示最后一個維度E已經被切分為X * Y * Z個分區，其中 BLE分別表示：Batch、Sequence Length、Model Embed；此外F表示MLP前向傳播中的維度。

（2）后綴“partialsum-x”：表示一個指定的Tensor在每個Chip上被本地合并（contracted/summed）了，但仍需要在TPU x軸上對芯片進行求和才是最終結果。

（3）all-reduce(x)：表示在x這個維度上對  這樣的切片進行合并，最終輸出  ，all-reduce = reduce-scatter + all-gather；

（4）reduce-scatter(x)：對x維度上芯片的張量進行求和，并再該維度上對求和后的張量進行重新分片，相當于消除掉partialsum-x并在任意維度增加一個下標x；

（5）all-gather(x)：可以直接理解為concat，可以消除掉某一個維度的下標；

（6）all-to-all(x)：相當于轉置，可以實現  ->  這樣的操作；

▲ 圖 2：常見的分布式通訊方法 

2）FFN的并行

隨著在更多芯片上并行計算，內存延遲和計算延遲通常會呈近線性下降。然而，通信延遲依然會保持大致不變，因為對于每一對矩陣乘法運算來說，整個激活矩陣是跨芯片聚合的。因此通訊延遲會成為超大規模分布式訓練的瓶頸；

其中FFN的并行方法大致可以分為三類：

▲ 圖 3：FFN分布式并行方法

▲ 圖 4：不同batch下不同并行策略的消耗

3.注意力層的并行

目前注意力的并行機制主要包括了MHA以及MQA（https://arxiv.org/pdf/1911.02150.pdf，19年Google提出的一種新的Attention機制，可以加快decoder的生成）的兩種Attention Layer的并行，其主要并行方式如下：

▲ 圖 5：Attention 并行方式

▲ 圖 6：MHA&MQA并行表示

MHA的切分和FFN Partion相似，不同之處在于將F維度替換為H維度進行切分。如圖6(a)所示，可以選擇僅切分H維度。在處理器數不能整除注意力頭數的情況下，特殊處理是必要的。對于處理器數量大于注意力頭數的情況，注意力頭的部分會進行部分復制（partially replicated），與FFN類似，這等效于一次AllReduce=ReduceScatter+Allgather。

MQA的切分相對較為復雜。如果仍然選擇在H維度上進行切分，如圖6(b)方案所示，K和V張量在所有注意力頭之間是共享的，但它們必須在每個芯片上進行復制，這會導致MQA的內存成本節省失效。因此，作者采用圖5(c)切分，即對B維度進行切分。然而，這也要求任務能夠組織成較大的batch，以便有足夠的切分余地。具體而言，Q、K和V矩陣在批次B維度上被分成N個分區，這降低了每個芯片加載KV緩存的內存成本，也因此減少了相應的訪存。與MHA的并行策略相比，MQA需要使用all-to-all對進行輸入輸出激活張量resharding，從而帶來額外的通信成本。值得注意的是，使用較大的Batch Size時，MQA可以減少內存訪問成本，但同時會增加并行通信開銷。

總的來說，MQA允許使用更大的批處理大小和序列長度，從而在減少內存時間的同時降低延遲，提高吞吐量（訓練更快）。

4.其他優化策略

1） 使用PaLM中并行的FFN以及Attention層：可以將兩層layerNorm優化成一層，并且對于FFN層的輸入矩陣、KV的投影矩陣、輸出矩陣都可以實現復用，從而提高模型整體的吞吐率；

2）Low-level優化：主要指的是使用了Looped CollectiveEinsum（https://dl.acm.org/doi/abs/10.1145/3567955.3567959）的技術，可以降低通訊成本；

4.總結

本文系統性地闡述了稠密LLM模型并行推理的目標、方法和規律。針對稠密LLM模型（主要指的是Transformer類的模型）并行推理的目標、方法和規律該論文進行詳細的總結。

首先，在Prefill和Decoding階段的并行策略，在模型推理過程中具有不同的特性和需求。Prefill主要發生在CPU上，而Decoding則主要在GPU上進行處理。因此，針對不同的計算資源和優化需求，對這兩個階段的并行化需要采用不同的策略。

其次，強調了在Low Latency（低延遲）和High Throughput（高吞吐量）的場景下，并行策略也呈現出不同的特征。在追求低延遲的情境下，需要特別關注推理速度和實時性。而在高吞吐量的情境下，更側重于系統整體性能的提升。因此，針對不同的應用場景，需要靈活選擇并實施不同的并行化策略以滿足特定需求。

最后，對于MQA（Mixed Quantization Attention）在并行Attention和Feed-Forward Network（FFN）方面的優勢進行了重點強調。MQA的采用對于推理性能提升具有顯著裨益，特別是在處理注意力機制和前饋網絡時。這進一步強調了在不同任務和模型結構下，采用差異化的并行化策略對于性能優化至關重要。

綜合而言，稠密LLM模型的高效并行推理具備一定的復雜性，不同階段、不同場景下的并行策略有不同的指導原則。MQA的引入更進一步豐富了并行計算的工具箱，為提升模型推理效率提供了有力支持。

參考文獻

[1] Pope R, Douglas S, Chowdhery A, et al. Efficiently scaling transformer inference[J]. Proceedings of Machine Learning and Systems, 2023, 5.

[2]? https://zhuanlan.zhihu.com/p/634236135

[3]? https://www.zhihu.com/question/591646269/answer/3246726206

[4] Wang S, Wei J, Sabne A, et al. Overlap communication with dependent computation via decomposition in large deep learning models[C]//Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1. 2022: 93-106.

[5] Shazeer N. Fast transformer decoding: One write-head is all you need[J]. arXiv preprint arXiv:1911.02150, 2019.

本文轉載自 ??AI遇見云??，作者： 楊希