AI News HubLIVE
站内改写

PyTorch 效能分析（第一部分）：torch.profiler 入門指南

本文是 PyTorch 效能分析系列的第一篇，從最簡單的矩陣乘加操作開始，引導讀者學習如何使用 torch.profiler 進行效能分析，包括設定分析器、解讀分析表和追蹤資料，以及理解 CPU 和 GPU 活動之間的時間關係。文章還討論了預熱和最佳化開銷等問題。

文章情報

投資人進階

要點

  • torch.profiler 可以生成效能分析表和時間線追蹤，幫助識別熱點和瓶頸。
  • 小矩陣乘法容易導致開銷受限，增大矩陣規模可轉為計算受限。
  • 預熱步驟對於消除啟動開銷至關重要，確保分析結果準確。
  • CPU 和 GPU 活動之間存在時間偏移，反映了核心啟動和同步的延遲。

為什麼重要

這條新聞值得關注，因為torch.profiler 可以生成效能分析表和時間線追蹤，幫助識別熱點和瓶頸。

技術影響

可能影響模型選型、推理成本、產品能力和評測基準。

在最佳化 PyTorch 模型時，效能分析是不可或缺的一步。無論你是想提高大語言模型（LLM）的每秒 token 數，縮短推理時間，還是僅僅想了解訓練迴圈為何比預期慢，最終都需要透過效能分析來定位問題。然而，效能分析有著較高的學習門檻：追蹤資料是密集的彩色矩形，事件名稱令人望而生畏，大多數教程假設讀者已經具備讀取能力。因此，即使我們知道應該進行效能分析，開啟追蹤檔案也常常被推遲或交給他人處理。本文旨在降低這一門檻。

本系列文章將從 PyTorch 效能分析的基礎開始，逐步培養讀者讀取分析追蹤資料並用其指導最佳化的能力。本部分（第一部分）從最簡單的操作——矩陣乘法加偏置——開始，教你如何理解分析器返回的資訊。後續部分將擴充套件到 nn.Linear 和小型 MLP，並最終應用於大語言模型中的 Transformer。

準備工作：核心概念

在開始之前，需要了解兩個定義：

  • GPU 核心（kernel）：是在 GPU 上並行執行的程式。
  • CPU 排程並啟動這些核心。

當使用 PyTorch 操作時，它會被自動轉換為一個或多個在 GPU 上執行的核心。

建立效能分析程式碼

我們從最簡單的操作開始：一個矩陣乘法加一個偏置加法。程式碼如下：

def fn(x, w, b):
    return torch.add(torch.matmul(x, w), b)

使用 torch.profiler 進行效能分析的步驟如下：

  1. 準備好要分析的程式碼（如上）。
  2. 用 torch.profiler.record_function 標註演算法（可選但推薦）。
  3. 將程式碼包裝在 torch.profiler.profile 上下文管理器中。
  4. 匯出分析結果。

分析器產生兩個不同的產物：

  • 分析表：提供事件的統計摘要，回答“什麼佔用了最多時間”。
  • 追蹤資料：提供時間維度的執行檢視，回答“操作何時發生以及為何發生”。

解讀分析表

執行指令碼時，會生成包含分析表的 .txt 檔案和包含追蹤資料的 .json 檔案。分析表的第一列是觸發的事件，其他列是 CPU、GPU 等裝置上的時間。透過觀察事件消耗的時間，可以直觀地瞭解哪些是熱點。注意“自 CPU/CUDA 時間”與“總 CPU/CUDA 時間”的區別：前者僅包括事件本身的時間，後者包括其所有子事件。

例如，對於 64×64 的矩陣，自 CUDA 時間僅為 23.104 微秒，而自 CPU 時間為 2.314 毫秒，說明 GPU 大部分時間空閒，這是典型開銷受限的情況。透過將矩陣大小增加到 4096×4096，自 CUDA 時間變為 4.495 毫秒，CPU 時間為 4.908 毫秒，GPU 時間顯著增加，說明從開銷受限轉向了計算受限。

解讀追蹤資料

追蹤資料可以用 Perfetto UI 開啟。CPU 和 GPU 有各自的泳道，條形寬度表示事件持續時間，垂直巢狀表示呼叫層次。在 64×64 的追蹤中，可以觀察到以下現象：

  • ProfileStep#2 耗時明顯更長：這是因為沒有預熱 GPU，導致啟動開銷被記錄。預熱步驟（先執行幾次程式碼再進行分析）可以消除這種一次性開銷。
  • CPU 和 GPU 泳道之間存在約 2.5 毫秒的偏移：這是 CPU 啟動核心後，GPU 實際開始執行之前的延遲，可能涉及記憶體傳輸、核心排隊等。

總結與展望

本部分介紹了 torch.profiler 的基本用法以及如何解讀分析表和追蹤資料。透過調整矩陣大小和預熱，可以觀察到從開銷受限到計算受限的轉變。在後續部分中，我們將擴充套件這些概念到更復雜的網路，並利用分析結果指導最佳化。

透過本部分的學習，你應該能夠：

  • 設定 torch.profiler 並理解其輸出。
  • 讀取分析表和時間線追蹤（CPU 泳道、GPU 泳道以及兩者之間的間隙）。
  • 跟蹤從 Python 呼叫到 CUDA 核心的完整事件鏈。
  • 理解 torch.compile 帶來的變化。