你的模型在AWS Trainium上运行吗？查一查便知

AWS Trainium 上到底能跑哪些模型，现在有了一个可以直接查询的公开数据库。Hugging Face 博客作者 David Corvoysier 发布文章介绍，团队基于此前推出的 TorchNeuron（AWS Trainium 的原生 PyTorch 后端，即 torch-neuronx 包）搭建了一个模型支持查询 Space，专门回答“这个模型能不能在 Trainium 上原样跑起来”这个问题。用户可以在页面上按下载量或热度浏览 Hub 上的热门模型，每一行会显示测试结果、模型被切分到多少个 Trainium 核心、包含编译时间在内的实际运行耗时，以及峰值显存占用。如果某个模型被有意跳过测试，页面也会说明原因——多数情况是优先级不够，少数是确实不兼容。目前该工具只覆盖文本生成和图文转文本两类 Transformers 模型，团队计划后续扩展到 diffusers、sentence transformers 等其他任务。 这件事的背景是，Trainium 作为 AWS 自研的 AI 训练芯片，长期面临一个现实问题：硬件有了，但开发者不清楚手里的模型能不能直接跑。TorchNeuron 解决的是“能不能跑”的底层能力问题，而这个查询工具解决的是“哪些能跑”的信息透明度问题。对使用 AWS 云资源做推理或微调的中国出海团队来说，这意味着选型前可以先查再动手，避免在编译和适配环节浪费算力成本。具体测试方式上，每个模型会被加载到 trn2 实例上，按能容纳其权重的最小核心数切分，多数模型采用张量并行，没有张量并行方案的混合专家模型则只用专家并行。测试生成 32 个 token 并校验结果：单核运行时与 CPU 结果逐 token 对比，切分运行时则用一个小型评判模型检查输出是否为流畅英文。需要注意的是，目前只测试能放进 4 台 Trainium2 设备（最多 16 核）的模型，且 PASS 结果不代表吞吐性能，也不代表长文本生成的表现。 所有结果都来自公开数据集 aws-neuron/torch-neuronx-model-support，按 neuronx-cc 编译器版本和库版本组合存放为 Parquet 文件，开发者可以直接用 pandas 查询某个模型在各版本栈下的表现。页面还会给出完整的复现脚本 run.py，遵循 TorchNeuron 发布时的配方：用 from_pretrained 加载模型、按需切分、配合静态缓存和 torch.compile(backend=“neuron”) 生成。在 trn2 实例上装好依赖后直接执行 python3 run.py 即可，不需要 torchrun，脚本会自行拉起分布式进程。首次运行需要编译计算图，想跳过这一步可以从公开的 NEFF 缓存桶同步对应编译器版本的图缓存。对正在评估 Trainium 性价比的团队而言，这个工具的价值在于把“听说能跑”变成了“查过能跑”，但也要清楚它的边界：能跑不等于跑得快，短输出验证也不等于生产环境下的稳定表现。