Entrenamiento de un modelo de atención lineal en AWS Trainium2 con PyTorch nativo

Investigadores de i14 y del Centro de Innovación en IA Generativa de AWS han logrado entrenar y servir con éxito el modelo Laplace 3B utilizando aceleradores AWS Trainium2. Publicado en octubre de 2026, este trabajo demuestra cómo las arquitecturas de atención lineal pueden optimizarse para hardware especializado sin depender de kernels personalizados de GPU.

La colaboración destaca un cambio respecto a la atención softmax tradicional de los transformers, que sufre de escalado cuadrático durante el entrenamiento y costes de memoria crecientes durante la inferencia. Al aprovechar el soporte nativo de PyTorch mediante TorchNeuron Native, el equipo logró métricas de rendimiento significativas manteniendo una huella de memoria fija por secuencia.

Qué ocurrió

Los modelos de lenguaje grandes suelen basarse en arquitecturas transformer que utilizan atención softmax. Durante la inferencia, estos modelos mantienen una caché clave-valor (KV) que crece linealmente con la longitud de la conversación o documento. Este crecimiento limita la concurrencia, aumenta los costes de servicio y a menudo obliga a las aplicaciones a descartar el contexto histórico. Para abordar estas limitaciones, i14 desarrolló Laplace 3B, un modelo basado en principios de atención lineal.

A diferencia de los transformers estándar, Laplace 3B utiliza un estado recurrente que permanece constante en 29,5 MB por secuencia, independientemente de la longitud del contexto. El equipo trabajó con AWS para adaptar esta arquitectura al hardware Trainium2. Utilizaron TorchNeuron Native, una función beta que permite ejecutar código nativo de PyTorch directamente sobre los aceleradores AWS Trainium. Este enfoque preservó el diseño de Atención Lineal Compuertada (GLA) del modelo y su uso del optimizador Muon.

Los resultados mostraron ganancias sustanciales de eficiencia. En una sola instancia trn2.48xlarge, el modelo se entrenó a aproximadamente 120.000 tokens por segundo. Durante la inferencia, sirvió 11.544 tokens de salida por segundo a través de 1.280 secuencias concurrentes. Estas cifras demuestran que los modelos de atención lineal pueden alcanzar un alto rendimiento en hardware especializado cuando la pila de software está correctamente alineada con la arquitectura subyacente.

Cómo funciona

Los modelos de atención lineal reorganizan la computación alrededor de un estado recurrente en lugar de comparaciones de pares de tokens. En los transformers tradicionales, las puntuaciones de atención crecen cuadráticamente con la longitud de la secuencia durante el entrenamiento. Laplace 3B sustituye esto por un mecanismo que actualiza un estado de matriz de tamaño fijo para cada token. Esto cambia la complejidad del entrenamiento de O(T^2) a O(T) y mantiene la memoria de inferencia constante. El modelo comprime el historial en este estado, lo que requiere una evaluación cuidadosa de la calidad de recuperación para cargas de trabajo específicas.

Adaptar este modelo para Trainium2 requirió tres cambios de ingeniería específicos. Primero, el equipo sustituyó los kernels Triton optimizados para GPU por operadores estándar de PyTorch. Estructuraron el kernel GLA para usar 16 fragmentos de 128 tokens, alineándose con las 128 particiones del búfer en chip NeuronCore. Segundo, rediseñaron la comunicación para el optimizador Muon. Al utilizar dos colectivos all-to-all en lugar de métodos de gather-and-broadcast, redujeron el tráfico por rango en 24 veces y aumentaron el rendimiento en 1,63 veces. Tercero, expandieron estáticamente el paso hacia atrás secuencial para ajustarse a los límites de tamaño de grafo y memoria del compilador.

La arquitectura Trainium2 jugó un papel crucial en estas optimizaciones. Cada NeuronCore presenta un Búfer de Estado (SBUF) gestionado por software de 28 MiB con 128 particiones. A diferencia de una caché, el compilador programa cada transferencia de datos por adelantado, asegurando que no ocurran desalojos ni fallos. Al mapear cada fragmento de 128 tokens a una partición, el equipo logró un aumento del 14 por ciento en el rendimiento medido. Esta alineación entre el tamaño de fragmento del modelo y las particiones físicas del hardware fue clave para desbloquear el rendimiento.

Detalles clave

• Laplace 3B mantiene un estado recurrente fijo de 29,5 MB por secuencia, independiente de la longitud del contexto.
• El entrenamiento en una instancia trn2.48xlarge alcanzó aproximadamente 120.000 tokens por segundo.
• La inferencia sirvió 11.544 tokens de salida por segundo a través de 1.280 secuencias concurrentes.
• La implementación utilizó TorchNeuron Native para ejecutar operadores estándar de PyTorch en Trainium2.
• El tráfico de comunicación para el optimizador Muon se redujo 24 veces utilizando colectivos all-to-all.
• El modelo usa 16 fragmentos de 128 tokens para alinearse con las 128 particiones SBUF en los NeuronCores.

Por qué importa

Para los ingenieros que construyen productos de IA, el creciente coste de las cachés KV es un cuello de botella importante. A medida que los contextos se alargan, el uso de memoria se dispara, limitando el número de usuarios concurrentes que un servidor puede manejar. Los modelos de atención lineal como Laplace 3B ofrecen una vía hacia costes de servicio predecibles. Dado que el tamaño del estado no crece con el contexto, las aplicaciones pueden retener historiales largos sin aumentos proporcionales en la asignación de memoria. Esto es particularmente valioso para flujos de trabajo agénticos y análisis de documentos grandes donde la memoria persistente es crítica.

Además, este trabajo demuestra que el hardware especializado como Trainium2 puede soportar arquitecturas novedosas sin requerir reescrituras de kernels de bajo nivel. Al usar PyTorch nativo a través de TorchNeuron Native, los desarrolladores pueden aprovechar optimizaciones específicas del hardware manteniéndose dentro de un framework familiar. El éxito de la estrategia de comunicación all-to-all para el optimizador Muon también proporciona una hoja de ruta para el entrenamiento distribuido en sistemas fragmentados. Muestra que reducir la sobrecarga de comunicación puede generar ganancias significativas de rendimiento, incluso cuando los FLOPs computacionales permanecen inalterados.

Qué puedes hacer

• Evalúa modelos de atención lineal para casos de uso que requieran ventanas de contexto largas y uso de memoria predecible.
• Prueba TorchNeuron Native si estás desplegando arquitecturas personalizadas en instancias AWS Trainium.
• Alinea los tamaños de fragmento de datos con las particiones del búfer de hardware para maximizar el rendimiento en aceleradores especializados.
• Considera patrones de comunicación all-to-all para optimizadores distribuidos para reducir cuellos de botella de tráfico de red.
• Monitorea la calidad de recuperación al cambiar de atención softmax a mecanismos de compresión de estado.
• Perfilá tu motor de inferencia para identificar si el crecimiento de la caché KV está limitando tus niveles de concurrencia.