Does Your Model Run on AWS Trainium? Look It Up

您的模型可以在 AWS Trainium 上运行吗？查找一下

返回文章列表

[-1

]

[0

发布于 2026 年 10 月 23 日

GitHub 上的更新

点赞

6

[

David Corvoysier

dacorvo

关注

在之前的博客文章中，我们介绍了 TorchNeuron（torch-neuronx 包），这是 AWS Trainium 的原生 PyTorch 后端。通过它，Transformers 可以像在其他设备上一样，在 Trainium 上运行自己的建模代码，因此 Hugging Face 模型库中的任何 Transformers 模型都可能运行。哪些模型确实可以运行？

torch-neuronx 模型支持空间针对模型库中最受欢迎的模型回答了这个问题。它会告诉您模型是否可以在 Trainium 上无修改地运行，并提供运行它的脚本。

查找模型

该空间默认以“精选视图”打开：我们根据下载量和趋势评分从模型库中挑选出的模型。每一行显示结果、模型分片所使用的 Trainium 核心数量、运行的墙时（包含编译时间）以及峰值设备内存。

要查找特定模型，切换到“趋势”或“下载量”视图并输入其名称的一部分。这些视图遵循模型库本身的排名，按任务逐一显示，无论模型是否经过测试。当模型被有意排除时，该空间会说明原因（大多数情况下是未被优先处理，有时则是真正不兼容）。您还可以选择 Neuron 编译器（neuronx-cc）和 Transformers 的版本进行查看。

截至撰写本文时，仅测试了 Transformers 模型，我们的选择仅限于文本生成和图像-文本到文本模型。我们计划扩展到扩散模型、句子嵌入模型和其他 Transformers 任务。

自行运行

选择模型以查看如何运行。该空间列出了模型测试所使用的 Neuron 堆栈、需要在其上安装的软件包以及运行的 run.py 脚本。该脚本遵循 TorchNeuron 宣布的配方：使用 from_pretrained 加载模型，需要时进行分片，然后使用静态缓存和 torch.compile(backend="neuron") 进行生成。

例如，在带有该堆栈的 Trainium2（trn2）实例上，如 AWS Neuron 深度学习容器中，安装软件包后运行：

python3 run.py

无需使用 torchrun：当模型被分片时，脚本会自行启动 torch.distributed 的 rank。

首次运行会先编译模型的图。要跳过此步骤，可从公共 NEFF 缓存存储桶中复制测试期间编译器生成的图。该空间会提供命令和正确的编译器版本：

hf buckets
sync
hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache

“运行”的含义

每次测试都会在 trn2 实例上加载模型，使用尽可能少的 Trainium 核心来容纳其权重。大多数模型使用张量并行进行分片。没有张量并行计划的专家混合模型仅使用专家并行。测试会生成 32 个 token 并进行检查。在单个核心上，这些 token 会与相同检查点的 CPU 运行结果进行比较。当模型被分片时，一个小的判断模型会检查文本是否为流畅的英文。

截至撰写本文时，仅测试了可以适应 4 个 Trainium2 设备（最多 16 个核心）的模型。通过测试并不意味着吞吐量或更长生成的性能。

结果是一个数据集

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 上运行时，可以查阅该数据集。

本文提到的数据集 1

本文提到的 Space 1

更多博客文章

open-source

torch

neuron

TorchNeuron：在 AWS Trainium 上运行 HF 模型从未如此简单

• +3

30

2026年10月19日

tpu

google

gcp

Hugging Face ❤️ TorchTPU：原生支持 TPUs

• +1

32

2026年10月20日

社区

编辑

预览

通过拖拽、粘贴或点击此处上传图片、音频和视频。

点击此处上传图片

评论

· 注册或登录以发表评论