在之前的文章中，我们介绍了 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 Deep Learning Container 中，安装这些包，然后运行：
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 上运行时，查一下。