BurstAttention: Un Marco de Atención Distribuida Eficiente para Secuencias Extremadamente Largas

BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences

March 14, 2024

Autores: Sun Ao, Weilin Zhao, Xu Han, Cheng Yang, Zhiyuan Liu, Chuan Shi, Maosong Sun, Shengnan Wang, Teng Su

cs.AI

Resumen

Los módulos de atención efectivos han desempeñado un papel crucial en el éxito de los modelos de lenguaje grandes (LLMs) basados en Transformers, pero las complejidades cuadráticas en tiempo y memoria de estos módulos de atención también representan un desafío al procesar secuencias largas. Una solución potencial para el problema de las secuencias largas es utilizar clústeres distribuidos para paralelizar el cálculo de los módulos de atención en múltiples dispositivos (por ejemplo, GPUs). Sin embargo, adoptar un enfoque distribuido introduce inevitablemente sobrecargas adicionales de memoria para almacenar los resultados locales de atención y genera costos adicionales de comunicación para agregar los resultados locales en resultados globales. En este artículo, proponemos un marco de atención distribuida llamado "BurstAttention" para optimizar el acceso a la memoria y las operaciones de comunicación tanto a nivel del clúster global como de los dispositivos locales. En nuestros experimentos, comparamos BurstAttention con otras soluciones competitivas de atención distribuida para el procesamiento de secuencias largas. Los resultados experimentales bajo diferentes configuraciones de longitud demuestran que BurstAttention ofrece ventajas significativas para procesar secuencias largas en comparación con estas líneas base competitivas, reduciendo un 40% la sobrecarga de comunicación y logrando una aceleración de 2X durante el entrenamiento de secuencias de 32K en 8 X A100.

English

Effective attention modules have played a crucial role in the success of Transformer-based large language models (LLMs), but the quadratic time and memory complexities of these attention modules also pose a challenge when processing long sequences. One potential solution for the long sequence problem is to utilize distributed clusters to parallelize the computation of attention modules across multiple devices (e.g., GPUs). However, adopting a distributed approach inevitably introduces extra memory overheads to store local attention results and incurs additional communication costs to aggregate local results into global ones. In this paper, we propose a distributed attention framework named ``BurstAttention'' to optimize memory access and communication operations at both the global cluster and local device levels. In our experiments, we compare BurstAttention with other competitive distributed attention solutions for long sequence processing. The experimental results under different length settings demonstrate that BurstAttention offers significant advantages for processing long sequences compared with these competitive baselines, reducing 40% communication overheads and achieving 2 X speedup during training 32K sequence length on 8 X A100.