DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training

Dacheng Li, Rulin Shao, Anze Xie, Eric P. Xing, Xuezhe Ma +3 more
10/5/2023
cs.LGcs.AIcs.DC

Official code: RulinShao/LightSeq (222 stars, last commit 2024-08-19).

Abstract

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.

View on arXivView PDF

Code Implementations(1)

RulinShao/LightSeqOfficial100%

Official repository for DistFlashAttn: Distributed Memory-efficient Attention for Long-context LLMs Training

22210PythonOct 5, 20232 years ago

Similar Papers

Titan-Llama: Retrofitting Pre-trained LLMs with External Neural Memory Modules

90% match1 repo2026

Long-Context Attention Benchmark: From Kernel Efficiency to Distributed Context Parallelism

90% match2025

Zebra: Extending Context Window with Layerwise Grouped Local-Global Attention

90% match2023

A Training-free Sub-quadratic Cost Transformer Model Serving Framework With Hierarchically Pruned Attention

89% match2024

BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences

88% match1 repo242024

Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention

88% match3 repos3742024

Cite this paper

@article{li2023distflashattn,
  title  = {DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training},
  author = {Dacheng Li and Rulin Shao and Anze Xie and Eric P. Xing and Xuezhe Ma and Ion Stoica and Joseph E. Gonzalez and Hao Zhang},
  year   = {2023},
  eprint = {2310.03294},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2310.03294v2}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2310.03294.svg)](https://codeofpaper.com/paper/2310.03294)

Discussion