Hugging Face Blog

Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.

6.9内容质量

TL;DR · AI 摘要

Profiling TorchNeuron Part 1 : System-level profiling on AWS Trainium. Back to Articles 0 -1 Enterprise Article Publishe...

核心要点

  • 主题聚焦：Profiling TorchNeuron Part 1 : System-level prof
  • 来源：Hugging Face Blog，建议结合原文判断细节。
  • AI 分析暂不可用，本条为保底评分与摘要。
#AI#编程#后端#安全#开源
打开原文

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

返回文章列表

[0

[-1

企业

]

文章

发布于 2026 年 10 月 21 日

点赞

6

[

Josh Longenecker

jlongen-aws

关注

amazon

许多机器学习实践者对如何让模型运行得更快有一些直觉，但很少有人知道如何分析模型运行缓慢的原因。在本系列的第一部分中，你将学习如何系统地分析推理性能，并识别时间消耗在何处。

本文从零开始，假设读者没有任何性能分析经验。本教程使用 AWS Trainium，这是一种专为深度学习训练和推理工作负载设计的 AI 芯片，以及其开源原生 PyTorch 后端 TorchNeuron。TorchNeuron 支持在 Trainium 上使用标准 PyTorch 工作流程，包括即时执行、torch.compile 和标准 PyTorch API，只需进行最小的代码修改。此外，TorchNeuron 与 Neuron Kernel Interface（NKI）集成，以启用自定义操作符并优化 Trainium 硬件的性能。到文章结束时，你应该了解：

  • 如何使用 Neuron 后端配置 torch.profiler 并捕获清晰的跟踪记录。
  • 如何通过 key_averages() 表识别消耗 Neuron 设备时间的操作，并比较它们的 Self CPU 时间。
  • 如何通过 Neuron Explorer 跟踪一个操作，从框架调度到运行时调度，直至在 NeuronCore 上的硬件执行。
  • 这些信号如何在相同操作的三个不同实现之间发生变化。

Neuron 作为 PyTorch 后端

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

Neuron 后端使用 PyTorch 的 PrivateUse1 扩展点。这种机制允许加速器在不添加专用调度密钥（如 cuda）的情况下与框架集成。torch_neuronx 包与上游 PyTorch 协作，将 Neuron 注册为 PyTorch 代码的原生后端。有关此扩展机制的详细信息，请参阅通过 PrivateUse1 实现新后端集成。

设置分析器

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

  • NeuronConfig 控制要捕获的 Neuron 数据流和 NeuronCores，以及输出的写入位置。你将其作为 experimental_config 传递给 torch.profiler.profile。
  • NeuronProfiler 接收一个 NeuronConfig，并提供一个 export_trace 回调。该回调将框架跟踪记录（trace.json）与 Neuron 系统跟踪记录一起写入。
code
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 参数选择要跟踪的 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 文件。神经元可执行文件格式（NEFF）包含编译后的神经元可执行文件，而对应的 NTFF 文件包含其硬件指令跟踪记录。这两个文件共享一个数字哈希值，用于标识它们是一对文件。要在神经元资源管理器中打开分析结果，请上传完整输出目录以进行系统或组合分析。对于仅设备的分析，请分别上传匹配的 .neff 和 .ntff 文件。有关 UI 和命令行操作的说明，请参见步骤 5：上传您的分析结果。

分析 RMSNorm

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

您将对同一操作的三种实现方式进行分析：

  • 使用单独 PyTorch 操作的朴素 RMSNorm：乘法、均值、加法、平方根倒数和乘法。
  • torch.rms_norm()，PyTorch 内置的已降低为单个融合神经元内核的操作。
  • 使用 torch.compile 的朴素 RMSNorm，与第一种实现方式使用相同的操作。

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

所有用于重现本工作的脚本均可在 torch-neuronx 示例目录中找到。

分析 1：朴素 PyTorch 操作

首先使用单独的 PyTorch 操作实现 RMSNorm。每一行都对应一个或多个 ATen 操作，并可能生成中间张量。

python
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
# 归一化 + 缩放

使用两个参数调用分析器以启用神经元后端：

  • 包含 ProfilerActivity.PrivateUse1 以收集神经元设备的活动数据。
  • 将 NeuronConfig 作为 experimental_config 传递以启用神经元特定的 Kineto 功能。

调度器在四次 prof.step() 调用中使用 wait=2、warmup=1 和 active=1。这会记录一次不包含初始设置开销的预热迭代。

在分析器循环中调用 torch.neuron.synchronize()。如果没有此同步，分析器可能在设备完成前导出数据，并导致神经元时间列为空。

python
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 操作。完整的 key_averages() 输出还包括 CPU 侧事件，如内存分配和主机到设备传输。

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

关键平均值表格显示，三个 aten::mul 调用占据了神经元设备时间的最大比例。

表格显示了哪些 ATen 操作被执行、执行频率以及它们在 CPU 和 Neuron 设备上消耗的时间。aten::mul 的总设备时间最高，因为该函数执行了三次乘法运算。

名称

调用次数

Self CPU

Self NEURON

% NEURON

aten::mul

3

60.9 微秒

1,164.5 微秒

49.5%

aten::add

1

26.6 微秒

434.6 微秒

18.5%

aten::rsqrt

15.6 微秒

415.0 微秒

17.6%

aten::mean

21.1 微秒

340.8 微秒

14.5%

Self NEURON 总计：2.355 毫秒

每一行直接对应到函数。aten::mul 分别用于 x * x、x * rms 以及最终的权重乘法。aten::mean 对最后一个维度进行降维。aten::add 添加 epsilon 值并触发主机到设备的拷贝，因为标量初始位于 CPU。aten::rsqrt 计算平方倒数。

Neuron Explorer 跟踪了从框架调度到运行时调度，再到 Neuron 硬件通道执行的 aten::mul 操作。

在 Explorer 中阅读追踪

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

系统追踪视图将事件组织到不同通道中。在默认分组模式下，标签遵循 instance/PID/subsystem/identifier 的格式：

通道

标签后缀

显示内容

框架调度

code
framework/TID:<id>

操作调度事件（PrivateUse1 运行时），按线程 ID 分组。

框架设备

code
framework/Stream:<n>

设备交互事件（PrivateUse1 驱动），按流编号分组。

Neuron 运行时

code
neuron_rt/NC:<n>

Neuron 运行时 API 调用，按逻辑 NeuronCore 分组。

Neuron 硬件

code
neuron_hw/NC<n>:<id>

硬件执行事件，按物理 NeuronCore 分组。

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

依赖链查看器突出了一个操作在框架、运行时和硬件通道中的生命周期。

分析 2：torch.rms_norm

将单个操作替换为 PyTorch 内置的 rms_norm 函数。Neuron 将该操作降低为一个融合内核：

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

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

aten::_fused_rms_norm

47.6 微秒

403.3 微秒

100%

Self NEURON 总计：403.3 微秒

融合的 RMSNorm 采用从框架到硬件的单一路径，而不是单独调度每个算术操作。

对比降低路径

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

分析 3：使用 torch.compile 的朴素 PyTorch

对于最后一个示例，使用 torch.compile 编译原始的朴素 RMSNorm 函数。

使用 @torch.compile(backend="neuron") 时，TorchDynamo 会将 Python 字节码捕获为 FX 图，然后 AOTAutograd 将其转换为 ATen 操作。Neuron 后端接收完整的图，而不是分别调度每个操作。它会融合操作、优化内存访问，并将图编译为单个 NEFF 软件包。该软件包作为一次硬件调度运行。

code
@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。

Torch-Compiled Region: 0/0

98.4 微秒

375.5 微秒

Self NEURON 总计：375.5 微秒

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

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

并排摘要

变体

NEURON 时间

Self CPU 时间

操作调用

简单操作

2,355 微秒

~125 微秒

torch.rms_norm

403 微秒

~48 微秒

简单 + torch.compile

376 微秒

~98 微秒

torch.rms_norm 和 torch.compile 都将 RMSNorm 融合为一个内核调度。在这些运行中，两种方法消耗的设备时间都约为简单实现的六分之一。融合消除了硬件执行之间的间隙，避免在分别调度的操作之间生成中间值（如 x_sq 和 mean_sq）。

结论

key_averages() 显示哪些操作消耗设备时间，其完整输出还揭示了 CPU 端的设置和传输。Explorer 提供了对应的可视化追踪。您可以从框架调度跟踪到运行时调度，直至硬件执行，并查看融合消除的间隙。

您可以在更大的模型上使用相同的流程：在表格中找到热点，在追踪中验证，然后确定融合操作或 torch.compile 是否能降低其成本。

在第二部分中，我们将进入设备级分析。您将打开 NEFF，检查 TensorEngine、VectorEngine 和直接内存访问（DMA）的分解，并将内核放置在性能极限线上，以确定它是计算受限还是内存受限。

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

本文提到的数据集 1

更多来自该作者的内容

通过 Strands Agents、LeRobot 和 Hugging Face 存储桶从一个地方记录、训练和部署

35

2026 年 8 月 13 日

一键从 Hugging Face 迁移到 Amazon SageMaker Studio

21

2026 年 7 月 7 日

社区

编辑

预览

通过在文本输入中拖放、粘贴或

点击此处

上传图片、音频和视频。

轻点或粘贴此处上传图片

评论

· 注册或登录以发表评论