TorchNeuron 性能剖析（一）：如何在 AWS Trainium 上做系统级性能分析

Hugging Face 社区博客·2026-10-22 00:04· 2 小时前AI 评分39

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

AI 导读

本教程演示如何用 torch.profiler 配合 NeuronConfig 与 NeuronProfiler 在 AWS Trainium 上抓取干净的推理性能 trace，并通过 key_averages() 表定位消耗 Neuron 设备时间的算子。

正文 · 原文

许多机器学习从业者对如何让模型跑得快有一些直觉，但很少有人知道如何分析模型为什么慢。在本系列的第 1 部分中，你将学习如何系统地分析推理性能，并找出时间花在了哪里。

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

• 如何为 torch.profiler 配置 Neuron 后端并抓取干净的 trace。
• 如何阅读 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 trace.json 与 Neuron 系统 trace 一起写出。

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 参数选择要抓取的数据流，例如设备指令和运行时 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() 以数值形式报告的内容，并检查调度之间发生了什么。

性能分析完成后，输出目录中包含系统跟踪记录；当你捕获设备数据时，还会包含匹配的 .neff 和 .ntff 文件。Neuron 可执行文件格式（NEFF）包含编译后的 Neuron 可执行文件，而匹配的 NTFF 包含其硬件指令跟踪。这些文件共享一个数字哈希值，用于标识它们是一对。要在 Neuron Explorer 中打开性能分析，请上传系统或组合性能分析的完整输出目录。对于仅设备性能分析，请分别上传匹配的 .neff 和 .ntff 文件。有关 UI 和命令行说明，请参阅步骤 5：上传你的性能分析。

分析 RMSNorm

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

你对同一操作的三种实现进行性能分析：

1. 使用单个 PyTorch 操作的朴素 RMSNorm：乘法、均值、加法、倒数平方根和乘法。
2. torch.rms_norm()，即 PyTorch 内置实现，被降级为单个融合的 Neuron 内核。
3. 使用 torch.compile 的朴素 RMSNorm，使用与第一个实现相同的操作。

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

用于复现这项工作的所有脚本都位于 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 后端所需的两个参数调用性能分析器：

1. 包含 ProfilerActivity.PrivateUse1 以收集来自 Neuron 设备的活动。
2. 将 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 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 以及最后与 weight 的乘法。aten::mean 对最后一维进行归约。aten::add 添加 epsilon 值并触发主机到设备复制，因为标量从 CPU 开始。aten::rsqrt 计算倒数平方根。

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

在 Explorer 中读取跟踪

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

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

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

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

Profile 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 调用覆盖完整计算。跟踪中包含一次调度和一个硬件事件，组件操作之间没有间隙。

Profile 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。

自身 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 提供了相应的可视化追踪。你可以跟踪一个操作从框架调度、运行时调度到硬件执行的整个过程，并看到融合所消除的间隙。

你可以在更大的模型上使用同样的工作流程：在表格中找到一个热点，在追踪中验证它，然后判断融合操作或 torch.compile 能否降低其成本。

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

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

来源：Hugging Face 社区博客 · huggingface.co