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

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

本文从零开始，不假设你有任何性能分析经验。本演练使用 AWS Trainium——一款专为深度学习训练和推理工作负载打造的 AI 芯片，以及 TorchNeuron，它是其开源的原生 PyTorch 后端。TorchNeuron 只需极少的代码更改，即可在 Trainium 上支持标准 PyTorch 工作流，包括 eager execution、 torch.compile以及标准 PyTorch API。此外，TorchNeuron 还与 Neuron Kernel Interface (NKI) 集成，以便为 Trainium 硬件实现自定义算子和性能优化。读完本文后，你应该了解：

• 如何配置 torch.profiler 并使用 Neuron 后端捕获干净的 trace。
• 如何阅读 key_averages() 表格，以识别哪些操作消耗 Neuron 设备时间，并比较它们的 Self CPU time。
• 如何在 Neuron Explorer 中跟踪一个操作，从框架分发到运行时调度，再到 NeuronCore 上的硬件执行。
• 这些信号在同一操作的三个不同实现中如何变化。

Neuron 作为 PyTorch 后端

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

Neuron 后端使用 PyTorch 的 PrivateUse1 扩展点。这一机制使加速器无需添加专用的 dispatch key（如 cuda）即可与框架集成。 torch_neuronx 包与上游 PyTorch 协同工作，将 Neuron 注册为 PyTorch 代码的原生后端。有关此扩展机制的详细信息，请参阅 Facilitating New Backend Integration by PrivateUse1.

设置 Profiler

你使用两个对象来分析 Neuron 工作负载：

• NeuronConfig 控制要捕获哪些 Neuron 数据流和 NeuronCore，以及输出的写入位置。你将其作为 torch.profiler.profile 传递给 experimental_config.
• 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)

The modes 参数选择要捕获的数据流，例如设备指令和运行时跟踪。The capture_enabled_for_nc 参数选择要跟踪的 NeuronCores。设置 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 Executable File Format（NEFF）包含编译后的 Neuron 可执行文件，而匹配的 NTFF 包含其硬件指令跟踪。这些文件共享一个数字哈希值，用于标识它们是一对。要在 Neuron Explorer 中打开分析结果，请上传系统分析或组合分析的完整输出目录。对于仅设备分析，请分别上传匹配的 .neff 和 .ntff 文件。有关 UI 和命令行说明，请参阅 Step 5: Upload your profile 。

分析 RMSNorm

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

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

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

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

复现本研究的所有脚本均可在 torch-neuronx examples 目录.

配置 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 功能。

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

Self 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/子系统/标识符 的模式：

在 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 在关键平均值表中显示为一个设备操作。

Self NEURON 总计：403.3 μs

融合后的 RMSNorm 遵循单一的框架到硬件路径，而不是单独分发每个算术操作。

比较 Lowering 路径

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

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

在最后一个示例中，使用以下代码编译原始的朴素 RMSNorm 函数： torch.compile.

使用 @torch.compile(backend="neuron")后，TorchDynamo 将 Python 字节码捕获为 FX 图，AOTAutograd 将其 lowers 为 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 总计: 375.5 μs

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

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

并排总结

Both 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 以及其他参与审阅的人!