Hugging Face 推出 TorchNeuron 模型支持查询 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 上，查一下。