Memory-efficient tensor parallelism for long-sequence Transformer training

Abstract

Transformer-based models like large language models (LLMs) have attracted significant attention in recent years due to their superior performance. A long sequence of input tokens is essential for industrial LLMs to provide better user services. However, memory consumption increases quadratically with the increase of sequence length, posing challenges for scaling up long-sequence training. Current parallelism methods produce duplicated tensors during execution, leaving space for improving memory efficiency. Additionally, tensor parallelism (TP) cannot achieve effective overlap between computation and communication. To solve these weaknesses, we propose a general parallelism method called memory-efficient tensor parallelism (METP), designed for the computation of two consecutive matrix multiplications and a possible function between them (O = f(AB)C), which is the kernel computation component in Transformer training. METP distributes subtasks of computing O to multiple devices and uses send/recv instead of collective communication to exchange submatrices for finishing the computation, avoiding producing duplicated tensors. We also apply the double buffering technique to achieve better overlap between computation and communication. We present the theoretical condition of full overlap to help instruct the long-sequence training of Transformers. Suppose the parallel degree is p; through theoretical analysis, we prove that METP provides O(1/p^3) memory overhead when not using FlashAttention to compute attention and could save at least 41.7% memory compared to TP when using FlashAttention to compute multi-head self-attention. Our experimental results demonstrate that METP can increase the sequence length by 2.38–2.99 times compared to other methods when using eight A100 graphics processing units (GPUs).

1. Introduction

Training language models with long sequences is crucial for handling long-context tasks such as document summarization, dialogue systems, and technical writing (Beltagy et al., 2020; Kaddour et al., 2023; Tarassow, 2023). This is because the context length determines the maximum word count a language model can handle, and longer sequences improve contextual understanding and semantic coherence, allowing models to generate more coherent and contextually appropriate text (Achiam et al., 2023; Huang et al., 2023). However, it requires O(s^2) memory for a sequence with s tokens to complete the attention computation in Transformer models. This memory limitation makes it challenging for language models to calculate extremely long sequences during training. Additionally, training with long sequences costs O(s^2) time for attention computation, bringing computation challenges to a single device.

To address these challenges, researchers proposed various techniques such as tensor parallelism (TP) (Narayanan et al., 2021a), fully sharded data parallelism (FSDP) (Zhao et al., 2023), sequence parallelism (SP) (Li SG et al., 2023; Jacobs et al., 2024), and pipeline parallelism (PP) (Huang et al., 2019; Narayanan et al., 2021b; Liu ZM et al., 2023) to distribute computation and memory costs across multiple devices. These techniques can reduce the memory cost of training large language models (LLMs) to O(1/p), where p denotes the parallel degree. Since these methods are orthogonal to each other, using multiple parallelism methods simultaneously, such as 3D parallelism (Narayanan et al., 2021a; Lai et al., 2023), can reduce the memory footprint to O(1/(pt pd pp)), where pt, pd, and pp represent the parallel degrees of TP, data parallelism (DP), and PP, respectively. However, 3D parallelism could only reduce memory footprint to Ω(1/#device). This work aims to find a 1D parallelism method that can achieve memory usage lower than Ω(1/#device).

FlashAttention (Dao et al., 2022) achieves high memory efficiency by employing tiling techniques to optimize the attention computation process, which comprises two matrix multiplications interleaved with a Softmax operation. This method divides the query, key, and value matrices into submatrices and calculates attention scores for each query submatrix by iterating over all the key and value submatrices, avoiding the significant memory overhead of storing the query@key. As a result, it reduces the memory footprint of computing attention from O(s^2) to O(s), where s represents the sequence length.