模型能跑在 AWS Trainium 上吗？查询工具来了

在前一篇博文中，我们介绍了 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 进程。

首次运行时，程序会先编译模型的计算图。若想跳过此步骤，可从公开的 NEFF 缓存 存储桶中，复制测试阶段相同编译器版本生成的计算图。Space 页面提供了对应的同步命令，并附带正确的编译器版本：

hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache

“可运行”的定义

每项测试均在 trn2 实例上加载模型，并将其分片到能容纳模型权重的最少 Trainium 核心上。大多数模型采用张量并行方式进行分片。缺乏张量并行方案的混合专家（MoE）模型则仅使用专家并行。测试阶段生成 32 个 token 并进行校验。在单核心场景下，生成的 token 会与同一检查点在 CPU 上运行生成的 token 进行比对；当模型采用多核心分片时，则使用一个小模型作为裁判，检查生成的英文文本是否通顺。

截至本文撰写时，仅测试能适配 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 上运行时，直接查一下就知道了。