bibtext: @inproceedings{
  li2024distflashattn,
  title={{DISTFLASHATTN}: Distributed Memory-efficient Attention for Long-context {LLM}s 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},
  booktitle={First Conference on Language Modeling},
  year={2024},
  url={https://openreview.net/forum?id=pUEDkZyPDl}
  }
abstract: FlashAttention 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 DistFlashAttention, 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 DistFlashAttention on Llama-7B and variants with sequence lengths from 32K to 512K. DistFlashAttention 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. Codes are available at https://github.com/RulinShao/LightSeq.
title: DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training
authors: Dacheng Li Rulin Shao Anze Xie Eric P. Xing Xuezhe Ma Ion Stoica Joseph E. Gonzalez Hao Zhang