Hugging Face 上线 Space，可查询模型能否在 AWS Trainium 上运行

查找模型

该 Space 打开时显示“精选”视图：这是我们自己挑选的模型，选自 Hub 上按下载量和趋势得分排名靠前的模型。每一行给出结果、模型分片所用的 Trainium 核心数、运行的实际耗时（含编译）以及峰值设备内存。

要查找某个模型，请切换到“趋势”或“下载量”视图，并输入其名称的一部分。这些视图遵循 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 中，安装软件包，然后运行：

无需 torchrun：当模型分片时，脚本会自行启动 torch.distributed 进程。

首次运行会先编译模型的图。要跳过该步骤，请从公共 NEFF 缓存桶中复制同一编译器在测试期间构建的图。该 Space 会给出命令，并附带正确的编译器版本：

“运行”的含义

每项测试都会在 trn2 实例上加载模型，并按能容纳其权重的最少 Trainium 核心数进行分片。大多数模型使用张量并行进行分片。没有张量并行方案的混合专家模型仅使用专家并行。随后测试会生成 32 个 token 并进行校验。在单核上，这些 token 会与同一检查点的 CPU 运行结果进行比较。当模型分片时，由一个小型评判模型检查文本是否为流利的英文。

在撰写本文时，仅测试能放入 4 台 Trainium2 设备的模型，即最多 16 个核心。PASS 并不代表吞吐量或更长生成内容的表现。

结果就是一个数据集

该 Space 显示的所有内容都来自公开的 aws-neuron/torch-neuronx-model-support 数据集，每对 neuronx-cc 和库版本对应一个 Parquet 文件。

每一行包含模型的结果、失败时的错误、整个栈的版本以及实际运行的 run.py。你可以直接查询它：

数据集卡片给出了每种结果的确切含义。

下次你想知道某个模型能否在 Trainium 上运行时，去查一查吧。

本文提及的数据集 1

aws-neuron/torch-neuronx-model-support

本文提及的 Space 1

torch-neuronx model support

哪些 Hub 模型可在 AWS Trainium 上通过 torch-neuronx 运行

来自我们博客的更多文章

TorchNeuron：在 AWS Trainium 上运行 HF 模型比以往更容易

Hugging Face ❤️ TorchTPU：原生支持 TPU