作者｜kipply

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

本文詳細闡述了大型語言模型推理性能的幾個基本原理，不含任何實驗數據或復雜的數學公式，旨在加深讀者對相關原理的理解。此外，作者還提出了一種極其簡單的推理時延模型，該模型與實證結果擬合度高，可更好地預測和解釋Transformer模型的推理過程。

為了更好地閱讀本文，讀者需了解一些Transformer模型的相關先驗知識，比如《圖解Transformer》的大部分內容。另外，了解與本文相關的參數計數文章也能更好地幫助讀者理解本文內容。本文主要包括以下內容：

• kv 緩存 (kv cache) 解釋了在推理過程中緩存自注意力向量所帶來的性能優化效果，以及可能導致的權衡（tradeoff）以及容量成本問題。

• 容量（capacity）考慮了kv緩存的存儲成本以及模型權重的存儲成本之間的聯系，并解釋了容量大小對模型性能的影響。

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

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

• 批大小（batch size）對性能的影響以及最優批大小為多少。

• 通過transformer blocks執行flops（每秒浮點運算次數）計數操作，可以識別對flops速度有實質性貢獻的操作。

• 中間內存成本涵蓋了激活（即激活函數的輸出結果）占用額外內存，以及一些真實基準測試中的內存帶寬成本。

• 對比真實基準測試是指將計算出的內容與英偉達 FasterTransformer基準測試結果進行對比，并確定其中的差異。

（本文經授權后由OneFlow編譯發布，譯文轉載請聯系OneFlow獲得授權。原文：https://kipp.ly/blog/transformer-inference-arithmetic/）

1

kv?緩存

采樣時，Transformer模型會以給定的prompt/context作為初始輸入進行推理（可以并行處理），隨后逐一生成額外的token來繼續完善生成的序列（體現了模型的自回歸性質）。在采樣過程中，Transformer會執行自注意力操作，為此需要給當前序列中的每個項目（無論是prompt/context還是生成的token）提取鍵值（kv）向量。這些向量存儲在一個矩陣中，通常被稱為kv緩存或者past緩存（開源GPT-2的實現稱其為past緩存）。past緩存通常表示為：[batch, 2, num_heads, seq_len, features]。

kv緩存是為了避免每次采樣token時重新計算鍵值向量。利用預先計算好的k值和v值，可以節省大量計算時間，盡管這會占用一定的存儲空間。每個token所存儲的字節數為：

?

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

我們乘以token嵌入（token embeddings）得到的權重為，其中每個token嵌入為。這樣，我們就可以算出所有層的k和v需進行的浮點運算次數為：

將乘以需要進行次浮點運算。另一個2表示我們需要重復兩次這樣的操作，一次用于計算k和一次用于計算v，然后再重復所有層數。

矩陣乘法（matmul）中的浮點運算次數為多少？

矩陣-向量（matrix-vector）乘法的計算公式是，其中，。對于矩陣-矩陣（matrix-matrix）乘法，計算公式是 ，其中，。因子十分重要，因為它反映了矩陣乘法中由乘法和加法組成的組合方式，即“乘法(1)-加法(2) 操作組合”。更多內容見講義（lecture notes）。

這意味著對于一個520億參數的模型來說 （以Anthropic中的模型為例，?，），其浮點運算次數為：

假設有一個A100 GPU，其每秒可執行的浮點運算次數為，其內存帶寬可達字節/秒。以下數字僅涉及kv權重及計算的數值：

Flops vs 內存有界性（Boundedness）

英偉達使用了數學帶寬這個術語，我覺得這個術語真的很可愛。從技術上講，這種描述存在于每個內核中，但可以抽象為操作組。?

Flops vs 內存有界性是Transformer推理和深度學習優化的常見問題。為了完成所需計算，通常需要加載權重，加載過程會占用內存帶寬。假設通過加載權重已經得到了很好的優化，那么我們可以在加載權重的同時開始計算。

在這種情況下，flop bound意味著一段時間內存中沒有任何數據傳輸；memory bound則意味著沒有進行任何計算操作。

英偉達使用數學帶寬（math bandwidth）來描述該情況，我覺得相當有意思。從技術上講，這種劃分通常是指每個內核（kernel）中的計算量受限，但也可以指這些操作組的計算量受限，將每一組視為抽象意義上的單元。

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

下圖的交點是208，不過實際上內存線（memory line）會有一些傾斜，這是因為中間計算（intermediate calculation）存在內存成本（上一節討論過）。

對于擁有520億參數的模型來說，一次完整的前向傳遞需要毫秒，這是處理208個token所需的時間（實際上我們會使用四個GPU進行并行處理，因此實際需要的時間約為17毫秒，后續章節將做詳細介紹）。如果語言環境存在416個token（即雙倍token），那么處理時間將翻倍，而處理312個token所需的時間是處理208個token的1.5倍。

計算kv緩存的token時，一個token所需的計算成本為模型中傳遞該token計算成本的1/6。總的來說，這些前向傳遞（獲取logits、嵌入和訓練時我們深有體會）非常便宜，因為可以進行并行計算。相比之下，采樣的成本要高得多，因為它需要強制讀取每個token的所有權重，進行自回歸預測。

但這并不意味著時間節省了1/6！假設出現了flops bound，在每個采樣步驟中，我們可以少進行次浮點運算，而解碼步驟（step）需要進行次浮點運算。

因此，在每個步驟中，我們節省的每秒浮點運算次數是序列中每個token每秒浮點運算次數的1/6（很大！），而且該數值會隨著采樣token數量的增加而增加。在沒有kv緩存的情況下，隨著token數量的增加，采樣的時間復雜度（time complexity）將以平方級增加。

考慮到存儲緩存相關的開銷和權衡（tradeoffs），以上說法并不全面。如果我們進行小批量設置，可能會出現內存帶寬受限，而非flops bound。在這種情況下，我們可能不會使用過去的緩存，而是傾向于重新計算，這會消耗flops（因為我們已經支付了采樣的內存成本）。

2

容量

我們對于GPU中存儲的kv緩存和權重有了一定認識，并且了解到GPU容量確實對Transformer模型的推理性能有著重要影響，也就有了充分的理由對其進行評估。

一般來說，Nvidia A100 GPU是用于推理的最佳GPU，其容量標準為40GB。雖然有一些GPU的容量高達80GB，并且具有更高的內存帶寬（為 2e12 而非 1.5e12），但它們尚未被任何大型云服務提供商采用，因此，于我而言它們缺乏實際價值。

將給定的參數計數（parameter count）乘以2，我們就可以獲取相應字節數，進而算出擁有520億參數的模型的權重大小。

?

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

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

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

同時使用四個GPU，就能處理個token。要想batch size更大，處理更多的數據，我們肯定會選擇四個GPU。若只用一個GPU，會降低2/3的效率，這顯然不是明智的做法。這不僅是batch size的問題，如果有大量的數據需要處理，那應該使用多個模型實例。我們的目標應該是，盡可能讓每個實例都可以通過更大的batch size處理，因為無法避免支付存儲權重的成本。

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

3

模型并行

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

模型并行的最終結果：通過內存和flops傳輸的所有權重成本都被分攤到使用的加速器數量上。

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

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

A100 GPU的通信帶寬為300GB/s，而文檔將其標記為600GB/s，這是因為NVIDIA 在每顆芯片上疊加了300GB/s的帶寬，同時向外輸出300GB/s的帶寬，而沒有使用雙向數值（對于計算，這種方法更為直觀）。

首先，我們在圖中黃色區域將token嵌入（token embedding）插入到模型底部。紫色盒子描述了權重在加速器上的分配情況，此處我們使用的是一個非常小的框架，以便能夠按比例繪制所有內容。

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

注意力機制通常涉及多頭（head），因此采取并行計算非常直觀。幾乎不需要通信就能走完大部分注意力層，因為注意力頭（attention head）是連接在一起的，需要乘以權重進行計算。在乘以后，我們需要將結果乘以shard的，以得到?。

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

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

我們最終的通信量為字節。GPU將對整個kv緩存進行劃分，并將其分配到各個頭部。

4

時延計算

我們已經較全面地討論了容量（capacity）、模型并行中的通信，以及一般的計算步驟。接下來我們將其構建到估計時延的方程中！

時延計算大多與flops和內存有界性相關。如果每個參數需要進行的乘法運算很少，那么我們可能會受到內存帶寬的限制。浮點運算量增加取決于批處理大小和參數數量，而內存只受參數數量的影響。

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

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

針對小batch（當batch size=1時，可以忽略batch因素）的方程如下（其中N為加速器數量，P為參數數量，b表示字節單位）：

因為我們需要通過內存傳遞所有參數，且每個參數都是2字節，所以這里是2*P。是加速器內存帶寬，成本在加速器之間分攤。每層有個通信請求，每個請求的時延較小，通常可以忽略。通信還有吞吐成本，但也可以忽略。

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

如前所述，內存實際上并非保持不變，每個批次會使用一些額外的內存來存儲中間激活值。我們之所以不計算這些額外內存，是因為其數量深受軟件堆棧、編譯器優化等因素的影響，很難精確計算。

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

其中，表示加速器的浮點運算能力，表示處理單元之間的通信帶寬。公式中做了2*P次浮點運算，這從對所有參數進行矩陣相乘也可看出。如前所述，因為，所以矩陣向量乘法是2mn。

在模型并行部分，每個層需要進行四次（N-1個因子，四舍五入為至N）大小向量的通信。考慮到時延可以忽略不計，這里使用吞吐量來計算通信時間。將需要傳輸的總數據量除以通信帶寬，可以得到通信所需時間。

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

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

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

這些數值明顯比“沙箱環境”中獲得的值要低得多（譯者注：“沙箱環境”即一個獨立于操作系統和其他程序運行的安全環境。在這個環境中，程序只能使用被授權的資源，無法對系統或其他程序產生影響），因為它假設了最佳的硬件使用，沒有考慮softmax，假設沒有通信時延，并忽略了許多其他較小的因素。盡管如此，這些數學推理仍有助于思考如何優化性能以及未來優化所帶來的變化。

5

Batch sizes

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

我們在前面部分進行了兩個計算，用于確定何時是內存帶寬bound，何時是flops bound。為了確定主導因素，我們可以對這些數值進行對比。

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

很容易計算出容量中的主流何時從kv緩存轉變為權重，但這之間并沒有明顯的二進制分界點（當kv緩存開始占用更多內存時并不會有特別的變化）。此外，也沒有特別重要的通信因素。隨著batch size增加，吞吐量開始超過時延，所以我們不再考慮時延因素。正如之前觀察到的，時延變得不重要的時間相對較晚（例如，52B通信成本上的512批大小仍有11%的時延）。

將其簡化一下，即通信需要在四個不同的步驟中進行，這意味著我們不僅希望計算時間比通信時間長，而且在每個步驟中都是如此（如果我們可以同時進行計算和通信）。為了達到這個目標，我們使用一個更奇特的比率，即每個字節的通信量需要多少次浮點運算。下圖是一個很好的計算表，后續內容也將會用到。

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

API負載較低時，會出現較小的batch sizes，這時會考慮放棄kv緩存等決策。當API負載較高時，為了優化每個請求的時延，會選擇提供最低batch size，以達到flop bound，即使仍有容量剩余。在AlphaCode等大規模推斷任務中，我們通常會盡可能多地插入芯片，然后利用這些容量執行最大的batch操作。雖然我經常使用“可能”一詞，但這些情況都是絕對存在的。

6

Flops計算

之前：

從所有參數的matmul操作來看，我們的flops運算為2*P次。

這是很合理的推理，我們還可以檢查transformer步驟來分解推理，以核查是否能得出2P這個結論。

以下是對token和layer的計算。準確來說，應該表述為，其中i最大為。為了計算時延，我簡化了，以包含所有heads。

>qkv計算

? ?>將乘以

? ?>Flop次數：

>計算z

? ?>公式為：

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

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

? ?>將乘以

? ?>Flop次數：

>前饋

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

? ?>Flop次數：

>其他

? ?>通常每次attention后都會運行layernorm，其中權重為長度向量。

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

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

將所有flops相加！

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

103079215104除以2約等于515億。我們得到的結果為515億，比520億略少一些，這是因為token（un）embeddings的參數接近了10億，相比，更適合用來計算時延，但這兩者的差異不到2%。

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

7

中間內存成本

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

那么我們該如何計算矩陣、layernorms等的時延呢？我們硬件上報告的flops是專門針對乘加運算的，因此即使我們可以計算flops，其結果也不對。幸運的是，這只是為了占用內存來進行softmax讀/寫，因為這有利于帶寬和flops比率。這是已知的時延因素！

在這里我將拋開第一性原則，討論?Data Movement Is All You Need一文中的表格A.1。我們發現softmax的時延時間略長于qkv的計算時間（softmax時延時間是qkv操作時間的三倍）。這有點令人擔憂，可能會影響整個神經網絡的性能。

出于同樣的原因，softmax會受到內存限制，所以qk、ReLU和dropout的乘法（multiplication）操作也相當昂貴。

GPU內核融合（Fusion）

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

通過計算所需的讀寫次數，我們發現這里的softmax沒有完美融合。理論上它可以是一次讀取和一次寫入（標準次數為四次）。qk是兩次讀取和一次寫入（兩次讀取可能可以保存）。三比一的比率表示softmax執行的內存傳遞量多于最佳值。我這樣說是因為，這表明了計算的軟件依賴程度，并且需要通過實驗來驗證，因為從理論上來說成本可能為0。

值得注意的是，隨著模型大小的增加，操作所花費的時間百分比會迅速降低，因為每層內存將以增加，flops將以 增加。本文為336M參數模型，。

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

這些memory bound中間操作的持續時間將延長8倍，因為這些操作是長度的向量。但是，flops次數將增加64倍，這意味著flop時間會延長64倍。

因此，若不考慮時延，使用?Data Movement Is All You Need中的優化方法，520億參數模型的推理時延大約會是中間計算的 5%。

8

實際基準對比

我在語言建模公司工作，我們公司有自己的基礎設施和基準，但IP是公司面臨的一大難題。可悲的是缺乏可用于模型并行推理的公共基準。目前我只知道Nvidia的FasterTransformer和Microsoft的Deepspeed，可能其他我不了解的論文也提出了一些基準。無論如何，我們可以根據真實基準來驗證計算！

因為我只想用2個GPU，所以使用了FasterTransformer來運行了一個130億參數的模型。該模型執行一連串內核融合，并提供張量并行功能。模型有40層，每層有40個頭，每個頭的維度為128，總維度大小為5120。這里是配置文件截圖，這些截圖顯示了許多有趣的細節，可能值得單獨撰寫一篇文章。

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

在這里我應該將內存寬帶（mem bandwidth）設置為1.555，而不是1.5。

實測結果表明，1個GPU應該是22.0ms，這意味著，我們的猜測的準確率為76%。可以確定，其中一部分時間花在了中間激活操作上，從理論上來說，我們可以獲得100%的內存帶寬，但實際上并沒有。

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

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

實驗結果顯示，使用2個GPU時，總時間為13.5毫秒。相比一個GPU，這次我們只占了內存寬帶的76%，離目標還有很大的差距。為了解決這個問題，我們重新檢查了配置文件，發現內存帶寬稍微差了一些，因為張量較小獲取的帶寬也比較少。

經過計算，寬帶沒有達到90%，只有87%，總時間為9.5毫秒，中間的激活時間需要大約2毫秒，這使得總時間為11.7毫秒。剩余的1.5毫秒需要找出通信問題。但是這個問題很容易解決，因為我們之前計算的1毫秒通信時間沒有被并行化。根據配置文件的數據，每個層的通信時間為40-50微秒，總共約為1.7毫秒，這就是很好的證明。

上述兩種操作的中間激活計數都比應有的要高一些，因為配置文件提供的時延始終略高于原始基準測試運行。基準測試運行的輸出為180.86ms（上下文時間：45.45ms）和 283.60ms（上下文時間：63.17ms）。

在前向傳遞過程中，模型需要將所有token發送到每個GPU上，然后每個GPU都會對其進行自己的注意力頭計算并存儲kv。由于這個過程需要發送大量的數據并進行并行計算，因此預計前向傳遞需要的時間會比解碼步驟長num_tokens/flops_to_bw_ratio倍。

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

單個GPU只缺失了部分kv儲存時間。查看配置文件，我們發現kv儲存時間每層為18微秒，總共為0.7毫秒。Memsets占了0.2毫秒。我們預計其中一個MLP乘法的flop時間（flop bound）為512x4x5120^2x2/312e12 = 344微秒。實際上，最低浮點運算時間應該是476微秒，也就是說，我們得到了預期flops的72%。

對于attention中的投影（projection），我們期望是512x5120^2x2/312e12 =86微秒。但是在配置文件中，我們發現最低是159微秒。請參閱本文的圖14，其中512x4000x4000的最終結果低于150TFLOPs/s。

9

練習

1.在給定batch size、上下文長度和next_n的情況下，我們該如何計算kv緩存節約的寬帶？

2.kv緩存會增加內存時間開銷嗎？

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

4.在GPU超過容量所需的情況下，我們應該進行怎樣的權衡和計算？例如，一個520億參數大小的模型（擁有8或16個GPU）。

5.如果我們有計算預測一個token時間的公式。我們應該如何計算執行整個樣本的時間呢? 是否應該先在上下文上進行前向傳遞，再預測所有的請求token？

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

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

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

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

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

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

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

其他人都在看

• “ChatGPT們”的淘金時代

• GPT-4，大增長時代的序幕

• GPT-4創造者：第二次改變AI浪潮的方向

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

• 比快更快，開源Stable Diffusion刷新作圖速度

• OneEmbedding:單卡訓練TB級推薦模型不是夢

• GLM訓練加速：性能最高提升3倍，顯存節省1/3

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

總結

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

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。