AWS 博客：用 torch.profiler 分析 Trainium 推理性能

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

本文从零开始，假设读者没有性能分析经验。演练使用 AWS Trainium，这是一款专为深度学习训练和推理工作负载打造的 AI 芯片，以及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，以及将输出写入何处。你将它作为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)

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

---

工具：Neuron Explorer

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

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

---

分析 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                                    # 平方
    mean_sq = x_sq.mean(dim=-1, keepdim=True)       # 平方的均值
    rms = torch.rsqrt(mean_sq + 1e-6)               # 1 / sqrt(mean_sq + eps)
    return t * rms * weight                         # 归一化 + 缩放

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

1. 包含ProfilerActivity.PrivateUse1以收集来自 Neuron 设备的活动。
2. 将NeuronConfig作为experimental_config传入，以启用 Neuron 特定的 Kineto 功能。

调度使用wait=2，warmup=1，和active=1跨越四个prof.step()调用。这记录了一次预热迭代，不包含初始设置开销。

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

使用 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),
) 作为 prof:
    用于 step 在 范围(4):
        使用 torch.no_grad():
            _ = rms_norm_fn(x)
        torch.neuron.synchronize()
        prof.step()

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

打印(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 中读取 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 将该操作降低为一个融合内核：

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

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

自身 NEURON 总计：403.3 μs

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

比较降低路径

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

---

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

在最后一个示例中，用以下方式编译原始的朴素 RMSNorm 函数torch.compile。

使用@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 以及其他人的审阅！