SemiAnalysis 实测 Trainium3 的 PyTorch 训练与 Neuron Explorer 剖析器

亚马逊的Trainium向训练王者发起冲击

Trainium最适合推理？但现在也擅长训练了：原生Trainium PyTorch、高质量CI、TorchTitan、GPU Impossible剖析器 Kimbo Chen与Dylan Patel 2026年10月23日∙付费

Trainium已在推理领域证明了自己。Anthropic正在Trainium2和Trainium3上运行生产推理和RL rollout。从历史上看，Trainium在推理方面表现出色，但并非研究训练的最佳选择。在推理中，模型架构的变化速度比训练慢得多；在训练中，研究人员会对模型架构进行大量不同的实验和消融，需要一个覆盖广泛API表面的技术栈。过去，这正是Trainium的短板所在。

AWS Neuron团队正积极努力改变这一缺陷，最大的转变是迈向一流的原生PyTorch训练体验。在训练循环中到处撒xm.step()、与定制的XLA式执行模型较劲的日子已经一去不复返。Trainium现在以PyTorch用户所期望的方式运行PyTorch——eager模式，配合torch.compile和标准API。

SemiAnalysis工程团队过去几周一直在试用Trainium3机架，发现体验很棒——我们很喜欢它整体的发展方向。需要说明的是，它还不完美，我们自己也遇到了几个bug。它远未达到CUDA生态的成熟度，但方向是对的。他们的团队借助智能体编码正在快速推进。

此外，与其他加速器公司不同——我们曾不得不费尽周折才能让他们的工程师获得一个还算稳定的CI集群——Trainium最多分配80,000个物理核心用于自动化质量控制。

在本文中，我们将介绍我们在Trainium3 TorchTitan和TorchNeuron技术栈上的体验，以及我们对Trainium的Neuron Explorer剖析器的体验——它远优于Nvidia和AMD的剖析器。我们制作了一系列演示，展示Trainium的剖析器如何做到“GPU Impossible”。

AWS副总裁Nafea Bshara在Trainium剖析器中展示Trainium的细粒度剖析器。来源：SemiAnalysis

在即将发表的一篇文章中，我们将深入探讨编写光速级NKI内核以及在Trainium3上进行微基准测试。如需深入了解Trainium硬件，请参阅我们此前进行完整硬件拆解的文章。

来源：SemiAnalysis我们正在招聘

SemiAnalysis正在招聘多个技术员工（Member of Technical Staff）岗位。如果你对这类分析感兴趣，请在下方申请：

在此申请

我们寻找的是：

积极进取且技术娴熟的技术员工，加入我们不断壮大的工程团队。

技术员工将在开发推理基准测试与系统建模方面发挥关键作用

我们正在招聘各经验层级的个人贡献者。薪酬具有竞争力

作为面试流程的一部分，你将完成一项付费编程挑战，其设计旨在反映SemiAnalysis的典型日常工作

薪酬具有竞争力，视经验、技能、地点和业务需求而定

任职资格

积极主动、自我驱动，能够在全球团队中独立工作。

具备PyTorch或JAX等机器学习框架的实证经验，可通过职业经历、个人项目或个人Substack博客体现

对以下至少1项有扎实理解：Transformer架构、LLM并行策略和/或CUDA并行编程

出色的研究能力，能够综合来自不同来源的信息以提炼洞见。

要求英语熟练，优先考虑中文熟练者迈向一流的PyTorch体验

亚马逊对其软件战略做出了重大方向调整，我们相信这将显著拓宽Trainium3的采用。TorchNeuron代表了一种原生PyTorch体验，而非通过PyTorch/XLA项目将Trainium生硬地拼接到PyTorch上。这一新后端面向所有不是前沿实验室内核工程师的人。AWS甚至聘请了谷歌原生Torch on TPU项目的负责人来从事Trainium原生Torch的工作。他从谷歌该项目启动之初就担任TPU软件负责人，涵盖XLA、xprof、Cloud TPU等。

旧方法让来自CUDA背景的人觉得Trainium很陌生。PyTorch/XLA依赖惰性张量图捕获，没有一流的eager模式。它也不支持PyTorch的原生分布式API（torch.distributed.*）或其原生并行API（DTensor、SimpleFSDP、DDP）。用户不得不依赖树外的XLA SPMD API，例如torch_xla.experimental.spmd_fsdp和torch_xla.distributed.spmd。对于习惯在CUDA上使用原生PyTorch的研究人员来说，切换到Trainium意味着要重新学习如何编写训练循环。

来源：AWS

通过在各处添加XLA标记来转换GPU研究代码库的痛苦，对任何用过PyTorch/XLA的人来说都是巨大的劝退因素。

来源：AWS

从高层来看，唯一需要做的改动就是将.to("cuda")替换为.to("neuron")，理论上绝大多数PyTorch ATen算子都能开箱即用地良好运行。

来源：AWS

理论上，借助PyTorch的torch.accelerator硬件无关API，像bf16 autocast这样的代码行甚至无需更改，尽管PyTorch中的一些API和一些研究代码库尚未迁移到torch.accelerator命名空间，仍直接调用torch.cuda.*命名空间。

由于机器人、生物技术、世界模型和自动驾驶领域的大多数学者和研究人员主要使用BF16，Trainium Torch在把BF16训练体验做好之前，不会在第0天支持TorchAO MXFP8。我们并不反对这一顺序，因为大多数进行MXFP8训练的人是大型一级实验室，他们有足够的内核工程师为大型训练任务编写自定义内核。Trainium Torch生态系统

PyTorch之所以流行，不仅因为核心PyTorch，还因为整个库生态系统，包括Hugging Face Transformers、PEFT、diffusers、TRL、Accelerate、TorchTitan等。我们惊讶地看到，Trainium Torch已经支持生态系统中大量此类库，包括Hugging Face Accelerate，其间接层数多得惊人。Trainium支持最受欢迎、最简洁的开源训练代码库之一TorchTitan，许多二级实验室将其分叉作为内部训练代码库的基础。Trainium的新PyTorch体验也完全支持这一点。

我们看到Google如何成功与Inferact/RadixArk合作，将TPU引入vLLM/SGLang生态系统，并且还让Meta成为TPU的生产客户。此外，通过将TPU引入vLLM/SGLang生态系统，TPU已经能够自动获得对Muse Spark Open的第0天支持，以及对即将推出的热门模型的第0天TPU支持的承诺。

虽然agentic coding的增加确实使引擎代码本身变得不那么重要，但vLLM/SGLang带来的是生态系统。无论开源模型是什么，在第0天，全球ML社区都信任vLLM/SGLang作为推理的首选引擎。没有任何vibe-coded引擎具有这种覆盖范围。

我们开始看到Trainium通过硬件访问以及在基于原生PyTorch构建的SGLang上（TorchNeuron）提前访问Trainium，迅速投资于vLLM/SGLang社区，类似于TorchTPU方法。

来源：SemiAnalysisTorchTitan体验

TorchTitan是最受欢迎的开源预训练框架之一，我们测试了Neuron的支持，以评估研究人员使用Trainium训练一个非平凡模型的用户体验。具体来说，我们在16颗Trainium2芯片上训练了一个Qwen3-32B，总共64个LNC=2逻辑核心，我们还尝试在具有64颗物理芯片的机架级Trainium3 UltraServer上进行训练。

总体而言，我们喜欢这种用户体验，尽管存在一些粗糙之处，Amazon团队已经在其路线图上进行修复。

虽然我们将在下面为完整性提供Trainium2 MFU，但需要关注的部分是，对于ML研究消融实验，最重要的不是光速优化的堆栈，而是良好的性能和对模型架构修改的广泛支持，这样研究人员就可以迭代而无需与堆栈作斗争。

使用Neuron官方的TorchTitan镜像，我们能够毫无问题地启动运行。我们在eager模式下测得平均训练吞吐量为每秒11,099个token，MFU为23.89%。对于compiled模式，我们看到每秒14,156个token和30.47% MFU，比eager模式提高了1.28倍吞吐量。我们尝试在建模代码中稍微更改模型架构，它也运行得很好。

Neuron还为更现代的模型架构提供训练配方，因此我们测试了其gpt-oss 20B训练配方。Neuron通过与GraphTrainer（TorchTitan的实验性训练器模块）集成来支持gpt-oss训练。具体来说，Neuron添加了SimpleFSDP的Trainium兼容配置、用于PyTorch计算图捕获的编译后端，以及用于注意力和MoE操作的自定义NKI内核。就用户体验而言，该配方开箱即用，一切顺利。

尚未使用上游API实现的一件事是TokenSwitch。在MoE时代，对于推理和训练，MoE combine和MoE dispatch是最重要的通信操作之一，占通信量的60-70%，而TokenSwitch是PyTorch为它们提供的一流API。Trainium尚不支持它，而是依赖一个名为NeuronEPTokenDispatcher的自定义TorchTitan类。将其迁移到TokenSwitch API看起来是一个简单的修复，AWS团队已承诺与Meta PyTorch团队合作修改TokenSwitch ABC，使其也能支持Trainium。

稍显令人失望的部分是 Neuron 编译器编译耗时很长，导致第一步等待时间很久。具体来说，第一步在 eager 模式下耗时 184.6 秒（compiled 模式下为 255.7 秒），比排除第一步后的平均步时长高出 30-55 倍。对于 gpt-oss，这次训练运行的编译时间高达 27 分 17 秒，这带来了糟糕的用户体验，尤其是对于尝试新模型架构的研究人员而言。造成这一现象的原因是，从历史上看，Trainium 编译器栈主要聚焦于推理，在推理场景下编译时间无关紧要，因为编译是离线完成的，并且每个镜像只缓存一次，不会影响生产推理；而在训练场景下则不然。虽然这在当时对推理而言是正确的选择，但对于迭代模型架构的研究人员来说，这是一个明显的摩擦点。

Neuron 团队已经在直接解决这个问题。即将发布的 TorchNeuron 版本将引入新的编译器后端，其中包含大量预编译的 NKI 内核，以及为快速编译而重新架构的编译栈，目标是在 eager 和 compiled 两种模式下将首步时间缩短最多两个数量级，同时性能回退极小或没有回退。如果他们兑现这一承诺，编译时间将不再是讨论的话题。训练 MatMul 性能

对于实验性训练而言，最重要的事情之一是在广泛的形状范围内具备不错的性能，因为 ML 研究人员会尝试许多不同的东西。我们在 PyTorch nn.functional.linear API 上对单个 Trainium3 逻辑核心（LNC=2）进行了基准测试，以了解其矩阵乘法性能，并作为参考与 GB300 和 B200 进行了比较。我们遍历了主要开放权重模型中投影操作的矩阵形状，并将 M 维度设置为 1、128 和 2048，以代表不同的批大小场景。

来源：SemiAnalysis

在我们的基准测试中，Trainium3 从矩阵乘法执行 1 GFLOP 计算的中等区间开始，就取得了明显高于 GB300 的 MFU。直到 100 GFLOP 之后，GB300 和 B200 才开始超过 Trainium3，而在 1 TFLOP 之后 Trainium3 再次胜出。FlexAttention API

FlexAttention 是一个高层 API，允许使用 PyTorch 实现不同的注意力变体，同时提供反向传播自动生成和合理的性能。良好的 FlexAttention 支持体现了硬件对研究人员的灵活性和易用性。我们在 Trainium3 上运行了 FlexAttention 基准测试套件，以了解其 FlexAttention 支持情况。

来源：SemiAnalysis

总体而言，Trainium3 在 96 个测试用例中通过了 89 个，通过率为 93%。所有 7 个失败用例都出现在大批大小（4）和序列长度（8192）的情况下，且全部是因编译超出时间预算而失败，而非数值错误。这些问题将在后续版本的 TorchNeuron 中修复。

此外，我们还尝试为 Trainium3 实现 FlexAttention 基准测试套件中未包含的自定义注意力变体，其数值结果与 Nvidia GPU 一致。

来源：SemiAnalysis 自动化质量保证

SemiAnalysis 在全球机器学习社区中以“CI 凯伦”著称，它为所有硬件厂商设定了高标准，并不惜一切代价促使 AI 芯片公司的高层真正为自家工程师提供足够的资源和 CI 集群，以维持高标准的自动化质量控制。

来源：SemiAnalysis

最近取得的一项成果，是与 AMD 高层斗争了 9 个多月，才为其工程师争取到足够多的稳定节点，从而在门禁测试组上达到 93% 的对等通过率。

来源：vLLM Dashboard

推理是一个范围狭窄的问题。一个实验室在服务少量生产模型时，需要的是前向传播、基本固定的架构以及一小组热点内核。

训练则是一个通用性问题。数千个算子中的每一个都需要反向传播，而研究代码库有时会意外地以更高精度运行算子，这也必须得到支持。为了确保 Trainium 能在广泛的模型架构上工作，并且能很好地支持 eager 模式和编译模式，AWS 收集了一份包含 10 万多个 Hugging Face 热门模型的庞大清单，这些模型涵盖了各种不同的归一化、MLP、MoE 布局和注意力类型。他们不仅测试它不会崩溃，还测试性能、哪些算子使用了 CPU 回退且尚未获得原生支持、图断点出现在哪里，最重要的是，他们针对 CPU 和 Nvidia 基线测试数值。他们还测试张量级对等性和损失曲线发散情况。在 10 分钟内完成一场考试很容易，但如果你考了 0 分，完成得多快都无关紧要。训练栈也是同样的道理。获得不错的 MFU% 很容易，但在不引入数值问题的同时获得不错的 MFU% 却很难。

来源：AWS

尽管通过率尚未达到 100%，但重要的是趋势。AWS 正在快速攀升，智能体集群正在系统性地修复剩余的失败项。我们认为，AWS 的内部汇总表——也就是高层看到的那张表——透明地展示差距而不是隐藏它们，这是一种优势。

我们这样说是有经验的。当 SemiAnalysis 和上游维护者试图在其他硬件厂商那里修复质量控制文化时，我们发现高层在看到真实的对等通过率时确实感到震惊。尤其是在某一家厂商，中层管理层的过滤如此严重，以至于当 CI 数据传到高层手中时显示为 99.99%，而当时实际的门禁对等通过率低于 5%。Neuron Explorer——优于 Nvidia 和 AMD 的剖析器

Trainium 的剖析器目前遥遥领先于业界，处于 SOTA 水平。在此我们分享初步印象，并计划在未来的 Trainium 内核文章中分享完整的用户体验。

Trainium 的剖析器深受 Anthropic 所有性能工程师的喜爱。Anthropic 的性能负责人（曾任职于 Jane 和 Finch Street）公开表示，在 Trainium 上进行性能分析远优于在 Nvidia GPU 上进行性能分析（显然也优于在 AMD GPU 上进行）。

来源：Tristan Hume 在 Twitter 上

We agree with Tristan Hume. In our experience, the profiling experience was amazing, after getting over the initial hump of the profiling setup being slightly confusing (but not impossible). There are multiple environment variables to set, different file types to save, and various commands to capture and view the profile traces. This is only the kernel profiling setup, and the PyTorch profiling setup is similar but with slight differences, making things even more confusing. With the help of Codex (since Claude refuses to write Trainium kernels), our best effort working setup is as follows:

Launch a persistent Neuron Explorer server

Profile the kernel with neuron-explorer capture command

Upload the profile trace as an artifact to the Neuron Explorer server as a web request

We shared feedback with AWS that we wished the documentation could be clearer on the workflow, explaining what environment variables need to be set for what reasons, and what purpose each file type serves.

Once the setup is done, the profiler user experience is great. Neuron Explorer is co-designed with the hardware. We particularly like the timeline separating hardware units in different rows, so that bubbles are visible at a glance. In contrast, GPU programmers have to roll their own setup for intra-SM profiling as the official Nsight tools don’t support that.

Source: SemiAnalysis

We also like how the viewer visualizes dependencies. When a block is selected, its dependencies and dependents light up in different colors, which helps us reason about work overlapping, e.g. whether computation sufficiently overlaps with memory loading. The image below shows that the selected white block is:

Flow-dependent on the red block

Order-dependent on the orange block

A flow dependency of the blue block

An output dependency of the purple block

Neuron Explorer dependency visualization. Source: SemiAnalysis

Neuron Explorer code editor mapping. Source: SemiAnalysis

Neuron Explorer: Event information when hovering on an event. Source: SemiAnalysis

Neuron Explorer memory bank utilization visualization. Source: SemiAnalysis

Neuron Explorer microsecond-level memory throughput visualization. Source: SemiAnalysis

Trainium Profiler Demo - Pixel Art

Neuron Explorer has many more features, including code editor mapping, MFU estimation, memory usage, and even creating pixel art due to how precise the profiler is. The images below are all from profiler traces of running on actual Trainium hardware, not fake profile reports. We were mainly inspired by the Bad Apple post from Daniel Liu (Anthropic Performance Engineer) and have created our own agentic coding SKILL.md for creating images. We did not have the time yet to create one to automatically generate videos frame by frame but will in the future and open source this SKILL.md for everyone to use. We mainly generate this pixel art by using the 16 DMA engines per physical core by “coloring in” regions by turning the DMA engines active, and the width of the shaded area is determined by how much data we have each DMA engine move.

Bad Apple on Neuron Explorer. Source: Daniel Liu

Visualizing the SemiAnalysis logo. Source: SemiAnalysis

Visualizing Jeff Bezos (pre looksmaxxing). Source: SemiAnalysis

Visualizing Dylan Patel (pre looksmaxxing). Source: SemiAnalysis

Since AWS is a partner of Nvidia, Jensen should tell his engineers to book a meeting with the AWS Trainium team on the hardware features needed to do accurate fine-grained profiling such that Nvidia can learn how to add this capability to Feynman or Feynman-Next GPUs.

Visualizing Jensen on Trainium3. Source: SemiAnalysis

Nvidia’s intra-kernel profiler Nsight Compute employs a different design. It offers aggregate metrics and detailed information on performance properties, rather than an event-based timeline viewer. The closest thing we found was Tri Dao’s QuACK intra-kernel profiler, which uses Nvidia’s In-Kernel Event Tracing API. While QuACK requires custom design to make Perfetto visualize accordingly, it’s great to see that AWS made this a well-integrated, first-class experience in Neuron Explorer.

QuACK’s intra-kernel profiler: Visualization of a QuACK GEMM kernel. Source: QuACK

Nsight Compute view of FlashAttention 4 workload analysis. Source: Nvidia Nsight