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

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（trace.json）与 Neuron 系统 trace 一并写出。

modes 参数用于选择要捕获的数据流，例如设备指令和运行时 trace。capture_enabled_for_nc 参数用于选择要跟踪的 NeuronCore。设置 correlate_device_time=True 会将设备时间从 trace 传播到 key_averages() 表中，你可以在其中比较 CPU 和 Neuron 设备时间。有关完整的参数参考和可用的 ProfileMode 值，请参阅 NeuronConfig Profile Modes 文档。

工具：Neuron Explorer

Neuron Explorer 是一个基于浏览器的查看器，用于查看系统 trace、设备时间线和各引擎分解。你可以通过 SSH 隧道或 Visual Studio Code 扩展访问它。在本文中，你使用 Explorer 来确认 key_averages() 以数值形式报告的内容，并检查各次调度之间发生了什么。

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

剖析 RMSNorm

在第一次分析中，你剖析 RMSNorm，这是一个小型操作，具有清晰的算术和归约操作序列。

你剖析同一操作的三种实现：

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

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

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

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

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

剖析 1：朴素 PyTorch 操作

首先使用单个 PyTorch 操作编写 RMSNorm。每一行映射到一个或多个 ATen 操作，并可能创建中间张量。

调用剖析器时，需要传入 Neuron 后端所需的两个参数：

包含 ProfilerActivity.PrivateUse1 以收集来自 Neuron 设备的活动。

将 NeuronConfig 作为 experimental_config 传入，以启用 Neuron 特有的 Kineto 功能。

该调度在四次 prof.step() 调用中使用 wait=2、warmup=1 和 active=1。这会记录一次预热迭代，而不包含初始设置开销。

在剖析器循环内调用 torch.neuron.synchronize()。如果没有此同步，剖析器可能会在设备完成之前导出，导致 Neuron 时间列为空。

剖析完成后，检查 key_averages() 表。按 self_privateuse1_time_total 排序，将设备开销最高的操作排在最前面。此处显示的表已筛选为报告 Neuron 设备时间的 RMSNorm 操作。完整的 key_averages() 输出还包括 CPU 侧事件，例如内存分配和主机到设备传输。

关键平均值表显示，三个 aten::mul 调用占 Neuron 设备时间的最大份额。

该表显示了哪些 ATen 操作运行过、运行频率以及它们在 CPU 和 Neuron 设备上花费的时间。aten::mul 的总设备时间最高，因为该函数执行了三次乘法。

自身 NEURON 总计：2.355 ms

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

Neuron Explorer 跟踪一个 aten::mul 操作，从框架调度经过运行时调度，直到在 Neuron 硬件通道上执行。

在 Explorer 中读取 Trace

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

系统 trace 视图将事件组织到各个通道中。在 Default 分组模式下，标签遵循 instance/PID/subsystem/identifier 的模式：

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

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

Profile 2：torch.rms_norm

将各个独立操作替换为 PyTorch 内置的 rms_norm 函数。Neuron 会将该操作降级为一个融合内核：

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

自身 NEURON 总计：403.3 μs

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

比较降级路径

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

Profile 3：使用 torch.compile 的朴素 PyTorch

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

使用 @torch.compile(backend="neuron") 时，TorchDynamo 将 Python 字节码捕获为 FX 图，AOTAutograd 将其降级为 ATen 操作。Neuron 后端接收完整的图，而不是分别调度每个操作。它会融合这些操作、优化内存访问，并将该图编译为单个 NEFF 产物。该产物作为一次硬件调度运行。

编译后的函数在关键平均值表中显示为一个 Torch-Compiled Region。

自身 NEURON 总计：375.5 μs

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

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

并排总结

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

结论

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

你可以在更大的模型上使用相同的工作流程：在表中找到热点，在 trace 中验证它，然后确定融合操作或 torch.compile 能否降低其开销。

在第 2 部分中，我们将进入设备级剖析。你将打开 NEFF，检查 TensorEngine、VectorEngine 和直接内存访问（DMA）的分解情况，并将内核放置在 roofline 上，以确定它是受计算限制还是受内存限制。

特别感谢 Aritra Roy Gosthipaty 启发我开启本系列，并感谢我的同事 Pinak Panigrahi、Samir Araujo、Akhil Azhikodan、Esha Lakhotia 以及其他同事的审阅！

本文提到的数据集 1

jlongen-aws/blog-assets

来自该作者的更多内容

使用 Strands Agents、LeRobot 和 Hugging Face Storage Buckets 在一个地方进行记录、训练和部署

从 Hugging Face 一键进入 Amazon SageMaker Studio

本文提到的数据集 1

jlongen-aws/blog-assets