作者｜kipply

翻譯｜楊婷、徐佳渝、賈川??

本文詳細(xì)闡述了大型語(yǔ)言模型推理性能的幾個(gè)基本原理，不含任何實(shí)驗(yàn)數(shù)據(jù)或復(fù)雜的數(shù)學(xué)公式，旨在加深讀者對(duì)相關(guān)原理的理解。此外，作者還提出了一種極其簡(jiǎn)單的推理時(shí)延模型，該模型與實(shí)證結(jié)果擬合度高，可更好地預(yù)測(cè)和解釋Transformer模型的推理過(guò)程。

為了更好地閱讀本文，讀者需了解一些Transformer模型的相關(guān)先驗(yàn)知識(shí)，比如《圖解Transformer》的大部分內(nèi)容。另外，了解與本文相關(guān)的參數(shù)計(jì)數(shù)文章也能更好地幫助讀者理解本文內(nèi)容。本文主要包括以下內(nèi)容：

• kv 緩存 (kv cache) 解釋了在推理過(guò)程中緩存自注意力向量所帶來(lái)的性能優(yōu)化效果，以及可能導(dǎo)致的權(quán)衡（tradeoff）以及容量成本問(wèn)題。

• 容量（capacity）考慮了kv緩存的存儲(chǔ)成本以及模型權(quán)重的存儲(chǔ)成本之間的聯(lián)系，并解釋了容量大小對(duì)模型性能的影響。

• 模型并行可幫助我們理解張量并行，以明確通信成本。

• 時(shí)延計(jì)算需要從其他概念中獲得理解，并創(chuàng)建用于確定推理速度底線（floorline）的方程。

• 批大小（batch size）對(duì)性能的影響以及最優(yōu)批大小為多少。

• 通過(guò)transformer blocks執(zhí)行flops（每秒浮點(diǎn)運(yùn)算次數(shù)）計(jì)數(shù)操作，可以識(shí)別對(duì)flops速度有實(shí)質(zhì)性貢獻(xiàn)的操作。

• 中間內(nèi)存成本涵蓋了激活（即激活函數(shù)的輸出結(jié)果）占用額外內(nèi)存，以及一些真實(shí)基準(zhǔn)測(cè)試中的內(nèi)存帶寬成本。

• 對(duì)比真實(shí)基準(zhǔn)測(cè)試是指將計(jì)算出的內(nèi)容與英偉達(dá) FasterTransformer基準(zhǔn)測(cè)試結(jié)果進(jìn)行對(duì)比，并確定其中的差異。

（本文經(jīng)授權(quán)后由OneFlow編譯發(fā)布，譯文轉(zhuǎn)載請(qǐng)聯(lián)系OneFlow獲得授權(quán)。原文：https://kipp.ly/blog/transformer-inference-arithmetic/）

1

kv?緩存

采樣時(shí)，Transformer模型會(huì)以給定的prompt/context作為初始輸入進(jìn)行推理（可以并行處理），隨后逐一生成額外的token來(lái)繼續(xù)完善生成的序列（體現(xiàn)了模型的自回歸性質(zhì)）。在采樣過(guò)程中，Transformer會(huì)執(zhí)行自注意力操作，為此需要給當(dāng)前序列中的每個(gè)項(xiàng)目（無(wú)論是prompt/context還是生成的token）提取鍵值（kv）向量。這些向量存儲(chǔ)在一個(gè)矩陣中，通常被稱(chēng)為kv緩存或者past緩存（開(kāi)源GPT-2的實(shí)現(xiàn)稱(chēng)其為past緩存）。past緩存通常表示為：[batch, 2, num_heads, seq_len, features]。

kv緩存是為了避免每次采樣token時(shí)重新計(jì)算鍵值向量。利用預(yù)先計(jì)算好的k值和v值，可以節(jié)省大量計(jì)算時(shí)間，盡管這會(huì)占用一定的存儲(chǔ)空間。每個(gè)token所存儲(chǔ)的字節(jié)數(shù)為：

?

第一個(gè)因子2表示k和v這兩個(gè)向量。在每一層中我們都要存儲(chǔ)這些k，v向量，每個(gè)值都為一個(gè)矩陣。然后再乘以2，以計(jì)算每個(gè)向量所需的字節(jié)數(shù)（在本文中，我們假設(shè)采用16位格式）。

我們乘以token嵌入（token embeddings）得到的權(quán)重為，其中每個(gè)token嵌入為。這樣，我們就可以算出所有層的k和v需進(jìn)行的浮點(diǎn)運(yùn)算次數(shù)為：

將乘以需要進(jìn)行次浮點(diǎn)運(yùn)算。另一個(gè)2表示我們需要重復(fù)兩次這樣的操作，一次用于計(jì)算k和一次用于計(jì)算v，然后再重復(fù)所有層數(shù)。

矩陣乘法（matmul）中的浮點(diǎn)運(yùn)算次數(shù)為多少？

矩陣-向量（matrix-vector）乘法的計(jì)算公式是，其中，。對(duì)于矩陣-矩陣（matrix-matrix）乘法，計(jì)算公式是 ，其中，。因子十分重要，因?yàn)樗从沉司仃嚦朔ㄖ杏沙朔ê图臃ńM成的組合方式，即“乘法(1)-加法(2) 操作組合”。更多內(nèi)容見(jiàn)講義（lecture notes）。

這意味著對(duì)于一個(gè)520億參數(shù)的模型來(lái)說(shuō) （以Anthropic中的模型為例，?，），其浮點(diǎn)運(yùn)算次數(shù)為：

假設(shè)有一個(gè)A100 GPU，其每秒可執(zhí)行的浮點(diǎn)運(yùn)算次數(shù)為，其內(nèi)存帶寬可達(dá)字節(jié)/秒。以下數(shù)字僅涉及kv權(quán)重及計(jì)算的數(shù)值：

Flops vs 內(nèi)存有界性（Boundedness）

英偉達(dá)使用了數(shù)學(xué)帶寬這個(gè)術(shù)語(yǔ)，我覺(jué)得這個(gè)術(shù)語(yǔ)真的很可愛(ài)。從技術(shù)上講，這種描述存在于每個(gè)內(nèi)核中，但可以抽象為操作組。?

Flops vs 內(nèi)存有界性是Transformer推理和深度學(xué)習(xí)優(yōu)化的常見(jiàn)問(wèn)題。為了完成所需計(jì)算，通常需要加載權(quán)重，加載過(guò)程會(huì)占用內(nèi)存帶寬。假設(shè)通過(guò)加載權(quán)重已經(jīng)得到了很好的優(yōu)化，那么我們可以在加載權(quán)重的同時(shí)開(kāi)始計(jì)算。

在這種情況下，flop bound意味著一段時(shí)間內(nèi)存中沒(méi)有任何數(shù)據(jù)傳輸；memory bound則意味著沒(méi)有進(jìn)行任何計(jì)算操作。

英偉達(dá)使用數(shù)學(xué)帶寬（math bandwidth）來(lái)描述該情況，我覺(jué)得相當(dāng)有意思。從技術(shù)上講，這種劃分通常是指每個(gè)內(nèi)核（kernel）中的計(jì)算量受限，但也可以指這些操作組的計(jì)算量受限，將每一組視為抽象意義上的單元。

現(xiàn)在模型架構(gòu)不再重要了——在給定硬件規(guī)格的情況下，我們得到了一個(gè)明顯的比率208。這意味著，我們計(jì)算一個(gè)token的kv所需的時(shí)間，與處理208個(gè)token的時(shí)間相同。若低于該值，會(huì)出現(xiàn)內(nèi)存帶寬限制，若高于該值，會(huì)出現(xiàn)flops限制。如果我們使用剩余的權(quán)重來(lái)完成完整的前向傳遞（即運(yùn)行剩余的transformer），那么結(jié)果仍然是208（分子和分母各乘以6）。這一點(diǎn)我們會(huì)在后面的章節(jié)詳細(xì)介紹。

下圖的交點(diǎn)是208，不過(guò)實(shí)際上內(nèi)存線（memory line）會(huì)有一些傾斜，這是因?yàn)橹虚g計(jì)算（intermediate calculation）存在內(nèi)存成本（上一節(jié)討論過(guò)）。

對(duì)于擁有520億參數(shù)的模型來(lái)說(shuō)，一次完整的前向傳遞需要毫秒，這是處理208個(gè)token所需的時(shí)間（實(shí)際上我們會(huì)使用四個(gè)GPU進(jìn)行并行處理，因此實(shí)際需要的時(shí)間約為17毫秒，后續(xù)章節(jié)將做詳細(xì)介紹）。如果語(yǔ)言環(huán)境存在416個(gè)token（即雙倍token），那么處理時(shí)間將翻倍，而處理312個(gè)token所需的時(shí)間是處理208個(gè)token的1.5倍。

計(jì)算kv緩存的token時(shí)，一個(gè)token所需的計(jì)算成本為模型中傳遞該token計(jì)算成本的1/6。總的來(lái)說(shuō)，這些前向傳遞（獲取logits、嵌入和訓(xùn)練時(shí)我們深有體會(huì)）非常便宜，因?yàn)榭梢赃M(jìn)行并行計(jì)算。相比之下，采樣的成本要高得多，因?yàn)樗枰獜?qiáng)制讀取每個(gè)token的所有權(quán)重，進(jìn)行自回歸預(yù)測(cè)。

但這并不意味著時(shí)間節(jié)省了1/6！假設(shè)出現(xiàn)了flops bound，在每個(gè)采樣步驟中，我們可以少進(jìn)行次浮點(diǎn)運(yùn)算，而解碼步驟（step）需要進(jìn)行次浮點(diǎn)運(yùn)算。

因此，在每個(gè)步驟中，我們節(jié)省的每秒浮點(diǎn)運(yùn)算次數(shù)是序列中每個(gè)token每秒浮點(diǎn)運(yùn)算次數(shù)的1/6（很大！），而且該數(shù)值會(huì)隨著采樣token數(shù)量的增加而增加。在沒(méi)有kv緩存的情況下，隨著token數(shù)量的增加，采樣的時(shí)間復(fù)雜度（time complexity）將以平方級(jí)增加。

考慮到存儲(chǔ)緩存相關(guān)的開(kāi)銷(xiāo)和權(quán)衡（tradeoffs），以上說(shuō)法并不全面。如果我們進(jìn)行小批量設(shè)置，可能會(huì)出現(xiàn)內(nèi)存帶寬受限，而非flops bound。在這種情況下，我們可能不會(huì)使用過(guò)去的緩存，而是傾向于重新計(jì)算，這會(huì)消耗flops（因?yàn)槲覀円呀?jīng)支付了采樣的內(nèi)存成本）。

2

容量

我們對(duì)于GPU中存儲(chǔ)的kv緩存和權(quán)重有了一定認(rèn)識(shí)，并且了解到GPU容量確實(shí)對(duì)Transformer模型的推理性能有著重要影響，也就有了充分的理由對(duì)其進(jìn)行評(píng)估。

一般來(lái)說(shuō)，Nvidia A100 GPU是用于推理的最佳GPU，其容量標(biāo)準(zhǔn)為40GB。雖然有一些GPU的容量高達(dá)80GB，并且具有更高的內(nèi)存帶寬（為 2e12 而非 1.5e12），但它們尚未被任何大型云服務(wù)提供商采用，因此，于我而言它們?nèi)狈?shí)際價(jià)值。

將給定的參數(shù)計(jì)數(shù)（parameter count）乘以2，我們就可以獲取相應(yīng)字節(jié)數(shù)，進(jìn)而算出擁有520億參數(shù)的模型的權(quán)重大小。

?

不過(guò)這顯然不能在一個(gè)GPU上完成，我們至少需要三個(gè)GPU才能加載所有權(quán)重（稍后將討論如何進(jìn)行分區(qū)（sharding））。但這樣就只剩下?可用kv緩存了，這足夠嗎？讓我們回到kv緩存內(nèi)存中每個(gè)token的方程式，再次使用520億參數(shù)大小的模型運(yùn)行。

使用這種GPU設(shè)置，我們可以將個(gè)token存儲(chǔ)在kv緩存中。或者我們可以將batch size設(shè)置為4，其中每個(gè)請(qǐng)求最多包含2048個(gè)token（token越少所需batch size越大）。

難辦的是，我們想要做更高的batch size，但卻受到容量限制。batch size越大，GPU處理相同請(qǐng)求所需的時(shí)間就越短，就能更有效地利用GPU資源。然而，batch size小，內(nèi)存又會(huì)受到限制。在這種情況下，應(yīng)該放棄使用kv緩存，選擇支付flops成本。

同時(shí)使用四個(gè)GPU，就能處理個(gè)token。要想batch size更大，處理更多的數(shù)據(jù)，我們肯定會(huì)選擇四個(gè)GPU。若只用一個(gè)GPU，會(huì)降低2/3的效率，這顯然不是明智的做法。這不僅是batch size的問(wèn)題，如果有大量的數(shù)據(jù)需要處理，那應(yīng)該使用多個(gè)模型實(shí)例。我們的目標(biāo)應(yīng)該是，盡可能讓每個(gè)實(shí)例都可以通過(guò)更大的batch size處理，因?yàn)闊o(wú)法避免支付存儲(chǔ)權(quán)重的成本。

中間計(jì)算步驟會(huì)占用一些額外空間，但這些空間可以忽略不計(jì)。

3

模型并行

這方面已有許多相關(guān)介紹，所以我就不再詳細(xì)介紹模型并行（model parallelism）及其實(shí)現(xiàn)細(xì)節(jié)了，僅提供一些有用信息，以幫助讀者做性能決策并計(jì)算通信成本。

模型并行的最終結(jié)果：通過(guò)內(nèi)存和flops傳輸?shù)乃袡?quán)重成本都被分?jǐn)偟绞褂玫募铀倨鲾?shù)量上。

我們將采用張量并行（一種模型并行），將模型的中間進(jìn)行劃分。每個(gè)加速器將使用其權(quán)重分片（shards）盡可能多地執(zhí)行操作，并在需要同步時(shí)進(jìn)行通信。相比之下，流水并行（pipeline parallel）更為簡(jiǎn)單，其中每個(gè)GPU將保留模型的一部分層。雖然這種方法的確平衡了權(quán)重加載成本，但存在明顯缺陷：只有一個(gè)GPU運(yùn)轉(zhuǎn)，其他的都被閑置！

在訓(xùn)練過(guò)程中，你可以采用流水并行（第一批數(shù)據(jù)移向下一個(gè)GPU時(shí)，新一批數(shù)據(jù)重新于第一個(gè)GPU開(kāi)始），以便有效利用各個(gè)GPU。雖然處理多個(gè)樣本請(qǐng)求時(shí)該方法十分有用，但在處理單個(gè)樣本請(qǐng)求時(shí)，這種方法就不太奏效了。此外，無(wú)論flops是否受限，流水線并行都無(wú)法充分利用內(nèi)存帶寬。簡(jiǎn)而言之，流水并行適用于通信，而模型并行適用于計(jì)算。流水并行會(huì)在每個(gè)加速器之間進(jìn)行次通信，而模型并行會(huì)在每個(gè)層內(nèi)進(jìn)行次通信，其中是加速器的數(shù)量。

A100 GPU的通信帶寬為300GB/s，而文檔將其標(biāo)記為600GB/s，這是因?yàn)镹VIDIA 在每顆芯片上疊加了300GB/s的帶寬，同時(shí)向外輸出300GB/s的帶寬，而沒(méi)有使用雙向數(shù)值（對(duì)于計(jì)算，這種方法更為直觀）。

首先，我們?cè)趫D中黃色區(qū)域?qū)oken嵌入（token embedding）插入到模型底部。紫色盒子描述了權(quán)重在加速器上的分配情況，此處我們使用的是一個(gè)非常小的框架，以便能夠按比例繪制所有內(nèi)容。

普遍的做法是，我們可以將兩個(gè)矩陣和進(jìn)行劃分，并計(jì)算shards的乘積，但這并不能完成 的矩陣乘法（matmul）。換言之，如果我們只是簡(jiǎn)單地將shards的乘積連接在一起，那么得到的矩陣就會(huì)太大，無(wú)法存儲(chǔ)或進(jìn)行計(jì)算。相反，如果我們想要進(jìn)行傳輸，可以計(jì)算出shard sum，并將shard sum進(jìn)行通信，然后進(jìn)行連接到輸出。

注意力機(jī)制通常涉及多頭（head），因此采取并行計(jì)算非常直觀。幾乎不需要通信就能走完大部分注意力層，因?yàn)樽⒁饬︻^（attention head）是連接在一起的，需要乘以權(quán)重進(jìn)行計(jì)算。在乘以后，我們需要將結(jié)果乘以shard的，以得到?。

然后，每個(gè)加速器會(huì)將自身的shard傳遞給其他加速器，其他加速器也會(huì)將其shard返回，其通信成本為。每個(gè)加速器平等分配相加的shard，以獲得輸出投影（output projection，譯者注：輸出層的輸出被稱(chēng)為輸出投影,它是一種將神經(jīng)網(wǎng)絡(luò)的輸出轉(zhuǎn)換為特定格式的過(guò)程。）然后，他們會(huì)重復(fù)上次進(jìn)行的通信，各個(gè)主機(jī)進(jìn)行連接（近似瞬間就能完成）。

其本質(zhì)與MLP（Multi-Layer Perceptron）層類(lèi)似！就像我們使用權(quán)重將多頭注意力（multi-headed attention）的結(jié)果投射回長(zhǎng)度為的向量一樣。同樣，我們也使用和將向量維度擴(kuò)大4倍，并將其重新投射回原始值大小。因此，在MLP的末端完成了兩次同樣的通信。

我們最終的通信量為字節(jié)。GPU將對(duì)整個(gè)kv緩存進(jìn)行劃分，并將其分配到各個(gè)頭部。

4

時(shí)延計(jì)算

我們已經(jīng)較全面地討論了容量（capacity）、模型并行中的通信，以及一般的計(jì)算步驟。接下來(lái)我們將其構(gòu)建到估計(jì)時(shí)延的方程中！

時(shí)延計(jì)算大多與flops和內(nèi)存有界性相關(guān)。如果每個(gè)參數(shù)需要進(jìn)行的乘法運(yùn)算很少，那么我們可能會(huì)受到內(nèi)存帶寬的限制。浮點(diǎn)運(yùn)算量增加取決于批處理大小和參數(shù)數(shù)量，而內(nèi)存只受參數(shù)數(shù)量的影響。

通信方面，關(guān)鍵不是有界性問(wèn)題，而是增加時(shí)延項(xiàng)和吞吐量項(xiàng)（即300GB/s）。由于時(shí)延方面的數(shù)據(jù)不透明，最理想化的估計(jì)也是每條消息發(fā)送大約需要8微秒。該數(shù)據(jù)源于Citadel的一篇論文，但該論文針對(duì)的是V100 NVLink。

受計(jì)算因素的影響，計(jì)算單個(gè)token解碼步驟的時(shí)延時(shí)間需要兩個(gè)公式：一個(gè)用于內(nèi)存帶寬bound（小batch），另一個(gè)用于 flops bound（大batch）。處理大batch數(shù)據(jù)時(shí)，我們通常會(huì)忽略通信的時(shí)延因素。

針對(duì)小batch（當(dāng)batch size=1時(shí)，可以忽略batch因素）的方程如下（其中N為加速器數(shù)量，P為參數(shù)數(shù)量，b表示字節(jié)單位）：

因?yàn)槲覀冃枰ㄟ^(guò)內(nèi)存?zhèn)鬟f所有參數(shù)，且每個(gè)參數(shù)都是2字節(jié)，所以這里是2*P。是加速器內(nèi)存帶寬，成本在加速器之間分?jǐn)偂Ｃ繉佑袀€(gè)通信請(qǐng)求，每個(gè)請(qǐng)求的時(shí)延較小，通常可以忽略。通信還有吞吐成本，但也可以忽略。

有時(shí)讀取kv緩存的時(shí)間也會(huì)對(duì)計(jì)算產(chǎn)生重要影響。不過(guò)，由于該因素取決于前后token的數(shù)量，而數(shù)量在同一批次內(nèi)會(huì)有所變化，我們要采樣的總token數(shù)量也有所不同，因此暫時(shí)將其排除在計(jì)算之外。這個(gè)時(shí)間會(huì)被計(jì)算為內(nèi)存帶寬時(shí)間。另一個(gè)缺失的內(nèi)存帶寬時(shí)間是讀取unembeddings，以計(jì)算每個(gè)采樣步驟的logits，即。

如前所述，內(nèi)存實(shí)際上并非保持不變，每個(gè)批次會(huì)使用一些額外的內(nèi)存來(lái)存儲(chǔ)中間激活值。我們之所以不計(jì)算這些額外內(nèi)存，是因?yàn)槠鋽?shù)量深受軟件堆棧、編譯器優(yōu)化等因素的影響，很難精確計(jì)算。

針對(duì)大batch（batch size=512）的方程如下（其中 B 是批量大小）：

其中，表示加速器的浮點(diǎn)運(yùn)算能力，表示處理單元之間的通信帶寬。公式中做了2*P次浮點(diǎn)運(yùn)算，這從對(duì)所有參數(shù)進(jìn)行矩陣相乘也可看出。如前所述，因?yàn)?#xff0c;所以矩陣向量乘法是2mn。

在模型并行部分，每個(gè)層需要進(jìn)行四次（N-1個(gè)因子，四舍五入為至N）大小向量的通信。考慮到時(shí)延可以忽略不計(jì)，這里使用吞吐量來(lái)計(jì)算通信時(shí)間。將需要傳輸?shù)目倲?shù)據(jù)量除以通信帶寬，可以得到通信所需時(shí)間。

接下來(lái)，在16個(gè)GPU上使用2600億參數(shù)的Gopher模型來(lái)進(jìn)行推理。使用小batch推理時(shí)，生成一個(gè)token需要22毫秒的時(shí)間。通過(guò)大batch公式計(jì)算出來(lái)的通信吞吐量成本約為35微秒，因此可以放心地降低該成本。

對(duì)于512的大batch，生成每個(gè)token所需的時(shí)間為53毫秒（即在62毫秒內(nèi)生成512個(gè)token）。通信的時(shí)延成本也為3毫秒（由于消息可以一起準(zhǔn)備，因此延時(shí)延不會(huì)隨著batch的增加而增加），這對(duì)于減少時(shí)延來(lái)說(shuō)有一定的意義，但如果假設(shè)通信和計(jì)算是并行的，那么這也可以接受。

在假定并行處理的情況下，我們會(huì)選擇計(jì)算和通信中較大的值。因此，我們希望避免通信時(shí)間大于計(jì)算時(shí)間（這是防止隨著芯片數(shù)量的增加趨近于零時(shí)延的機(jī)制，最終通信時(shí)間將占用越來(lái)越多的時(shí)間）。但并不能保證所有系統(tǒng)都能夠完美地并行處理。

這些數(shù)值明顯比“沙箱環(huán)境”中獲得的值要低得多（譯者注：“沙箱環(huán)境”即一個(gè)獨(dú)立于操作系統(tǒng)和其他程序運(yùn)行的安全環(huán)境。在這個(gè)環(huán)境中，程序只能使用被授權(quán)的資源，無(wú)法對(duì)系統(tǒng)或其他程序產(chǎn)生影響），因?yàn)樗僭O(shè)了最佳的硬件使用，沒(méi)有考慮softmax，假設(shè)沒(méi)有通信時(shí)延，并忽略了許多其他較小的因素。盡管如此，這些數(shù)學(xué)推理仍有助于思考如何優(yōu)化性能以及未來(lái)優(yōu)化所帶來(lái)的變化。

5

Batch sizes

（Batch sizes）批大小是性能的一個(gè)重要因素，尤其是對(duì)特定用途性能的理解。

我們?cè)谇懊娌糠诌M(jìn)行了兩個(gè)計(jì)算，用于確定何時(shí)是內(nèi)存帶寬bound，何時(shí)是flops bound。為了確定主導(dǎo)因素，我們可以對(duì)這些數(shù)值進(jìn)行對(duì)比。

我們正在處理與kv緩存部分相同的比率。內(nèi)存帶寬bound的最小batch size為。該比率非常有用！在足夠負(fù)載的情況下，我們更傾向于flops bound，因?yàn)檫@樣計(jì)算效率更高。但是，如果是flops bound，增大batch size并不會(huì)提高計(jì)算速度。

很容易計(jì)算出容量中的主流何時(shí)從kv緩存轉(zhuǎn)變?yōu)闄?quán)重，但這之間并沒(méi)有明顯的二進(jìn)制分界點(diǎn)（當(dāng)kv緩存開(kāi)始占用更多內(nèi)存時(shí)并不會(huì)有特別的變化）。此外，也沒(méi)有特別重要的通信因素。隨著batch size增加，吞吐量開(kāi)始超過(guò)時(shí)延，所以我們不再考慮時(shí)延因素。正如之前觀察到的，時(shí)延變得不重要的時(shí)間相對(duì)較晚（例如，52B通信成本上的512批大小仍有11%的時(shí)延）。

將其簡(jiǎn)化一下，即通信需要在四個(gè)不同的步驟中進(jìn)行，這意味著我們不僅希望計(jì)算時(shí)間比通信時(shí)間長(zhǎng)，而且在每個(gè)步驟中都是如此（如果我們可以同時(shí)進(jìn)行計(jì)算和通信）。為了達(dá)到這個(gè)目標(biāo)，我們使用一個(gè)更奇特的比率，即每個(gè)字節(jié)的通信量需要多少次浮點(diǎn)運(yùn)算。下圖是一個(gè)很好的計(jì)算表，后續(xù)內(nèi)容也將會(huì)用到。

A100芯片每字節(jié)通信的浮點(diǎn)運(yùn)算次數(shù)為(312e12/300e9)=1040次。我們希望最后一行的值大于硬件每字節(jié)的浮點(diǎn)運(yùn)算次數(shù)，以保持flops bound（假設(shè)沒(méi)有memory bound）。對(duì)于任何embedding維度超過(guò)1024（每個(gè)芯片）的模型，相對(duì)比較安全！但對(duì)于512維度的模型而言，情況就有些棘手。

API負(fù)載較低時(shí)，會(huì)出現(xiàn)較小的batch sizes，這時(shí)會(huì)考慮放棄kv緩存等決策。當(dāng)API負(fù)載較高時(shí)，為了優(yōu)化每個(gè)請(qǐng)求的時(shí)延，會(huì)選擇提供最低batch size，以達(dá)到flop bound，即使仍有容量剩余。在AlphaCode等大規(guī)模推斷任務(wù)中，我們通常會(huì)盡可能多地插入芯片，然后利用這些容量執(zhí)行最大的batch操作。雖然我經(jīng)常使用“可能”一詞，但這些情況都是絕對(duì)存在的。

6

Flops計(jì)算

之前：

從所有參數(shù)的matmul操作來(lái)看，我們的flops運(yùn)算為2*P次。

這是很合理的推理，我們還可以檢查transformer步驟來(lái)分解推理，以核查是否能得出2P這個(gè)結(jié)論。

以下是對(duì)token和layer的計(jì)算。準(zhǔn)確來(lái)說(shuō)，應(yīng)該表述為，其中i最大為。為了計(jì)算時(shí)延，我簡(jiǎn)化了，以包含所有heads。

>qkv計(jì)算

? ?>將乘以

? ?>Flop次數(shù)：

>計(jì)算z

? ?>公式為：

? ?>這里不需要矩陣乘，flops為的某個(gè)因子。

>乘以輸出投影矩陣（projection matrix）

? ?>將乘以

? ?>Flop次數(shù)：

>前饋

? ?>我們有用于兩個(gè)線性變換的MLP權(quán)重和(中間有一個(gè)小的ReLU)。

? ?>Flop次數(shù)：

>其他

? ?>通常每次attention后都會(huì)運(yùn)行l(wèi)ayernorm，其中權(quán)重為長(zhǎng)度向量。

? ?>這里還有另一個(gè)線性層，然后是位于頂部的softmax，softmax就是我們的輸出（token）嵌入（embedding）、解嵌（unembedding?）、反嵌（de-embedding）或嵌入（?embedding）。

? ?>原始transformer有余弦絕對(duì)位置編碼（cosine absolute positional encoding scheme），它是token嵌入的加法操作。

將所有flops相加！

在8192模型中，flops大約為100B；

103079215104除以2約等于515億。我們得到的結(jié)果為515億，比520億略少一些，這是因?yàn)閠oken（un）embeddings的參數(shù)接近了10億，相比，更適合用來(lái)計(jì)算時(shí)延，但這兩者的差異不到2%。

該如何計(jì)算Z和其他步驟呢？這些都是向量到向量（甚至向量到標(biāo)量）的運(yùn)算，所以它們都是圍繞因子建造起來(lái)的，而不是因子。即使每一層（layer)要運(yùn)行100次這樣的操作，最終也會(huì)有一億次flops, 占已計(jì)算flops的0.1%。

7

中間內(nèi)存成本

Data Movement Is All You Need（https://arxiv.org/pdf/2007.00072.pdf，主要內(nèi)容是優(yōu)化transformers的低層級(jí)數(shù)據(jù)移動(dòng)，與本文內(nèi)容不太相關(guān)）一文提出了一個(gè)很好的分類(lèi)操作方法。首先，在大矩陣（包括線性層）當(dāng)中，張量縮并（tensor contractions）是最重要的，統(tǒng)計(jì)歸一化（包括softmax和layernorm）次之，最后是逐元素操作，比如偏置（ biases）、dropouts和激活（activations）等。

那么我們?cè)撊绾斡?jì)算矩陣、layernorms等的時(shí)延呢？我們硬件上報(bào)告的flops是專(zhuān)門(mén)針對(duì)乘加運(yùn)算的，因此即使我們可以計(jì)算flops，其結(jié)果也不對(duì)。幸運(yùn)的是，這只是為了占用內(nèi)存來(lái)進(jìn)行softmax讀/寫(xiě)，因?yàn)檫@有利于帶寬和flops比率。這是已知的時(shí)延因素！

在這里我將拋開(kāi)第一性原則，討論?Data Movement Is All You Need一文中的表格A.1。我們發(fā)現(xiàn)softmax的時(shí)延時(shí)間略長(zhǎng)于qkv的計(jì)算時(shí)間（softmax時(shí)延時(shí)間是qkv操作時(shí)間的三倍）。這有點(diǎn)令人擔(dān)憂，可能會(huì)影響整個(gè)神經(jīng)網(wǎng)絡(luò)的性能。

出于同樣的原因，softmax會(huì)受到內(nèi)存限制，所以qk、ReLU和dropout的乘法（multiplication）操作也相當(dāng)昂貴。

GPU內(nèi)核融合（Fusion）

GPU 以“內(nèi)核”為單位執(zhí)行操作。內(nèi)核融合意味著2個(gè)內(nèi)核可以融合為一個(gè)，這樣我們就可以在內(nèi)存中重復(fù)利用負(fù)載，減少冗余負(fù)載和存儲(chǔ)。例如，一個(gè)乘加（multiply-add）是一個(gè)內(nèi)核。但是如果一個(gè)乘加有兩個(gè)內(nèi)核，一個(gè)內(nèi)核負(fù)責(zé)加載+加法+存儲(chǔ)，另一個(gè)內(nèi)核負(fù)責(zé)加載+乘法+存儲(chǔ)。內(nèi)核融合以后，我們可以運(yùn)行加載+加法+乘法+存儲(chǔ)，以簡(jiǎn)化步驟。

通過(guò)計(jì)算所需的讀寫(xiě)次數(shù)，我們發(fā)現(xiàn)這里的softmax沒(méi)有完美融合。理論上它可以是一次讀取和一次寫(xiě)入（標(biāo)準(zhǔn)次數(shù)為四次）。qk是兩次讀取和一次寫(xiě)入（兩次讀取可能可以保存）。三比一的比率表示softmax執(zhí)行的內(nèi)存?zhèn)鬟f量多于最佳值。我這樣說(shuō)是因?yàn)?#xff0c;這表明了計(jì)算的軟件依賴(lài)程度，并且需要通過(guò)實(shí)驗(yàn)來(lái)驗(yàn)證，因?yàn)閺睦碚撋蟻?lái)說(shuō)成本可能為0。

值得注意的是，隨著模型大小的增加，操作所花費(fèi)的時(shí)間百分比會(huì)迅速降低，因?yàn)槊繉觾?nèi)存將以增加，flops將以 增加。本文為336M參數(shù)模型，。

我將“Ours”列中內(nèi)存限制的所有值的時(shí)延相加（包括 element-wise操作）。結(jié)果顯示，中間步驟占用了43%的時(shí)間。可以看到，這些操作在大小為520億（ 是這個(gè)模型的八倍）的模型中沒(méi)有那么重要。

這些memory bound中間操作的持續(xù)時(shí)間將延長(zhǎng)8倍，因?yàn)檫@些操作是長(zhǎng)度的向量。但是，flops次數(shù)將增加64倍，這意味著flop時(shí)間會(huì)延長(zhǎng)64倍。

因此，若不考慮時(shí)延，使用?Data Movement Is All You Need中的優(yōu)化方法，520億參數(shù)模型的推理時(shí)延大約會(huì)是中間計(jì)算的 5%。

8

實(shí)際基準(zhǔn)對(duì)比

我在語(yǔ)言建模公司工作，我們公司有自己的基礎(chǔ)設(shè)施和基準(zhǔn)，但I(xiàn)P是公司面臨的一大難題。可悲的是缺乏可用于模型并行推理的公共基準(zhǔn)。目前我只知道Nvidia的FasterTransformer和Microsoft的Deepspeed，可能其他我不了解的論文也提出了一些基準(zhǔn)。無(wú)論如何，我們可以根據(jù)真實(shí)基準(zhǔn)來(lái)驗(yàn)證計(jì)算！

因?yàn)槲抑幌胗?個(gè)GPU，所以使用了FasterTransformer來(lái)運(yùn)行了一個(gè)130億參數(shù)的模型。該模型執(zhí)行一連串內(nèi)核融合，并提供張量并行功能。模型有40層，每層有40個(gè)頭，每個(gè)頭的維度為128，總維度大小為5120。這里是配置文件截圖，這些截圖顯示了許多有趣的細(xì)節(jié)，可能值得單獨(dú)撰寫(xiě)一篇文章。

首先是輸出的512個(gè)上下文長(zhǎng)度（context length）、batch size為1和10個(gè)token。對(duì)于2個(gè)GPU上的一個(gè)小batch的token，我們預(yù)計(jì)為8.4毫秒，大約1毫秒的通信，那么一個(gè)GPU則是16.8毫秒，0毫秒的通信。（2x40x12x5120^2/1.5e12）

在這里我應(yīng)該將內(nèi)存寬帶（mem bandwidth）設(shè)置為1.555，而不是1.5。

實(shí)測(cè)結(jié)果表明，1個(gè)GPU應(yīng)該是22.0ms，這意味著，我們的猜測(cè)的準(zhǔn)確率為76%。可以確定，其中一部分時(shí)間花在了中間激活操作上，從理論上來(lái)說(shuō)，我們可以獲得100%的內(nèi)存帶寬，但實(shí)際上并沒(méi)有。

對(duì)于這些維度，測(cè)試表明我們可以獲得高達(dá)90%的內(nèi)存帶寬利用率（我們將矩陣乘法的預(yù)期成本與單個(gè)矩陣乘法內(nèi)核的持續(xù)時(shí)間進(jìn)行比較，由于加載的張量不同，帶寬利用率會(huì)有所變化）。考慮到這一點(diǎn)，我們預(yù)計(jì)需要18.5毫秒。加上中間激活操作的成本后（我們可以從測(cè)試結(jié)果中獲得），需要額外的2.2毫秒，總共需要20.7毫秒！

為了解釋剩下的1.4毫秒，我們考慮到了一些其他的亞毫秒級(jí)操作，比如token嵌入、top-(k|p)操作、少于90%的帶寬利用率（不想計(jì)算平均值，直接取了我能找到的最高帶寬利用率），或者是內(nèi)核啟動(dòng)時(shí)間。

實(shí)驗(yàn)結(jié)果顯示，使用2個(gè)GPU時(shí)，總時(shí)間為13.5毫秒。相比一個(gè)GPU，這次我們只占了內(nèi)存寬帶的76%，離目標(biāo)還有很大的差距。為了解決這個(gè)問(wèn)題，我們重新檢查了配置文件，發(fā)現(xiàn)內(nèi)存帶寬稍微差了一些，因?yàn)閺埩枯^小獲取的帶寬也比較少。

經(jīng)過(guò)計(jì)算，寬帶沒(méi)有達(dá)到90%，只有87%，總時(shí)間為9.5毫秒，中間的激活時(shí)間需要大約2毫秒，這使得總時(shí)間為11.7毫秒。剩余的1.5毫秒需要找出通信問(wèn)題。但是這個(gè)問(wèn)題很容易解決，因?yàn)槲覀冎坝?jì)算的1毫秒通信時(shí)間沒(méi)有被并行化。根據(jù)配置文件的數(shù)據(jù)，每個(gè)層的通信時(shí)間為40-50微秒，總共約為1.7毫秒，這就是很好的證明。

上述兩種操作的中間激活計(jì)數(shù)都比應(yīng)有的要高一些，因?yàn)榕渲梦募峁┑臅r(shí)延始終略高于原始基準(zhǔn)測(cè)試運(yùn)行。基準(zhǔn)測(cè)試運(yùn)行的輸出為180.86ms（上下文時(shí)間：45.45ms）和 283.60ms（上下文時(shí)間：63.17ms）。

在前向傳遞過(guò)程中，模型需要將所有token發(fā)送到每個(gè)GPU上，然后每個(gè)GPU都會(huì)對(duì)其進(jìn)行自己的注意力頭計(jì)算并存儲(chǔ)kv。由于這個(gè)過(guò)程需要發(fā)送大量的數(shù)據(jù)并進(jìn)行并行計(jì)算，因此預(yù)計(jì)前向傳遞需要的時(shí)間會(huì)比解碼步驟長(zhǎng)num_tokens/flops_to_bw_ratio倍。

更新的內(nèi)存帶寬為：312e12/（1.5e12x0.9）=231。在1個(gè)GPU的設(shè)置中，22ms是我們預(yù)期的解碼步驟，我們可以得到22*（512/231）= 48，并不是63。在2個(gè)GPU設(shè)置中，我們通過(guò)計(jì)算得到了更加糟糕的結(jié)果：13.5*（512/231）=30ms。

單個(gè)GPU只缺失了部分kv儲(chǔ)存時(shí)間。查看配置文件，我們發(fā)現(xiàn)kv儲(chǔ)存時(shí)間每層為18微秒，總共為0.7毫秒。Memsets占了0.2毫秒。我們預(yù)計(jì)其中一個(gè)MLP乘法的flop時(shí)間（flop bound）為512x4x5120^2x2/312e12 = 344微秒。實(shí)際上，最低浮點(diǎn)運(yùn)算時(shí)間應(yīng)該是476微秒，也就是說(shuō)，我們得到了預(yù)期flops的72%。

對(duì)于attention中的投影（projection），我們期望是512x5120^2x2/312e12 =86微秒。但是在配置文件中，我們發(fā)現(xiàn)最低是159微秒。請(qǐng)參閱本文的圖14，其中512x4000x4000的最終結(jié)果低于150TFLOPs/s。

9

練習(xí)

1.在給定batch size、上下文長(zhǎng)度和next_n的情況下，我們?cè)撊绾斡?jì)算kv緩存節(jié)約的寬帶？

2.kv緩存會(huì)增加內(nèi)存時(shí)間開(kāi)銷(xiāo)嗎？

3.我們是否可以在前向傳遞時(shí)選擇memory bound，在采樣步驟中選擇flops bound？

4.在GPU超過(guò)容量所需的情況下，我們應(yīng)該進(jìn)行怎樣的權(quán)衡和計(jì)算？例如，一個(gè)520億參數(shù)大小的模型（擁有8或16個(gè)GPU）。

5.如果我們有計(jì)算預(yù)測(cè)一個(gè)token時(shí)間的公式。我們應(yīng)該如何計(jì)算執(zhí)行整個(gè)樣本的時(shí)間呢? 是否應(yīng)該先在上下文上進(jìn)行前向傳遞，再預(yù)測(cè)所有的請(qǐng)求token？

6.在容量部分，我曾提到中間計(jì)算的內(nèi)存可以忽略不計(jì)。那么這些內(nèi)存到底有多小呢？

7.在batch size部分，我們討論了每字節(jié)通信的token數(shù)。如果我們的嵌入維度為512，我們需要做怎樣的權(quán)衡？

8.假設(shè)GPU都連接到了同一主機(jī)，但可以像訓(xùn)練那樣在主機(jī)之間進(jìn)行GPU通信。AWS 有400GB/s。這種情況該怎么辦呢？

9.在模型并行部分，我們可以實(shí)際溝通所有分片（shards），然后讓每個(gè)加速器執(zhí)行所有添加（不只是其中一部分添加）。這部分會(huì)對(duì)時(shí)延產(chǎn)生怎樣的影響呢？

10.在batch size為256的4xGPUs上計(jì)算520億的大批量速度。計(jì)算約為21毫秒，通信約為4毫秒。

11.從最后一層中取出向量，將其乘以未嵌入矩陣，存儲(chǔ)logits，然后進(jìn)行top-k或top-p采樣（需要排序）。對(duì)于一個(gè)包含520億參數(shù)的模型，這個(gè)過(guò)程需要多長(zhǎng)時(shí)間，我們可以在這里并行化什么操作？

12.如何進(jìn)行分片token嵌入？在輸入token嵌入和未嵌入token之間，是否需要以不同的方式劃分分片和使用層歸一化（Layernorms），這會(huì)引起額外的通信開(kāi)銷(xiāo)嗎？

其他人都在看

• “ChatGPT們”的淘金時(shí)代

• GPT-4，大增長(zhǎng)時(shí)代的序幕

• GPT-4創(chuàng)造者：第二次改變AI浪潮的方向

• ChatGPT作者Schulman：我們成功的秘密武器

• 比快更快，開(kāi)源Stable Diffusion刷新作圖速度

• OneEmbedding:單卡訓(xùn)練TB級(jí)推薦模型不是夢(mèng)

• GLM訓(xùn)練加速：性能最高提升3倍，顯存節(jié)省1/3

試用OneFlow: github.com/Oneflow-Inc/oneflow/http://github.com/Oneflow-Inc/oneflow/

總結(jié)

以上是生活随笔為你收集整理的大型语言模型的推理演算的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。