Hugging Face lanza un buscador para saber qué modelos de IA funcionan en los chips Trainium de Amazon
La plataforma ha publicado una herramienta que permite consultar si los modelos más populares de inteligencia artificial se ejecutan directamente en los procesadores especializados de AWS, con instrucciones para probarlo uno mismo.
Hugging Face ha publicado un buscador que indica si los modelos de IA más populares de su repositorio funcionan en AWS Trainium (el chip especializado de Amazon para redes neuronales). La herramienta muestra resultados de pruebas reales y proporciona el código Python necesario para ejecutar cada modelo compatible, incluyendo instrucciones para descargar grafos precompilados y ahorrar tiempo de compilación.
Hugging Face (una plataforma donde desarrolladores comparten y descargan modelos de IA ya entrenados) ha lanzado un buscador público que permite consultar si un modelo de inteligencia artificial funciona en AWS Trainium, el chip especializado de Amazon para entrenar y ejecutar redes neuronales. La herramienta, llamada torch-neuronx model support Space, muestra los resultados de pruebas sobre los modelos más populares del repositorio y proporciona el código exacto para ejecutarlos.
Según informó Hugging Face en su blog oficial, el buscador ofrece tres vistas: una selección propia de modelos destacados, una lista por descargas y otra por popularidad reciente. Para cada modelo probado, la herramienta indica si funciona sin modificaciones en Trainium, en cuántos núcleos del chip debe repartirse (lo que se conoce como sharding, o fragmentación del modelo para distribuir su carga) y cuánto tiempo tarda en ejecutarse, incluyendo la compilación inicial. Cuando un modelo no ha sido probado, el sistema explica el motivo, que suele ser falta de priorización o incompatibilidad técnica.
La utilidad no solo informa del resultado: también ofrece un script de Python listo para ejecutar en instancias Trainium2 (la segunda generación del chip de Amazon). El código sigue una receta estándar: carga el modelo con la función from_pretrained de la biblioteca Transformers, lo fragmenta si es necesario y genera texto usando torch.compile con el backend neuron (el sistema que adapta PyTorch para que funcione en Trainium). Para ahorrar tiempo, el buscador proporciona un comando que descarga los grafos precompilados del modelo desde un bucket público de Amazon S3, evitando así la compilación inicial.
Actualmente, las pruebas se limitan a modelos de Transformers para generación de texto e imagen-texto combinados, ejecutados en un máximo de cuatro dispositivos Trainium2, lo que equivale a 16 núcleos. Los modelos se fragmentan usando paralelismo tensorial (una técnica que divide las operaciones matemáticas entre varios procesadores) o, en el caso de modelos mixture-of-experts (arquitecturas que activan solo parte de sus parámetros según la entrada), paralelismo entre expertos. Hugging Face planea ampliar las pruebas a modelos de difusión (como Stable Diffusion), transformadores de frases y otras tareas.
Todos los datos del buscador provienen de un dataset público llamado aws-neuron/torch-neuronx-model-support, con archivos en formato Parquet (un formato de almacenamiento columnar muy eficiente) que se pueden consultar directamente con bibliotecas como pandas. Cada fila del dataset incluye el resultado de la prueba, las versiones exactas del compilador Neuron y de las bibliotecas usadas, y el script que se ejecutó. Según la documentación, un resultado PASS indica que el modelo carga y genera 32 tokens correctos, pero no garantiza buen rendimiento en generaciones más largas.
