MEMO: Fine-grained Tensor Management For Ultra-long Context LLM Training

TL;DR

MEMO tackles memory and fragmentation bottlenecks in ultra-long context LLM training by combining fine-grained activation memory management with a bi-level memory planning approach. It offloads skeletal activations to CPU memory and uses token-wise recomputation and selective swapping to overlap computation with data movement, while solving memory allocation via offline, hierarchical MIP planning. The method yields substantial MFU improvements over Megatron-LM and DeepSpeed (average ~1.97x and ~1.80x) and enables training of 7B LLMs with 1 million-token contexts on 8 A800 GPUs, achieving around 52% MFU. This approach reduces memory fragmentation, minimizes recomputation and communication, and demonstrates practical gains for long-context LLM training with realistic hardware constraints.

Abstract

Nowadays, Large Language Models (LLMs) have been trained using extended context lengths to foster more creative applications. However, long context training poses great challenges considering the constraint of GPU memory. It not only leads to substantial activation memory consumption during training, but also incurs considerable memory fragmentation. To facilitate long context training, existing frameworks have adopted strategies such as recomputation and various forms of parallelisms. Nevertheless, these techniques rely on redundant computation or extensive communication, resulting in low Model FLOPS Utilization (MFU). In this paper, we propose MEMO, a novel LLM training framework designed for fine-grained activation memory management. Given the quadratic scaling of computation and linear scaling of memory with sequence lengths when using FlashAttention, we offload memory-consuming activations to CPU memory after each layer's forward pass and fetch them during the backward pass. To maximize the swapping of activations without hindering computation, and to avoid exhausting limited CPU memory, we implement a token-wise activation recomputation and swapping mechanism. Furthermore, we tackle the memory fragmentation issue by employing a bi-level Mixed Integer Programming (MIP) approach, optimizing memory reuse across transformer layers. Empirical results demonstrate that MEMO achieves an average of 1.97x and 1.80x MFU compared to Megatron-LM and DeepSpeed, respectively. This improvement is attributed to MEMO's ability to minimize memory fragmentation, reduce recomputation and intensive communication, and circumvent the delays associated with the memory reorganization process due to fragmentation. By leveraging fine-grained activation memory management, MEMO facilitates efficient training of 7B LLM with 1 million sequence length on just 8 A800 GPUs, achieving an MFU of 52.30%.

Figures (11)

• Figure 1: The left figure, generated using PyTorch's snapshot API torch_mem_utils, shows the allocated and reserved GPU memory of PyTorch when training a 7B GPT model with sequence length 512K. The right figure shows the time consumption of FlashAttention computation, one transformer layer forward computation, and one-layer full activation offloading when training a 7B GPT on 8 A800 GPUs with a TP size of 8.
• Figure 2: An overview of Memo. We devise a fine-grained recomputation and swapping mechanism to manage the skeletal activations for backward propagation, and leverage a bi-level memory planning method to reuse the memory space of transient activations across the transformer layers.
• Figure 3: (a): The architecture of a typical LLM. (b): An example memory request sequence of a transformer layer's forward and backward pass. Tensors 15 and 16 are skeletal tensors, while the others are transient tensors.
• Figure 4: Illustration of the transformer layer architecture. The sizes of skeletal activations are provided in the brackets.
• Figure 5: Forward and backward propagation with rounding buffers for token-wise recomputation/swapping. During forward propagation, the darker part in the rounding buffers is offloaded to CPU, while the lighter part is discarded; during backward propagation, the darker part in the rounding buffers is prefetched from CPU, while the lighter part is recomputed.