HyScale acelera el entrenamiento con atención dispersa con un millón de tokens
HyScale: Scaling Long-Context Training for Learned Sparse Attention · arXiv:2610.32057
En palabras sencillas
Algunos sistemas de IA solo eligen unas pocas palabras anteriores a las que prestar atención, pero encontrar esas palabras puede llevar mucho tiempo cuando el texto es muy largo. HyScale es un sistema de entrenamiento que encuentra las palabras elegidas mientras las puntúa y guarda esas selecciones para no tener que buscarlas de nuevo durante el entrenamiento.
En la configuración evaluada, HyScale entrenó GLM-5.1 con secuencias de un millón de tokens 2.07 veces más rápido que Megatron-LM. Ese resultado corresponde al modelo y a las 512 GPU H20 indicados.
La investigación
Si un sistema de atención usa solo unas pocas entradas de una secuencia muy larga, ¿por qué encontrar esas entradas puede consumir tanto tiempo de entrenamiento?
En la configuración de GLM-5.1 medida, la selección ocupó casi la mitad de un paso de entrenamiento, lo que señala una oportunidad concreta para reducir el costo del entrenamiento con contexto largo.
HyScale combina la selección sobre la marcha con su repetición durante la recomputación para reducir el costo de entrenar la atención dispersa aprendida con secuencias largas.
HyScale introduce dos cambios en la selección: elige entradas mientras las puntúa y luego reutiliza los índices elegidos durante la recomputación. Así evita escribir una tabla completa de puntuaciones en la memoria de la GPU y repetir todo el examen de selección. El sistema también incluye optimizaciones para entrenar con contextos largos.
En una comparación del entrenamiento completo de GLM-5.1 a partir de los pesos publicados, HyScale es 2.07 veces más rápido que Megatron-LM con secuencias de un millón de tokens y 512 GPU H20.
Una ablación independiente de una capa atribuye la mayor parte del ahorro medido en tiempo de capa a los dos métodos de selección, lo que respalda la conclusión de que reducir el trabajo de selección contribuye de manera sustancial en las configuraciones evaluadas.
Estos resultados respaldan un entrenamiento más rápido y una ejecución más rápida de las capas en las configuraciones evaluadas. No demuestran que se obtengan las mismas aceleraciones ni la misma calidad de entrenamiento con otros modelos, longitudes de secuencia o equipos.
Qué muestran las pruebas
- La selección es el cuello de botella
HyScale es un sistema de entrenamiento con contexto largo para la atención dispersa aprendida, que selecciona unas pocas entradas anteriores para cada consulta. La pregunta es si el entrenamiento puede evitar dedicar más tiempo a seleccionar entradas que a usarlas.
En la configuración de Megatron-LM medida, con un millón de tokens, la selección ocupó el 47% de un paso de entrenamiento de GLM-5.1, sin contar la espera de la canalización, y tardó 6.6 veces más que la atención dispersa.
- Elegir entradas mientras se puntúan
El método HyTopK de HyScale estima un umbral usando claves espaciadas uniformemente dentro del contexto de cada consulta. Filtra las puntuaciones mientras las calcula, evitando una tabla completa de puntuaciones en la memoria de la GPU, y luego elige las entradas top-k exactas entre las candidatas.
Como la estimación puede fallar, el sistema detecta las consultas con muy pocas candidatas o con desbordamiento y vuelve a examinar en la GPU todo su contexto.
- Reutilizar las selecciones durante el entrenamiento
La repetición de la selección aborda otro costo: la recomputación de activaciones puede repetir todo el examen de selección. HyScale guarda los índices elegidos durante el paso hacia adelante y los usa para volver a puntuar solo esas entradas durante la recomputación. Esto se ajusta a las necesidades del paso hacia atrás, que requiere las entradas seleccionadas y sus puntuaciones, no una nueva búsqueda por todo el contexto.
- Las mejoras de capa se reflejan en el entrenamiento completo
En una ablación acumulativa de una capa de atención de GLM-5.1 con un millón de tokens, HyTopK y la repetición de la selección explican el 84% de la reducción medida del 66% en el tiempo de capa de Megatron-LM; las optimizaciones de implementación representan el otro 11% del tiempo de referencia de la capa.
En el entrenamiento completo de GLM-5.1 con 512 GPU H20, HyScale es 2.07 veces más rápido con un millón de tokens. La ablación de la capa y la comparación del entrenamiento completo miden alcances distintos.
- Pruebas de otras dos capas
Los autores también comparan capas de atención individuales de dos arquitecturas por bloques con sus implementaciones de Megatron-LM. Con un millón de tokens y 32 GPU H800, HyScale ejecuta la capa CSA de DeepSeek-V4-Flash 2.01 veces más rápido y la capa QSA de Qwen3.8-Flash-Next 1.57 veces más rápido.
Estos resultados muestran mejoras medidas en esas pruebas por capa, no aceleraciones del entrenamiento completo de los dos modelos.
Detalles
- Autores del paper original
- Yangsheng Deng, Xinjing Huang, Yulong Wang, Yunzhuo Liu, Jun Hu
- Enviado a arXiv
- Licencia
- arXiv non-exclusive