ATTENTION2D: Communication Efficient Distributed Self-Attention Mechanism
Transformer-based models have emerged as a leading architecture for natural language processing, natural language generation, and image generation tasks. A fundamental element of the transformer architecture is self-attention, which allows the model to capture intricate dependencies within the data. However, the self-attention mechanism also incurs significant computational and memory costs, particularly for long sequences. In this paper, we introduce ATTENTION2D, a novel approach that exploits parallelism along two dimensions - query and key/value - of the self-attention operation. This method enables efficient distribution and parallelization of computations across multiple devices. Our approach facilitates asymptotically faster training and inference phases compared to previous methods, without relying on approximations or incurring additional computational or memory overheads. Furthermore, unlike existing techniques that struggle to scale with an increasing number of processing units, our approach effectively scales with additional processing units. Our experimental results confirm the effectiveness of our method in improving communication efficiency and scalability. Compared to Ring Attention, our approach demonstrated up to a 5x performance boost on a GPT-3-like model using 64 NVIDIA A100 GPUs across 16 nodes, and up to a 9.4x performance boost on 64 NVIDIA H100 GPUs across 64 nodes.
Code (0)
등록된 구현이 없습니다.
Tasks
Image GenerationText GenerationMethods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training
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 distribute…
GPUHomogeneous Learning: Self-Attention Decentralized Deep Learning
Federated learning (FL) has been facilitating privacy-preserving deep learning in many walks of life such as medical image classification, network intrusion detection, and so forth. Whereas it necessitates a central para…
Deep LearningFederated Learningimage-classificationImage Classification+5LV-XAttn: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
Cross-attention is commonly adopted in multimodal large language models (MLLMs) for integrating visual information into the language backbone. However, in applications with large visual inputs, such as video understandin…
GPUVideo UnderstandingFederated Attention: A Distributed Paradigm for Collaborative LLM Inference over Edge Networks
Large language models (LLMs) are proliferating rapidly at the edge, delivering intelligent capabilities across diverse application scenarios. However, their practical deployment in collaborative scenarios confronts funda…
Computational EfficiencyPolicy Gradient with Self-Attention for Model-Free Distributed Nonlinear Multi-Agent Games
Multi-agent games in dynamic nonlinear settings are challenging due to the time-varying interactions among the agents and the non-stationarity of the (potential) Nash equilibria. In this paper we consider model-free game…