剖析 TorchNeuron（第 1 部分）：在 AWS Trainium 上进行系统级性能分析

许多机器学习从业者对于如何让模型运行得更快有一定直觉，但较少人知道如何分析模型为何缓慢。在本系列的第1部分中，你将学习如何系统地分析推理性能并识别时间消耗所在。

本文从零开始，假设读者没有先前的性能分析经验。演练使用AWS Trainium（专为深度学习训练和推理工作负载打造的AI芯片）及其开源原生PyTorch后端TorchNeuron。TorchNeuron支持在Trainium上使用标准PyTorch工作流，包括即时执行、torch.compile和标准PyTorch API，只需最少的代码改动。此外，TorchNeuron与Neuron Kernel Interface（NKI）集成，以支持自定义算子和针对Trainium硬件的性能优化。读完本文后，你应了解：

如何使用Neuron后端配置torch.profiler并捕获干净的跟踪。

如何阅读key_averages()表格，以识别哪些操作消耗Neuron设备时间，并比较它们的Self CPU时间。

如何通过Neuron Explorer跟踪一个操作，从框架分发经运行时调度到NeuronCore上的硬件执行。

这些信号如何在同一操作的三个实现之间发生变化。

Neuron作为PyTorch后端

在分析工作负载之前，了解Neuron如何与PyTorch集成会有帮助。

Neuron后端使用PyTorch的PrivateUse1扩展点。该机制允许加速器与框架集成，而无需添加如cuda这样的专用分发键。torch_neuronx包与上游PyTorch配合，将Neuron注册为PyTorch代码的原生后端。有关此扩展机制的详细信息，请参阅Facilitating New Backend Integration by 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

Neuron Explorer是基于浏览器的查看器，用于系统跟踪、设备时间线和每个引擎的分解。你可以通过SSH隧道或Visual Studio Code扩展访问它。在本文中，你使用Explorer来确认key_averages()在数值上报告的内容，并检查分发之间发生的情况。

分析后，输出目录包含系统跟踪，并且在捕获设备数据时包含匹配的.neff和.ntff文件。Neuron可执行文件格式（NEFF）包含已编译的Neuron可执行文件，而匹配的NTFF包含其硬件指令跟踪。这些文件共享一个数字哈希，将它们标识为一对。要在Neuron Explorer中打开配置文件，请上传系统或组合配置文件的完整输出目录。对于仅设备配置文件，请单独上传匹配的.neff和.ntff文件。有关UI和命令行说明，请参阅Step 5: Upload your profile。

分析RMSNorm

对于首次分析，你分析RMSNorm，这是一个具有清晰算术和归约操作序列的小型操作。

你分析同一操作的三个实现：

使用单独PyTorch操作的朴素RMSNorm：乘法、均值、加法、倒数平方根和乘法。

torch.rms_norm()，即降低为单个融合Neuron内核的PyTorch内置函数。

使用与第一个实现相同操作的朴素 RMSNorm，配合 torch.compile。

这些脚本使用相同的输入、分析器配置和调度，以便结果可以直接比较。

重现此工作的所有脚本均可在 torch-neuronx 示例目录中找到。

分析 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 # 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

使用 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 的总设备时间最高，因为该函数执行了三次乘法。

自身 NEURON 总计：2.355 ms

每一行直接映射到函数。aten::mul 运行于 x * x、x * rms 以及最终与权重的乘法。aten::mean 对最后一维进行归约。aten::add 添加 epsilon 值，并触发主机到设备的复制，因为标量从 CPU 开始。aten::rsqrt 计算倒数平方根。

Neuron Explorer 跟踪 aten::mul 操作，从框架分发通过运行时调度到 Neuron 硬件通道上的执行。

在 Explorer 中读取跟踪

在 Explorer 中，时间从左到右流动，事件生命周期从上到下流动。每个 ATen 操作遵循相同的模式：在框架通道中创建，在运行时通道中作为 nrta_execute_schedule 调度，并在硬件通道中作为 nc_exec_running 进行设备执行。

系统跟踪视图将事件组织成通道。在默认分组模式下，标签遵循 instance/PID/subsystem/identifier 模式：

在 eager 模式下，每个操作独立运行。使用依赖链查看器从框架层通过运行时调度跟踪 aten::mul.Tensor 到设备执行。

依赖链查看器突出显示一个操作在框架、运行时和硬件通道中的生命周期。

分析 2：torch.rms_norm

用 PyTorch 内置的 rms_norm 函数替换单独的操作。Neuron 将该操作降低为一个融合内核：

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

融合的 RMSNorm 在关键平均值表中显示为一个设备操作。

自身 NEURON 总计：403.3 μs

融合的 RMSNorm 遵循一条从框架到硬件的路径，而不是分别调度每个算术运算。

比较降低路径

在朴素实现中，每个操作独立跨越调度边界，并遵循各自的从框架到运行时到硬件的路径。在融合实现中，一次 Python 调用覆盖完整计算。跟踪包含一次调度和一次硬件事件，组件操作之间没有间隙。

配置文件 3：带有 torch.compile 的朴素 PyTorch

对于最后一个示例，使用 torch.compile 编译原始的朴素 RMSNorm 函数。

使用 @torch.compile(backend="neuron")，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。

Self NEURON total: 375.5 μs

编译后的图通过运行时调度和硬件执行遵循一条路径。

该表显示一次设备执行，与配置文件 2 中相同，但代码采用不同的路径。编译器不是调用专门构建的融合操作，而是从原始 Python 函数中发现融合。在此运行中，Neuron 时间与 torch.rms_norm 相当，为 375 μs，而后者为 403 μs。

并排摘要

torch.rms_norm 和 torch.compile 都将 RMSNorm 融合为一次内核调度。在这些运行中，任一方法大约占用朴素实现设备时间的六分之一。融合消除了硬件执行之间的间隙，并避免了在单独调度的操作之间物化中间值（例如 x_sq 和 mean_sq）。

结论

key_averages() 显示哪些操作消耗设备时间，而其完整输出还揭示了 CPU 端的设置和传输。Explorer 提供相应的可视化跟踪。您可以跟踪一个操作从框架调度到运行时调度再到硬件执行，并看到融合消除的间隙。

您可以对更大的模型使用相同的工作流程：在表中找到热点，在跟踪中验证它，然后确定融合操作或 torch.compile 是否可以降低其成本。

在第 2 部分中，我们将进入设备级配置文件。您打开 NEFF，检查 TensorEngine、VectorEngine 和直接内存访问（DMA）分解，并将内核放置在屋顶线上，以确定它是计算受限还是内存受限。

特别感谢 Aritra Roy Gosthipaty 启发我开始这个系列，以及我的同行 Pinak Panigrahi、Samir Araujo、Akhil Azhikodan、Esha Lakhotia 和其他人进行审阅！