TorchNeuron：改一个设置，就能把现有 PyTorch 代码跑在 AWS Trainium 上

「不用改代码」到底承诺了什么

Amazon 的演示刻意做得很无聊，而这正是硬件移植里你最想要的东西。工程师把同一份训练脚本分别跑在常规芯片和 Trainium 上，唯一的差别是一个告诉代码该用哪块芯片的设置。Amazon 说模型以同样的速度学会了同样的东西。没有重写的数据加载器，没有换掉的优化器，也没有 fork 出来的训练循环。

机制上它成立，是因为 TorchNeuron 通过 PyTorch 内置的「新硬件自我注册」机制注册成后端。这件事比听起来重要：当设备是通过标准路径接入、而不是私有 fork，那么挂在 dispatcher 上的一切——自动求导、checkpoint、profiling 钩子、分布式封装——都会继续工作。你原本用来测试、调试和扩模型的那套工具，还是那套工具。

这条说法诚实的作用域是「可移植」，不是「等价」。可移植说的是你的代码能跑；等价说的是你的数字对得上。公告强力支持前者，后者只由 Amazon 自己的测试支撑。所以你该用哪个测试来卡住迁移，就很清楚了。

# TorchNeuron is announced on October 19, 2026 and open-sourced under Apache 2.0
# at PyTorch Conference North America on October 20, 2026. Do not install a
# "torchneuron" package you found in a search result. Read the official repo first.

# The repository named in Amazon's announcement:
git clone https://github.com/aws-neuron/torch-neuronx
cd torch-neuronx
git log -1 --format='%H %ci'      # pin the exact commit you evaluated
cat LICENSE                        # confirm Apache 2.0 yourself

# Two facts worth writing down before you start, because they decide scope:
#   1. supported silicon is listed as Trainium2 and Trainium3
#   2. the backend registers through PyTorch's standard mechanism for new
#      hardware, so tooling on the dispatcher keeps working

# Do not trust a blog post (including this one) for the device string.
# Read the release notes in the repo, then write it down once:
echo 'export TORCH_DEVICE=<backend-name-from-release-notes>' >> ~/.neuron_env

第一天唯一重要的数字是损失对齐

人的本能是先搬最大的训练任务，因为省钱在那儿。别这么干。先搬你真正在乎的最小任务，在把昂贵硬件预订下来之前，证明损失曲线是对得上的。

对齐要同时满足三件事：相同的随机种子和数据顺序；一个你在看结果之前就写下来的损失容差；以及第一个开始分叉的 step 记录。第三件事最常被跳过。早期出现一个很小、之后一直平的差距，通常是浮点累加顺序变了；而逐步扩大的差距是真的数值或 kernel 问题。这两种情况的处理方式完全不同。

下面这个对齐脚本故意写得很笨，而在这里「笨」是优点。它通过一个设备开关把同一份脚本跑两遍，按固定 step 导出损失，一旦分叉就大声失败。有个细节：不要从任何博客（包括这篇）里抄设备字符串。去仓库的 release notes 里读，然后把值放进一个环境变量，这样改错只需要改一处，而不是四十处。

# Parity harness. Run this BEFORE you move a big job. Three assertions:
#   1. same seed and same data order on both devices
#   2. loss stays inside a tolerance you wrote down in advance
#   3. you record the divergence point, not just pass/fail

import json, os, subprocess, sys

TOLERANCE = 2e-3          # decide this before you look at the numbers
STEPS_TO_COMPARE = [10, 50, 100, 200]

def run(device: str, out_path: str) -> None:
    env = dict(os.environ, TORCH_DEVICE=device, LOSS_DUMP=out_path)
    subprocess.run([sys.executable, "train.py", "--seed", "1337"], check=True, env=env)

def load(path: str) -> dict:
    return {int(k): float(v) for k, v in json.load(open(path)).items()}

run("cpu", "/tmp/loss_cpu.json")
run(os.environ["TORCH_DEVICE"], "/tmp/loss_hw.json")

a, b = load("/tmp/loss_cpu.json"), load("/tmp/loss_hw.json")
for step in STEPS_TO_COMPARE:
    delta = abs(a[step] - b[step])
    status = "ok" if delta <= TOLERANCE else "FAIL"
    print(f"step {step:>4} cpu={a[step]:.6f} hw={b[step]:.6f} delta={delta:.2e} {status}")

# A failed step is information, not a blocker: it tells you whether the gap is
# floating-point ordering (small, early, doesn't grow) or a real kernel issue.

工程周到底花在哪

Amazon 说它每次构建都对几十万个 AI 模型跑测试，好让问题在客户负载跑起来之前暴露。这是一个关于覆盖面的有效信号，但它不是关于你的模型的陈述。你的模型有私有的数据加载器、有两个没人写文档的自定义算子，还有那段只在你的序列长度下才起作用的梯度裁剪。

发布时的集成面值得记一下，因为它告诉你哪些路已经铺好：Hugging Face 已提供原生 Trainium 支持；TorchTitan 配 FSDP 可以做分布式训练；Amazon 向 vLLM 提交了重要改动；它还把整套性能 kernel 开源，给要自己写 kernel 的团队当起点。

公告里有个容易被忽略的二阶项：Amazon 说它在向上游 PyTorch 提交改动，去掉那些把框架绑死在特定芯片上的代码。如果这真的落地，移植税对所有人都会下降，不只是 Trainium 用户。押注某个后端是厂商决策；押注标准路径是工程决策，而且这是更容易在评审里站住的那个。

# Distributed training stays on standard PyTorch: TorchTitan + FSDP. The launcher
# is torchrun, unchanged. What changes is the device backend, and that is the point
# of the announcement. Keep every experiment reproducible by pinning all three:
# torch version, neuron backend commit, and the launcher config.

torchrun --nproc_per_node=8 train.py \
  --model llama3_8b \
  --data configs/data.json \
  --steps 500 \
  --fsdp full_shard \
  --checkpoint-dir s3://your-bucket/parity/run-$(date +%Y%m%d-%H%M) \
  2>&1 | tee logs/parity-$(date +%s).log

# Then diff the artifacts, not the vibes:
#   sha256sum of the final checkpoint from both runs, and a loss log diff.
#   If the hashes differ, that is expected with different kernels; what must
#   match is the loss trajectory inside your tolerance.

损失曲线要在事先定好的容差里比，而不是看完数字再定标准。

早期用户，以及该问他们的三个问题

公告点名了两个早期用户。做对话 AI 的 Cartesia 说它把标准推理流水线跑到 Trainium 上，对 PyTorch 代码只做了最小改动。Lieber 脑发育研究所用生物、化学和语言模型做药物靶点发现，那里的分子动力学贵到「用学习模型换取仿真」本身就是能算出钱的事。

两个例子都来自厂商，所以把它们当存在性证明，不要当基准测试。如果你能问到参考客户，问三个案例里不会写的问题：你的训练代码实际上改了多少比例？你因此失去了什么能力？你们的对齐容差是多少，怎么量的？

比任何一个演示都更有力的，是结构性的论点：同一套硬件既做训练也做推理，于是低谷时段的推理算力可以转去跑训练，而不是空转。这是利用率论点，而云账单实际上是由利用率决定的，不是由每芯片小时的标价决定的。

# The decision is rarely "is the chip cheaper". It is "does the engineering time
# pay back inside the project's lifetime". Put real numbers in, not hopes.

def payback(port_weeks: float, hours_per_week: float, loaded_hourly: float,
            monthly_compute_now: float, monthly_compute_after: float) -> dict:
    port_cost = port_weeks * hours_per_week * loaded_hourly
    monthly_saving = monthly_compute_now - monthly_compute_after
    if monthly_saving <= 0:
        return {"verdict": "no savings on these inputs", "port_cost": port_cost}
    months = port_cost / monthly_saving
    return {
        "port_cost": round(port_cost),
        "monthly_saving": round(monthly_saving),
        "breakeven_months": round(months, 1),
        "verdict": "worth a pilot" if months <= 6 else "not this quarter",
    }

# Example: 3 engineer-weeks at 35 h/week and 120 currency-units/hour is 12,600 up
# front. A 4,000/month saving breaks even in 3.2 months. A 900/month saving does not.
print(payback(3, 35, 120, 14000, 10000))
print(payback(3, 35, 120, 14000, 13100))

一份能在评审里站住的迁移清单

TorchNeuron 支持 Trainium2 和 Trainium3，并在 10 月 20 日以 Apache 2.0 发布。在这之前，先回答一个问题并写下来：你们团队这个季度能不能吃下一个新后端。一条命令就能回滚的试点叫实验；没有回滚路径的迁移叫下注，而下注该写在战略文档里，不该塞进一个迭代。

三种情况应该让你早点停下。如果你维护手写的设备专用 kernel，逐个验证，而不是假设标准路径都覆盖。如果你的依赖图钉住的框架版本还没被后端支持，就等。如果团队里没人能半夜对着不熟悉的硅片处理训练故障，那省下来的钱是纸面上的。

如果这三条都不成立，就把清单跑成脚本，把输出贴到评审里：torch 版本、后端 commit、设备是否可见、哪些算子缺失、以及那条能把你送回老路径的命令。这份文档才是把厂商说法变成团队能自己承担的决定的东西。

#!/usr/bin/env bash
# Pre-flight checklist. The output is the artifact you attach to the review.
set -euo pipefail

echo "== versions =="
python -c "import torch; print('torch', torch.__version__)"
git -C torch-neuronx log -1 --format='backend %H'

echo "== device visible? =="
python - <<'PY'
import os, torch
dev = os.environ.get("TORCH_DEVICE", "")
print("requested:", dev or "<unset>")
print("available:", [d for d in (torch.accelerator.current_accelerator(),) if d])
PY

echo "== unsupported-op probe (fails loudly, that is the point) =="
python tools/probe_ops.py --model ./your-model-config --exit-on-missing || \
  echo "MISSING OPS FOUND -> treat as a real port task, not a config flip"

echo "== rollback path =="
test -n "${CPU_FALLBACK_COMMAND:-}" || echo "no fallback documented -> stop here"

# Reversible experiment beats a heroic migration. If rollback is not one command,
# you are not running a pilot, you are running a bet.

最有力的论点是利用率：低谷时段的推理算力可以转去跑训练。