¿Se ejecuta tu modelo en AWS Trainium? Consúltalo

En un publicación anterior, presentamos TorchNeuron, el backend nativo de PyTorch para AWS Trainium (el torch-neuronx paquete). Con él, Transformers ejecuta su propio código de modelado en Trainium, tal como lo hace en cualquier otro dispositivo, por lo que cualquier modelo de Transformers del Hub puede ejecutarse. ¿Cuáles lo hacen realmente?

El compatibilidad de modelos de torch-neuronx Space responde a esa pregunta para los modelos más populares del Hub. Te indica si un modelo se ejecuta sin modificaciones en Trainium y te proporciona el script para ejecutarlo tú mismo.

Buscar un modelo

El Space abre en el Seleccionado vista: nuestra propia selección de modelos elegidos entre los mejores modelos del Hub por descargas y por puntuación de tendencia. Cada fila muestra el resultado, el número de núcleos Trainium sobre los que está fragmentado el modelo, el tiempo total de ejecución (compilación incluida) y la memoria máxima del dispositivo.

Para consultar un modelo determinado, cambie a la Tendencias o Descargas ver y escribir parte de su nombre. Estas vistas siguen los propios rankings del Hub, una tarea a la vez, independientemente de si un modelo fue evaluado o no. Cuando un modelo fue excluido a propósito, el Space te indica por qué (la mayoría de las veces no fue priorizado, a veces es verdaderamente incompatible). También puedes elegir las versiones del compilador de Neuron (neuronx-cc) y de Transformers que mirar.

En el momento de escribir este artículo, solo se han probado los modelos de Transformers, y nuestra selección se limita a los modelos text-generation y image-text-to-text . Planeamos ampliar a diffusers, sentence transformers y otras tareas de transformers.

Ejecútalo tú mismo

Selecciona un modelo para ver cómo ejecutarlo. El Space enumera el stack de Neuron con el que se probó el modelo, los paquetes que se deben instalar además de él y el run.py script que se ejecutó. El script sigue la receta del anuncio de TorchNeuron: carga el modelo con from_pretrained, fragmentándolo (sharded) cuando es necesario, y luego genera con una caché estática y torch.compile(backend="neuron").

En una instancia Trainium2 (trn2) con ese stack, por ejemplo en el AWS Neuron Deep Learning Container, instala los paquetes y luego ejecuta:

python3 run.py

No hace falta torchrun: el script lanza sus propios torch.distributed ranks cuando el modelo está particionado.

Una primera ejecución comienza compilando los grafos del modelo. Para saltar ese paso, copia los grafos que el mismo compilador construyó durante la prueba, desde el bucket público NEFF cache . El Space te da el comando, con la versión correcta del compilador:

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

Qué significa "runs"

Cada prueba carga el modelo en una trn2 instancia, particionado sobre el menor número de núcleos Trainium que puedan contener sus pesos. La mayoría de los modelos se particionan con paralelismo de tensores. Los modelos de mezcla de expertos sin un plan de paralelismo de tensores usan solo paralelismo de expertos. La prueba genera entonces 32 tokens y los verifica. En un solo núcleo, los tokens se comparan con los de una ejecución en CPU del mismo checkpoint. Cuando el modelo está particionado, un pequeño modelo juez comprueba que el texto sea inglés fluido.

En el momento de escribir esto, solo se prueban modelos que caben en 4 dispositivos Trainium2, lo que significa 16 núcleos como máximo. Un PASS no dice nada sobre el rendimiento ni sobre generaciones más largas.

Los resultados son un dataset

Todo lo que el Space muestra proviene del dataset público aws-neuron/torch-neuronx-model-support , con un archivo Parquet por cada par de versiones de neuronx-cc y librería.

Cada fila contiene el resultado de un modelo, el error si falló, las versiones de todo el stack y el run.py que se ejecutó. Puedes consultarlo directamente:

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"]]

La dataset card da el significado exacto de cada resultado.

La próxima vez que te preguntes si un modelo funciona en Trainium, consúltalo.