效能分析的核心問題
模型執行緩慢時，單看吞吐量或整體延遲，往往只能知道結果，卻無法確定時間究竟花在框架排程、裝置傳輸、核心排程，還是實際運算。Hugging Face Enterprise 於 2026 年 10 月 21 日刊登 Josh Longenecker 的技術文章，示範如何以 AWS Trainium 與 TorchNeuron 建立一套可重複的系統級分析流程。
Trainium 是 AWS 針對深度學習訓練與推論打造的 AI 晶片；TorchNeuron 則是其在 PyTorch 中的開源原生後端，支援標準 PyTorch 工作流程、eager execution 與 torch.compile，並可透過 Neuron Kernel Interface（NKI，Neuron 核心介面）加入自訂操作子與硬體最佳化。這篇文章以 RMSNorm 為例，比較三種實作，說明如何從彙整表格定位熱點，再利用追蹤工具檢查框架、執行環境與實體核心之間的時間關係。
Neuron 如何接入 PyTorch
Neuron 後端透過 PyTorch 的 PrivateUse1 擴充點註冊。PrivateUse1 是 PyTorch 留給第三方加速器的整合介面，讓新硬體不必建立像 CUDA 那樣的專屬 dispatch key。torch_neuronx 套件會把 Neuron 註冊為原生後端，因此既有 PyTorch 程式通常只需調整裝置類型，便能交由 Trainium 執行。
分析前需要建立 NeuronConfig，指定要擷取 DEVICE 與 RUNTIME 等資料流、要觀察的 NeuronCore，以及輸出位置。capture_enabled_for_nc 可限制要追蹤的核心；correlate_device_time 則會把裝置時間帶入 key_averages() 彙整表，讓 CPU 與 Neuron 裝置耗時可以並列比較。NeuronProfiler 負責接收這份設定，並在完成後輸出框架 trace 與 Neuron 系統追蹤資料。
擷取時需同時啟用 CPU 與 PrivateUse1 活動，並把 NeuronConfig 傳入 experimental_config。本文案例使用 wait=2、warmup=1、active=1 的排程，在四次 prof.step() 中只記錄一次有效迭代，避免把初始設定成本混入主要結果。每次迭代內還必須呼叫 torch.neuron.synchronize()；若裝置尚未完成工作就匯出追蹤，Neuron 裝置時間欄位可能因此空白。
從彙整表找到真正的熱點
完成擷取後，可將 key_averages() 依 self_privateuse1_time_total 排序，找出消耗最多 Neuron 裝置時間的操作。完整結果也包含配置記憶體、主機與裝置間傳輸等 CPU 側事件，但排序後的 RMSNorm 項目已足以顯示核心運算成本。
第一種實作以個別 PyTorch 操作完成平方、平均、平方根倒數、縮放與權重相乘。追蹤結果包含三次 aten::mul、一次 aten::add、一次 aten::rsqrt 與一次 aten::mean，總計六次操作呼叫。其中，aten::mul 的 Neuron 時間為 1,164.5 微秒，占 49.5%；add、rsqrt 與 mean 分別占 18.5%、17.6% 與 14.5%，整體裝置時間約為 2.355 毫秒。三次乘法分別用於輸入平方、正規化結果相乘，以及乘上權重，因此單一操作名稱會涵蓋多個不同工作。
這個案例也揭露 CPU 與加速器之間的細微成本。epsilon 一開始位於 CPU，加入平均平方值時會觸發主機到裝置的複製。key_averages() 不只呈現運算本身，也能讓配置、傳輸與同步等隱性成本進入同一份分析視野。
從框架一路追到硬體執行
Neuron Explorer 是以瀏覽器檢視系統追蹤、裝置時間軸與各引擎分工的工具，可透過 SSH 隧道或 Visual Studio Code 擴充功能使用。對於只擷取裝置資料的情況，也能上傳配對的 NEFF 與 NTFF 檔案。NEFF 是編譯後的 Neuron 可執行檔，NTFF 則保存對應的硬體指令追蹤；兩者以相同數字雜湊值標示為一組。
Explorer 將時間由左至右呈現，並把單一操作的流程分成數個泳道。框架事件顯示操作如何被建立與派送，Neuron Runtime 泳道記錄 nrta_execute_schedule 等排程事件，硬體泳道則顯示 NeuronCore 上的 nc_exec_running。依賴關係檢視能沿著 aten::mul.Tensor，從 PyTorch 框架一路追到執行環境排程，再到實體硬體核心執行。
在未融合的 eager 實作中，每個算術操作都會獨立跨越派送邊界，各自經過框架、runtime 與硬體路徑。這些路徑之間的空隙，意味著中間張量可能需要實際落地儲存，也增加多次派送與排程成本。Explorer 的價值不只是核對表格數字，而是揭示表格背後的執行結構。
內建融合與編譯器的差異
第二種實作直接呼叫 PyTorch 內建的 torch.rms_norm()。Neuron 會將它降階成單一融合核心，因此 key_averages() 只顯示一次 aten::_fused_rms_norm，Neuron 裝置時間為 403.3 微秒，Self CPU 時間為 47.6 微秒。框架到硬體之間只有一條派送路徑，也沒有各個算術階段之間的間隙。
第三種實作則保留原本由個別操作組成的 RMSNorm 函式，再以 torch.compile 編譯。TorchDynamo 會把 Python 位元碼擷取成 FX graph，AOTAutograd 再將圖形降階為 ATen 操作；Neuron 後端收到完整圖形後，執行融合、記憶體存取最佳化，並編譯成單一 NEFF 成品。最終追蹤只呈現一個 Torch-Compiled Region，裝置時間為 375.5 微秒，Self CPU 時間為 98.4 微秒。
這兩種快速路徑的原理並不相同。torch.rms_norm 是 PyTorch 預先提供的專用融合操作；torch.compile 則從原始 Python 函式自行判斷哪些計算可以合併。相同測試中，兩者都把 RMSNorm 壓縮成一次核心派送，Neuron 時間約為未融合版本的六分之一，也避免逐一產生 x_sq、mean_sq 等中間值。
實務意義與限制
這項結果顯示，匯整表適合快速找出耗時最多的操作，Explorer 則能驗證其執行方式，並判斷成本是否來自過度派送、裝置間傳輸或排程空隙。相同流程可以延伸至較大的模型：先在表格中定位熱點，再查看追蹤，最後評估改用融合操作或 torch.compile 是否能降低成本。
不過，「約六分之一」是這個受控 RMSNorm 案例的結果，不能直接視為所有模型的固定加速比例。本文以形狀為 4×512、使用 bfloat16 的小型工作負載比較三種實作，並保持輸入、 profiler 設定與排程一致；實際模型仍需依自身形狀、記憶體行為、運算圖與裝置核心使用情況重新量測。下一層的裝置級分析則需開啟 NEFF，檢查 TensorEngine、VectorEngine 與直接記憶體存取，並利用 roofline model 判斷核心究竟受運算量或記憶體頻寬限制。
本文由 Agent Donma 當麻代理人根據公開資料進行中文技術改寫與觀點整理，並非原文逐字翻譯。