每個人都想讓模型訓練得更快，但是你真的找對方法了嗎？在康奈爾大學本科生、曾在 PyTorch 團隊實習的 Horace He 看來，這個問題應該分幾步解決：首先，你要知道為什么你的訓練會慢，也就是說瓶頸在哪兒，其次才是尋找對應的解決辦法。在沒有了解基本原理（第一性原理）之前就胡亂嘗試是一種浪費時間的行為。

在這篇文章中，Horace He 從三個角度分析可能存在的瓶頸：計算、內存帶寬和額外開銷，并提供了一些方式去判斷當前處于哪一個瓶頸，有助于我們更加有針對性地加速系統。這篇文章得到了陳天奇等多位資深研究者、開發者的贊賞。

以下是原文內容：

怎樣才能提高深度學習模型的性能？一般人都會選擇網上博客中總結的一些隨機技巧，比如「使用系統內置的運算算子，把梯度設置為 0，使用 PyTorch1.10.0 版本而不是 1.10.1 版本……」

在這一領域，當代（特別是深度學習）系統給人的感覺不像是科學，反而更像煉丹，因此不難理解用戶為什么傾向于采用這種隨機的方法。即便如此，這一領域也有些第一性原理可以遵循，我們可以據此排除大量方法，從而使得問題更加容易解決。

比如，如果你的訓練損失遠低于測試損失，那么你可能遇到了「過擬合」問題，而嘗試著增加模型容量就是在浪費時間。再比如，如果你的訓練損失和你的驗證損失是一致的，那對模型正則化就顯得不明智了。

類似地，你也可以把高效深度學習的問題劃分為以下三個不同的組成部分：

1. 計算：GPU 計算實際浮點運算（FLOPS）所花費的時間；
2. 內存：在 GPU 內傳輸張量所花費的時間；
3. 額外開銷：花在其它部分的時間。

在訓練機器學習模型的時候，知道你遇到的是哪類問題非常關鍵，使模型高效的問題也是如此。例如，當模型花費大量時間進行內存到 GPU 的轉移的時候（也就是內存帶寬緊張的時候），增加 GPU 的 FLOPS 就不管用。另一方面，如果你正在運行大量的矩陣乘法運算（也就是計算緊張的時候），將你的程序重寫成 C++ 去減輕額外開銷就不會管用。

所以，如果你想讓 GPU 絲滑運行，以上三個方面的討論和研究就是必不可少的。

慘痛教訓的背后有大量工程師保持 GPU 高效運行。

注意：這個博客中的大多數內容是基于 GPU 和 PyTorch 舉例子的，但這些原則基本是跨硬件和跨框架通用的。

計算

優化深度學習系統的一個方面在于我們想要最大化用于計算的時間。你花錢買了 312 萬億次浮點數運算，那你肯定希望這些都能用到計算上。但是，為了讓你的錢從你昂貴的矩陣乘法中得到回報，你需要減少花費在其他部分的時間。

但為什么這里的重點是最大化計算，而不是最大化內存的帶寬？原因很簡單 —— 你可以減少額外開銷或者內存消耗，但如果不去改變真正的運算，你幾乎無法減少計算量。

與內存帶寬相比，計算的增長速度增加了最大化計算利用率的難度。下表顯示了 CPU 的 FLOPS 翻倍和內存帶寬翻倍的時間 （重點關注黃色一欄）。

一種理解計算的方式是把它想象成工廠。我們把指令傳達給我們的工廠（額外消耗），把原始材料送給它（內存帶寬），所有這些都是為了讓工廠運行得更加高效（計算）。

所以，如果工廠容量擴展的速度高于我們提供給它原材料的速度，它就很難達到一個頂峰效率。

即使我們工廠容量（FLOP）翻倍，但帶寬跟不上，我們的性能也不能翻倍。

關于 FLOPS 還有一點要說，越來越多的機器學習加速器都有專門針對矩陣乘法的硬件配置，例如英偉達的「Tensor Cores」。

所以，你要是不做矩陣乘法的話，你只能達到 19.5 萬億次運算，而不是 312 萬億次。注意，并不是只有 GPU 這么特殊，事實上 TPU 是比 GPU 更加專門化的計算模塊。

除了矩陣乘法以外，GPU 處理其他運算時都比較慢，這一現象乍看上去似乎有問題：比如像是層歸一化或者激活函數的其它算子怎么辦呢？事實上，這些算子在 FLOPS 上僅僅像是矩陣乘法的舍入誤差一樣。例如，看看下表對于 BERT 中的不同算子類型占用的 FLOP 數，其中的「Tensor Contraction」就是指矩陣乘法。

可以看到，非矩陣乘法運算僅僅占所有運算的 0.2%，所以即使它們的速度僅為矩陣乘法的 1/15 也沒什么問題。

事實上，歸一化運算和逐點（pointwise）運算使用的 FLOPS 僅為矩陣乘法的 1/250 和 1/700。那為什么非矩陣乘法運算會遠比它們應該使用的運行時間更多呢？

回到前文「工廠」的類比，罪魁禍首經常還是如何將原始材料運到以及運出工廠，換句話說，也就是「內存帶寬」。

帶寬

帶寬消耗本質上是把數據從一個地方運送到另一個地方的花費，這可能是指把數據從 CPU 移動到 GPU，從一個節點移動到另一個節點，甚至從 CUDA 的全局內存移動到 CUDA 的共享內存。最后一個是本文討論的重點，我們一般稱其為「帶寬消耗」或者「內存帶寬消耗」。前兩者一般叫「數據運輸消耗」或者「網絡消耗」，不在本文敘述范圍之內。

還是回到「工廠」的類比。雖然我們在工廠中從事實際的工作，但它并不適合大規模的存儲。我們要保證它的存儲是足夠高效的，并且能夠很快去使用（SRAM），而不是以量取勝。

那么我們在哪里存儲實際的結果和「原材料」呢？一般我們要有一個倉庫，那兒的地足夠便宜，并且有大量的空間（DRAM）。之后我們就可以在它和工廠之間運送東西了（內存帶寬）。

這種在計算單元之間移動東西的成本就是所謂的「內存帶寬」成本。事實上，nvidia-smi 命令中出現的那個「內存」就是 DRAM，而經常讓人抓狂的「CUDA out of memory」說的就是這個 DRAM。

值得注意的是：我們每執行一次 GPU 核運算都需要把數據運出和運回到我們的倉庫 ——DRAM。

現在想象一下，當我們執行一個一元運算（如 torch.cos）的時候，我們需要把數據從倉庫（DRAM）運送到工廠（SRAM），然后在工廠中執行一小步計算，之后再把結果運送回倉庫。運輸是相當耗時的，這種情況下，我們幾乎把所有的時間都花在了運輸數據，而不是真正的計算上。

因為我們正把所有的時間都花費在內存帶寬上，這種運算也被稱作內存限制運算（memory-bound operation），它意味著我們沒有把大量時間花費在計算上。

顯然，這并不是我們想要的。那我們能做什么呢？讓我們來看看算子序列長什么樣子。

一個逐點算子序列可能的樣子。

在全局內存和計算單元之間來回傳輸數據的做法顯然不是最佳的。一種更優的方式是：在數據工廠中一次性執行完全部運算再把數據傳回。

這就是算子融合（operator fusion）—— 深度學習編譯器中最重要的優化。簡單地說，這種方法不會為了再次讀取而將數據寫入全局內存，而是通過一次執行多個計算來避免額外的內存訪問。

例如，執行 x.cos ().cos () 運算，寫入內存的方式需要 4 次全局讀寫。

x1 = x.cos() # Read from x in global memory, write to x1
x2 = x1.cos() # Read from x1 in global memory, write to x2

而算子融合只需要 2 次全局內存讀寫，這樣就實現了 2 倍加速。

x2 = x.cos().cos() # Read from x in global memory, write to x2

但是這種做法也并不容易，需要一些條件。首先，GPU 需要知道執行完當前運算后下一步會發生什么，因此無法在 PyTorch 的 Eager 模式（一次運行一個運算符）下進行此優化。其次，我們需要編寫 CUDA 代碼，這也不是一件簡單的事。

并不是所有的算子融合都像逐點算子那樣簡單。你可以將逐點算子融合到歸約（reduction）或矩陣乘法上。甚至矩陣乘法本身也可以被認為是一種融合了廣播乘法（broadcasting multiply）和歸約的運算。

任何 2 個 PyTorch 算子都可以被融合，從而節省了讀取 / 寫入全局內存的內存帶寬成本。此外，許多現有編譯器通?？梢詧绦小负唵巍沟娜诤希ɡ?NVFuser 和 XLA）。然而，更復雜的融合仍然需要人們手動編寫，因此如果你想嘗試自己編寫自定義 CUDA 內核，Triton 是一個很好的起點。

令人驚訝的是，融合后的 x.cos ().cos () 運算將花費幾乎與單獨調用 x.cos () 相同的時間。這就是為什么激活函數的成本幾乎是一樣的，盡管 gelu 顯然比 relu 包含更多的運算。

因此，重新實現 / 激活檢查點會產生一些有趣的結果。從本質上講，進行額外的重新計算可能會導致更少的內存帶寬，從而減少運行時間。因此，我們可以通過重新實現來減少內存占用和運行時間，并在 AOTAutograd 中構建一個簡潔的 min-cut 優化通道。

推理內存帶寬成本

對于簡單的運算，直接推理內存帶寬是可行的。例如，A100 具有 1.5 TB / 秒的全局內存帶寬，可以執行 19.5 teraflops / 秒的計算。因此，如果使用 32 位浮點數（即 4 字節），你可以在 GPU 執行 20 萬億次運算的同時加載 4000 億個數字。

此外，執行簡單的一元運算（例如將張量 x2）實際上需要將張量寫回全局內存。

因此直到執行大約一百個一元運算之前，更多的時間是花在了內存訪問而不是實際計算上。

如果你執行下面這個 PyTorch 函數：

def f(x: Tensor[N]):
   for _ in range(repeat):
       x = x * 2
   return x

并使用融合編譯器對其進行基準測試，就可以計算每個 repeat 值的 FLOPS 和內存帶寬。增大 repeat 值是在不增加內存訪問的情況下增加計算量的簡單方法 - 這也稱為增加計算強度 （compute intensity)。

具體來說，假設我們對這段代碼進行基準測試，首先要找出每秒執行的迭代次數；然后執行 2N（N 是張量大小）次內存訪問和 N *repeat FLOP。因此，內存帶寬將是 bytes_per_elem * 2 * N /itrs_per_second，而 FLOPS 是 N * repeat /itrs_per_second。

現在，讓我們繪制計算強度的 3 個函數圖象：運行時間、flops 和內存帶寬。 

請注意，在執行 64 次乘法之前，運行時間根本不會顯著增加。這意味著在此之前主要受內存帶寬的限制，而計算大多處于空閑狀態。

一開始 FLOPS 的值是 0.2 teraflops。當我們將計算強度加倍時，這個數字會線性增長，直到接近 9.75 teraflops 的峰值，一旦接近峰值 teraflops 就被認為是「計算受限的」。

最后，可以看到內存帶寬從峰值附近開始，隨著我們增加計算強度開始下降。這正是我們所期待的，因為這說明執行實際計算的時間越來越多，而不是訪問內存。

在這種情況下，很容易看出何時受計算限制以及何時受內存限制。repeat< 32 時，內存帶寬接近飽和，而未進行充分的計算；repeat> 64 時，計算接近飽和（即接近峰值 FLOPS），而內存帶寬開始下降。

對于較大的系統，通常很難說是受計算限制還是內存帶寬限制，因為它們通常包含計算限制和內存限制兩方面的綜合原因。衡量計算受限程度的一種常用方法是計算實際 FLOPS 與峰值 FLOPS 的百分比。 

然而，除了內存帶寬成本之外，還有一件事可能會導致 GPU 無法絲滑運行。

額外開銷

當代碼把時間花費在傳輸張量或計算之外的其他事情上時，額外開銷（overhead）就產生了，例如在 Python 解釋器中花費的時間、在 PyTorch 框架上花費的時間、啟動 CUDA 內核（但不執行）所花費的時間， 這些都是間接開銷。

額外開銷顯得重要的原因是現代 GPU 的運算速度非常快。A100 每秒可以執行 312 萬億次浮點運算（312TeraFLOPS）。相比之下 Python 實在是太慢了 ——Python 在一秒內約執行 3200 萬次加法。

這意味著 Python 執行單次 FLOP 的時間，A100 可能已經運行了 975 萬次 FLOPS。

更糟糕的是，Python 解釋器甚至不是唯一的間接開銷來源，像 PyTorch 這樣的框架到達 actual kernel 之前也有很多層調度。PyTorch 每秒大約能執行 28 萬次運算。如果使用微型張量（例如用于科學計算），你可能會發現 PyTorch 與 C++ 相比非常慢。

例如在下圖中，使用 PyTorch 執行單次添加，僅有一小塊圖是實際執行計算的內容，其他的部分都是純粹的額外開銷。

鑒于此，你可能會對 PyTorch 成為主流框架的現象感到不解，而這是因為現代深度學習模型通常執行大規模運算。此外，像 PyTorch 這樣的框架是異步執行的。因此，大部分框架開銷可以完全忽略。

如果我們的 GPU 算子足夠大，那么 CPU 可以跑在 GPU 之前（因此 CPU 開銷是無關緊要的）。另一方面，如果 GPU 算子太小，那么 GPU 將在 paperweight 上花費大部分時間。

那么，如何判斷你是否處于這個問題中？由于額外開銷通常不會隨著問題的規模變化而變化（而計算和內存會），所以最簡單的判斷方法是簡單地增加數據的大小。如果運行時間不是按比例增加，應該可以說遇到了開銷限制。例如，如果將批大小翻倍，但運行時間僅增加 10%，則可能會受到開銷限制。

另一種方法是使用 PyTorch 分析器。如下圖，粉紅色塊顯示了 CPU 內核與 GPU 內核的匹配情況。

CPU 運行地比 GPU 更超前。

另一方面，nvidia-smi 中的「GPU-Util」（不是「Volatile GPU-Util」）入口會測量實際運行的 GPU 內核的百分占比，所以這是另一種觀察是否遇到開銷限制的好方法。這種開銷是 PyTorch 等所有靈活的框架所具有的，本質上都需要花費大量時間來「弄清楚要做什么」。

這可能來自 Python（查找屬性或調度到正確的函數）或 PyTorch 中的代碼。例如，當你執行 a + b 時，需要執行以下步驟：

1. Python 需要在 a 上查找__add__調度到的內容。
2. PyTorch 需要確定張量的很多屬性（比如 dtype、device、是否需要 autograd）來決定調用哪個內核。
3. PyTorch 需要實際啟動內核。

從根本上說，這種開銷來自能夠在每個步驟中執行不同運算的靈活性。如果不需要這種靈活性，解決這種靈活性的一種方法是跟蹤它，例如使用 jit.trace、FX 或 jax.jit?；蛘?，可以換用 CUDA Graphs 之類的東西在更低的級別上執行此運算。

不幸的是，這是以失去靈活性為代價的。一種兩全其美的方法是，通過在 VM 級別進行 introspect 來編寫更多符合「真實」的 JIT 的內容。有關更多信息，可參閱 TorchDynamo (https://dev-discuss.pytorch.org/t/torchdynamo-an-experiment-in-dynamic-python-bytecode-transformation/361)。

總結

如果你想加速深度學習系統，最重要的是了解模型中的瓶頸是什么，因為瓶頸決定了適合加速該系統的方法是什么。

很多時候，我看到研究人員和其他對加速 PyTorch 代碼感興趣的人，會在不了解所處問題的情況下盲目嘗試。

當然，另一方面，如果用戶需要考慮這些東西，也反映了框架的部分失敗。盡管 PyTorch 是一個活躍的關注領域，但 PyTorch 的編譯器或配置文件 API 并不是最容易使用的。

總而言之，我發現對系統基本原理的理解幾乎總是有用的，希望這對你也有用。