Does Your Model Run on AWS Trainium? Look It Up

Does Your Model Run on AWS Trainium? Look It Up

在之前的文章中，我们介绍了 TorchNeuron，这是 AWS Trainium 的原生 PyTorch 后端（即 torch-neuronx 软件包）。有了它，Transformers 库可以像在任何其他设备上一样，在 Trainium 上运行其自身的建模代码，因此 Hub 上的任何 Transformers 模型理论上都可以运行。那么，哪些模型真正做到了呢？torch-neuronx 模型支持空间（Space）回答了 Hub 上最热门模型的相关问题。它会告诉你模型是否可以在未经修改的情况下在 Trainium 上运行，并提供让你亲自运行它的脚本。

Look up a model

该空间（Space）默认打开“Selected”视图：这是我们根据下载量和趋势评分，从 Hub 的顶级模型中挑选出的精选模型。每一行都显示了运行结果、模型分片所使用的 Trainium 核心数量、运行耗时（包含编译时间）以及峰值设备内存。要查找特定模型，请切换到“Trending”或“Downloads”视图并输入名称的一部分。这些视图遵循 Hub 自身的排名，一次显示一个任务，并标注模型是否经过测试。当某个模型被刻意排除时，空间会说明原因（大多数情况是因为优先级不高，有时则是因为确实不兼容）。你还可以选择要查看的 Neuron 编译器（neuronx-cc）和 Transformers 版本。在撰写本文时，仅测试了 Transformers 模型，且我们的选择仅限于文本生成和图文生成模型。我们计划扩展到 diffusers、sentence transformers 以及其他 transformers 任务。

Run it yourself

选择一个模型即可查看如何运行它。该空间列出了模型测试所使用的 Neuron 堆栈、需要额外安装的软件包，以及运行时的 run.py 脚本。该脚本遵循 TorchNeuron 发布时的方案：它使用 from_pretrained 加载模型（必要时进行分片），然后使用静态缓存和 torch.compile(backend="neuron") 进行生成。例如，在配备该堆栈的 Trainium2 (trn2) 实例上（如 AWS Neuron 深度学习容器中），安装软件包后运行：python3 run.py。无需使用 torchrun：当模型进行分片时，脚本会自动启动自己的 torch.distributed 进程。首次运行会从编译模型的图开始。若要跳过此步骤，可以从公共 NEFF 缓存存储桶中复制编译器在测试期间构建的图。空间会提供带有正确编译器版本的命令：hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache。

What “runs” means

每次测试都会在 trn2 实例上加载模型，并将其分片到能够容纳其权重的最少 Trainium 核心上。大多数模型使用张量并行进行分片。没有张量并行计划的混合专家模型仅使用专家并行。随后，测试会生成 32 个 token 并进行检查。在单核心上，这些 token 会与同一检查点在 CPU 上运行的结果进行比较。当模型进行分片时，会由一个小型的判别模型检查生成的文本是否为流畅的英语。在撰写本文时，仅测试了能够适配 4 个 Trainium2 设备（即最多 16 个核心）的模型。通过（PASS）测试并不代表吞吐量或更长生成的表现。

The results are a dataset

该空间显示的所有内容均来自公共数据集 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 上运行时，直接去查一下即可。