2024年05月07日08:31:13 科技 5856

在本系列文章《》中，作者對transformer解碼器的文本生成算法進行了高層次概述，着重介紹了兩個階段：提示的處理階段和逐個生成補全詞元的多步生成階段。在上一篇文章《》中，深入探討了kv cache優化。

本文將轉變方向，探索可能影響機器學習模型速度的不同性能瓶頸。本文所詳細介紹的概念廣泛適用於任何ml模型，無論是用於訓練還是推理，不過提供的示例主要聚焦於llm推理設置。

作者 | pierre lienhart

oneflow編譯

翻譯｜張雪聃

四種性能瓶頸

如果你對模型性能不太滿意，並且打算繼續改進，那麼第一步就是確定性能瓶頸的類型。瓶頸主要有四類——其中三類與硬件限制相關，另一類與軟件相關。

我們先來看看硬件瓶頸。每一個都有其對應的特定操作模式。

計算受限模式：大部分處理時間（即時延）被用於執行算術運算（見圖1）。與其他模式相比，由於主要成本是計算，計算受限模式是最具成本效益的，因此我們應該以此為目標。

圖1：計算受限過程圖（計算和數據傳輸時間分別用黃色和藍色表示）

內存帶寬受限模式：大部分處理時間花在了芯片內存和處理器之間搬運數據，如權重矩陣、中間計算等（見圖2）。

圖2：內存帶寬受限過程圖（計算和數據傳輸時間分別用黃色和藍色表示）

通信受限模式（引文[1]中未涵蓋）：僅適用於計算和數據分布在多個芯片之間的情況。大部分處理時間花在了芯片之間的網絡數據傳輸（見圖3）。

圖3：通信受限過程圖（計算、數據傳輸和網絡通信時間分別用黃色、藍色和綠色表示）

注意：我用“芯片”這個詞是因為這些概念適用於任何類型的芯片：cpu、gpu、定製化芯片（如google的tpu、aws的neuron cores等）。

需注意的是，現代硬件和框架經過了高度優化且計算和數據傳輸任務常常出現部分重疊的現象（見圖4）。簡單起見，我們在本文中將繼續假設它們按順序進行。

圖4：出現重疊數據傳輸的通信受限過程

最後一種模式稱為計算開銷受限，與軟件引起的限制相關。在這種模式下，大部分處理時間花在了調度工作並將其提交給硬件上 —— 基本上，我們花更多的時間去弄清楚要做什麼，而不是在硬件上執行實際操作（見圖5）。計算開銷受限的情況在使用非常靈活的語言（例如python）或框架（例如pytorch）時更有可能出現，這些語言或框架不需要在運行時明確指定全部所需的信息（如張量數據類型、目標設備、要調用的kernel等）。這些缺失的信息必須在運行時推理出來，相應的cpu周期稱為計算開銷。現在的加速硬件與cpu相比速度更快，以至於計算開銷影響硬件利用率，進而影響成本效益的情況很可能發生 —— 基本上，有時硬件處於空閑狀態，等待軟件提交下一個工作項。

圖5：計算開銷受限過程圖（計算、數據傳輸和軟件開銷時間分別用黃色、藍色和紫色表示）

執行模型的反向或正向傳播涉及運行多kernel的執行（~ gpu函數調用）。所有kernel都在相同的模式下運行是不可能實現的。問題的關鍵在於確定大部分執行時間都花在哪個模式上。由此，優先事項就變成針對主要瓶頸進行優化，找出下一個最重要的瓶頸，依此類推。

正確地判斷瓶頸類型至關重要。每個問題需要不同的解決方案。如果判斷錯誤，可能會浪費不少時間，即使實施這種優化確實有所幫助，但結果可能依舊令人失望。

判斷限制因素

這篇文章中我們不會深入討論細節問題，但博文[1]強調了當計算開銷受限制時，運行時不會隨着增加的計算或數據傳輸成比例地擴展。換句話說，如果你將計算或數據傳輸容量加倍，但運行時未相應增加，那麼你的程序很可能受到計算開銷限制。否則，你可能受到硬件限制，但區分計算與內存帶寬等瓶頸需要訪問諸如flop計數和數據傳輸量等指標，即使用性能分析器（profiler）。

當我們重新把llm納入討論範圍時，請記住，訓練和推理的預填充階段通常是計算受限，而推理解碼階段通常在大多數硬件上是內存帶寬受限。因此，主要用於訓練的優化（例如低精度矩陣乘法）如果應用於減少由解碼時延為主的總推理時延，可能就不會那麼有幫助。

基於瓶頸類型優化以降低時延

讓我們看看針對每種瓶頸類型的優化。

如果處於計算受限模式下，你可以：

升級到更強大、更昂貴、峰值flops更高的芯片。
針對特定操作，如矩陣乘法，可以使用專門的、更快的核心，如nvidia tensor cores。例如，nvidia h100 pcie [2]使用通用cuda核心的峰值計算性能為51 tflops，而使用專門的tensor cores（全精度）可達到378 tflops。
減少所需的操作數。更具體地說，對於ml模型，這可能意味着，使用更少的參數來實現相同的結果。剪枝或知識蒸餾等技術可以幫助實現這一點。
使用較低精度和更快速的數據類型進行計算。例如，對於相同的nvidia h100 pcie，8位tensor cores的峰值flops（1,513 tflops）是16位峰值flops（756 tflops）的兩倍，是32位峰值flops（378 tflops）的兩倍。但這需要量化所有的輸入（例如，權重矩陣和激活函數），並使用專門的低精度kernel。

如果處於內存帶寬受限模式下，你可以：

升級到更強大、更昂貴、有更高的內存帶寬的芯片
減少所搬運數據的大小，例如使用模型壓縮技術，如量化或不太流行的剪枝和知識蒸餾。就llm而言，數據大小問題主要通過僅包含權重的量化技術（例如gtpq [5] 和awq [6] 量化算法）與kv cache量化來解決。
減少內存操作次數。gpu上的運行任務可以歸結為執行一系列kernel（gpu函數調用）的有向圖。對於每個kernel，必須從內存中獲取輸入並將輸出寫入其中。kernel融合，即將最初分散在多個kernel中的操作作為單個kernel調用執行，可以減少內存操作次數。算子融合（見圖6）可以由編譯器自動執行，也可以通過編寫自己的kernel來手動執行（這更難，但對於複雜的融合是必要的）。

對於transformer模型而言，針對注意力層開發高效融合的kernel仍然是一個十分活躍的領域。許多優化的kernel都基於流行flashattention算法[7]。transformer融合kernel庫包括flashattention、meta的xformers和現已棄用的nvidia fastertransformer（已合併到nvidia tensorrt-llm中）。

圖6：應用於cnn的水平和垂直層（操作）融合的示例，初始狀態（上方）和最終狀態（下方）[8]

如果處於通信受限模式下，你可以：

升級到更強大、更昂貴、具有更高的網絡帶寬的芯片
通過選擇更高效的分區和集合通信策略來減少通信量。例如，通過引入新的張量並行策略，[9]擴展了來自[10]中用於transformer模型的流行張量並行布局，使通信時間能夠更好地擴展（即防止它們成為瓶頸），以適應大量芯片或批大小。

例如，[10]中的張量並行布局保持權重分片不變，而激活分片在芯片之間移動。例如，在預填充階段和對於非常大的序列批次，例如，[9]指出激活可能比權重更大。因此，從通信的角度來看，保持激活不變並將權重分片移動會更有效，就像他們所採用的“權重聚合（weight-gathered）”分區策略那樣。

如果處於計算開銷受限模式下，你可以：

通過使用不太靈活但更高效的語言，如c++，來換取更少的計算開銷。
將kernel分組提交以分攤跨多個kernel提交的計算開銷，而不是按每個kernel都算一次計算開銷。
這在需要多次提交相同的“短命（short-lived）”kernel組時特別有益（例如在迭代工作負載中）。cuda graph（自pytorch 1.10版以來已集成）通過提供工具捕獲由代碼塊引起的所有gpu活動，形成一個kernel啟動的有向圖，實現一次性提交，從而服務於這一目的。

你可以通過使用模型編譯器進一步優化圖。

在任何情況下，你再次通過犧牲靈活性來減少計算開銷，因為追蹤/編譯需要參數（如張量大小、類型等）是靜態的，因此運行時（runtime）保持不變。控制流結構，如if-else，在這一過程中通常也會丟失。

在對靈活性有需求但同時與aot編譯不兼容的情況中（例如，動態張量形狀、控制流等），即時（jit）編譯器能通過在模型代碼在執行之前進行動態優化來提供幫助（雖然不如aot編譯器徹底）。例如，pytorch 2.x 引入了一個名為torchdynamo的jit編譯器。由於無需修改你的pytorch程序即可使用它，因此你可以使用jit編譯器減少計算開銷，同時保持python開發體驗的易用性。

附註：模型優化器和（aot）編譯器之間是否有區別？在我看來，這個區分比較模糊。我在概念上對這兩個術語的區分如下。

首先，它們都是提前執行的。典型的aot編譯器工作流程是：跟蹤來自支持框架（如pytorch、tensorflow、mxnet等）的代碼，將計算圖提取到中間表示（ir）中，應用與硬件無關的優化（如代數重寫、循環展開loop unrolling、算子融合等）生成優化後的圖形，最後為目標硬件創建可部署的工件，包括硬件感知優化（選擇最合適的內核、數據移動優化等）。一些aot模型編譯器的示例包括pytorch的torchinductor、xla和meta的glow。

模型優化器是一種工具，包含aot編譯，但通常針對特定硬件（例如，openvino針對英特爾硬件，tensorrt和tensorrt-llm針對nvidia硬件），能夠執行額外的後訓練優化，如量化或剪枝。

到現在，我們只聚焦於時延（處理單個請求所需的時間），接下來讓我們通過深入探討計算和內存帶寬受限模式，重新引入吞吐量（單位時間內可以處理的請求數量）。

瓶頸 = f（硬件，算術強度）

有趣的是，處理相同輸入的相同算法可以是計算受限模式，也可以是內存帶寬受限，這取決於所使用的硬件。適用的模式由算法的算術強度確定，即每個訪問的字節執行的算術操作數量。

我們想要的是能使我們處於或更具成本效益的計算受限模式的強度值。正如我們將看到的那樣，較高的強度與更好的吞吐量和成本效率相關。然而，一些強度驅動因素可能會降低延遲。我們不得不在時延和吞吐量之間進行權衡。

設b為每次運行從內存傳輸的數據字節數，p為每次運行執行的浮點運算次數（flops）。設bw_mem（以tb/s表示）為硬件的內存帶寬，bw_math（以tflops表示）為數學帶寬，也稱為峰值flops。設 t_mem 為移動數據字節所花費的時間，t_math 為算術運算所花費的時間。

處於計算受限模式意味着在算術運算上花費的時間比傳輸數據多（見圖7）。

圖7：計算受限與內存帶寬受限模式對比圖（計算和數據傳輸時間分別用黃色和藍色表示）

因此，當以下條件滿足時，我們處於計算受限狀態：

a是算法的算術強度，其維度為flop/byte。每字節傳輸的數據進行的算術運算越多，算術強度就越高。

如方程所示，要使算法處於計算受限狀態，其算術強度必須超過峰值flops與內存帶寬的比率有關的硬件相關比值。相反，內存帶寬受限意味着操作強度低於相同帶寬比率（見圖8）。

圖8：內存帶寬 / 計算受限邊界

讓我們來看一些關於在nvidia硬件上進行半精度矩陣乘法（即使用張量核心）的實際數據（表1）：

表1：用於訓練或為llm提供服務的常用nvidia數據中心gpu的規格

這意味着什麼？以nvidia a10為例，帶寬比為208意味着在該特定硬件上搬運一字節的數據與執行208 flops的速度相同。因此，如果在nvidia a10上運行的算法每傳輸一字節不能執行至少208次flops（或等效地每傳輸半精度數執行416次flops），那麼它必然會花費更多的時間在數據傳輸上而不是在計算上，即它是內存帶寬受限。換句話說，算法的算術強度低於硬件帶寬比率的情況下，就是內存帶寬受限。證明完畢。

已知llm推理過程的解碼階段具有較低的算術強度（在下一篇博文中詳細介紹），這導致它在大多數硬件上成為內存帶寬受限。與nvidia h100相比，nvidia h200針對這種低強度工作負載具有更有利的帶寬比率。這解釋了nvidia將h200宣傳為“為生成式ai推理提速”的原因，因為其硬件設計針對的是這種內存受限情況。

現在我們將算術強度與時延和吞吐量聯繫起來：

注意：這裡的吞吐量表示為tflops而不是每秒請求數，但兩者是直接成比例的。另外，吞吐量以tflops表示突顯了它與硬件利用率的聯繫，因此與成本效率有關。為了使這種聯繫更明顯，吞吐量更精確地表示為每個芯片-秒（chip-second）的請求數量，每個請求的芯片秒數越低，即吞吐量越高，成本效率越高（參見[9]第4節）。

如果我們將算術強度作為自變量繪製在x軸上，將（最大可實現的）吞吐量作為因變量繪製在y軸上，我們就得到了所謂的（簡單）屋頂線模型[12]（圖9）。

圖9：屋頂線模型圖

讓我們進行一個小小的思維實驗，以更好地理解為什麼此圖中的吞吐量值是最大可實現的水平。顯而易見，在計算受限模式下：沒有什麼可以阻止我們利用全部計算能力，我們只受到硬件峰值容量的限制。在內存帶寬受限模式下，我們在1秒內可以獲取的最大數據量由硬件的內存帶寬bw_mem確定。考慮到算法的算術強度a，因此我們在1秒內可以實現的最大flop數量是bw_mem乘以a。

增加算法的算術強度會產生什麼影響？我們可以考慮三種情況（圖10）：

圖10：算術強度增加的三種情況

情景1：算術強度的增加太小，無法擺脫內存帶寬受限的狀態，但吞吐量線性增長。系統仍然受到內存帶寬的限制，因此對時延的影響取決於更大強度如何影響該特定算法的數據傳輸。
情景2：算術強度的增加使系統轉換到計算受限模式。吞吐量增加到硬件的峰值吞吐量。現在處於計算受限狀態，時延的影響取決於更高強度如何改變該算法的總操作量。

情景3：由於已經處於計算受限狀態並且達到了峰值吞吐量，增加強度不會提供吞吐量增益。時延的影響仍取決於更高強度如何影響該算法的總計算量。

如何具體增加算術強度？這完全取決於算法的具體情況。在下一篇文章中，我們將考查控制transformer解碼器塊算術強度的關鍵參數。例如，我們會看到如何通過提高批大小來增加某些操作的算術強度。

一些已經討論過的優化方法也可以增加算術強度，從而提高吞吐量和資源利用率。對於transformer模型（其解碼階段受內存帶寬限制），通過操作融合和數據（權重矩陣、kv cache）量化來減少數據傳輸的數量和大小，來改善算術強度。

目前為止，我們進行了一個關鍵的假設，即算法的實現充分利用了硬件資源。例如，在傳輸數據時，假設算法的實現使用了硬件的理論內存帶寬的100%。顯然，在實踐中並非總是如此（儘管一些實現達到了接近最佳資源利用率），那麼次優的資源使用對分析有何影響呢？

很簡單：上述帶寬數字必須被實際實現的數字所替代。次優系統位於它自己的屋頂線曲線上，位於最優化的屋頂線曲線下方（見圖11）。現在有兩個提高吞吐量的自由度：增加算術強度和/或改善算法的實現以更好地利用硬件資源。

圖11：具有次優資源利用的屋頂線模型

讓我們通過一個實際的改進實現的示例來總結。在版本2.2之前，flashattention kernel的實現在應用於推理的解碼階段時可能會變得非常不理想。先前的數據加載實現使得kernel在解碼階段利用內存帶寬方面的效率較低。更糟的是，隨着批大小的增加，帶寬利用實際上進一步降低；因此，對於由於內存限制需要較小批次的長序列，性能受到的影響最大。flashattention團隊通過解決這個問題（主要是通過在kv cache序列長度維度上並行加載數據）發布了一個名為flashdecoding的優化解碼階段kernel，為長序列長度實現了顯著的時延改進。[13]

總結

通過本文，我們了解了影響模型時延的四種瓶頸類型。識別模型時延主要成因的類別至關重要，因為每種類別都需要特定的緩解策略。

不考慮分布式設置以簡化問題，實際的硬件操作要麼是計算受限，要麼是內存帶寬受限。kernel的算術強度決定了綁定的模式。在較低強度的內存帶寬受限模式下，最大可實現的吞吐量與算術強度呈線性關係。相反，在計算受限模式下，吞吐量受到峰值硬件flops的限制。根據影響強度的因素，我們可能能夠增加強度以提高最大吞吐量，甚至可能達到計算受限的性能。然而，這種強度增益可能會增加時延。

在下一篇文章中，我們將運用這些新知識來研究llm，具體探討transformer解碼器塊的算術強度。

參考文獻（請上下滑動）