返回文章 無法分析的事物，就無法優化。無論您是想從大型語言模型 (LLM) 中榨取更多每秒詞元、縮短推論時間數毫秒，還是僅僅想了解為何您的訓練迴圈運行速度不如規格表所承諾，最終都必須透過效能分析來解決。問題在於，效能分析的學習曲線相當陡峭。
追蹤圖是密密麻麻的彩色矩形牆，事件名稱令人望而生畏。大多數教學都假設您已經會解讀它們。因此，即使我們知道應該進行效能分析，打開追蹤圖仍可能感覺像是一件最好留到以後（或交給別人）的苦差事。本文及其開創的系列，正是我們試圖降低這個門檻的嘗試。這是「PyTorch 效能分析」系列的首篇文章，我們將在此系列中逐步培養解讀分析器追蹤圖的技能，並利用這些技能來推動優化。
我們的計畫是：第一部分（本文）：從最簡單的操作開始，即矩陣乘法後接偏置加法，學習如何解讀分析器回傳的資訊。第二部分：擴展到 nn.Linear 和一個小型多層感知器 (MLP)，利用追蹤圖來激發優化，並一窺底層的 GPU 核心。第三部分：將所有知識應用於使用 Transformer 的大型語言模型。
我們從初學者的角度記錄這段旅程。除了基本的 PyTorch 知識外，沒有其他先決條件。請將其視為一次輕鬆的閱讀，並期待一些「原來如此！」的時刻。本文的結構特意以問題導向：我們打開一個追蹤圖，問「等等，為什麼會這樣？」，然後追尋答案直到豁然開朗。
讀完本文，您應該會了解：如何設定 torch.profiler 以及它實際回傳了什麼，如何解讀分析器表格和追蹤圖（CPU 執行緒、GPU 執行緒以及兩者之間可疑的間隙），從 Python 呼叫一直到 CUDA 核心的事件鏈，當您套用 torch.compile 時，哪些會改變（以及更有趣的是，哪些不會改變）。
在我們開始之前，先提供兩個定義，這將有助於您更好地理解下文：GPU 核心 (GPU kernel) 是一個在 GPU 多個執行緒上平行運行的程式。CPU 負責排程並啟動這些核心。您通常不需要自己編寫 GPU 核心；當您使用 PyTorch 操作時，它會自動轉換為一個或多個在 GPU 上執行任務的核心。
有了這兩個概念，我們就開始提問吧。這是我們用於本文的完整腳本：01_matmul_add.py。我們建議您在單獨的分頁中打開此腳本，並逐步查看程式碼。我們使用 NVIDIA A100-SXM4-80GB GPU 來運行這些腳本。矩陣乘法與加法操作正如 Sara Hooker 博士精闢地指出，就像我們主要由水組成一樣，深度神經網路主要由矩陣乘法組成。
它們是如此基礎，若不以此開始我們的效能分析之旅，那將會是個遺憾。def fn(x, w, b): return torch.add(torch.matmul(x, w), b)矩陣加法與矩陣乘法一同模擬了權重和偏置在神經元中的交互方式。這個「加法」（雙關語）將幫助我們理解它如何為本文後續的編譯鋪路。
為了進行效能分析，我們將使用 torch.profiler 模組。涉及的步驟包括：準備好要分析的程式碼（這裡的 def fn 封裝了矩陣乘法和矩陣加法）註解演算法。雖然這完全是可選的，但我們建議這樣做。record_function 將我們的函數註解為 matmul_add，這將有助於在追蹤圖中輕鬆導航（正如我們稍後會提到的）。
def step(): with torch.profiler.record_function("matmul_add"): return fn(x, w, b)使用 torch.profiler.profile 上下文管理器包裝程式碼with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, # CPU 活動 torch.profiler.ProfilerActivity.CUDA, # GPU 活動 ], ) as prof: # 建議多次運行事件以預熱 GPU for _ in range(5): step() prof.step()匯出分析結果# 分析器表格 prof.key_averages().table(sort_by="cuda_time_total", row_limit=15) # 分析器追蹤圖 prof.export_chrome_trace(trace_path)分析器匯出兩種不同的產物：分析器表格：提供演算法的統計摘要。
它回答了「什麼佔用了最多的時間」。這對於找出熱點非常有幫助。熱點可能是耗時最長的事件、管線的瓶頸，或是被觸發多次的事件。分析器追蹤圖：提供時間執行視圖。它回答了「操作何時以及為何發生」，描繪了 CPU 和 GPU 上發生的活動。這對於我們調查啟動的核心、啟動它們的任何延遲、CPU 和 GPU 活動之間的任何重疊等情況時非常有用。
讓我們透過第一次執行來看看這兩者的實際應用。（這裡有完整的 01_matmul_add.py 腳本）建議在配備 GPU 的機器上運行此腳本。uv run 01_matmul_add.py --size 64如果您運行上述腳本（在 GPU 機器上），您將在 traces/01_matmul_add 資料夾中找到這兩個產物：64_bf16_cold_eager.json 64_bf16_cold_eager.txt圖 1：64 大小矩陣的矩陣乘法加法分析器表格.txt 檔案包含分析器表格。
如圖 1 所示，打開檔案後，您會看到一個大表格，第一欄列出了在分析範圍內觸發的事件。其他欄位則與事件在 CPU、GPU 或 torch.profiler.profile 中指定活動的任何其他裝置上花費的時間有關。請查看哪些事件佔用了最多的時間，並嘗試直觀地理解該事件是否確實應該花費那麼多時間。
同時，查看「# of Calls」（呼叫次數）欄位也很重要，它指示了事件被觸發的次數。順帶一提，我們也來談談「Self CPU/CUDA」與「CPU/CUDA total」的區別。「Self」欄位測量的是僅在事件本身內部花費的時間，不包括其子事件。
「total」欄位則包括事件本身及其所有子事件的總和。因此，如果您查看 matmul_add 的「CPU total」，它包含了事件本身花費的時間加上它觸發的子事件。這是一個重要的細微差別，需要注意。如果您查看表格的最後兩行，您會注意到分析器告訴我們：Self CPU time total: 2.314ms Self CUDA time total: 23.104usCPU 時間以毫秒 (ms) 為單位，而 GPU 時間以微秒 (us) 為單位。
從這個角度來看，GPU 上花費的時間（核心 ampere_bf16_s16816gemm...）不到 CPU 上花費時間（matmul_add 操作）的 1%。GPU 大部分時間處於閒置狀態，這是一個立即的警訊。發生這種情況的原因是 GPU 可以非常快速地計算小型矩陣乘法，因此我們的程式碼大部分時間都花費在準備核心、在 GPU 上啟動它們、發送數據進行乘法以及收集結果上。
這個概念被稱為「開銷限制型演算法」(overhead-bound algorithm)。擺脫這種情況最簡單的方法是使用更大的矩陣乘法。uv run 01_matmul_add.py --size 4096圖 2：4096 大小矩陣的矩陣乘法加法分析器表格圖 2 中的最後兩行是：Self CPU time total: 4.908ms Self CUDA time total: 4.495ms兩個時間都以毫秒 (ms) 為單位，這意味著我們僅透過增加矩陣乘法的大小，就顯著增加了 GPU 的運行時間。
如果您查看圖 2，您還會注意到現在大部分的 CUDA 時間都由 GPU 核心 (ampere_bf16_s16816gemm_..) 佔用，而不是由啟動它的 CPU 操作 (matmul_add) 佔用。這表示我們確實能夠從「開銷限制型」轉變為「計算限制型」(compute bound)。
現在我們將進入調度鏈的可視化，它存在於 .json 產物中。您可以將它們上傳到 Perfetto UI 查看追蹤圖，或者直接使用 uvx trace-util traces -b traces 來生成 Perfetto 連結。64x64 追蹤圖圖 3：64 大小矩陣的矩陣乘法與加法分析器追蹤圖在圖 3 中，我們看到了矩陣乘法和加法的分析器追蹤圖。
其中，長條的寬度表示事件的持續時間，垂直巢狀結構是呼叫層級，CPU 執行緒表示在 CPU 上發生的事件，而 GPU 執行緒則顯示實際的核心執行。您可能還會注意到空白區域，這些是等待或閒置時間。該腳本以預設配置運行，包括：size 64：輸入、權重和偏置的大小為 (64, 64)dtype bf16：資料型別為 bfloat16no compile：我們尚未編譯 torch 操作no warmup：我們在分析前沒有預熱 GPU使用 Perfetto 時，我們建議使用鍵盤來更快地存取追蹤圖。
您可以使用「W A S D」來導航追蹤圖。圖 4：PyTorch 分析器追蹤圖的 CPU 和 GPU 執行緒圖 4 中有兩個執行緒，一個用於 CPU 活動，一個用於 GPU 活動。在 CPU 執行緒中，您會注意到三個分析步驟（從 ProfilerStep#2 開始）。
這來自於排程設定。schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)wait 跳過了嘈雜的初始化（ProfilerStep#0），warmup 在不記錄的情況下運行分析器（ProfilerStep#1），而 active 則是在追蹤圖中顯示的部分。
您可以在這裡找到腳本中使用的排程設定。讓我們戴上偵探帽，調查追蹤圖並提出一些問題。為什麼 ProfilerStep#2 花費這麼長時間？圖 5：ProfileStep#2 明顯比後續步驟寬在圖 5 中，我們注意到 ProfileStep#2 比其他步驟花費了更多的時間，仔細觀察，您也會在 matmul_add 註解中看到類似的模式。
問題的癥結點在於註解內部，而不是註解本身：步驟 matmul_add 開始 aten::matmul 開始 間隙 #2 138.736 366.493 227.757 µs #3 517.926 523.447 5.521 µs #4 610.039 614.527 4.488 µs圖 6：record_function("matmul_add") 與 aten::matmul 之間約 228 微秒的閒置期圖 6 中顯示的約 228 微秒是進入 record_function("matmul_add") 和 PyTorch 實際調度 aten::matmul 之間的「閒置期」。
這可能由於多種原因造成，包括工作區記憶體分配、cuBLAS（NVIDIA 專有的 GPU 加速函式庫，用於執行基本線性代數運算）的啟發式演算法，或延遲模組載入。我們可以選擇忽略它，或者在分析前運行更多的暖身步驟（這是標準做法）。在效能分析方面，「暖身」(warmup) 是指在實際分析事件之前，先運行幾次事件。
GPU 所做的預備工作（包括上述幾點）是一次性的努力，我們不希望將其納入分析。在我們的範例中，我們有兩個暖身階段：一個是在進入分析器之前實際循環運行函數，另一個則是在分析器內部透過 warmup 參數實現。在本節中，我們已經啟用了實際的迭代以及排程。
uv run 01_matmul_add.py --warmup帶有暖身功能的 64x64 Perfetto 追蹤圖圖 7：暖身後，每個分析步驟花費的時間相似在圖 7 中，我們看到每個分析步驟花費的時間相似，但這並不意味著我們能夠優化一次性開銷。
我們進行暖身運行是為了不將這些開銷納入分析。我們認為如果突然結束本節而不提供解決方案的提示，將對讀者不公平，因此這裡提供一個連結，供您閱讀有關進一步優化啟動開銷的資訊。為什麼 CPU 和 GPU 執行緒之間有約 2.5 毫秒的偏移？圖 8：CPU 和 GPU 執行緒之間約 2.5 毫秒的偏移在圖 8 中，我們看到 CPU 和 GPU 執行緒之間存在約 2.5 毫秒的偏移：這是 CPU 提交 CUDA 核心後到實際執行之間的時間延遲。