聽名字就知道，這個(gè)會(huì)代表了美國(guó) MLSys 的最核心圈子，會(huì)議的發(fā)起人 Jeff Dean、Fei-Fei Li、Eric Xing、Alex Smola 都來自的 Google、Meta 這些硅谷科技巨頭和 UCB、Stanford、CMU 這些老牌系統(tǒng)強(qiáng)校的巨佬。鑒于現(xiàn)在大模型技術(shù)還是美國(guó)主導(dǎo)的，MLSys 一定程度上也代表著先進(jìn)生產(chǎn)力的前進(jìn)方向和最新的思潮發(fā)展方向。

今天我們一起來學(xué)習(xí)一下 MLSys 中關(guān)于大模型推理的重量級(jí)文章。鏈接如下，它也是三篇 Outstanding Paper Award 之一，Jeff Dean 尾作署名。

論文鏈接：

去年 Dean 署名的文章是 Pathways，也是 22 年的 Outstanding Paper之一。我感覺 Dean 署名論文投稿 MLSys 就類似《求是》雜志發(fā)表咱們大領(lǐng)導(dǎo)的文章，起到給全美的 MLSys 工作定調(diào)子的作用。言外之意類似：2022 年一整年，經(jīng)過全美各企業(yè)和科研院校的辛勤努力，大模型訓(xùn)練系統(tǒng)研發(fā)已取得重要成果。如今，備受矚目的焦點(diǎn)已轉(zhuǎn)向推理優(yōu)化。在新的一年里，希望全美各界再接再厲，踔厲奮發(fā)，進(jìn)一步提升推理效率，努力譜寫通用人工智能技術(shù)的發(fā)展新篇章！

推理優(yōu)化不學(xué)此雄文，十倍加速也枉然?？墒钱?dāng)你通讀此文后，也許會(huì)些許失望。文章不就是講怎么做模型并行么，說的我好像早就知道，沒什么花哨的想法，也沒什么神奇的結(jié)果。話說，金庸小說《天龍八部》有一個(gè)片段，天下英雄齊聚聚賢莊和喬峰比武，喬峰打出一套太祖長(zhǎng)拳，就折服武林眾高手。太祖長(zhǎng)拳別成“大宋廣播體操”，在宋朝人人皆會(huì)，太祖拳法的精要，可說是無人不知?？墒菃谭逄骈L(zhǎng)拳的一招打出，各路高手都還是情不自禁的喝了一聲采！同理，這篇文章的效果就好比喬幫主的太祖長(zhǎng)拳。

這篇文章好就好在它的樸素直接。常言道，有錢人的生活就是這么樸實(shí)無華且無趣，當(dāng)其他文章爭(zhēng)奇斗艷，拿錘找釘，欲蓋彌彰地討好審稿人博出位的時(shí)候，一篇文章只告訴你最樸素的道理，只給你最直接結(jié)論，如此存粹的文風(fēng)就像一縷清風(fēng)，一股清泉讓此時(shí)此刻的你在這爾虞我詐的花花世界中得到了片刻的安寧。這也是這篇文章美學(xué)藝術(shù)之所在。

下面咱們進(jìn)入正式解讀部分。

背景

對(duì)于 100B+ 規(guī)模的稠密大模型（LLM），即非 MoE 結(jié)構(gòu)的 Transformer 模型，其推理實(shí)現(xiàn)是極其具有挑戰(zhàn)的。比如，在線服務(wù) 540B 的 GaLM 的成本是極其高，如果不能把它的成本控制在一定范圍內(nèi)，即使 Pathways 訓(xùn)練出模型也沒法普惠到客戶，這是 Google 內(nèi)部遇到的具體問題。

LLM 推理應(yīng)用的需求有可分兩類：1. 在線推理，比如聊天機(jī)器人交互式的，需要非常嚴(yán)格的延遲約束，為了這個(gè)目標(biāo)放松成本限制。2. 離線推理，包括用于 Rewards Model 或蒸餾的離線推理任務(wù)。它強(qiáng)調(diào)高吞吐量和低成本，延遲反而不重要。

大模型推理可以分解為處理 prompt 的的階段（prefill）和自回歸生成輸出標(biāo)記解階段（Decoding）。二者性質(zhì)不同，前者是計(jì)算密集型的，后者是訪存密集型的。尤其是，Decoding 階段以難以優(yōu)化而臭名昭著。因?yàn)槊看沃荒芙獯a一個(gè) token，計(jì)算量本來就很小，還需要訪問一個(gè)隨著解碼步數(shù)增加而不斷累計(jì) KVCache，因此嚴(yán)重受限于加速卡的片上帶寬。

由于參數(shù)巨大，單個(gè)加速卡內(nèi)存肯定存不下 100B 級(jí)別的模型，因此并行推理成為必選項(xiàng)。尤其是對(duì)于第一類任務(wù)，希望滿足推理延遲，必須增加加速卡數(shù)目，來提升整體算力和帶寬。

推理工程優(yōu)化是一個(gè)多層次、模塊化的系統(tǒng)工程。需要算子優(yōu)化、量化加速、并行策略、內(nèi)存管理、Batching 策略、Decoding 方法、甚至結(jié)合模型改動(dòng)一起，鏈路非常長(zhǎng)。而本文關(guān)注并行策略部分。

它提供一套 engineering principles，來指導(dǎo) Transformer 結(jié)構(gòu)推理并行策略的優(yōu)化，從而幫助我們找到在模型大小和特定應(yīng)用需求下經(jīng)驗(yàn)性地推到出最佳模型的并行策略。簡(jiǎn)單來說，將大家腦子中對(duì)并行這塊多多少少有的模糊概念落筆成文，進(jìn)行了系統(tǒng)性地闡述，并給出了 Google 一手的實(shí)驗(yàn)數(shù)據(jù)。

方法

在 LLM 推理應(yīng)用時(shí)，我們需要考慮三個(gè)指標(biāo)，延遲、吞吐和成本。延遲就是端到端完成一次推理時(shí)間，包括 prefill+decoding兩個(gè)階段的延遲總和。吞吐就是單位時(shí)間可以處理多少 tokens。成本就是用多少機(jī)時(shí)，大家常說的 MFU（Model Flops Utilization）指標(biāo)和成本掛鉤。這些指標(biāo)往往互相矛盾，比如要達(dá)到低延遲，可能導(dǎo)致不能打出大 Batch 從而造成低吞吐，也可能需要加芯片導(dǎo)致高成本。

面對(duì)這些實(shí)際問題，在不同的 Batch Size，模型尺寸，芯片數(shù)目情況下，最優(yōu)的模型切分策略應(yīng)該是不同的，本文會(huì)給一些經(jīng)驗(yàn)性的結(jié)論，比如下面這幅圖 Figure 1。圖中每條線上的點(diǎn)表示成本與延遲二元目標(biāo)優(yōu)化的 Pareto frontier，不同顏色表示不同大小的模型。因?yàn)?Decoding 和 Prefilling 的計(jì)算性質(zhì)不同，需要分開觀察。左圖：Decoding 生成 64 個(gè) token 平均延遲情況。右圖：處理 2048 個(gè)輸入 token Prefilling 的總延遲。Chip count 是 C, batch size 是 B。

這幅圖就像手冊(cè)一樣，比如我們可以根據(jù)延遲需求，查找最少成本的 B、C 的值和模型類別，從而配置相應(yīng)地去配置我們的推理服務(wù)。一些應(yīng)用實(shí)例：對(duì)于 540B 參數(shù)的 LLM，在 64 個(gè) TPU v4 芯片上運(yùn)行時(shí)，我們?cè)谏蛇^程中實(shí)現(xiàn)了每個(gè)標(biāo)記 29 毫秒的low-batch-size延遲，并在 large-batch-size 處理輸入標(biāo)記時(shí)實(shí)現(xiàn)了 76% 的 MFU。實(shí)現(xiàn)了離線在線應(yīng)用的兩開花的目的。

而想做出上面這兩幅圖，還需要一樣法寶，如何對(duì)給定 B、C 和模型類別，選擇最優(yōu)的并行策略實(shí)現(xiàn)。這也是本文大篇幅討論的內(nèi)容。本文提出一個(gè) Partitioning Framework 解決此事。

1. 符號(hào)化抽象

Partitioning 就是 Tensor Parallelism。首先，對(duì)多個(gè) TPU 組成的硬件集群，可以抽象成一個(gè) N-D 的 device mesh，和 UCB alpa 論文 [1] 里的 device mesh一 樣。稍微特殊一點(diǎn)在于，在 TPU v4 集群里，由于網(wǎng)絡(luò)硬件是 3D torus topology ，非常適合抽象成 X×Y×Z 形式的 3D mesh。

本文最出彩的地方表示并行策略的抽象符號(hào)體系。符號(hào)化抽象對(duì)一套理論的普及至關(guān)重要，一個(gè)例子是牛頓和萊布尼茨同時(shí)發(fā)明了微積分，他們獨(dú)立地發(fā)展了微積分的基本原理和符號(hào)體系。牛頓使用的是差分法和流數(shù)法，而萊布尼茨使用的是微分法和積分法。萊布尼茨提出的微分法和積分法更加簡(jiǎn)潔、直觀且易于使用，所以成為現(xiàn)在的主流。

怎么表示模型并行呢？本質(zhì)是表示 tensor layout 和 device mesh 的映射，目前有如下幾種主流符號(hào)體系：

（1）Megatron-LM 式 [2]，把矩陣切分寫成行列式方式，[Y1, Y2] = [GeLU(XA1), GeLU(XA2)] 表達(dá) row-wise split。表達(dá) 1D device mesh 很適合，但是對(duì)更復(fù)雜的 device mesh 則捉襟見肘。不過 GPU 集群通常 NVLink+IB，3D device mesh屬于長(zhǎng)尾需求。

（2）oneflow 式 [3]，使用 SBP signature，用字母（數(shù)字）形式表達(dá)并行，字母表示 Tensor layout 方式，括號(hào)內(nèi)數(shù)字表達(dá) device mesh 的維度。比如 column-wise TP 矩陣乘表示：(S(0), B) (B, S(1))->(S(0), S(1))。這套體系對(duì) N-D tensor->N-D device 映射還不是很直觀。

（3）GSPMD 式 [4]，使用一個(gè)數(shù)組 dims_mapping 表示 tensor->device mesh 映射，devich mesh 不同維度在 dims_mapping 中應(yīng)該最多出現(xiàn)一次，數(shù)組中 -1 表示不存在的映射，它可以表示 tiled, partially tiled, and replicated sharding。盡管這套符號(hào)體系支持 N-D mesh 了，但是它真的非常非常抽象，戰(zhàn)鷹可以做它的代言人。曾經(jīng)有一位做 DL 框架 N 年資深巨佬看了半個(gè)小時(shí)也沒看明白原始論文圖 1 的意思。

在 GSPMD 基礎(chǔ)之上，本文改進(jìn)了一下它的符號(hào)體系，兼顧了直觀性和表達(dá)能力。通過以下規(guī)定表達(dá)一個(gè)張量的分布式的 layout：

（1）Shard 表示：使用下標(biāo)來指定被 shard 的張量維度。例如，符號(hào) 表示邏輯形狀為 BLE 的張量的最后一個(gè)維度 E 被分割成  X × Y × Z 個(gè)分區(qū)，其中 x、y 和 z 表示物理 TPU v4 軸，每個(gè)芯片的張量形狀為 [B, L, E/(X × Y × Z)]。

（2）Replicated 表示：如果一個(gè)張量在軸x上被replicated存儲(chǔ)，那么該軸將在符號(hào)中被省略。

（3）Partial 表示：后綴 “partialsum-x” 來表示給定的張量已經(jīng)在每個(gè)芯片上 locally 地進(jìn)行了處理（concate 或 sum），但仍需要在 TPU x 軸上對(duì)芯片進(jìn)行求和才是最終結(jié)果。

因?yàn)?reduce-scatter(x)，all-gather(x) 這些操作有對(duì)應(yīng)的通信量公式，可以方便計(jì)算出并行推理的通信開銷。有了這套符號(hào)體系，和 MPI 的運(yùn)算規(guī)則，咱們可以看懂論文中的圖了。

2. FFN并行方法

比如 FFN 的兩個(gè) Linear 層有三種 Partition 方法，如下圖所示，它們的名字分是 1D weight-stationary，2D weight-stationary 和 weight-gather。

1D weight 就是 Megatron-LM 中的 FFN 并行方法，用一個(gè) column-wise+row-wise TP，forward 只需要 Allgather+ReduceScatter 兩次通信。

后兩者大大家不常見，它對(duì) weight 的兩個(gè)維度都進(jìn)行 partition，對(duì) 2D 對(duì) activation 最低維度切分，weight-gather 則對(duì)最低和最高維度分別切分。值得注意的是，作者對(duì) L 維度沒有切分。

3. Attention Layer并行方法

作者設(shè)計(jì)了針對(duì)經(jīng)典的 MHA（multihead attention）和改進(jìn)的 MQA（multiquery attention）的并行方案，后者多頭共享參數(shù)，可以減少 KVCache 讀取帶寬需求，但是給并行策略設(shè)計(jì)造成一些麻煩，因?yàn)樯倭?head 維度切分的選擇。如 Figure 4 所示，MQA KV Cache 沒有 head 維度了。

MHA 的切分和 FFN Partion 類似，只是把 F 維度換成 H，參與切分。如 Figure 4（a）所示，可以只切 H 維度，和 Megatron-LM 類似。需要特殊處理的 H 不能整除處理器數(shù)的情況，對(duì)于處理器數(shù)量大于注意力頭數(shù)的情況，注意力頭部分 partially replicated。如下圖 Figure 5（a）是 MHA 的并行方案細(xì)節(jié)。和 FFN 同行類似，都是等價(jià)于一次 AllReduce=ReduceScatter+Allgather。

MQA 的切分比較棘手，如果還切 H 維度，就是如 Figure 4（b）所示方案，K 和 V 張量在所有頭部之間共享，但它們必須在每個(gè)芯片上市 replicated 的，這導(dǎo)致 MQA 的內(nèi)存成本節(jié)省將會(huì)喪失。所以作者采用 Figure 4（c）切分，對(duì) B 維度切分，不過這也要求任務(wù)能組比較大的 batch，有切分余地。

具體來說，Q、K 和 V 矩陣在批次 B 維度上分區(qū)為 N 個(gè)分區(qū)。Figure 4（c）顯示，這將使每個(gè)芯片加載 KV 緩存的內(nèi)存成本降低 N 倍，從而將訪存也降低相同倍數(shù)。具體并行細(xì)節(jié)參考 Figure 5（b），與Figure 5（a）中 MHA 并行策略相比，MQA 需要使用 all-to-all 對(duì)進(jìn)行輸入輸出激活張量 resharding，從而產(chǎn)生額外的通信成本。可見，MQA 減少了訪存，但是增加了并行通信開銷。

對(duì) Figure 2 和 Figue 5 的五個(gè) workflow，作者在原文中都給出了通信開銷。讓讀者可以直觀地感受到通信開銷和芯片數(shù)量、B、L、E 三個(gè)維度的變化關(guān)系。建議讀者使用時(shí)去查閱論文。

其他優(yōu)化

出了并行策略，文章還有介紹了其他對(duì)超大模型推理有裨益的優(yōu)化。

首先，模型結(jié)構(gòu)微調(diào)。和 GPT 結(jié)構(gòu)不同，PaLM 中 FFN 和 Attention Layer 是并行的。區(qū)別如下：

GPT: y = MLP(LayerNorm(x + Attn(LayerNorm(x))

PaLM: y = x + MLP(LayerNorm(x)) + Attn(LayerNorm(x)) 

PaLM 的方式有幾個(gè)對(duì)工程優(yōu)化的好處。第一，省掉一次 LayerNorm。第二， 有很多矩陣乘法融合增加 FLOPS 的機(jī)會(huì)了。FFN 的輸入矩陣可以與 Attention 的 融合，Attn 的 KV 矩陣 和 可以融合，F(xiàn)FN 的輸出矩陣可以與 Attn 的輸出投影矩陣 融合。第三，也是對(duì)于并行最重要的，它還消除了每個(gè) Transformer 層中用于 dff/nheads 并行性的兩個(gè)全局歸約操作中的一個(gè)，將沿著這個(gè)軸的通信時(shí)間減少了一半。

然后，還有一些 Low-level 優(yōu)化。

比較值得關(guān)注的是 Looped CollectiveEinsum 技術(shù)，使通信與計(jì)算同時(shí)進(jìn)行。這個(gè)技術(shù)發(fā)表在 ASPLOPS 23 上 [6]，值得 highlight 一下，簡(jiǎn)單來說就是把矩陣乘法和通信分塊流式起來，隱藏通信開銷。

Google 作者聲稱這項(xiàng)技術(shù)使我們能夠部分或完全隱藏 Figure 2和 Figure 5 中大部分 reduce-scatter 和 all-gather 操作的通信時(shí)間。因?yàn)橛?Looped CollectiveEinsum 技術(shù)，所以作者對(duì)于 Figure 2 和 Figure 5 中的所有 reduce-scatter 操作，可以選擇將其 reduce-scatter 到批次或序列維度（B 或 L）或隱藏維度（E 或 F），作者都選擇了后者，因?yàn)樗鼮檠h(huán) CollectiveEinsum 提供了更多有效的機(jī)會(huì)。

最后，還有內(nèi)存和量化優(yōu)化，這里按下不表了。

實(shí)驗(yàn)和結(jié)論

1. FFN如何并行結(jié)論：

在 Prefilling 階段，我們根據(jù) batch 中的 token 數(shù)量從 weight-stationary and weight-gathered layouts 中進(jìn)行選擇并行策略。在 Decoding 階段，我們選擇 2D weight-stationary，因?yàn)?token 的 batch 大小始終很小。

2. Attn如何并行結(jié)論：

在 Prefilling 階段，MQA 和 MHA 產(chǎn)生類似的推理延遲，因?yàn)槲覀儾⑿杏?jì)算許多 attention queries，并且在注意力矩陣乘法上的計(jì)算是受限于計(jì)算能力的。

在 Decoding 階段，MQA 只有一種并行方式了，所以也沒得選了。優(yōu)化的 MQA 布局 Figure 4（a）比 Figure 4（b）提升效率。當(dāng)上下文長(zhǎng)度較短時(shí)，速度提升較小，因?yàn)閹缀跛袝r(shí)間都花在了 FFN 上。隨著上下文長(zhǎng)度變得更長(zhǎng)，加載注意力層中的 KV 緩存所需的時(shí)間成為總體推理時(shí)間的很大一部分。MQA 可擴(kuò)展到 8192-32768 個(gè) token 的序列長(zhǎng)度（分別為批次大小 512 和 128），attention 僅占總運(yùn)行時(shí)間的 8-31%。如下圖所示，MAQ 比 MHA 在長(zhǎng)序列時(shí)候提升很多。

3. 模型結(jié)構(gòu)微調(diào)效果：

在 540B 規(guī)模下，Decoding 階段，串行公式每步的推理延遲比并行版本高 14%，這是主要因?yàn)榧せ畹耐ㄐ艜r(shí)間增加了。在 Prefilling 階段，這種差異會(huì)縮小，因?yàn)闄?quán)重聚集的布局會(huì)減少激活通信。

4. 端到端分析

回看 Figure 1，對(duì)于具體應(yīng)用，本質(zhì)上是在求解 Latency 和 Cost 作為目標(biāo)的多目標(biāo)優(yōu)化問題。圖中每一個(gè)點(diǎn)都是選擇最佳的 FFN 和 Attn 并行策略的結(jié)構(gòu)，我們抽出兩個(gè)點(diǎn)具體看看，如 Table 2,3 所示。對(duì)于低延遲、高吞吐兩種需求，Prefill 和 Decoding 兩個(gè)階段，最佳的并行策略都不同。

5. 和 FasterTransformer 對(duì)比

TPU+ 本文方法和 GPU+FasterTransformers，兩套軟硬件組合進(jìn)行了對(duì)比。如 Figure 9 所示，大部分情況下，紅色系線都優(yōu)于藍(lán)色系的線。

結(jié)論

本文系統(tǒng)性地闡述了稠密 LLM 模型并行推理的目標(biāo)、方法和規(guī)律。在方法部分，比較亮眼的是作者一套簡(jiǎn)單、有力的符號(hào)化抽象，描述了 FFN 和 Attn 不同切分策略，并推導(dǎo)出通信的開銷。作者給出了一些規(guī)律性的結(jié)論：1）Prefill 和 Decoding 的并行策略是不同的；2）Low Latency 和 High Throught場(chǎng)景下并行策略是不同的；3）MQA 和并行 Attn 和 FFN 的模型改進(jìn)都是大有裨益的。文章里的圖畫的很精美，信息量很大，細(xì)節(jié)拉滿，可謂一圖勝千言之典范。

作者的方法和 Google 內(nèi)部基建過于耦合，比如 3D Torus 網(wǎng)絡(luò)更適合 3-D Device Mesh，一些結(jié)論在 GPU 集群上還有改變。另外，本文的工作還需要非常強(qiáng)大的 Batching 實(shí)現(xiàn)配合才能在真實(shí)業(yè)務(wù)中發(fā)揮作用。

#投 稿 通 道#

 讓你的文字被更多人看到 

如何才能讓更多的優(yōu)質(zhì)內(nèi)容以更短路徑到達(dá)讀者群體，縮短讀者尋找優(yōu)質(zhì)內(nèi)容的成本呢？答案就是：你不認(rèn)識(shí)的人。

總有一些你不認(rèn)識(shí)的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學(xué)者和學(xué)術(shù)靈感相互碰撞，迸發(fā)出更多的可能性。 

PaperWeekly 鼓勵(lì)高校實(shí)驗(yàn)室或個(gè)人，在我們的平臺(tái)上分享各類優(yōu)質(zhì)內(nèi)容，可以是最新論文解讀，也可以是學(xué)術(shù)熱點(diǎn)剖析、科研心得或競(jìng)賽經(jīng)驗(yàn)講解等。我們的目的只有一個(gè)，讓知識(shí)真正流動(dòng)起來。

?? 稿件基本要求：

· 文章確系個(gè)人原創(chuàng)作品，未曾在公開渠道發(fā)表，如為其他平臺(tái)已發(fā)表或待發(fā)表的文章，請(qǐng)明確標(biāo)注 

· 稿件建議以 markdown 格式撰寫，文中配圖以附件形式發(fā)送，要求圖片清晰，無版權(quán)問題

· PaperWeekly 尊重原作者署名權(quán)，并將為每篇被采納的原創(chuàng)首發(fā)稿件，提供業(yè)內(nèi)具有競(jìng)爭(zhēng)力稿酬，具體依據(jù)文章閱讀量和文章質(zhì)量階梯制結(jié)算

?? 投稿通道：