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

許多機器學習從業者對於如何加速模型執行有所了解，但卻不清楚如何分析模型緩慢的原因。在本系列文章的第一部分，您將學習如何系統性地分析推論效能，並找出時間耗費的環節。

本文從零開始，假設讀者沒有任何效能分析經驗。教學將使用 AWS Trainium，這是一款專為深度學習訓練和推論工作負載設計的 AI 晶片，以及其開源的原生 PyTorch 後端 TorchNeuron。TorchNeuron 支援 Trainium 上的標準 PyTorch 工作流程，包括即時執行（eager execution）、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：上傳您的分析結果以獲取使用者介面和命令列說明。

對於第一次分析，我們將分析 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))

關鍵平均值表格顯示，三個 aten::mul 呼叫佔 Neuron 裝置時間的最大份額。表格顯示了哪些 ATen 運算執行了，執行頻率，以及它們在 CPU 和 Neuron 裝置上花費了多少時間。aten::mul 具有最高的總裝置時間，因為該函式執行了三次乘法。

表格中，aten::mul 呼叫了 3 次，自我 CPU 時間為 60.9 微秒，自我 NEURON 時間為 1,164.5 微秒，佔 NEURON 時間的 49.5%。aten::add 呼叫了 1 次，自我 CPU 時間為 26.6 微秒，自我 NEURON 時間為 434.6 微秒，佔 18.5%。

aten::rsqrt 呼叫了 1 次，自我 CPU 時間為 15.6 微秒，自我 NEURON 時間為 415.0 微秒，佔 17.6%。aten::mean 呼叫了 1 次，自我 CPU 時間為 21.1 微秒，自我 NEURON 時間為 340.8 微秒，佔 14.5%。總自我 NEURON 時間為 2.355 毫秒。

每一行都直接對應到函式。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 模式。框架分派通道顯示運算子分派事件（PrivateUse1 執行時），按執行緒 ID 分組。框架裝置通道顯示裝置互動事件（PrivateUse1 驅動程式），按串流編號分組。

Neuron 執行時通道顯示 Neuron 執行時 API 呼叫，按邏輯 NeuronCore 分組。Neuron 硬體通道顯示硬體執行事件，按實體 NeuronCore 分組。

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

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

將單獨的運算替換為 PyTorch 內建的 rms_norm 函式。Neuron 將該運算降級為一個融合的核心：

def rms_norm_fn(t):

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

融合的 RMSNorm 在關鍵平均值表格中顯示為一個裝置運算。表格中，aten::_fused_rms_norm 呼叫了 1 次，自我 CPU 時間為 47.6 微秒，自我 NEURON 時間為 403.3 微秒，佔 100%。總自我 NEURON 時間為 403.3 微秒。

融合的 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

編譯後的函式在關鍵平均值表格中顯示為一個 Torch-Compiled Region。表格中，Torch-Compiled Region: 0/0 呼叫了 1 次，自我 CPU 時間為 98.4 微秒，自我 NEURON 時間為 375.5 微秒，佔 100%。總自我 NEURON 時間為 375.5 微秒。

編譯後的圖遵循一條透過執行時排程和硬體執行的路徑。表格顯示了一次裝置執行，如範例 2 所示，但程式碼採用了不同的路徑。它不是呼叫專門設計的融合運算，而是編譯器從原始 Python 函式中發現了融合。在此次執行中，Neuron 時間與 torch.rms_norm 相當，為 375 微秒，而 torch.rms_norm 為 403 微秒。

以下是三種變體的並排摘要：

原始運算（Naive ops）的 NEURON 時間為 2,355 微秒，自我 CPU 時間約 125 微秒，運算子呼叫次數為 6 次。

torch.rms_norm 的 NEURON 時間為 403 微秒，自我 CPU 時間約 48 微秒，運算子呼叫次數為 1 次。

原始程式碼加上 torch.compile 的 NEURON 時間為 376 微秒，自我 CPU 時間約 98 微秒，運算子呼叫次數為 1 次。

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

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