Los modelos de la biblioteca Transformers pueden ejecutarse en los chips Trainium de AWS gracias a TorchNeuron, un backend nativo para PyTorch, la biblioteca de aprendizaje automático de código abierto. La cuestión es saber cuáles funcionan realmente. El Space torch-neuronx model support, publicado por Hugging Face, responde a esta pregunta para los modelos más populares del Hub.
- El Space torch-neuronx model support de Hugging Face indica qué modelos Transformers funcionan en AWS Trainium.
- Permite consultar resultados, elegir versiones de Neuron y ejecutar el script run.py en una instancia trn2.
- Solo prueba modelos de generación de texto y de texto a imagen que caben en dieciséis núcleos Trainium.
- Un resultado PASS no indica nada sobre el rendimiento ni sobre generaciones más largas.
Buscar un modelo
El Space se abre con la vista Selected, una selección de modelos elegidos entre los más descargados y tendencia del Hub. Cada línea indica el resultado de la prueba, el número de núcleos Trainium en los que se distribuye el modelo, la duración total de la ejecución, incluyendo la compilación, y la memoria máxima utilizada en el dispositivo.
Para buscar un modelo concreto, basta con cambiar a las vistas Trending o Downloads e introducir parte de su nombre. Estas vistas siguen la clasificación del Hub, tarea por tarea, independientemente de si el modelo ha sido probado o no. Cuando un modelo ha sido descartado, el Space especifica la razón: la mayoría de las veces no era prioritario, a veces es realmente incompatible. El usuario también puede elegir las versiones del compilador Neuron (neuronx-cc) y de Transformers a examinar.
Por el momento, solo se prueban los modelos Transformers, y la selección se limita a las tareas de generación de texto y de texto a imagen. El equipo planea ampliar el sistema a diffusers, a los sentence transformers y a otras tareas de Transformers.
Ejecutar un modelo uno mismo
Al seleccionar un modelo, se ve la pila de software Neuron en la que fue probado, los paquetes a instalar y el script run.py utilizado durante la prueba. Este script carga el modelo con from_pretrained, lo distribuye en varios núcleos si es necesario, y luego genera texto con una caché estática y torch.compile(backend= »neuron »).
En una instancia trn2 con la misma pila, por ejemplo en el AWS Neuron Deep Learning Container, basta con instalar los paquetes y luego ejecutar python3 run.py. No se necesita ningún torchrun, el script lanza por sí mismo los procesos de distribución cuando el modelo está distribuido. La primera ejecución comienza con la compilación de los grafos del modelo. Para evitarla, se pueden copiar los grafos producidos durante la prueba desde el bucket público NEFF cache, con el comando proporcionado por el Space, que tiene en cuenta la versión del compilador.
Qué significa « funciona »
Cada prueba carga el modelo en una instancia trn2, distribuido en el menor número de núcleos Trainium capaz de contener sus pesos. La mayoría de los modelos utilizan el paralelismo de tensores, que divide los cálculos entre varios núcleos. Los modelos de expertos mixtos, sin plan de paralelismo de tensores, solo utilizan el paralelismo de expertos. La prueba genera luego 32 tokens, es decir, unidades de texto, y los verifica. En un solo núcleo, estos tokens se comparan con los de una ejecución en CPU del mismo punto de control. Cuando el modelo está distribuido, un pequeño modelo juez verifica que el texto sea un inglés fluido.
Solo se prueban los modelos que caben en cuatro dispositivos Trainium2, es decir, un máximo de dieciséis núcleos. Un resultado PASS no dice nada sobre el rendimiento ni sobre generaciones más largas.
Todos los resultados provienen del conjunto de datos público aws-neuron/torch-neuronx-model-support, que contiene un archivo Parquet por cada par de versiones de neuronx-cc y de biblioteca. Cada línea proporciona el resultado de un modelo, el error eventual, las versiones de la pila y el run.py ejecutado. Es posible consultarlo directamente con pandas, por ejemplo para el modelo Qwen/Qwen3-0.6B. La ficha del conjunto de datos especifica el significado exacto de cada resultado.
