Efficiently Training 7B LLM with 1 Million Sequence Length on 8 GPUs

2024-07-16

arXiv.org

AI Summary

Training 7B-language models on ultra-long contexts (1M tokens) across 8×A800 GPUs suffers from severe activation memory explosion and GPU memory fragmentation, resulting in low model flops utilization (MFU). Method: We propose a fine-grained activation memory management framework featuring (i) token-level dynamic offloading and reloading, (ii) a two-tier mixed-integer programming (MIP) optimizer for CPU–GPU heterogeneous memory reuse under CPU memory constraints, and (iii) tight integration of FlashAttention, activation recomputation, and distributed training primitives. Contribution/Results: Our approach achieves 52.30% MFU—97% and 80% higher than Megatron-LM and DeepSpeed, respectively—while enabling stable, efficient training of million-token sequences on 8 A800 GPUs for the first time. The framework preserves computational continuity and respects system memory budgets without compromising throughput or scalability.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search

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%.

Problem

• Multi-GPU Training
• Large Language Models
• Memory Management

Innovation

• Memory Optimization
• GPU Training Efficiency
• Large Language Models