¿Su modelo corre en AWS Trainium? Así puede verificarlo y probarlo

Introducción

AWS Trainium es la infraestructura de Amazon diseñada para acelerar la inferencia y el entrenamiento con chips Trainium. Hugging Face, junto con el backend nativo para PyTorch llamado TorchNeuron (paquete torch-neuronx), permite ejecutar modelos de Transformers directamente sobre esa arquitectura. Pero ¿cómo saber si un modelo concreto en el Hub funciona sin modificaciones sobre Trainium? Hugging Face puso una Space pública que responde exactamente eso: qué modelos corren en Trainium y cómo reproducir la ejecución.

En este artículo describimos qué ofrece la Space, cómo buscar un modelo, cómo reproducir una ejecución en un nodo Trainium2 (trn2), qué significa realmente que un modelo “corra” y cómo consultar el dataset público que contiene los resultados. También añadimos contexto sobre la relevancia para equipos y tomadores de decisión en América Latina.

Qué ofrece la Space de soporte para modelos torch-neuronx

La Space lista los modelos más populares del Hub y prueba si funcionan con la pila Neuron (neuronx-cc y torch-neuronx). Para cada modelo muestra:

• El resultado del test (por ejemplo PASS o un fallo).
• El número de núcleos Trainium usados y la forma de sharding.
• El tiempo total de pared de la ejecución (incluye compilación).
• La memoria máxima usada en dispositivo.

La vista inicial (Selected) incluye una selección propia basada en descargas y tendencia. También puede cambiar entre vistas como Trending y Downloads para ver listas ordenadas igual que en el Hub, tarea por tarea. Si un modelo no fue probado por decisión explícita, la Space indica la razón; en otros casos la incompatibilidad es real.

Actualmente la prueba cubre modelos de Transformers orientados a text-generation y a tareas de imagen-text-to-text. Hugging Face planea ampliar la cobertura a diffusers, sentence transformers y otras tareas Transformer en el futuro.

Cómo buscar un modelo en la Space

Para localizar un modelo basta con cambiar a la vista Trending o Downloads y escribir parte del nombre del modelo. Cada entrada muestra la combinación de versiones del compilador Neuron y de la librería usada en la prueba, de modo que pueden comparar resultados entre distintas versiones.

La Space además facilita el script exacto run.py que se usó para cada prueba, así como las instrucciones sobre la pila Neuron y los paquetes necesarios.

Ejecutar un modelo en su propio entorno Trainium

Si quiere reproducir una prueba, la Space incluye el run.py que siguió la receta de TorchNeuron: el script hace from_pretrained para cargar el checkpoint, aplica sharding cuando corresponde, después hace generación usando cache estático y torch.compile(backend='neuron').

En una instancia Trainium2 (trn2) con la pila adecuada, por ejemplo usando el AWS Neuron Deep Learning Container, instale los paquetes listados y ejecute:

python3 run.py

No es necesario usar torchrun: el script crea internamente las réplicas distribuidas (torch.distributed) cuando el modelo está shardeado.

La primera ejecución inicia la compilación de los grafos del modelo. Si desea saltarse esa compilación puede descargar los grafos ya compilados usados en la prueba desde el bucket público NEFF. La Space proporciona el comando correcto con la versión del compilador, por ejemplo:

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

Reemplace <version> por la versión indicada en la Space para que coincida con la compilación usada en la prueba.

Qué significa exactamente que un modelo “corra” en Trainium

Cada test carga un modelo en una instancia trn2, shardeándolo en el número mínimo de núcleos Trainium necesarios para alojar los pesos. En la mayoría de casos el sharding usa paralelismo tensorial. Los modelos mixture-of-experts (MoE) que no tienen un plan de tensor-parallel usan solo expert parallelism.

La prueba genera 32 tokens y valida el resultado:

• Si el modelo se ejecutó en un solo núcleo, los 32 tokens generados se comparan con la ejecución en CPU del mismo checkpoint.
• Si el modelo está shardeado, un modelo juez evalúa que el texto sea fluido en inglés (fluency check).

Importante: en el momento de la prueba Hugging Face limita las pruebas a modelos que caben en 4 dispositivos Trainium2, es decir 16 núcleos como máximo. Un resultado PASS solo indica que la ejecución fue correcta según los criterios del test; no es una medición de throughput, latencia en cargas de trabajo reales o estabilidad en generaciones muy largas.

El dataset público con todos los resultados

Toda la información que muestra la Space proviene del dataset público aws-neuron/torch-neuronx-model-support. Los resultados están almacenados en Parquet, con un archivo por cada par de versiones de neuronx-cc y librería. Cada fila contiene el id del modelo, el resultado del test, el error si falló, las versiones de las herramientas y el run.py usado.

Puede consultar el dataset directamente desde Python, por ejemplo:

import pandas as pd

df = pd.read_parquet('hf://datasets/aws-neuron/torch-neuronx-model-support/data')
print(df[df.model_id == 'Qwen/Qwen3-0.6B'][['neuronx_cc', 'library', 'library_version', 'neuron_cores', 'result']])

La tarjeta del dataset (dataset card) describe con precisión qué significa cada resultado y cómo interpretar los metadatos.

Por qué esto importa para América Latina

Para equipos en América Latina que diseñan productos de ML o toman decisiones de infraestructura, esta Space y el dataset asociado son herramientas prácticas para evaluar compatibilidad con hardware especializado como Trainium. Antes de migrar cargas de trabajo a instancias con aceleradores específicos, conviene saber si el checkpoint que usan corre sin cambios, qué sharding se requiere y cuánto esfuerzo implicaría reproducir las compilaciones.

Además, la posibilidad de usar caches públicos de grafos compilados acelera pruebas y prototipos, algo útil para equipos con presupuesto limitado que quieran experimentar con infraestructuras en la nube.

Conclusión

Si se preguntan si un modelo del Hub funciona en AWS Trainium, no hace falta suponerlo: la Space de Hugging Face y el dataset aws-neuron/torch-neuronx-model-support permiten verificarlo, obtener el script exacto de ejecución y, en muchos casos, reproducir la corrida en una instancia trn2. Recuerden que pasar un test de compatibilidad no equivale a garantizar rendimiento de producción, pero sí reduce la incertidumbre técnica al evaluar migraciones a hardware Trainium.

La recomendación para equipos en América Latina: antes de comprometerse con una arquitectura específica, consulten la Space para su modelo, prueben localmente con los scripts y caches provistos, y utilicen los resultados del dataset para planear validaciones de rendimiento más profundas.