从 PyTorch 分派到 Trainium 执行：TorchNeuron 性能分析阅读

全文概览

分析模型为何慢，需要把 Python 算子、主机调度和设备执行连接起来。AWS 的 TorchNeuron 教程以一个小型 RMSNorm 运算为例，展示如何先用 PyTorch Profiler 找热点，再在 Neuron Explorer 中追踪同一操作跨越框架、运行时和 NeuronCore 的过程。

原文比较朴素 PyTorch 运算、内置融合算子和编译后的朴素实现。后两者在作者的单次示例中减少了设备侧累计时间，轨迹也显示分派路径发生变化。理解这些结果，需要同时看采样窗口、设备同步、算子语义和计时范围，不能仅凭一张耗时表推断整模型收益。

本文阅读 Josh Longenecker 在 AWS 官方账号发布的 Profiling TorchNeuron (Part 1)。原文页面日期为 2026 年 10 月 21 日，官方 RSS 时间为当日 16:04:43 UTC，即北京时间 10 月 22 日 00:04:43。先用下面的跨层轨迹建立整体认识，后文再解释采集、读图与优化验证。

图源为 Josh Longenecker / AWS 的原文 Profile 1 截图，用于解释性能诊断，版权归原作者。打开高清图。黄色关联线把框架事件连到运行时与硬件事件，横向是时间，纵向是不同职责的泳道。内嵌图用于识别路径，细小事件名称和时间刻度请打开高清图查看。

背景与约束

TorchNeuron 通过 PyTorch 的 PrivateUse1 扩展机制接入 Trainium，使代码使用 neuron 设备。原生后端接入后，eager 运算、融合算子与 torch.compile 会走不同的分派路径，同一数学表达式可能对应多个执行事件，也可能被编译为一个图。

RMSNorm，即均方根归一化，包含平方、均值、加 epsilon、倒平方根和带权缩放。教程选用形状为 4×512 的 BF16 输入，朴素表达式中有三次乘法、一次均值、一次加法和一次倒平方根。它足够小，能把多次分派与中间张量的代价暴露出来，工作负载仍与真实长上下文模型有明显区别。

这里的首要问题是计时对象。Self CPU 描述主机侧事件耗时，Neuron 设备时间描述对应设备活动，端到端请求延迟还包括其他步骤。主机和设备可以异步工作，不能把表中两个 Self 时间直接相加，或把 Python 调用返回视为设备执行结束。

整体方案与关键变化

用同一套采样规则保留跨层关联

教程通过 NeuronConfig 选择 DEVICE 和 RUNTIME 数据流，并选择需要采集的 NeuronCore。将 correlate_device_time=True 打开后，设备活动才能关联回 key_averages() 中对应的框架事件。Profiler 活动同时包含 CPU 与 PrivateUse1，Neuron 配置通过 experimental_config 传入。

采样安排为 wait=2、warmup=1、active=1，对应四次 prof.step()。前面的步骤用于等待和预热，最后记录一个已预热迭代。这样的窗口避免把初次初始化完全混入稳态示例，但一个活跃迭代仍不足以说明重复运行的波动。

每轮中的 torch.neuron.synchronize() 用来等待设备完成。原文指出，缺少同步可能使导出发生在设备结束之前，导致设备时间列为空。这里的同步服务于采样完整性，在正式服务循环中插入同步则可能改变并行性，应单独衡量。

从热点表走到执行路径

第一步按 self_privateuse1_time_total 排序，找设备时间最多的操作。朴素 RMSNorm 中 aten::mul 调用三次，因此累计设备时间最多。这个排名帮助定位调查对象，尚不能证明乘法内核本身存在算法瓶颈。

第二步进入 Explorer。框架线程泳道展示算子分派，框架流泳道展示驱动侧交互，运行时中可以找到 nrta_execute_schedule，硬件泳道则包含 nc_exec_running。沿同一事件的依赖链查看，可以区分主机提交、运行时等待与实际硬件活动。截图中不同层事件的长度与数量并不需要一一相等。

导出文件也要保持完整。trace.json 记录框架轨迹，NEFF 是编译后的 Neuron 可执行文件，配对的 NTFF 是硬件指令轨迹。原文要求用相同数字 hash 对应的文件配对，系统或联合 profile 上传完整输出目录，设备 profile 则使用匹配的文件，避免将不同运行混在一起。

融合算子与图编译改变了什么

内置 RMSNorm 在 Neuron 后端降低为一个融合算子。相较逐个执行算术操作，它把完整运算交给一条较集中的分派路径，减少组件操作之间的间隙以及中间结果在独立算子间的物化。

图源为 Josh Longenecker / AWS 的原文 Profile 2 截图。打开高清图。这里应关注融合算子与下游调度、硬件事件之间的关联。两张截图的可视窗口与时间范围不同，不能直接用彩色条在页面上的像素长度计算加速倍数。

torch.compile(backend="neuron", fullgraph=True) 则从朴素 Python 函数捕获图，经 TorchDynamo 和 AOTAutograd 形成后端可处理的 ATen 图，再完成融合、内存访问优化与 NEFF 编译。作者报告的表中呈现一个 Torch-Compiled Region。它说明编译器可以从组合运算中找到融合机会，实际是否成功仍需查看目标函数的图与运行轨迹。

证据与工程代价

原文三组结果使用相同输入和采样配置。下表保留作者侧的计时范围，单位统一为微秒。CPU 数值为原文汇总的近似值。

这组数据支持了在该小算子示例中减少分派、融合计算的作用。编译版本设备时间更低，但其 Self CPU 高于内置融合版本，优化不同层会有不同代价。376 µs 与 403 µs 的差异来自教程报告的运行，缺少重复次数和误差范围时，不能据此建立稳定的优劣排名。

计时表与截图也不适合直接逐项相减。轨迹可能显示框架包装、调度、等待和具体硬件区间，Profiler 的聚合统计又有自身的归属规则。图中的单个硬件事件时长不必等于一行 Self NEURON 累计值，诊断时应先确认事件和度量的定义。

教程没有提供本文可据以统一横向比较的完整实例类型、Trainium 代际和全部软件版本组合，也没有给出整模型吞吐、首 token 延迟或请求分位数。本文未在 Trainium 上独立复现，以上数字均按作者示例解读。读者可从原文提供的 示例脚本目录开始，再记录自己的软硬件条件。

局限与适用边界

首先应对齐数值语义。原文朴素实现显式使用 1e-6，内置调用没有显式传入 eps。PyTorch 官方 RMSNorm 参考说明，eps=None 由计算类型决定，不能自动视为与显式常量相等。原文输入为 BF16，而权重创建时未显式指定 dtype，复现时还应检查权重、输出类型与累加精度。

下面仅示意如何把 epsilon 写入内置调用，其他实现也应保持相同语义。这段代码没有在本文环境中运行。

import torch.nn.functional as F

eps = 1e-6
y = F.rms_norm(x, (x.shape[-1],), weight, eps=eps)

显式 epsilon 仍不保证不同实现逐位一致，低精度中间值与融合累加可能引入数值差异。比较耗时前应使用与任务精度要求相符的容差检查输出，之后再保留同一份输入、预热和采样条件进行计时。这个复现建议来自对示例代码和 API 的核对，并非对作者设备结果的独立验证。

其次，Profiler 自身会增加开销，显式同步也会改变异步执行方式。对 4×512 的小输入观察到的分派收益，可能在大 batch、长序列或算力饱和时改变。教程提及的 TensorEngine、VectorEngine、DMA 与 roofline 分析属于后续设备级主题，本篇系统级轨迹不足以判断所有内核是计算受限还是带宽受限。

可验证的启发

以下是基于原文的后续尝试，本文未执行。先固定输入、权重 dtype 和 epsilon，对三种实现做输出一致性检查，再重复采集多个已预热迭代，同时记录设备时间、主机时间和不带 Profiler 的墙钟时间。如果聚合表改善而墙钟不改善，就继续检查同步、提交与采样开销，暂不宣称服务变快。

随后可以只改变 batch 和归一化宽度，观察朴素与融合路径的差距如何变化，同时保存分派次数和中间张量相关事件。若加速主要集中在很小输入，结果更支持主机与分派开销解释，仍需要设备级证据判断较大输入的瓶颈。

迁移到完整模型时，可以先选真实 trace 中占比明确的一个热点做改动，再回测首 token 延迟、解码吞吐和请求分位数。若热点自身占比很小，单算子的大幅改善也可能只带来有限的整链路收益。保留最小改动与完整 trace，才能判断优化是否改变了真正影响用户等待的路径。

参考资料

• Josh Longenecker / AWS：Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.，2026 年 10 月 21 日发布，重点为 Profiler 配置、Profile 1–3、跨层事件与汇总表。

• PyTorch RMSNorm 官方参考，用于核对 epsilon 与归一化定义，当前 main 文档不代表教程未写明的安装版本。

• 原文公开示例目录，保留复现入口，本文未执行其中代码。

• 两张 Neuron Explorer 截图均来自 AWS 作者原文，仅选择解释跨层关联与融合路径所需的截图作技术评论，保留作者、出处和高清原始地址。