Enregistré dans:
| Auteurs principaux: | , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2023
|
| Sujets: | |
| Accès en ligne: | https://arxiv.org/abs/2310.03294 |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _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 |