AWS Trainium推理性能瓶颈分析与TorchNeuron实践

系统级分析：在 AWS Trainium 上识别推理性能瓶颈

许多机器学习从业者对如何提升模型运行速度有一定直觉，但往往缺乏深入分析“为何运行缓慢”的能力。本文作为系列教程的第一部分，旨在介绍如何在 AWS Trainium 硬件平台上，利用 TorchNeuron 进行系统级的推断性能分析，精准定位时间消耗点。

本指南不假设读者具备先前的分析经验。所涉及的硬件为专门用于深度学习训练和推理的 AI 芯片 AWS Trainium，软件栈则基于开源原生 PyTorch 后端 TorchNeuron。TorchNeuron 支持标准 PyTorch 工作流程（包括急切执行、torch.compile 及常规 API），仅需最小化代码更改即可适配 Trainium。此外，通过与神经元内核接口 (NKI) 集成，TorchNeuron 还支持针对特定算子和 Trainium 硬件的性能优化。

通过本文的学习，你将掌握以下技能：

• 配置 torch.profiler 以配合神经元后端工作，并捕获清晰的性能痕迹。
• 解读 key_averages 表，识别消耗神经元设备时间的操作，并对比其 Self CPU 时间。
• 使用 Neuron Explorer 追踪单个操作从框架发送、运行时调度到 NeuronCore 硬件执行的完整生命周期。
• 理解不同实现方式下，上述信号的变化规律。

神经元作为 PyTorch 后端的集成机制

在深入具体工作量之前，需了解 Neuron 与 PyTorch 的集成原理。神经元后端利用 PyTorch 的 PrivateUse1 扩展点，允许加速器在不添加专用发送密钥（如 cuda）的情况下与框架集成。torch_neuronx 包与上游 PyTorch 协同工作，将 Neuron 注册为本地后端。关于此扩展机制的细节，可参考“PrivateUse1 促进新的后端集成”相关文档。

配置简介器 (Profiler)

描述一个神经元工作负载主要依赖两个对象：

• NeuronConfig：控制要捕获的神经元数据流和核心，以及输出写入位置。该配置需作为 experimental_config传递给 torch.profiler.profile。
• NeuronProfiler：使用 NeuronConfig 并提供 export_trace 回调。调用该回调会生成包含框架跟踪 (trace.json) 和神经系统跟踪的输出文件。

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 指定要跟踪的神经元核心。设置 correlate_device_time=True 可将设备时间传递至 key_averages() 表中，以便比较 CPU 与 Neuron 设备时间。完整的参数引用及可用的 ProfileMode 值请参阅 NeuronConfig 文档。

工具：Neuron Explorer

Neuron Explorer 是一个基于浏览器的工具，用于查看系统跟踪、设备时间表及每个引擎的分析数据。用户可通过 SSH 隧道或 Visual Studio Code 扩展访问它。在本文中，Explorer 用于验证 key_averages 报告的数值，并检查发送之间的底层行为。

输出目录包含系统跟踪文件，若捕获了设备数据，还会生成匹配的 .neff 和 .ntff 文件。NEFF（神经元可执行文件格式）包含编译后的神经元可执行文件，而 NTFF 则包含对应的硬件指令跟踪。这两个文件共享相同的数字哈希标识为一对。要在 Neuron Explorer 中打开配置文件，需上传完整的输出目录；对于仅设备的配置文件，则单独上传匹配的 .neff 和 .ntff 文件。

案例分析：RMSNorm 操作的三种实现

RMSNorm 是一个简单的算术运算序列，适合用于演示性能分析。本文对比了同一操作的三种实现方式：

1. 朴素 PyTorch 操作：由乘法、平均值、加法、平方根倒数和乘法等单个 ATen 操作组成。
2. 内置融合内核：使用 PyTorch 内置的 torch.rms_norm，降低为单一的融合神经元内核。
3. torch.compile 实现：使用 torch.compile 编译第一个朴素实现。

所有脚本使用相同的输入、配置器设置和时间表，确保结果可直接比较。示例脚本可在 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                                    
    mean_sq = x_sq.mean(dim=-1, keepdim=True)       
    rms = torch.rsqrt(mean_sq + 1e-6)               
    return t * rms * weight                         

为 Neuron 后端配置分析器需要两个关键步骤：

1. 包含 ProfilerActivity.PrivateUse1 以收集来自神经元设备的活动。
2. 通过 experimental_config 传入 NeuronConfig 以启用神经元特定的 Kineto 功能。

时间表设置为 wait=2, warmup=1, active=1，并在四个 prof.step() 调用中使用，以记录一个热迭代并排除初始设置开销。在循环中必须调用 torch.neuron.synchronize，否则分析器可能在设备完成前导出，导致神经元时间列为空。

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 排序可将设备成本最高的操作置于首位。下表展示了消耗神经元设备时间的 RMSNorm 操作（完整输出还包含分配和主机到设备传输等 CPU 端事件）。

print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))

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

Self NEURON 总计：2.355 ms

表中每行直接对应函数逻辑：aten::mul 执行三次乘法（x * x, x * rms, 最后乘以权重）；aten::mean 减少最后一个维度；aten::add 添加 epsilon 值（因标量始于 CPU，触发主机到设备复制）；aten::rsqrt 计算平方根的倒数。

Neuron Explorer 追踪了一个 aten::mul 操作，展示其从框架发送到运行时调度，再到 Neuron 硬件通道执行的过程。

在 Explorer 中阅读痕迹

在 Explorer 中，时间从左向右流动，事件生命周期从上向下流动。每个 ATen 操作遵循相同模式：在框架路径中创建，在运行时路径中安排为 nrta_execute_schedule，在硬件路径中执行设备为 nc_exec_running。

系统跟踪视图将事件组织成车道。默认分组模式下，标签遵循 instance/PID/subsystem/id 模式：

在急切模式下，每个操作独立运行。使用依赖链查看器可以追踪 aten::mul.Tensor 从框架层经过运行时调度到设备执行的生命周期。

The dependency chain viewer highlights the lifecycle of one operation across the framework, runtime, and hardware lanes.

配置文件 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 遵循一条从框架到硬件的路径，而不是分别调度每个算术操作。

比较降低路径

在朴素实现中，每个操作独立跨越调度边界，遵循各自的框架到运行时再到硬件的路径。而在融合实现中，一次 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。

Self NEURON 总计：375.5 μs

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

表格显示了一次设备执行，与配置文件 2 类似，但代码采取了不同的路径。编译器不是调用专门的融合操作，而是从原始 Python 函数中发现融合。在本次运行中，Neuron 时间与 torch.rms_norm 相当，为 375 μs 对比 403 μs。

并列总结