Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.

本文为原文前 6,000 字符的节选翻译，完整内容请查看原文。

Back to Articles Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium. Enterprise Article Published October 21, 2026 Upvote 6 Josh Longenecker jlongen-aws Follow amazon Many ML practitioners have some intuition about how to make models run fast, but fewer know how to analyze why a model is slow. In Part 1 of this series, you will learn how to systematically analyze inference performance and identify where the time is spent. This post starts from scratch and assumes no prior profiling experience. The walkthrough uses AWS Trainium, a purpose-built AI chip for deep learning training and inference workloads, and TorchNeuron, its open source native PyTorch backend.

返回文章：Profiling TorchNeuron（第一部分）：AWS Trainium 上的系统级性能分析。企业文章，发布于 2026 年 10 月 21 日，点赞数 6，作者 Josh Longenecker (jlongen-aws)，关注 amazon。许多机器学习从业者对于如何让模型运行得更快有一些直觉，但很少有人知道如何分析模型为何运行缓慢。在本系列的第一部分中，你将学习如何系统地分析推理性能并确定时间消耗在哪里。本文从零开始，不假设你有任何性能分析经验。本指南使用 AWS Trainium（一种专为深度学习训练和推理工作负载打造的 AI 芯片）以及 TorchNeuron（其开源原生 PyTorch 后端）。

TorchNeuron supports standard PyTorch workflows on Trainium, including eager execution, torch.compile, and standard PyTorch APIs, with minimal code changes. Additionally, TorchNeuron integrates with Neuron Kernel Interface (NKI) to enable custom operators and performance optimization for Trainium hardware. By the end, you should know: How to configure torch.profiler with the Neuron backend and capture a clean trace. How to read the key_averages() table to identify which operations consume Neuron device time and compare their Self CPU time. How to follow one operation through Neuron Explorer, from framework dispatch through runtime scheduling to hardware execution on a NeuronCore. How these signals change across three implementations of the same operation.

TorchNeuron 支持 Trainium 上的标准 PyTorch 工作流，包括即时执行（eager execution）、torch.compile 和标准 PyTorch API，且只需极少的代码更改。此外，TorchNeuron 集成了神经元内核接口（NKI），以实现自定义算子和针对 Trainium 硬件的性能优化。读完本文，你应该了解：如何配置带有 Neuron 后端的 torch.profiler 并捕获清晰的跟踪记录；如何读取 key_averages() 表以识别哪些操作消耗了 Neuron 设备时间，并比较它们的自身 CPU 时间；如何通过 Neuron Explorer 跟踪一个操作，从框架调度到运行时调度，再到 NeuronCore 上的硬件执行；以及这些信号在同一操作的三种不同实现中是如何变化的。

Neuron as a PyTorch Backend Before you profile a workload, it helps to understand how Neuron integrates with PyTorch. The Neuron backend uses PyTorch’s PrivateUse1 extension point. This mechanism lets an accelerator integrate with the framework without adding a dedicated dispatch key such as cuda. The torch_neuronx package works with upstream PyTorch to register Neuron as a native backend for PyTorch code. For details about this extension mechanism, see Facilitating New Backend Integration by PrivateUse1.

作为 PyTorch 后端的 Neuron：在分析工作负载之前，了解 Neuron 如何与 PyTorch 集成会很有帮助。Neuron 后端使用 PyTorch 的 PrivateUse1 扩展点。这种机制允许加速器与框架集成，而无需添加像 cuda 这样的专用调度键。torch_neuronx 包与上游 PyTorch 协作，将 Neuron 注册为 PyTorch 代码的原生后端。有关此扩展机制的详细信息，请参阅“Facilitating New Backend Integration by PrivateUse1”。

Set Up the Profiler You use two objects to profile a Neuron workload: NeuronConfig controls which Neuron data streams and NeuronCores to capture and where to write the output. You pass it to torch.profiler.profile as experimental_config. NeuronProfiler takes a NeuronConfig and provides an export_trace callback. The callback writes the framework trace, trace.json, alongside the Neuron system traces. nc = NeuronConfig( modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME], profile_output_dir=OUTPUT_DIR, capture_enabled_for_nc=‘0’, correlate_device_time=True, ) exporter = NeuronProfiler(nc) The modes parameter selects the data streams to capture, such as device instructions and runtime traces. The capture_enabled_for_nc parameter selects the NeuronCores to trace. Setting correlate_device_time=True propagates device time from the trace into the key_averages() table, where you can compare CPU and Neuron device time. For the complete parameter reference and available ProfileMode values, see the NeuronConfig Profile Modes documentation.

设置性能分析器：你需要使用两个对象来分析 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 文档。

Tooling: Neuron Explorer Neuron Explorer is a browser-based viewer for the system trace, device timeline, and per-engine breakdown. You can access it through an SSH tunnel or the Visual Studio Code extension. In this post, you use Explorer to confirm what key_averages() reports numerically and inspect what happens between dispatches. After profiling, the output directory contains the system trace and, when you capture device data, matched .neff and .ntff files. The Neuron Executable File Format (NEFF) contains the compiled Neuron executable, while the matching NTFF contains its hardware instruction trace. The files share a numeric hash that identifies them as a pair. To open the profile in Neuron Explorer, upload the complete output directory for a system or combined profile. For a device-only profile, upload the matching .neff and .ntff files individually. See Step 5: Upload your profile for the UI and command-line instructions.

工具：Neuron Explorer：Neuron Explorer 是一个基于浏览器的查看器，用于查看系统跟踪、设备时间轴和各引擎的细分情况。你可以通过 SSH 隧道或 Visual Studio Code 扩展访问它。在本文中，你将使用 Explorer 来确认 key_averages() 在数值上报告的内容，并检查调度之间发生了什么。分析完成后，输出目录中包含系统跟踪文件；如果捕获了设备数据，还会包含匹配的 .neff 和 .ntff 文件。Neuron 可执行文件格式 (NEFF) 包含编译后的 Neuron 可执行文件，而匹配的 NTFF 包含其硬件指令跟踪。这些文件共享一个数字哈希值，用于将它们识别为一对。要在 Neuron Explorer 中打开分析结果，请上传系统或组合分析的完整输出目录。对于仅设备分析，请分别上传匹配的 .neff 和 .ntff 文件。请参阅“步骤 5：上传你的分析结果”以获取 UI 和命令行说明。

Profile RMSNorm For this first analysis, you profile RMSNorm, a small operation with a clear sequence of arithmetic and reduction operations. You profile three implementations of the same operation: Naive RMSNorm with individual PyTorch operations: multiply, mean, add, reciprocal square root, and multiply. torch.rms_norm(), the PyTorch built-in lowered to a single fused Neuron kernel. Naive RMSNorm with torch.compile, using the same operations as the first implementation. The scripts use the same inputs, profiler configuration, and schedule so the results are directly comparable. All scripts to reproduce this work are available in the torch-neuronx examples directory.

分析 RMSNorm：在第一次分析中，你将分析 RMSNorm，这是一个具有清晰算术和归约操作序列的小型操作。你将分析同一操作的三种实现：使用单个 PyTorch 操作（乘法、均值、加法、倒数平方根和乘法）的朴素 RMSNorm；torch.rms_norm()，即降级为单个融合 Neuron 内核的 PyTorch 内置函数；以及使用与第一种实现相同操作的、带有 torch.compile 的朴素 RMSNorm。这些脚本使用相同的输入、分析器配置和调度，因此结果具有直接可比性。所有用于重现此工作的脚本都可以在 torch-neuronx 示例目录中找到。

Profile 1: Naive PyTorch Operations Start by writing RMSNorm with individual PyTorch operations. Each line maps to one or more ATen operations and can create intermediate tensors. 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 # square mean_sq = x_sq.mean(dim=-1, keepdim=True) # mean of squares rms = torch.rsqrt(mean_sq + 1e-6) # 1 / sqrt(mean_sq + eps) return t * rms * weight # normalize + scale Call the profiler with two arguments required for the Neuron backend: Include ProfilerActivity.PrivateUse1 to collect activity from the Neuron device. Pass NeuronConfig as experimental_config to enable the Neuron-specific Kineto features. The schedule uses wait=2, warmup=1, and active=1 across four prof.step() calls. This records one warm iteration without including initial setup overhead. Call torch.neuron.synchronize() inside the profiler loop. Without this synchronization, the profiler can export before the device finishes and leave the Neuron time column empty. 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

分析 1：朴素 PyTorch 操作：首先编写使用单个 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 功能。调度在四次 prof.step() 调用中使用 wait=2、warmup=1 和 active=1。这记录了一次预热迭代，而不包含初始设置开销。在分析器循环内调用 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