設定 Profiler 以捕捉 Neuron 裝置活動

要分析 AWS Trainium 上的模型效能，開發者需使用 torch.profiler 並指定 Neuron 後端。這需要建立一個 NeuronConfig 物件，用於控制要捕捉的資料流（如裝置指令與執行期追蹤）以及輸出目錄。

關鍵設定包括將 modes 設為 ProfileMode.DEVICE 和 ProfileMode.RUNTIME，並啟用 correlate_device_time=True，以便將裝置時間傳播到 key_averages() 表格中進行比較。此外，必須在 ProfilerActivity 中包含 PrivateUse1，以收集來自 Neuron 裝置的活動數據。

為了確保追蹤資料的完整性，程式碼需在 profiler 迴圈內呼叫 torch.neuron.synchronize()。若缺少此同步步驟，Profiler 可能在裝置完成工作前就結束導出，導致 Neuron 時間欄位為空。

朴素實作的效能瓶頸：分散運算與中間張量

文章首先分析以個別 PyTorch 操作（如乘法、均值、加法）手動實作的 RMSNorm。這種方式會產生多個 ATen 操作，每個操作都獨立跨越分派邊界，並可能建立中間張量。

透過 key_averages() 表格按 self_privateuse1_time_total 排序後發現，aten::mul 佔據了最高的裝置時間比例（49.5%），因為函數中包含了三次乘法操作。其他如 aten::add、aten::rsqrt 和 aten::mean 也各自佔有顯著的時間份額。

總計而言，這種朴素實作的 Neuron 裝置執行時間達到 2,355 微秒。表格數據顯示，雖然 CPU 端的自耗時間較短（約 125 微秒），但主要的效能消耗發生在硬體執行階段，且由於操作分散，存在大量的排程開銷與中間值寫入記憶體的成本。

使用內建融合核心 torch.rms_norm 優化

第二種實作方式改用 PyTorch 內建的 torch.rms_norm 函數。Neuron 後端會將此操作降低為單一的融合核心（fused kernel），在 key_averages() 表格中顯示為 aten::_fused_rms_norm。

這種方法將原本的多個運算合併為單一裝置操作，消除了各組件操作之間的空隙，並避免了中間值（如 x_sq 和 mean_sq）的實體化。測試結果顯示，Neuron 裝置執行時間降至 403.3 微秒，僅為朴素實作的約六分之一。

在 Neuron Explorer 中觀察此流程，可以看到單一 Python 呼叫對應到一次框架分派與一次硬體事件，沒有中間的排程延遲或額外的記憶體存取開銷。

torch.compile 自動融合與編譯優化

第三種實作保留原始的朴素 Python 程式碼，但透過 torch.compile(backend='neuron') 進行編譯。TorchDynamo 會捕捉 Python 位元組碼並轉換為 FX 圖，AOTAutograd 將其降低為 ATen 操作，最後由 Neuron 後端接收完整圖形進行融合與優化。

編譯後的函數在 key_averages() 表格中顯示為 Torch-Compiled Region，同樣表現為單一的裝置執行事件。測試數據顯示其 Neuron 時間為 375.5 微秒，與使用內建 torch.rms_norm 的效能相當（403.3 微秒）。

雖然兩者最終都產生了單一的核心分派，但路徑不同：torch.rms_norm 是直接調用預先編譯好的融合操作，而 torch.compile 則是從原始 Python 函數中動態發現並執行融合。這顯示自動編譯工具在處理複雜或自定義邏輯時，也能達到接近手動優化的效能。

利用 Neuron Explorer 視覺化追蹤與除錯

Neuron Explorer 是一個基於瀏覽器的檢視工具，用於查看系統追蹤、裝置時間軸以及每個引擎的詳細分解。它可以透過 SSH 隧道或 Visual Studio Code 擴充功能存取。

在 Explorer 中，時間從左至右流動，事件生命週期從上至下展開。開發者可以追蹤單一 ATen 操作（如 aten::mul）從框架層的建立、執行期層排程（nrta_execute_schedule），到硬體層執行（nc_exec_running）的完整過程。

系統追蹤視圖將事件組織在不同的軌道（lanes）中，包括框架分派、框架裝置互動、Neuron 執行期 API 呼叫以及硬體執行事件。透過依賴鏈檢視器，開發者可以確認 key_averages() 中的數值報告，並檢查操作之間的空隙與排程行為。