在上一篇文章中，我们介绍了 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 模型，并且我们的选择仅限于文本生成和图像-文本-到-文本模型。我们计划扩展到 diffusers、sentence transformers 以及其他 transformers 任务。
自己运行
选择一个模型以查看如何运行它。Space 列出了模型测试所用的 Neuron 堆栈、需要在其之上安装的包，以及运行的 run.py 脚本。该脚本遵循 TorchNeuron 公告中的配方：它使用 from_pretrained 加载模型，在需要时进行分片，然后使用静态缓存和 torch.compile(backend="neuron") 进行生成。
在配备该堆栈的 Trainium2（trn2）实例上，例如在 AWS Neuron Deep Learning Container 中，安装这些包，然后运行：
python3 run.py无需使用 torchrun：当模型被分片时，脚本会启动自己的 torch.distributed ranks。
首次运行会从编译模型的图开始。要跳过该步骤，从公共 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 上运行时，查一查吧。