¿Tu modelo de Hugging Face funciona en AWS Trainium? Descúbrelo
Hugging Face lanza un Space que indica qué modelos de Transformers se ejecutan sin cambios en AWS Trainium y cómo probarlos.
Un nuevo recurso para usuarios de IA
Hugging Face ha publicado un Space llamado torch‑neuronx‑model‑support que permite consultar, de forma rápida y visual, si un modelo del Hub se ejecuta sin modificaciones en la arquitectura de AWS Trainium.
El Space muestra, para los modelos más populares, el resultado de la prueba (PASS/FAIL), el número de núcleos Trainium utilizados, el tiempo total de ejecución (incluida la compilación) y la memoria máxima del dispositivo. Se pueden filtrar por Selected (una selección propia), Trending o Downloads, y buscar cualquier modelo introduciendo parte de su nombre.
Cuando un modelo no ha sido probado, la herramienta explica el motivo, que suele ser falta de prioridad o incompatibilidad real. Además, permite elegir la versión del compilador Neuron (neuronx‑cc) y la versión de Transformers para reproducir los resultados.
Cómo funciona la prueba
La prueba carga el modelo en una instancia Trainium2 (trn2) dentro del contenedor AWS Neuron Deep Learning, instala los paquetes indicados y ejecuta un script run.py. El script sigue la receta anunciada por TorchNeuron: carga el modelo con from_pretrained, lo reparte en shards cuando es necesario, genera texto usando una caché estática y compila con torch.compile(backend="neuron"). Se generan 32 tokens y se comparan con una ejecución en CPU (para un solo núcleo) o con un modelo juez (para modelos shardados).
Solo se han probado modelos que caben en cuatro dispositivos Trainium2 (máximo 16 núcleos). Un resultado PASS no garantiza rendimiento de throughput ni generación prolongada.
Datos del dataset
Los resultados provienen del dataset público aws‑neuron/torch‑neuronx‑model‑support, disponible en formato Parquet por cada combinación de versiones del compilador y de la librería. Cada fila contiene el modelo, el resultado, el error (si lo hubo), las versiones y el script usado.
💡 Opinión iaPasión
Esta herramienta es ideal para desarrolladores y emprendedores que necesiten validar rápidamente la viabilidad de sus modelos en Trainium antes de invertir en infraestructura. No aportará mucho a usuarios que trabajen exclusivamente con CPU o GPU y no tengan acceso a AWS.