AWS Trainium模型运行查询

在之前的文章中，我们介绍了TorchNeuron，它是AWS Trainium的原生PyTorch后端（torch-neuronx包）。 有了它，Transformers可以在Trainium上运行自己的模型代码，就像在任何其他设备上一样，因此Hub上的任何Transformers模型都可以运行。 哪些模型实际上可以呢？ torch-neuronx模型支持Space回答了这个问题，它告诉你哪些模型可以在未经修改的情况下在Trainium上运行，并为你提供运行它的脚本。

查找模型Space在Selected视图中打开：我们自己的选择，从Hub上下载量和趋势分数最高的模型中挑选。 每一行都给出结果，模型在Trainium核心上分片的数量，运行时间（包括编译），以及峰值设备内存。 要查找特定模型，切换到Trending或Downloads视图并输入其名称的一部分。 这些视图遵循Hub自己的排名，一次一个任务，无论模型是否经过测试。 如果一个模型是故意被排除的，Space会告诉你原因（大多数时候是因为没有优先考虑，有时它确实不兼容）。

你也可以选择Neuron编译器（neuronx-cc）和Transformers的版本来查看。 在撰写本文时，仅测试了Transformers模型，并且我们的选择仅限于文本生成和图像文本到文本模型。 我们计划扩展到diffusers、句子转换器和其他Transformers任务。 自己运行它选择一个模型以查看如何运行它。 Space列出了模型测试的Neuron堆栈，要在其上安装的软件包，以及运行的run.py脚本。

该脚本遵循TorchNeuron公告的配方：它使用from_pretrained加载模型，需要时进行分片，然后使用静态缓存和torch.compile(backend="neuron")进行生成。 例如，在Trainium2（trn2）实例上，在AWS Neuron深度学习容器中安装软件包，然后运行：python3 run.py不需要torchrun：脚本在模型分片时启动自己的torch.distributed ranks。 第一次运行时，会编译模型的图。

要跳过这一步，请从在测试期间由同一编译器构建的公共NEFF缓存桶中复制图。 Space为你提供了正确的编译器版本的命令：hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache “运行”的含义每个测试都在trn2实例上加载模型，分片到可以容纳其权重的最少数量的Trainium核心。 大多数模型使用张量并行进行分片。 没有张量并行计划的混合专家模型仅使用专家并行。

然后测试生成32个标记并检查它们。 在单个核心上，将标记与相同检查点的CPU运行进行比较。 当模型分片时，一个小型裁判模型检查文本是否流畅的英语。 在撰写本文时，仅测试了可以放在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" ]] 数据集卡给出了每个结果的准确含义。

下次当你想知道模型是否在Trainium上运行时，请查找它。