在介绍 TorchNeuron——AWS Trainium 的原生 PyTorch 后端之后,一个核心问题随之而来:哪些 Hugging Face Hub 上的模型真正能在 Trainium 上运行?

AWS Neuron 团队推出的 torch-neuronx-model-support 空间为此提供了答案。该工具旨在展示未经修改即可在 Trainium 上运行的模型,并直接提供相应的运行脚本。
如何查找与筛选模型
该空间默认以“精选”视图打开,展示了根据下载量和趋势分数筛选出的顶级模型。每一行数据详细列出了模型所需的 Trainium 核心数量、包含编译时间的墙钟时间(Wall Time)以及峰值设备内存占用。
用户可切换至“趋势”或“下载”视图,通过输入名称片段来检索特定模型。这些视图遵循 Hub 的排名逻辑,无论模型是否经过测试均会显示。对于未通过测试的模型,空间会说明原因(通常为优先级不足或不兼容)。此外,用户还可指定 Neuron 编译器(neuronx-cc)和 Transformers 的版本进行查询。
目前仅针对 Transformers 模型进行测试,范围限于文本生成和图像转文本任务。未来计划扩展至 Diffusers、Sentence Transformers 等其他任务领域。
本地部署与运行指南
选定模型后,空间将列出其测试所用的 Neuron 堆栈、安装的软件包版本以及 run.py 脚本。该脚本遵循 TorchNeuron 的标准配方:按需加载分片模型,并使用静态缓存配合 torch.compile(backend="neuron") 进行推理。
在 AWS Trainium2 (trn2) 实例(如 AWS Neuron 深度学习容器)中,安装必要包后即可执行:
python3 run.py
若模型被分片,脚本会自动启动 torch.distributed 进程组。
首次运行需经历图表编译过程。为跳过此步骤,可从公共 NEFF 缓存桶复制预编译图表。使用正确的编译器版本时,空间提供的命令如下:
hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache
测试标准与局限性
测试流程定义为:在 trn2 实例上加载模型,使用能承载其权重的最少 Trainium 核心数。大多数模型采用张量并行策略;无专家并行图的 MoE 模型则仅使用专家并行。随后生成 32 个 token 并进行验证:单核运行时,token 会与同一检查点在 CPU 上的输出进行比对;多核场景下,由一个小法官模型检查英语文本的流畅性。
目前的测试范围仅限于适配 4 个 Trainium2 设备(即最多 16 个核心)的模型。关于 Flux 或更长上下文长度的支持情况尚无定论。
基于数据集的结果查询
空间展示的所有信息均源自公开的 aws-neuron/torch-neuronx-model-support 数据集,其中每个 Parquet 文件对应一组 neuronx-cc 和库版本。
每行记录包含模型结果、失败时的错误信息、完整堆栈版本及运行脚本。用户可直接通过代码查询数据:
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 上的兼容性时,可参考此资源。





