PyTorch 性能分析：torch.profiler 初學者指南

PyTorch 性能分析：torch.profiler 初學者指南

了解 torch.profiler 以進行效能優化

效能分析是優化大型語言模型（LLM）和深度學習管線的首要步驟，因為它讓開發者能辨識模型是受 overhead 限制還是受 compute 限制。torch.profiler 模組提供兩個主要的診斷資產：profiler table，提供「什麼」佔用最多時間的統計摘要；以及 profiler trace，提供「何時」與「為何」操作在 CPU 與 GPU 通道中發生的時間視圖。

辨識瓶頸：Overhead 限制 vs. Compute 限制

效能瓶頸常透過比較 profiler table 中 CPU 與 GPU 所花費的時間來顯示。

  • Overhead-Bound: 當 Self CPU time total 明顯高於 Self CUDA time total（例如 CPU 時間為毫秒級，而 GPU 時間為微秒級）時，演算法屬於 overhead 限制。這表示 CPU 花費較多時間在準備與啟動 kernel，而 GPU 的執行時間較少。此情況常見於小規模矩陣乘法。
  • Compute-Bound: 當 Self CPU time total 與 Self CUDA time total 均在毫秒範圍且相近時，演算法屬於 compute 限制。這是高效能運算的理想狀態，GPU 成為主要瓶頸。

增大工作負載（例如從 64x64 矩陣升至 4096x4096 矩陣）通常會使模型從 overhead 限制的狀態轉變為 compute 限制。

分析 CPU 與 GPU 的分派鏈

PyTorch 操作遵循從 Python 呼叫到 CUDA kernel 的特定分派鏈。典型的矩陣乘法與加法序列如下：

ProfileStep → record_function (user annotation) → aten::matmul (ATen-level dispatch) → aten::mm (2D matrix-matrix multiply backend) → cudaLaunchKernel.

CPU 通道關鍵指標

  • Cold-Start Overhead（冷啟動開銷）：首個 ProfileStep 通常比後續步驟寬，原因是工作空間分配、cuBLAS 啟發式策略以及延遲模組載入。可透過熱身迭代來緩解。
  • CUDA Occupancy Queries（CUDA 占用率查詢）：在 cudaLaunchKernel 前出現 cudaOccupancyMaxActiveBlocksPerMultiprocessor 表示此為「重量級」kernel（如 GEMM 或卷積）。CPU 正在向驅動程式查詢，以根據硬體容量決定最佳 block 大小。
  • Resource-Light Kernels（資源輕量 kernel）：逐元素或 reduction kernel 通常不會有占用率查詢，因為其資源佔用固定且較小。
  • Synchronization（同步）：追蹤結尾的長時間 cudaDeviceSynchronize 常是 profiler 刷新事件；其持續時間反映 CPU 等待 GPU 完成未完成工作的時間。

GPU 通道關鍵指標

  • Activity Buffer Requests（活動緩衝請求）：kernel 之間的間隙或 CPU 與 GPU 通道之間的初始偏移，通常是因 profiler 分配或重新填充其事件緩衝區所致。
  • Runtime Variance（執行時間變異）：相同的 kernel 在不同步驟可能呈現不同的執行時間，原因可能是 GPU 時脈波動、熱度、電源管理或驅動程式維護。

torch.compile 的影響

使用 torch.compile 會透過 TorchDynamo、AOTAutograd 與 TorchInductor，將即時執行的 PyTorch 程式碼轉換為最佳化圖形。對已編譯函式進行效能分析會顯示出多項架構變化：

分派層級融合

對於 torch.add(torch.matmul(x, w), b) 之類的操作，torch.compile 會在圖層級執行算子融合，將分別的 aten::add 與 aten::mm 呼叫合併為單一的 aten::addmm 呼叫。然而，這僅是 分派層級 的融合，非 kernel 層級。GPU 仍會先執行 Memcpy DtoD（將偏差寫入目標緩衝區），再執行帶有 bias 加法尾端的 GEMM kernel。

執行時架構與開銷

編譯區塊會在 CPU 端引入特定的層級結構：

  1. TorchDynamo 快取查詢：驗證輸入的形狀、資料型別與裝置是否與快取的編譯相符。此步驟在每次呼叫時皆會執行。
  2. Torch-Compiled Region：進入編譯版本的包裝器。
  3. AOTDispatcher 執行時包裝器前置：處理張量的元資料與視圖追蹤。
  4. Call CompiledFxGraph：執行產生的程式碼（以內容雜湊識別）。

對於極小的操作，torch.compile 可能會增加 CPU 開銷，因為遍歷 Dynamo $\rightarrow$ AOTAutograd $\rightarrow$ Inductor 堆疊的成本超過融合帶來的節省。隨著模型中操作數量的增加，這些開銷會被攤平。

追蹤閱讀摘要表

觀測 可能含義
Self CPU $\gg$ Self CUDA Overhead 限制；增加 batch 大小或融合操作
cudaOccupancy... before launch 重量級、適應性啟動的 kernel（GEMM/Conv）
aten::matmul $\rightarrow$ aten::bmm 在 3D+ 張量上的批次矩陣乘法
Torch-Compiled Region 執行於 torch.compile 最佳化的區塊內
Memcpy DtoD before GEMM 編譯模式下 addmm 尾端的 bias 複製

SUMMARY: Hugging Face 提供了完整的指南，說明如何使用 torch.profiler 來辨識瓶頸、了解 CPU-GPU 分派鏈，並分析 torch.compile 對 kernel 執行的影響。

TITLE: PyTorch 性能分析：torch.profiler 初學者指南

Sources