TorchNeuron 效能分析（一）：AWS Trainium 系統層級剖析

許多機器學習實踐者對於如何加速模型執行有所了解，但卻很少人知道如何分析模型為何緩慢。在本系列文章的第一部分，您將學習如何系統性地分析推論效能，並找出時間花費在哪裡。

這篇文章將從零開始，假設讀者沒有任何剖析經驗。本文將使用 AWS Trainium，這是一款專為深度學習訓練與推論工作負載設計的 AI 晶片，以及其開源原生 PyTorch 後端 TorchNeuron。TorchNeuron 支援 Trainium 上的標準 PyTorch 工作流程，包括即時執行、torch.compile 和標準 PyTorch API，且只需最少的程式碼修改。

此外，TorchNeuron 還與 Neuron Kernel Interface (NKI) 整合，以實現客製化運算子和 Trainium 硬體的效能最佳化。學完本文後，您應該會了解：

如何使用 Neuron 後端配置 torch.profiler 並擷取清晰的追蹤記錄。

如何解讀 key_averages() 表格，以識別哪些運算消耗 Neuron 裝置時間，並比較它們的自我 CPU 時間。

如何透過 Neuron Explorer 追蹤一個運算，從框架分派、執行期排程到 NeuronCore 上的硬體執行。

這些訊號在同一運算的三種不同實作方式中如何變化。

在剖析工作負載之前，了解 Neuron 如何與 PyTorch 整合會很有幫助。

Neuron 後端使用 PyTorch 的 PrivateUse1 擴充點。這種機制讓加速器能夠與框架整合，而無需添加專用的分派鍵，例如 cuda。torch_neuronx 套件與上游 PyTorch 協同工作，將 Neuron 註冊為 PyTorch 程式碼的原生後端。有關此擴充機制的詳細資訊，請參閱「透過 PrivateUse1 促進新後端整合」。

您使用兩個物件來剖析 Neuron 工作負載：

NeuronConfig 控制要擷取哪些 Neuron 資料流和 NeuronCore，以及將輸出寫入何處。您將其作為 experimental_config 傳遞給 torch.profiler.profile。

NeuronProfiler 接受 NeuronConfig 並提供 export_trace 回呼。此回呼將框架追蹤記錄 trace.json 與 Neuron 系統追蹤記錄一起寫入。

nc = NeuronConfig(

modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],

profile_output_dir=OUTPUT_DIR,

capture_enabled_for_nc='0',

correlate_device_time=True,

)

exporter = NeuronProfiler(nc)

modes 參數選擇要擷取的資料流，例如裝置指令和執行期追蹤記錄。capture_enabled_for_nc 參數選擇要追蹤的 NeuronCore。設定 correlate_device_time=True 會將裝置時間從追蹤記錄傳播到 key_averages() 表格中，您可以在其中比較 CPU 和 Neuron 裝置時間。

有關完整的參數參考和可用的 ProfileMode 值，請參閱 NeuronConfig Profile Modes 文件。

Neuron Explorer 是一個基於瀏覽器的檢視器，用於查看系統追蹤記錄、裝置時間軸和每個引擎的細目分析。您可以透過 SSH 通道或 Visual Studio Code 擴充功能存取它。在本文中，您將使用 Explorer 來確認 key_averages() 數值報告的內容，並檢查分派之間發生了什麼。

剖析完成後，輸出目錄包含系統追蹤記錄，當您擷取裝置資料時，還會有匹配的 .neff 和 .ntff 檔案。Neuron 可執行檔格式 (NEFF) 包含編譯後的 Neuron 可執行檔，而匹配的 NTFF 則包含其硬體指令追蹤記錄。這些檔案共享一個數值雜湊，以識別它們為一對。

要在 Neuron Explorer 中開啟剖析結果，請上傳系統或組合剖析的完整輸出目錄。對於僅裝置剖析，請單獨上傳匹配的 .neff 和 .ntff 檔案。請參閱「步驟 5：上傳您的剖析結果」以獲取 UI 和命令列說明。

對於首次分析，您將剖析 RMSNorm，這是一個具有清晰算術和歸約運算序列的小型運算。

您將剖析同一運算的三種實作方式：

使用個別 PyTorch 運算的樸素 RMSNorm：乘法、平均、加法、倒數平方根和乘法。

torch.rms_norm()，PyTorch 內建函數，降低為單一融合的 Neuron 核心。

使用 torch.compile 的樸素 RMSNorm，使用與第一種實作相同的運算。

這些腳本使用相同的輸入、剖析器配置和排程，因此結果可以直接比較。所有重現此工作的腳本都可以在 torch-neuronx 範例目錄中找到。

首先，使用個別的 PyTorch 運算編寫 RMSNorm。每行對應一個或多個 ATen 運算，並可以建立中間張量。

device = torch.device('neuron')

x = torch.randn(4, 512, dtype=torch.bfloat16, device=device)

weight = torch.ones(512, dtype=torch.bfloat16, device=device)

def rms_norm_fn(t):

x_sq = t * t # 平方

mean_sq = x_sq.mean(dim=-1, keepdim=True) # 平方平均

rms = torch.rsqrt(mean_sq + 1e-6) # 1 / sqrt(mean_sq + eps)

return t * rms * weight # 正規化 + 縮放

使用 Neuron 後端所需的兩個參數呼叫剖析器：

包含 ProfilerActivity.PrivateUse1 以從 Neuron 裝置收集活動。

將 NeuronConfig 作為 experimental_config 傳遞，以啟用 Neuron 特定的 Kineto 功能。

排程使用 wait=2、warmup=1 和 active=1，跨四個 prof.step() 呼叫。這會記錄一次暖機迭代，而不包含初始設定開銷。在剖析器迴圈內呼叫 torch.neuron.synchronize()。如果沒有此同步，剖析器可能會在裝置完成之前匯出，並導致 Neuron 時間欄位為空。

with torch.profiler.profile(

activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],

experimental_config=nc,

record_shapes=False,

on_trace_ready=exporter.export_trace,

schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),

) as prof:

for step in range(4):

with torch.no_grad():

_ = rms_norm_fn(x)

torch.neuron.synchronize()

prof.step()

剖析完成後，檢查 key_averages() 表格。按 self_privateuse1_time_total 排序，將裝置成本最高的運算放在最前面。此處顯示的表格已篩選為報告 Neuron 裝置時間的 RMSNorm 運算。完整的 key_averages() 輸出還包括 CPU 端的事件，例如記憶體分配和主機到裝置的傳輸。

print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))

key averages 表格顯示三個 aten::mul 呼叫佔 Neuron 裝置時間的最大份額。

該表格顯示了哪些 ATen 運算運行、它們運行的頻率以及它們在 CPU 和 Neuron 裝置上花費了多少時間。aten::mul 具有最高的總裝置時間，因為該函數執行了三次乘法。

| Name | # Calls | Self CPU | Self NEURON | % NEURON |

|------------------|---------|-----------|-------------|----------|

| aten::mul | 3 | 60.9 μs | 1,164.5 μs | 49.5% |

| aten::add | 1 | 26.6 μs | 434.6 μs | 18.5% |

| aten::rsqrt | 1 | 15.6 μs | 415.0 μs | 17.6% |

| aten::mean | 1 | 21.1 μs | 340.8 μs | 14.5% |

Self NEURON total: 2.355 ms

每一行都直接對應到函數。aten::mul 運行 x * x、x * rms 和最終的權重乘法。aten::mean 歸約最後一個維度。aten::add 添加 epsilon 值並觸發主機到裝置的複製，因為純量始於 CPU。aten::rsqrt 計算倒數平方根。

Neuron Explorer 追蹤 aten::mul 運算，從框架分派、執行期排程到 Neuron 硬體通道上的執行。

在 Explorer 中，時間從左到右流動，事件生命週期從上到下流動。每個 ATen 運算都遵循相同的模式：在框架通道中建立、在執行期通道中排程為 nrta_execute_schedule，以及在硬體通道中執行為 nc_exec_running。

系統追蹤檢視將事件組織成通道。在預設分組模式下，標籤遵循 instance/PID/subsystem/identifier 模式：

| Lane | Label suffix | What it shows |

|------------------|--------------------|------------------------------------------------------|

| Framework dispatch | framework/TID:<id> | 運算子分派事件 (PrivateUse1 執行期)，按執行緒 ID 分組。 |

| Framework device | framework/Stream:<n> | 裝置互動事件 (PrivateUse1 驅動程式)，按串流編號分組。 |

| Neuron Runtime | neuron_rt/NC:<n> | Neuron 執行期 API 呼叫，按邏輯 NeuronCore 分組。 |

| Neuron Hardware | neuron_hw/NC<n>:<id> | 硬體執行事件，按實體 NeuronCore 分組。 |

在即時模式下，每個運算獨立運行。使用依賴鏈檢視器追蹤 aten::mul.Tensor，從框架層到執行期排程再到裝置執行。

依賴鏈檢視器突顯了一個運算在框架、執行期和硬體通道中的生命週期。

將個別運算替換為 PyTorch 內建的 rms_norm 函數。Neuron 將該運算降低為一個融合的核心：

def rms_norm_fn(t):

return torch.rms_norm(t, [512], weight)

融合的 RMSNorm 在 key averages 表格中顯示為一個裝置運算。

| Name | # Calls | Self CPU | Self NEURON | % NEURON |

|----------------------|---------|----------|-------------|----------|

| aten::_fused_rms_norm | 1 | 47.6 μs | 403.3 μs | 100% |

Self NEURON total: 403.3 μs

融合的 RMSNorm 遵循一個從框架到硬體的路徑，而不是單獨分派每個算術運算。

在樸素實作中，每個運算獨立跨越分派邊界，並遵循自己的從框架到執行期再到硬體的路徑。在融合實作中，一個 Python 呼叫涵蓋了完整的計算。追蹤記錄包含一個分派和一個硬體事件，組件運算之間沒有間隙。

對於最後一個範例，使用 torch.compile 編譯原始的樸素 RMSNorm 函數。

使用 @torch.compile(backend="neuron", fullgraph=True)，TorchDynamo 將 Python 位元碼擷取為 FX 圖，而 AOTAutograd 將其降低為 ATen 運算。Neuron 後端接收完整的圖，而不是單獨分派每個運算。它融合運算、最佳化記憶體存取，並將圖編譯成單一 NEFF 成品。該成品作為一個硬體分派運行。

@torch.compile(backend='neuron', fullgraph=True)

def naive_compiled(t):

x_sq = t * t

mean_sq = x_sq.mean(dim=-1, keepdim=True)

rms = torch.rsqrt(mean_sq + 1e-6)

return t * rms * weight

編譯後的函數在 key averages 表格中顯示為一個 Torch 編譯區域。

| Name | # Calls | Self CPU | Self NEURON | % NEURON |

|--------------------------|---------|----------|-------------|----------|

| Torch-Compiled Region: 0/0 | 1 | 98.4 μs | 375.5 μs | 100% |

Self NEURON total: 375.5 μs

編譯後的圖遵循一個通過執行期排程和硬體執行的路徑。

該表格顯示一次裝置執行，與剖析 2 相同，但程式碼採用不同的路徑。它不是呼叫專門設計的融合運算，而是編譯器從原始 Python 函數中發現融合。在此運行中，Neuron 時間與 torch.rms_norm 相當，為 375 μs，而 torch.rms_norm 為 403 μs。

| Variant | NEURON time | Self CPU time | Operator calls |

|-----------------------|-------------|---------------|----------------|

| Naive ops | 2,355 µs | ~125 µs | 6 |

| torch.rms_norm | 403 µs | ~48 µs | 1 |

| Naive + torch.compile | 376 µs | ~98 µs | 1 |

torch.rms_norm 和 torch.compile 都將 RMSNorm 融合到一個核心分派中。在這些運行中，這兩種方法所花費的裝置時間約為樸素實作的六分之一。融合消除了硬體執行之間的間隙，並避免在單獨分派的運算之間具體化中間值，例如 x_sq 和 mean_sq。

key_averages() 顯示哪些運算消耗裝置時間，而其完整輸出還揭示了 CPU 端的設定和傳輸。Explorer 提供了相應的視覺追蹤記錄。您可以追蹤一個運算從框架分派、執行期排程到硬體執行，並查看融合所消除的間隙。