1.背景

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

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

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

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

▲ 圖 1：《Efficiently Scaling Transformer Inference》

2.推理成本權(quán)衡

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

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

內(nèi)存成本：在prefill和decoder兩個階段都會存在，主要體現(xiàn)在將數(shù)據(jù)從HBM（high-bandwidth memory）傳輸?shù)接嬎愫诵牡某杀旧稀?/p>

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

3.分布式推理

1）符號表示

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

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

（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)：相當于轉(zhuǎn)置，可以實現(xiàn)  ->  這樣的操作；

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

2）FFN的并行

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

其中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維度。在處理器數(shù)不能整除注意力頭數(shù)的情況下，特殊處理是必要的。對于處理器數(shù)量大于注意力頭數(shù)的情況，注意力頭的部分會進行部分復(fù)制（partially replicated），與FFN類似，這等效于一次AllReduce=ReduceScatter+Allgather。

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

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

4.其他優(yōu)化策略

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

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

4.總結(jié)

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

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

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

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

綜合而言，稠密LLM模型的高效并行推理具備一定的復(fù)雜性，不同階段、不同場景下的并行策略有不同的指導(dǎo)原則。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.

本文轉(zhuǎn)載自 ??AI遇見云??，作者： 楊希