OFERRY正在进入矩阵

Amazon 开源 TorchNeuron：Trainium 变成 PyTorch 的原生设备，改动只有一行

10 月 19 日夜里，AWS 连着发了两份材料：PyTorch 官方博客挂出《Introducing TorchNeuron: Native PyTorch backend for AWS Trainium》，Amazon 的新闻页同步发了一篇《Amazon opens its AWS Trainium AI chips to every PyTorch developer, no rewrite needed》。两篇的落点一致——今天（10 月 20 日）在 PyTorch Conference North America 上，TorchNeuron 按 Apache 2.0 开源，代码在 github.com/aws-neuron/torch-neuronx，装机方式是 pip install torch-neuronx，同时以 public beta 上线，支持 Trainium2 和 Trainium3。

宣传语只有一句：不用重写代码就能换芯片。这句话在中文技术圈传得很快，但「不用重写」到底覆盖到哪一层、哪一层它自己承认还没做，官方文档其实写得很清楚。这篇把两边都拆开。

「不用重写」到底指什么#

PyTorch 博客给的样例是把整个训练循环改成 Trainium，改动是一行设备指针：

PYTHON
# Run a training loop on Trainium
model = MyModel().to('neuron')   # the only change
optimizer = torch.optim.AdamW(model.parameters())
for batch in dataloader:
    optimizer.zero_grad()
    loss = criterion(model(batch), targets)
    loss.backward()
    optimizer.step()
READ / EXECUTE / UNDERSTANDUTF-8

第二层是编译路径。如果你用 torch.compile，要么在第一次编译前声明一次默认后端，要么在每次编译时点名：

PYTHON
# Run a compiled training loop on Trainium
torch.compiler.set_default_backend("neuron")  # call once, before any compile
model = MyModel().to('neuron')
model.compile()  # no backend argument
READ / EXECUTE / UNDERSTANDUTF-8

第三层是加载权重。Hugging Face 侧的写法就是 device_map：

PYTHON
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    # Neuron hardware is automatically detected
    device_map="neuron"
)
READ / EXECUTE / UNDERSTANDUTF-8

官方说 device_map="auto" 也能自动发现，因为 Neuron 已经进了 PyTorch 的设备枚举。这三段都需要一台 trn2 或 trn3 实例才能真正跑起来——本机装个包是跑不通的，这一点先记住，后面还会说。

它是怎么接进 PyTorch 的：PrivateUse1#

这是全篇最值得看的一点。官方自述，TorchNeuron 站在 PyTorch 社区造的两个扩展点上面：PrivateUse1 device backend 和 torch.accelerator APIs（PyTorch 博客原文：builds directly on the accelerator-agnostic extension points, including the PrivateUse1 device backend and torch.accelerator APIs, created by the PyTorch community）。

PrivateUse1 是 PyTorch 给自研硬件留的官方位置——把设备注册进去之后，torch.nn、torch.tensor、torch.autocast、torch.distributed、torch.profiler 这些标准 API 就会按写好的样子作用在 neuron 设备上。换句话说，AWS 这次不是又发一套「Trainium 专用 API」，而是把自家硬件伪装成 PyTorch 的一个普通设备。差异在迁移成本上：专用 API 意味着训练/推理栈要分叉，标准设备意味着上游库（Transformers、vLLM、TorchTitan）不需要为它写分支。

代价是名字上的混乱。pip 包名沿用了 torch-neuronx——这个字符串和 Neuron SDK 里既有的 PyTorch 包重名，搜文档时很容易把新旧两条路径混着看。装之前先确认自己看的是哪一版文档。

torch.compile 的降级路径#

官方给出的编译链路是：TorchDynamo 抓图 → AOTAutograd 提前 trace 反向图 → Neuron 后端经 Torch-MLIR 降到 StableHLO → Neuron 编译器产出 NEFF 文件在硬件上执行。

两种执行模式都支持：eager 逐算子立即执行、每个算子各自编译一个 NEFF 并搬运中间结果，适合调试；torch.compile 负责稳定后的高性能路径，文档说可以只编译模型的一部分，并支持 FlexAttention。另外有一个实验性的 adaptive eager mode，目的是降低启动开销、在不改代码的前提下把算子融合得接近编译模式的性能。

基于同一套 torch.compile 支持，AWS 还搭了两个推理插件：vLLM Neuron（服务 LLM）和 vLLM Omni Neuron（服务图像/视频生成的扩散模型）。

官方自己列的取舍清单#

这部分是文章里最该抄下来的一段。PyTorch 博客写明，2026 年 8 月那次上游测试跑了约 154K 个测试结果，其中约 69K 被跳过：约 30K 是 PyTorch 自身的资格检查（CUDA-only、CPU-only 这类），与 Neuron 无关；剩下约 39K 是 Neuron 主动跳过，分三类：

  • 暂时还没有原生 Neuron lowering 的算子（约 22K，最大的一块）——逐元素、reduce、foreach、高级索引、高级线代算子的边角情形，以及融合的 LSTM/GRU cell。
  • CUDA 专属测试——CUDA graphs、NCCL 专用集合通信、CUDA 专属的 Dynamo 路径。
  • 首个版本明确选择不支持的功能：composite-operator 合规性、稀疏张量与稀疏算子、NestedTensor、前向模式自动微分（JVP）、原生多头注意力里的任意 attention mask（只支持无 mask 与 causal）、per-channel 量化、negative/conjugate view bits、torch.signal 的窗函数、bilinear/bicubic 上采样，以及动态与符号形状。

把这张表读懂，比读任何营销句都有用：如果你手上是稀疏模型、动态形状的序列、或者需要自定义 attention mask 的架构，这次迁移不是改一行，而是先在跳过清单里排队。

数值类型上，官方口径是 BF16、FP32、FP16 加整型与布尔都能用，torch.autocast 与混合精度代码按原样工作；TorchAO 和 FP8 目前是私有 Beta，下一个版本才转正式。

分布式这边覆盖面相当宽：DDP、FSDP2、SimpleFSDP、张量并行、流水线并行，以及 FSDP2+TP、ZeRO 这类组合，还有 DTensor 与 DeviceMesh；集合通信（all_gather、reduce_scatter、all_reduce）与 P2P（send、recv）都在。通信路径是硬件原生的——机内走 Trainium 专用的 collective core，机间走 EFA。

有一个语义细节容易被忽略，值得单独点出来：torch_neuronx.Stream / torch_neuronx.Event 映射了 PyTorch 的流语义，但非默认流是复用在同一条硬件队列上的，不会各自拿到一个独立引擎；优先级、顺序和事件语义都保留，只是没有真正的并行引擎。另外，当 host collectives 被关掉时，发在非默认流上的集合通信会保持顺序但跑在默认硬件队列上。写多流流水线的代码要按这个前提设计。

想自己写算子：NKI#

除了吃标准算子，官方还给了低层入口 Neuron Kernel Interface（NKI），可以写自定义算子并在 eager 与 torch.compile 下执行。PyTorch 博客给了一个按 register_autograd 风格注册的完整例子（节选）：

PYTHON
import torch
import nki
import nki.language as nl
from torch_neuronx import nki_op

@nki.jit
def sin_kernel(in_ptr0, out_ptr):
    input_tile = nl.load(in_ptr0[0:128])
    output_tile = nl.sin(input_tile)
    nl.store(out_ptr[0:128], value=output_tile)
    return out_ptr

@nki_op("mylib::pt_nki_sin", mutates_args={})
def pt_nki_sin(x: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    sin_kernel(x, out)
    return out

def pt_nki_sin_backward(ctx, grad):
    x, = ctx.saved_tensors
    return grad * pt_nki_cos(x)  # cos is derivative of sin
READ / EXECUTE / UNDERSTANDUTF-8

配套的 NKI-lib 也开源了，里面是 Trainium 优化过的高性能算子，覆盖 FlashAttention 与 MoE blockwise kernel。我顺手核了一下仓库：aws-neuron/nki-library 是 Apache-2.0，最近一次推送是 2026-10-19 20:56 UTC（发布会前一天），当前 73 stars。

性能怎么量#

Profiling 走标准入口，用 NeuronConfig 扩展 torch.profiler：

PYTHON
from torch_neuronx.profiling import NeuronConfig, ProfileMode, NeuronProfiler

nc = NeuronConfig(
   modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],
   profile_output_dir=OUTPUT_DIR,
   capture_enabled_for_nc='0',
   correlate_device_time=True,
)
with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
    experimental_config=nc,
    on_trace_ready=NeuronProfiler(nc).export_trace,
) as prof:
    ...
# 按设备侧耗时排序，直接看热点
print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))
READ / EXECUTE / UNDERSTANDUTF-8

trace 可以在 Neuron Explorer 里看，能从任一事件跳到源码行，也有 dependency chain viewer 看每个 stall 在等谁。另外它还能输出 Parquet（per-op、per-engine 指标），官方明说这份数据可以被 agent 拿去做自动瓶颈分析——这算是这条工具链里比较新的一个设计取向。

生态侧几条可核的事实#

  • TorchTitan 已经能跑。官方给了从源码装 torchtitan-neuron 并在单台 trn2.48xlarge 上用 Simple FSDP + torch.compile 训练 Llama3-8B 的完整命令，关键参数是 NEURON_RT_NUM_CORES=64、--compile.enable --compile.backend neuron。
  • Hugging Face 侧已在上游支持。官方口径：Trainium 是 Transformers、Diffusers、trl、kernels 里的受支持设备，Diffusers 的张量并行已上游合入，NKI kernel 通过 HF kernels 库提供。Hugging Face 的 Jeff Boudier 在官方稿里的说法是「TorchNeuron 让 Trainium 成为 Transformers 与 Diffusers 的原生后端」。
  • 上游在解耦硬件假设。AWS 说自己在往 PyTorch 上游提交，把「绑定特定芯片」的代码去掉，并附了 20 多个 PyTorch 主仓 PR 编号（179063、192365、193439、193441、193442 等）与 5 个 TorchTitan PR。这条是它和纯粹「再包一层」的区别：如果只是包一层，不会去动主仓的 CUDA 假设。
  • 客户口径要看清性质。Deepgram 的 CTO 在官方稿里给了一个数字：在它们评估的基准里，Trainium 每美元处理的输入音频秒数比其 GPU 基线多约 60%——请注意这是厂商自述、且没有公开基准细节，属于「官方确认发布、第三方未复现」；Cartesia 的说法是标准推理流水线「改动极小」就跑起来了，同样没有原始数据。
  • 仓库侧的事实（我用 GitHub API 查的）：aws-neuron/torch-neuronx 是 Apache-2.0、默认分支 main、描述写明支持 eager / torch.compile / 分布式 / NKI 自定义 kernel，创建时间 2025-09-24、最近一次推送 2026-10-17，发布当天 stars 数 54。仓库先于发布存在一年，说明这不是发布会当天临时搭的壳。
  • Amazon 的芯片底账：官方稿说自研芯片业务已超过 250 亿美元年化收入，2015 年收购 Annapurna Labs 起建了四代芯片家族。TechCrunch 的定性更直接——「Amazon 对 NVIDIA 最猛的一次出手」，这是第三方报道口径。

动手前先做一次算子自检#

官方把「22K 条没有原生 lowering」写在博客里，等于把「会不会踩到」这件事交给你自己判断。最省时间的做法不是先跑，而是先用 PyTorch 自己的 dispatcher 把你模型实际用到的算子列出来，再拿这份名单去对官方的算子覆盖表。

TorchDispatchMode 是 PyTorch 自带的钩子，能在不碰模型代码的前提下记录每一次算子调用：

PYTHON
# op_inventory.py —— 迁移前先看看你的模型到底碰了哪些 aten 算子
import torch, torch.nn as nn
from torch.utils._python_dispatch import TorchDispatchMode

class OpInventory(TorchDispatchMode):
    def __init__(self):
        self.seen = {}

    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        key = str(func)                       # 例如 aten.gelu.default
        self.seen[key] = self.seen.get(key, 0) + 1
        return func(*args, **(kwargs or {}))

model = nn.Sequential(nn.Linear(64, 64), nn.GELU(), nn.LayerNorm(64))
inv = OpInventory()
with inv:                                    # 这里的每一次算子调用都会被记录
    model(torch.randn(8, 64))

for name, n in sorted(inv.seen.items(), key=lambda kv: -kv[1]):
    print(f"{n:5d}  {name}")
READ / EXECUTE / UNDERSTANDUTF-8

跑一个小模型只会打出二三十行，跑真实模型会打出几百行。值得警惕的通常不是 aten.mm / aten.addmm 这类基础算子（它们在第一批支持名单里），而是三类：aten._sparse_*（稀疏）、带 _nested / NestedTensor 字样的调用、以及需要传任意 attn_mask 的 _scaled_dot_product_*。这三类正好压在官方「明确不支持」的第一版清单上。

顺便给一份按项目类型分开的判断，比逐条对清单更快：

  • 标准 Transformer 的训练与微调：主路径已经被 vLLM Neuron、TorchTitan 和 Hugging Face 上游覆盖，可以试。
  • 依赖稀疏 / 动态形状 / 自定义注意力 mask / per-channel 量化：先假设不可用，再逐个验证——这四项都在官方跳过清单上。
  • 重度依赖 CUDA graphs 或 NCCL 专有行为：需要改造，因为那些是 CUDA 专属测试里被跳过的部分。
  • 要自己写高性能算子：NKI + NKI-lib 是这次一起放出来的，属于「有路可走」，但要接受写硬件相关代码的成本。

这次发布里，其它自研芯片厂商值得抄的两点#

抛开 AWS 自己的得失，这次开源有两个动作是可以直接抄的：

一是先改上游，再发自家后端。 官方在稿子里列出了二十多个 PyTorch 主仓 PR 与 5 个 TorchTitan PR，内容是把「绑定特定芯片」的假设从上游代码里拆掉（例如让 torch.accelerator 的设备枚举不再隐含 CUDA）。如果只把自己的硬件包一层，这些改动是不需要的。先投入去解耦上游，等于把「生态以后要不要支持你」这件事变成默认选项而不是谈判项。

二是把不支持清单和盘托出。 22K 条没有 lowering、八项首版不支持——这些数字对市场部没有好处，但对工程团队是唯一的决策依据。对一个想让人把生产负载搬过来的平台，「我知道哪里不行」比「我哪里都行」更可信。

反过来说，这次发布也给所有自研加速器留了一个共同问题：PyTorch 已经证明它愿意为第三方硬件提供标准扩展位（PrivateUse1 与 torch.accelerator），那么剩下的门槛就不是框架，而是「你的汇编器/编译器能覆盖多少算子」和「云上买不买得到」。 前者是工程题，后者是商业题，跟「要不要写一套新 API」已经没关系了。

对国内开发者的三条现实#

第一，门槛不在代码，在账号和区域。 本文所有命令都需要 trn2/trn3 实例；Trainium 是 AWS 自有的云上加速器，自有服务器和国内云都买不到对应硬件。看到「改一行就能跑」时，要补一句「前提是你已经在 AWS 上有 Trainium 配额」。本次发布我没有在官方页面找到关于中国大陆区域可用性的说明，按未确认处理，不要按「国内能直接买到」做规划。

第二，public beta + 沿用旧包名，生产别直接压。 官方明写 10 月 20 日起是 public beta，pip 包名又和既有的 torch-neuronx 重名。API 在 beta 期变动是常态。

第三，别把「无感迁移」当成「全部可用」。 官方跳过清单里那几项——稀疏、NestedTensor、动态/符号形状、任意 attention mask、per-channel 量化——恰恰是 LLM 之外很常见的需求。迁移前先拿自己模型的算子名单对一遍 torch.compile 的支持面，比看完宣传语直接开跑省时间。

总结#

TorchNeuron 真正的价值不在「改一行」这个口号，而在两件事：一是把 Trainium 注册进 PyTorch 的 PrivateUse1，让上游生态不用为自己分叉；二是把 22K 条没有原生 lowering 的算子、以及首版明确不支持的八项功能写进了同一篇官方博客。前者决定了迁移成本，后者决定了迁移边界——两个都读，才知道这次开源对你手上的项目意味着什么。

要动手的话，顺序是：一台 trn2 实例 → pip install torch-neuronx → 把 .to('cuda') 换成 .to('neuron') → 跑 eager 确认能过 → 再上 torch.compile(backend="neuron") → 用 NeuronConfig 拉一份 trace 看有哪些算子在设备外。至于国内团队，先把「什么时候能用上 Trainium」这个前提解决，再谈迁移。

最后补一句时间口径：PyTorch 博客发布于 10 月 19 日，官方把开源节点写在 10 月 20 日（PyTorch Conference North America 当天），仓库最近一次推送是 10 月 17 日。看二手摘要时会同时遇到「已经开源了」和「即将开源」两种说法，它们指的是同一件事的不同环节——写这篇文章时，仓库已经是公开可访问的 Apache-2.0 仓库，而「正式发布」的仪式在会场上。

参考与致谢#

Amazon 开源 TorchNeuron：Trainium 变成 PyTorch 的原生设备，改动只有一行

https://www.oferry.com/posts/a482/
作者
摆渡
发布于
许可协议
CC BY-NC-SA 4.0

除另有声明外，本文采用上述许可协议。转载或改编请注明作者与原文地址，并遵守许可条款。