Hugging Face 上线 Trainium 模型兼容查询:Hub 模型能否不改代码跑在 Trainium 上,附可直接运行的脚本
10-20 那期报过 TorchNeuron 让 Transformers 直接跑在 Trainium 上,这次补上「具体哪个模型能跑」的查询表。
// 要点
- 新 Space「torch-neuronx model support」可按精选、热门、下载量查模型,每条结果给出所需核心数、含编译的运行时间和设备峰值显存。
- 目前只测 Transformers 里的 text-generation 和 image-text-to-text 模型,且只测能放进 4 张 Trainium2(最多 16 个核心)的模型;之后计划加入 diffusers、sentence transformers 等。
- 测试只生成 32 个 token:单核结果与 CPU 对比,分片的模型由一个小评判模型检查输出是否通顺。官方说明 PASS 不代表吞吐量或长文本生成的表现。
- 每个模型附 Neuron 版本、依赖和 run.py,可从公开的 NEFF 缓存同步预编译图跳过编译;底层数据集 aws-neuron/torch-neuronx-model-support 可直接用 pandas 查询。
开发者视角想把开源模型挪到 Trainium 省钱的,第一步不用再自己踩坑,先查这张表。但它只证明「能跑起来」:我会拿 PASS 的模型跑一遍自己的真实请求,测吞吐和长输出,再和现在的 GPU 账单比。