Hugging Face 推出 torch-neuronx model support Space，可查询哪些模型能在 AWS Trainium 上运行

在上篇文章中，我们介绍了 TorchNeuron，即 AWS Trainium 的原生 PyTorch 后端（torch-neuronx 包）。借助它，Transformers 可以像在其他任何设备上一样，在 Trainium 上运行自己的建模代码，因此 Hub 上的任何 Transformers 模型都可能运行。但哪些模型实际能运行呢？

torch-neuronx 模型支持 Space 为 Hub 上最流行的模型解答了这个问题。它会告诉你某个模型能否在 Trainium 上未经修改地运行，并提供让你自己运行它的脚本。

查找模型

该 Space 打开时显示精选（Selected）视图：这是我们根据下载量和热度分数从 Hub 上的热门模型中挑选出的模型。每一行给出运行结果、模型分片所用的 Trainium 核心数、运行耗时（包括编译）以及设备峰值内存。

要查找特定模型，请切换到热门（Trending）或下载（Downloads）视图并输入其名称的一部分。这些视图遵循 Hub 自身的排名，每次一个任务，无论模型是否经过测试。当模型被有意排除时，Space 会告诉你原因（大多数时候是因为未被优先安排，有时则是确实不兼容）。你还可以选择要查看的 Neuron 编译器（neuronx-cc）和 Transformers 的版本。

截至撰写本文时，仅测试了 Transformers 模型，且我们的选择仅限于 text-generation 和 image-text-to-text 模型。我们计划扩展到 diffusers、sentence transformers 以及其他 transformers 任务。

自己动手运行

选择一个模型即可查看如何运行它。该 Space 会列出测试模型时所用的 Neuron 软件栈、需要额外安装的包，以及运行所用的 run.py 脚本。该脚本遵循 TorchNeuron 发布公告中的方案：使用 from_pretrained 加载模型（在需要时进行分片），然后使用静态缓存和 torch.compile(backend="neuron") 进行生成。

在搭载该软件栈的 Trainium2（trn2）实例上，例如 AWS Neuron 深度学习容器中，安装这些包，然后运行：

python3 run.py

无需 torchrun：当模型需要分片时，脚本会自行启动 torch.distributed 进程。

首次运行会先编译模型的计算图。要跳过这一步，可以从公开的 NEFF 缓存存储桶中复制测试期间由同一编译器构建的计算图。Space 会为你提供带有正确编译器版本的命令：

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 个核心）的模型。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 上时，查一查。