AWS Trainium 模型支持：Hugging Face 推出 torch-neuronx 兼容性查询工具

TL;DR — Hugging Face 和 AWS 推出了 torch-neuronx 模型支持 Space，这是一个公开的查询工具，可以告诉你 Hugging Face Hub 上的某个 Transformers 模型是否能在 AWS Trainium 上无需修改直接运行，并提供用于测试的确切脚本和软件栈。它将“这个模型能在 Trainium 上运行吗？”从反复试验的问题变成了可查询的数据集。

该 Space 基于 TorchNeuron，这是 torch-neuronx 包中提供的 AWS Trainium 原生 PyTorch 后端，在之前的 Hugging Face 文章 中已有介绍。借助 TorchNeuron，Transformers 在 Trainium 上运行其自身的建模代码，就像在其他设备上一样——因此原则上 Hub 上的任何 Transformers 模型都可能运行。新的 Space 回答了哪些模型确实能运行，针对的是 Hub 上最流行的模型。

发布内容

aws-neuron/torch-neuronx-model-support Space 是一个公开的、可浏览的 AWS Trainium 兼容性表格。对于列出的每个模型，它报告：

  • 模型是否能在 Trainium 上无需修改运行（PASS 结果）
  • 模型分片所用的 Trainium 核心数
  • 运行的总时间（包括编译时间）
  • 峰值设备内存使用量

对于每个模型，Space 还公开了运行的脚本，因此结果可以被复现，而不是仅仅被信任。

如何查询模型

Space 打开时显示 Selected 视图，描述为作者从 Hub 上按下载量和趋势评分选出的顶级模型中的精选集合。

Trending 和 Downloads 视图遵循 Hub 自身的排名，一次一个任务。这些视图列出模型，无论它们是否经过测试，因此缺失的条目永远不会含糊：当模型被有意排除时，Space 会说明原因。

在撰写本文时，仅测试了 Transformers 模型，我们的选择仅限于 text-generation 和 image-text-to-text 模型。我们计划扩展到 diffusers、sentence transformers 和其他 transformers 任务。

Space 还允许你选择要检查的 Neuron 编译器（neuronx-cc）和 Transformers 的版本，因此结果可以按软件栈版本过滤，而不是作为单一的未注明日期的快照来阅读。

要查找特定模型，请切换到 Trending 或 Downloads 视图并输入其名称的一部分。Space 会报告模型被排除的原因——最常见的是因为它未被优先考虑，有时是因为它确实与 Trainium 不兼容。

如何自行运行模型

在 Space 中选择一个模型会显示如何运行它。Space 列出了模型测试所用的 Neuron 软件栈、在该软件栈之上需要安装的包，以及实际运行的 run.py 脚本。

该脚本遵循 TorchNeuron 公告中的方法：使用 from_pretrained 加载模型，在需要时进行分片，然后使用静态缓存和 torch.compile(backend="neuron") 进行生成。

在带有该软件栈的 Trainium2（trn2）实例上——例如在 AWS Neuron 深度学习容器 中——你安装列出的包，然后运行：

python3 run.py

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

使用 NEFF 缓存跳过编译

首次运行会先编译模型的图，这占据了报告的总时间的大部分。要跳过此步骤，你可以从公共的 NEFF 缓存 存储桶中复制测试期间同一编译器构建的图。Space 提供了带有正确编译器版本的命令：

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

使用缓存意味着只要使用与原始测试相同的 neuronx-cc 版本，复现就可以跳过重新编译。

“运行”的含义

兼容性声明是狭窄且具体的，Space 对此进行了精确记录。

每个测试在 trn2 实例上加载模型，分片到能容纳其权重的最少 Trainium 核心数。大多数模型使用张量并行分片；缺乏张量并行方案的混合专家模型仅使用专家并行。测试随后生成 32 个 token 并检查它们。

正确性检查取决于分片配置：

  • 在单核上，生成的 token 与同一检查点的 CPU 运行结果进行比较。
  • 当模型分片时，一个小型评判模型检查生成的文本是否为流利的英语。

PASS 的范围是刻意限定的。在撰写本文时，仅测试了适合 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"]]

数据集卡片 定义了每个结果值的精确含义，因此 PASS 或失败可以根据精确的方法论来解释，而不是推断。

覆盖范围限制和后续计划

四个约束限制了 Space 当前能证明的内容：

  1. 仅测试 Transformers 模型。
  2. 仅覆盖 text-generation 和 image-text-to-text 任务。
  3. 仅测试适合 4 个 Trainium2 设备（16 个核心）的模型。
  4. PASS 报告的是 32 个生成 token 的正确性，而不是吞吐量或长上下文行为。

已声明的计划是将覆盖范围扩展到 diffusers、sentence transformers 和其他 Transformers 任务。

为什么这很重要

硬件可移植性问题通常通过群体经验、临时问题或昂贵的试运行来回答。将 Trainium 兼容性作为带版本、可查询的数据集发布，并附上复现脚本，改变了这个问题的成本结构：答案是查询，证据是可检查的。

两个设计选择使该资源比排行榜更有用。首先，结果与特定的 neuronx-cc 和库版本相关联，因此可以根据你实际计划部署的软件栈来评估兼容性。其次，数据集记录了失败及其产生的错误，以及成功，这使得以前未记录的兼容性问题变得可搜索。

实际要点：在 AWS Trainium 上运行模型之前，请检查 aws-neuron/torch-neuronx-model-support Space——或直接查询其 Parquet 数据集——并重用已发布的 run.py 和 NEFF 缓存条目，在 trn2 实例上复现结果。

Sources