你的模型能跑在 AWS Trainium 上吗？查一下就知道
本文目录
作者：David Corvoysier
发布于 2026 年 10 月 23 日
在上一篇文章中，我们介绍了 TorchNeuron——AWS Trainium 的原生 PyTorch 后端（即 torch-neuronx 包）。借助它，Transformers 可以像在任何其他设备上一样，在 Trainium 上运行自己的建模代码，因此 Hub 上任何 Transformers 模型都可能跑得起来。但究竟哪些能跑？
这个 torch-neuronx 模型支持 Space 回答了这个问题，覆盖 Hub 上最热门的模型。它会告诉你某个模型是否能不做修改地在 Trainium 上运行，并给出你自己运行的脚本。
查找一个模型#
Space 打开时处于 Selected 视图：这是我们自己的模型选集，从 Hub 上下载量和趋势分最高的模型中挑选。每一行给出结果、模型分片（shard）所用的 Trainium 核数、运行的实际耗时（含编译），以及峰值显存。
要查某个特定模型，切换到 Trending 或 Downloads 视图，输入名称的一部分即可。这两个视图遵循 Hub 自己的排名，一次一个任务，无论该模型是否被测过。如果某个模型是被有意排除的，Space 会告诉你原因（多数情况是未被优先处理，有时是真的不兼容）。你还可以选择要查看的 Neuron 编译器（neuronx-cc）和 Transformers 版本。
截至撰写时，只有 Transformers 模型被测试，且我们的选集限于文本生成和图文（image-text-to-text）模型。我们计划扩展到 diffusers、sentence transformers 以及其他 transformers 任务。
自己运行#
选择一个模型，即可看到如何运行。Space 会列出该模型测试所用的 Neuron 软件栈、需要在其之上安装的包，以及实际运行的 run.py 脚本。该脚本遵循 TorchNeuron 发布时的配方：用 from_pretrained 加载模型（需要时分片），然后用静态缓存和 torch.compile(backend="neuron") 生成。
在配有该软件栈的 Trainium2（trn2）实例上（例如在 AWS Neuron Deep Learning Container 中），安装这些包，然后运行：
python3 run.py不需要 torchrun：当模型需要分片时，脚本会自己启动 torch.distributed 的各个 rank。
首次运行会先编译模型的计算图。要跳过这一步，可以从公开的 NEFF 缓存桶里复制同一编译器在测试期间构建的计算图。Space 会给出命令，并带上正确的编译器版本：
hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache"能跑"是什么意思#
每次测试都会在一个 trn2 实例上加载模型，并按能容纳其权重的最小 Trainium 核数分片。多数模型用张量并行（tensor parallelism）分片。没有张量并行方案的专家混合（mixture-of-experts）模型只使用专家并行（expert parallelism）。测试随后生成 32 个 token 并检查它们。在单核上，这些 token 会与同一检查点在 CPU 上运行的结果比较。当模型被分片时，则由一个小型裁判模型检查文本是否为通顺的英语。
截至撰写时，只测试能放进 4 台 Trainium2 设备的模型，也就是最多 16 个核。PASS 并不代表吞吐量，也不代表更长生成的情况。
结果就是一个数据集#
Space 展示的一切都来自公开的 aws-neuron/torch-neuronx-model-support 数据集，每一对 neuronx-cc 与库版本对应一个 Parquet 文件。
每一行保存一个模型的结果、失败时的错误信息、整套软件栈的版本，以及实际运行的 run.py。你可以直接查询它：
import pandas as pd
df = pd.read_parquet("hf://datasets/aws-neuron/torch-neuronx-model-support/data")
df[df.model_id == "Qwen/Qwen3-0.6B"][["neuronx_cc", "library", "library_version", "neuron_cores", "result"]]数据集卡给出了每种结果的确切含义。
下次你想知道某个模型能否在 Trainium 上运行时，查一下就好。