BurstAttention: Эффективная распределенная модель внимания для чрезвычайно длинных последовательностей

BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences

March 14, 2024

Авторы: Sun Ao, Weilin Zhao, Xu Han, Cheng Yang, Zhiyuan Liu, Chuan Shi, Maosong Sun, Shengnan Wang, Teng Su

cs.AI

Аннотация

Эффективные модули внимания сыграли ключевую роль в успехе крупных моделей языка на основе трансформера (LLM), однако квадратичная сложность по времени и памяти этих модулей внимания также представляет собой вызов при обработке длинных последовательностей. Одно из потенциальных решений проблемы длинных последовательностей - использовать распределенные кластеры для параллельного выполнения вычислений модулей внимания на нескольких устройствах (например, GPU). Однако принятие распределенного подхода неизбежно вносит дополнительные накладные расходы памяти для хранения локальных результатов внимания и вызывает дополнительные затраты на коммуникацию для агрегации локальных результатов в глобальные. В данной статье мы предлагаем распределенную архитектуру внимания под названием "BurstAttention" для оптимизации доступа к памяти и операций коммуникации как на уровне глобального кластера, так и на уровне локального устройства. В наших экспериментах мы сравниваем BurstAttention с другими конкурентоспособными распределенными решениями внимания для обработки длинных последовательностей. Экспериментальные результаты при различных настройках длины демонстрируют, что BurstAttention предлагает значительные преимущества для обработки длинных последовательностей по сравнению с этими конкурентными базовыми вариантами, сокращая накладные расходы на коммуникацию на 40% и достигая ускорения в 2 раза во время обучения последовательности длиной 32K на 8 устройствах 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.