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

许多机器学习从业者对如何让模型快速运行有一些直觉，但知道如何分析模型运行缓慢原因的人要少得多。在本系列的第 1 部分中，你将学习如何系统地分析推理性能，并确定时间花在了哪里。

本文从零开始讲解，假设你没有任何性能分析经验。本教程使用 AWS Trainium——一款专为深度学习训练和推理工作负载打造的 AI 芯片——以及 TorchNeuron，这是它的开源原生 PyTorch 后端。TorchNeuron 支持在 Trainium 上使用标准 PyTorch 工作流，包括即时执行、torch.compile 和标准 PyTorch API，且只需进行少量代码修改。此外，TorchNeuron 集成了 Neuron Kernel Interface (NKI)，以支持自定义算子和针对 Trainium 硬件的性能优化。完成本文后，你将了解：

• 如何使用 Neuron 后端配置 torch.profiler 并捕获干净的跟踪记录。
• 如何阅读 key_averages() 表格，找出哪些操作会占用 Neuron 设备时间，并比较它们的自身 CPU 时间。
• 如何在 Neuron Explorer 中跟踪一项操作，从框架分发到运行时调度，再到 NeuronCore 上的硬件执行。
• 同一操作的三种实现中，这些指标如何变化。

Neuron 作为 PyTorch 后端

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

Neuron 后端使用 PyTorch 的 PrivateUse1 扩展点。这一机制使加速器能够与框架集成，而无需添加专用的分发键，例如 cuda。torch_neuronx 软件包与上游 PyTorch 协作，将 Neuron 注册为 PyTorch 代码的原生后端。有关此扩展机制的详情，请参阅 通过 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 和命令行说明，请参阅 步骤 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 的形式执行。

系统跟踪视图将事件整理到各个通道中。在默认分组模式下，标签遵循以下格式：实例/PID/子系统/标识符：

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

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）的细分情况，并将内核标注在屋顶线图上，以判断它是受计算能力还是内存带宽限制。

特别感谢 Aritra Roy Gosthipaty 带来开启本系列的灵感，也感谢我的同事 Pinak Panigrahi、Samir Araujo、Akhil Azhikodan、Esha Lakhotia 等人帮忙审阅！