Hugging Face 上线 torch-neuronx 模型支持查询 Space

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

torch-neuronx 模型支持 Space 为 Hub 上最热门的模型回答了这个问题。它会告诉你某个模型是否能在 Trainium 上未经修改地运行，并给出你自己运行它的脚本。

查找模型

该 Space 打开时显示的是精选视图：这是我们自己从 Hub 上下载量和趋势得分最高的模型中挑选出的一组模型。每一行给出结果、模型分片所跨的 Trainium 核心数、运行的墙钟时间（包括编译）以及峰值设备内存。

要查找某个特定模型，请切换到趋势或下载视图，然后输入其名称的一部分。这些视图遵循 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 个 rank。

首次运行会先编译模型的图。要跳过该步骤，请从公共 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 上时，查一查。