DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2023
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866916185094750208 |
|---|---|
| author | Li, Dacheng Shao, Rulin Xie, Anze Xing, Eric P. Ma, Xuezhe Stoica, Ion Gonzalez, Joseph E. Zhang, Hao |
| author_facet | Li, Dacheng Shao, Rulin Xie, Anze Xing, Eric P. Ma, Xuezhe Stoica, Ion Gonzalez, Joseph E. Zhang, Hao |
| contents | FlashAttention (Dao, 2023) effectively reduces the quadratic peak memory usage to linear in training transformer-based large language models (LLMs) on a single GPU. In this paper, we introduce DISTFLASHATTN, a distributed memory-efficient attention mechanism optimized for long-context LLMs training. We propose three key techniques: token-level workload balancing, overlapping key-value communication, and a rematerialization-aware gradient checkpointing algorithm. We evaluate DISTFLASHATTN on Llama-7B and variants with sequence lengths from 32K to 512K. DISTFLASHATTN achieves 8x longer sequences, 4.45 - 5.64x speedup compared to Ring Self-Attention, 2 - 8x longer sequences, 1.24 - 2.01x speedup compared to Megatron-LM with FlashAttention. It achieves 1.67x and 1.26 - 1.88x speedup compared to recent Ring Attention and DeepSpeed-Ulysses. Code is available at https://github.com/RulinShao/LightSeq. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2310_03294 |
| institution | arXiv |
| publishDate | 2023 |
| record_format | arxiv |
| spellingShingle | DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training Li, Dacheng Shao, Rulin Xie, Anze Xing, Eric P. Ma, Xuezhe Stoica, Ion Gonzalez, Joseph E. Zhang, Hao Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing FlashAttention (Dao, 2023) effectively reduces the quadratic peak memory usage to linear in training transformer-based large language models (LLMs) on a single GPU. In this paper, we introduce DISTFLASHATTN, a distributed memory-efficient attention mechanism optimized for long-context LLMs training. We propose three key techniques: token-level workload balancing, overlapping key-value communication, and a rematerialization-aware gradient checkpointing algorithm. We evaluate DISTFLASHATTN on Llama-7B and variants with sequence lengths from 32K to 512K. DISTFLASHATTN achieves 8x longer sequences, 4.45 - 5.64x speedup compared to Ring Self-Attention, 2 - 8x longer sequences, 1.24 - 2.01x speedup compared to Megatron-LM with FlashAttention. It achieves 1.67x and 1.26 - 1.88x speedup compared to recent Ring Attention and DeepSpeed-Ulysses. Code is available at https://github.com/RulinShao/LightSeq. |
| title | DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training |
| topic | Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing |
| url | https://arxiv.org/abs/2310.03294 |