在上一篇文章中，我们介绍了 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个 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数据集，每对 对应一个 Parquet 文件neuronx-cc和库版本。
每一行保存一个模型的结果、失败时的错误、整个技术栈的版本以及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 上运行时，查一下它。