剖析 TorchNeuron 第一部分：AWS Trainium 上的系统级推理分析
许多机器学习从业者了解如何加速模型，但分析模型为何缓慢需要系统性的剖析。本指南将介绍如何使用 TorchNeuron 在 AWS Trainium 上分析推理性能，提供一种全面的方法来识别时间消耗热点，并展示内核融合的性能优势。
Neuron 作为 PyTorch 后端的集成
Neuron 后端通过 PrivateUse1 扩展点与 PyTorch 集成。该机制允许加速器在不使用专用分发键（如 cuda）的情况下与框架集成。torch_neuronx 包利用 PyTorch 的高级扩展能力，将 Neuron 注册为 PyTorch 代码的原生后端。
设置剖析器
要剖析 Neuron 工作负载，需要配置两个对象：NeuronConfig 和 NeuronProfiler。NeuronConfig 控制捕获哪些 Neuron 数据流和 NeuronCore，并指定输出位置。NeuronProfiler 接受此配置，并提供 export_trace 回调，将框架跟踪（trace.json）与 Neuron 系统跟踪一起写入。
关键配置参数包括：
modes：选择要捕获的数据流，例如设备指令和运行时跟踪。capture_enabled_for_nc：选择要跟踪的特定 NeuronCore。correlate_device_time=True：此设置将设备时间从跟踪传播到key_averages()表中，从而可以直接比较 CPU 和 Neuron 设备时间。
工具：Neuron Explorer
Neuron Explorer 是一个基于浏览器的查看器，用于可视化系统跟踪、设备时间线和每个引擎的细分。它允许用户确认 key_averages() 中的数值数据，并检查操作的详细生命周期。
剖析后，输出目录包含系统跟踪，以及匹配的 .neff（Neuron 可执行文件格式）和 .ntff（硬件指令跟踪）文件（如果捕获了设备数据）。要在 Neuron Explorer 中查看剖析结果，必须上传完整的输出目录以进行系统或组合剖析，或者单独上传匹配的 .neff 和 .ntff 文件以进行仅设备剖析。
案例研究：剖析 RMSNorm
选择 RMSNorm 进行分析，因为它是一个小操作，具有清晰的算术和归约操作序列。对 RMSNorm 的三种实现进行了剖析以比较性能：
- 朴素 RMSNorm： 使用单个 PyTorch 操作（乘法、均值、加法、倒数平方根、乘法）实现。
torch.rms_norm()： PyTorch 内置函数，被降级为单个融合的 Neuron 内核。- 朴素 RMSNorm +
torch.compile： 使用torch.compile和 Neuron 后端编译的原始朴素函数。
所有脚本使用相同的输入、剖析器配置和调度，以确保结果可直接比较。
剖析 1：朴素 PyTorch 操作
当 RMSNorm 使用单个 PyTorch 操作实现时，每行映射到一个或多个 ATen 操作，并可能创建中间张量。剖析结果显示，三个 aten::mul 调用占 Neuron 设备时间的最大份额。
| 名称 | 调用次数 | 自身 CPU | 自身 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% |
自身 NEURON 总计：2.355 ms
在这种朴素实现中，每个操作独立地跨越分发边界。例如，aten::add 操作会触发主机到设备的复制，因为标量 epsilon 从 CPU 开始。
剖析 2：torch.rms_norm()
通过用 PyTorch 的内置 rms_norm 函数替换单个操作，Neuron 将整个操作降级为一个融合内核。关键平均值表将此操作显示为单个设备操作。
| 名称 | 调用次数 | 自身 CPU | 自身 NEURON | NEURON 百分比 |
|---|---|---|---|---|
| aten::_fused_rms_norm | 1 | 47.6 μs | 403.3 μs | 100% |
自身 NEURON 总计：403.3 μs
这种融合实现遵循单一的框架到硬件路径，避免了单独分发每个算术操作的开销。
剖析 3：使用 torch.compile 的朴素 PyTorch
当使用 @torch.compile(backend='neuron', fullgraph=True) 编译原始朴素 RMSNorm 函数时，TorchDynamo 将 Python 字节码捕获为 FX 图。Neuron 后端接收此完整图，并将其融合并编译为单个 NEFF 工件。此工件作为一次硬件分发运行。
| 名称 | 调用次数 | 自身 CPU | 自身 NEURON | NEURON 百分比 |
|---|---|---|---|---|
| Torch-Compiled Region: 0/0 | 1 | 98.4 μs | 375.5 μs | 100% |
自身 NEURON 总计：375.5 μs
此编译图也遵循通过运行时调度和硬件执行的单一路径。此处实现的 Neuron 时间（375 μs）与专门构建的 torch.rms_norm（403 μs）相当。
并排性能摘要
| 变体 | NEURON 时间 | 自身 CPU 时间 | 操作调用次数 |
|---|---|---|---|
| 朴素操作 | 2,355 µs | ~125 µs | 6 |
| torch.rms_norm | 403 µs | ~48 µs | 1 |
| 朴素 + torch.compile | 376 µs | ~98 µs | 1 |
torch.rms_norm 和 torch.compile 都成功地将 RMSNorm 融合为一次内核分发。这些融合方法将设备时间减少到朴素实现设备时间的约六分之一。融合消除了硬件执行之间的间隙，并防止在单独分发的操作之间物化中间值，例如 x_sq 和 mean_sq。
结论
key_averages() 表识别哪些操作消耗设备时间，而完整输出提供 CPU 端设置和传输的上下文。Neuron Explorer 提供相应的可视化跟踪，允许用户从框架分发到运行时调度再到硬件执行跟踪操作，并观察融合消除的间隙。
此工作流程可应用于更大的模型：在表中识别性能热点，在跟踪中可视化验证，然后确定融合操作或 torch.compile 是否可以降低其成本。未来的分析将转向设备级剖析，检查 TensorEngine、VectorEngine 和 DMA 细分，以确定内核是计算密集型还是内存密集型。