選自horace部落格
作者：Horace He
機器之心編譯
編輯：Juniper
深度學習是門玄學？也不完全是。
每個人都想讓模型訓練得更快，但是你真的找對方法了嗎？在康奈爾大學大學生、曾在 PyTorch 團隊實習的 Horace He 看來，這個問題應該分幾步解決：首先，你要知道為什麼你的訓練會慢，也就是說瓶頸在哪兒，其次才是尋找對應的解決辦法。在沒有了解基本原理（第一性原理）之前就胡亂嘗試是一種浪費時間的行為。
在這篇文章中，Horace He 從三個角度分析可能存在的瓶頸：計算、記憶體帶寬和額外開銷，并提供了一些方式去判斷目前處于哪一個瓶頸，有助于我們更加有針對性地加速系統。這篇文章得到了陳天奇等多位資深研究者、開發者的贊賞。
以下是原文内容：
怎樣才能提高深度學習模型的性能？一般人都會選擇網上部落格中總結的一些随機技巧，比如「使用系統内置的運算算子，把梯度設定為 0，使用 PyTorch1.10.0 版本而不是 1.10.1 版本……」
在這一領域，當代（特别是深度學習）系統給人的感覺不像是科學，反而更像煉丹，是以不難了解使用者為什麼傾向于采用這種随機的方法。即便如此，這一領域也有些第一性原理可以遵循，我們可以據此排除大量方法，進而使得問題更加容易解決。
比如，如果你的訓練損失遠低于測試損失，那麼你可能遇到了「過拟合」問題，而嘗試着增加模型容量就是在浪費時間。再比如，如果你的訓練損失和你的驗證損失是一緻的，那對模型正則化就顯得不明智了。
類似地，你也可以把高效深度學習的問題劃分為以下三個不同的組成部分：
計算：GPU 計算實際浮點運算（FLOPS）所花費的時間；
記憶體：在 GPU 内傳輸張量所花費的時間；
額外開銷：花在其它部分的時間。
在訓練機器學習模型的時候，知道你遇到的是哪類問題非常關鍵，使模型高效的問題也是如此。例如，當模型花費大量時間進行記憶體到 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 次全局讀寫。
而算子融合隻需要 2 次全局記憶體讀寫，這樣就實作了 2 倍加速。
但是這種做法也并不容易，需要一些條件。首先，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 函數：
并使用融合編譯器對其進行基準測試，就可以計算每個 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 就被認為是「計算受限的」。
最後，可以看到記憶體帶寬從峰值附近開始，随着我們增加計算強度開始下降。這正是我們所期待的，因為這說明執行實際計算的時間越來越多，而不是通路記憶體。
在這種情況下，很容易看出何時受計算限制以及何時受記憶體限制。repeat64 時，計算接近飽和（即接近峰值 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 時，需要執行以下步驟：
Python 需要在 a 上查找__add__排程到的内容。
PyTorch 需要确定張量的很多屬性（比如 dtype、device、是否需要 autograd）來決定調用哪個核心。
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 并不是最容易使用的。
總而言之，我發現對系統基本原理的了解幾乎總是有用的，希望這對你也有用。