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 帶來的變化。