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 the reuse of memory across transformer layers. Empirical results demonstrate that MEMO achieves an average of 2.42x and 2.26x 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%.

Enabling Energy-Efficient Deployment of Large Language Models on Memristor Crossbar: A Synergy of Large and Small

Aging Aware Retraining for Memristor-based Neuromorphic Computing

Bayesian Inference Based Robust Computing on Memristor Crossbar

BRoCoM: A Bayesian Framework for Robust Computing on Memristor Crossbar.

Modeling and Analysis of a Large-Scale Memristor Crossbar Array for Neuromorphic Chip

Long short-term memory networks in memristor crossbars

Long short-term memory networks in memristor crossbar arrays

BETTER: Bayesian-Based Training and Lightweight Transfer Architecture for Reliable and High-Speed Memristor Neural Network Deployment

A Memristor Crossbar Based Computing Engine Optimized for High Speed and Accuracy

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

Reducing the Energy Dissipation of Large Language Models (llms) with Approximate Memories

Analogue Signal and Image Processing with Large Memristor Crossbars

Bulk-Switching Memristor-Based Compute-In-Memory Module for Deep Neural Network Training

Benchmarking the Performance of Large Language Models on the Cerebras Wafer Scale Engine

An Efficient Programming Framework for Memristor-based Neuromorphic Computing

Memristor Crossbar Architectures for Implementing Deep Neural Networks

Higher order neural processing with input-adaptive dynamic weights on MoS2 memtransistor crossbars

Mitigate IR-Drop Effect by Modulating Neuron Activation Functions for Implementing Neural Networks on Memristor Crossbar Arrays

Hecaton: Training Large Language Models with Scalable Chiplet Systems

AutoGMap: Learning to Map Large-Scale Sparse Graphs on Memristive Crossbars

Efficient and Economic Large Language Model Inference with Attention Offloading