SWattention: diseñando atención rápida y eficiente en memoria para una nueva supercomputadora Sunway

Resumen En los últimos años, los modelos de lenguaje grandes (LLM) basados en Transformer se han convertido en la tecnología dominante en una serie de aplicaciones. Para aumentar la longitud de secuencia del Transformer, se propone FlashAttention para calcular atención exacta con requisitos de memoria reducidos y ejecución más rápida. Sin embargo, implementar el algoritmo FlashAttention en la nueva generación de supercomputadoras Sunway enfrenta muchas limitaciones, como la arquitectura heterogénea única y el ancho de banda de memoria limitado. Este trabajo propone SWattention, un método altamente eficiente para calcular la atención exacta en el procesador SW26010pro. Para utilizar completamente los 6 grupos de núcleos (CG) y 64 núcleos por CG en el procesador, diseñamos una estrategia de partición de tareas paralelas en dos niveles. Se emplea acceso a memoria asincrónico para garantizar que el acceso a memoria se superponga con la computación. Además, se introduce una estrategia de subdivisión para determinar tamaños óptimos de bloques SRAM. En comparación con la atención estándar, SWattention logra alrededor de 2.0x de aceleración para entrenamiento FP32 y 2.5x de aceleración para entrenamiento de precisión mixta. Las longitudes de secuencia varían de 1k a 8k y escalan hasta 16k sin quedarse sin memoria. En cuanto al rendimiento de extremo a extremo, SWattention logra hasta 1.26x de aceleración para el entrenamiento de modelos estilo GPT, lo que demuestra que SWattention permite una mayor longitud de secuencia para el entrenamiento de LLM.