Hugging Face 發布 TorchNeuron 系統層級剖析教學

背景

TorchNeuron 是 AWS Trainium 的 PyTorch 原生後端，讓開發者能使用標準 PyTorch API 在 Trainium 硬體上執行模型。AWS Trainium 是專為深度學習訓練與推理設計的 AI 晶片。許多機器學習實務人員對如何讓模型跑得更快有直覺，但較少人知道如何分析模型變慢的原因。 官方原文・Hugging Face Blog 官方原文・github.com

發生什麼事

Hugging Face 於 2026 年 10 月 21 日發布了「剖析 TorchNeuron（第一部分）」部落格文章。該文章是系列教學的第一篇，假設讀者沒有先前的剖析經驗，從零開始介紹如何系統性地分析推理效能並識別時間消耗位置。 官方原文・Hugging Face Blog

重點內容

文章以 RMSNorm（一種包含算術與歸約操作的小型運算）為例，比較了三種實作方式的效能：1. 使用個別 PyTorch 操作的 Naive RMSNorm；2. 使用 PyTorch 內建 torch.rms_norm() 函式；3. 使用 torch.compile 編譯的 Naive RMSNorm。測試顯示，Naive 實作的 NEURON 裝置時間為 2,355 μs，而 torch.rms_norm 為 403.3 μs，torch.compile 版本為 375.5 μs。後兩者將運算融合為單一核心調用，裝置時間約為 Naive 實作的六分之一。 官方原文・Hugging Face Blog

文章詳細說明如何配置 torch.profiler 以捕捉 Neuron 後端資料，並使用 Neuron Explorer（基於瀏覽器的檢視工具）追蹤單一操作從框架分派、執行時排程到硬體執行的完整路徑。Neuron Explorer 可透過 SSH 隧道或 Visual Studio Code 擴充功能存取。剖析輸出包含系統追蹤檔及配對的 .neff（編譯後的 Neuron 可執行檔）與 .ntff（硬體指令追蹤）檔案。 官方原文・Hugging Face Blog

在 key_averages() 表格中，Naive 實作顯示 aten::mul 佔據最高的 NEURON 裝置時間（1,164.5 μs，佔比 49.5%），因為該函式執行三次乘法。相比之下，融合後的實作僅顯示單一操作項目。文章指出，融合消除了硬體執行之間的間隙，並避免了在獨立調用的操作之間物化中間值（如 x_sq 和 mean_sq）。 官方原文・Hugging Face Blog

對使用者的意義

開發者可以利用此工作流程在更大的模型上找到效能熱點，驗證追蹤結果，並決定是否使用融合操作或 torch.compile 來降低成本。所有用於重現此工作的腳本均可在 torch-neuronx 的 examples 目錄中找到。使用者需具備 Trn2 或 Trn3 實例，以及 Python 3.12、3.13 或 3.14 環境。 官方原文・Hugging Face Blog 官方原文・github.com

限制與未定事項

此文章為系列教學的第一部分，僅涵蓋系統層級剖析。第二部分將進入裝置層級剖析，包括檢查 TensorEngine、VectorEngine 和 DMA（直接記憶體存取）的分解，並利用 roofline 模型判斷核心是受計算限制還是記憶體限制。TorchNeuron 目前處於 Beta 狀態，支援 PyTorch 2.14.0 和 Neuron SDK 2.33。 官方原文・Hugging Face Blog 官方原文・github.com