Making Deep Learning Go Brrrr: A First-Principles Guide to GPU Performance
優化深度學習模型的效能往往感覺像是在進行煉金術。開發者經常訴諸於「各種技巧的組合」——切換 PyTorch 版本、將梯度設置為 None，或使用 in-place 操作——卻沒有清楚地理解為什麼這些改變有效。然而，透過從第一原理（first principles）進行推理，我們可以擺脫猜測，系統性地識別出阻礙效能的實際瓶頸。
要優化一個系統，你首先需要知道你處於哪種「模式」（regime）。就像訓練損失與驗證損失可以告訴你是否正在發生過擬合或欠擬合一樣，分析系統的資源利用率可以告訴你時間實際上花在了哪裡。在深度學習中，效率通常分為三個組成部分：Compute（計算）、Memory（記憶體）和 Overhead（開銷）。
The Three Pillars of Performance
1. Compute (The Factory)
計算是指 GPU 花費在執行實際浮點運算（FLOPS）的時間。大多數優化的目標是最大化處於計算受限（compute-bound）模式的時間。你支付了現代 GPU（例如 A100 的 312 TeraFLOPS）龐大的 TFLOPs，你當然希望能夠實際利用它們。
值得注意的是，這些峰值數字通常是指專門的硬體，例如 Tensor Cores，它們是專門為矩陣乘法設計的。如果你的操作不是矩陣乘法，你的效能將僅為該峰值的一小部分。然而，在大多數深度學習模型（如 BERT）中，非矩陣乘法操作（layer norm、activations）僅佔總 FLOPS 的極小部分，這意味著非矩陣乘法操作的低效率通常不是計算的主要瓶頸——但對於記憶體來說卻是。
2. Memory Bandwidth (The Warehouse)
記憶體頻寬是指將數據從一個地方移動到另一個地方的成本——特別是從 GPU 的 DRAM（「倉庫」）到計算單元/SRAM（「工廠」）的成本。
許多操作是**記憶體受限（memory-bound）**的，這意味著 GPU 花在搬運數據的時間比實際計算的時間還要多。一個簡單的一元操作（unary operation）如 torch.cos 就是一個典型的例子：GPU 從 DRAM 讀取數據，執行一個微小的計算，然後將其寫回。計算速度非常快，以至於 GPU 幾乎把所有時間都花在等待記憶體傳輸。
The Power of Operator Fusion
為了對抗記憶體瓶頸，我們使用算子融合（operator fusion）。與其將每個操作的結果寫回全域記憶體，然後在下一步驟中再次讀取它，融合技術將多個操作組合進一個單一的 GPU kernel。
例如，x.cos().cos() 通常需要四次全域記憶體訪問（兩次讀取，兩次寫入）。透過融合，它只需要兩次（一次讀取，一次寫入）。這就是為什麼像 GELU 這樣複雜的激活函數通常與像 ReLU 這樣簡單的函數成本相同；瓶頸在於記憶體訪問，而不是數學運算的數量。
3. Overhead (The Manager)
開銷是指除了計算和記憶體傳輸以外的所有內容。這包括 Python 解釋器、PyTorch 框架的調度邏輯，以及啟動 CUDA kernel 的時間。
現代 GPU 的速度非常快，以至於 Python 變成了巨大的瓶頸。在 Python 執行單個 FLOP 的時間內，A100 可以處理數百萬個。PyTorch 透過**異步（asynchronously）**執行 kernel 來緩解這一點；CPU 會「領先」於 GPU，排隊等待工作，以便 GPU 永不閒置。
然而，如果你的張量（tensors）太小，GPU 完成工作比 CPU 排隊下達下一個任務的速度還要快。在這種模式下，你的 GPU 就變成了一個「昂貴的鎮紙」。
Identifying Your Bottleneck
了解你處於哪種模式決定了你的解決方案。如果你將 batch size 翻倍，而運行時間幾乎沒有增加，你很可能處於開銷受限（overhead-bound）。如果你增加了操作的複雜度，但運行時間保持不變，你很可能處於頻寬受限（memory-bandwidth bound）。
| Performance Regime | Plausible Solutions |
|---|---|
| Overhead-Bound | Tracing (jit.trace, FX), CUDA Graphs, 或移動到像 TorchDynamo 這樣的 JIT 編譯器 |
| Bandwidth-Bound | Operator Fusion (Triton, NVFuser, XLA) |
| Compute-Bound | 利用 Tensor Cores, 升級硬體 |
Synthesis and Critical Perspectives
雖然第一原理法提供了一個清晰的思維模型，但實際應用可能很混亂。正如社群討論中所提到的，效能是很少能直接移植的。一個導出到 ONNX 的模型可能根據它是透過 ONNX Runtime 還是 TensorRT 運行，表現會有所不同，且結果會根據目標硬體和記憶體調優，而有所差異。
此外，「慘痛的教訓」（bitter lesson）指出，雖然人類在算子融合方面的智慧很有價值，but 隨著長期的趨轉向於大規模增加原始 TFLOPs 和頻寬。NVIDIA 的能力使其能夠在計算和互連技術方面保持指數級增長，這確保了硬體繼續推動可能性的邊界，即使軟體層在努力追趕時，硬體仍然是核心驅動力。
最終，目標是提高計算強度（compute intensity）——即計算與記憶體訪問的比例。透過減少開銷和融合算子，我們為 GPU 能夠做它最擅長的事情鋪平了道路：以峰值速度進行大規模矩陣乘法。