Training Ultra Long Context Language Model with Fully Pipelined Distributed Transformer

Training Ultra Long Context Language Model with Fully Pipelined Distributed Transformer

Speakers: Jinghan Yao, Ohio State University; Sam Ade Jacobs, Microsoft; Masahiro Tanaka, Microsoft; Olatunji Ruwase, Microsoft; Dhabaleswar Panda, Ohio State University

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

The rapid advancement of large language models (LLMs) has highlighted a critical bottleneck: the ability to process and train on ultra-long input sequences. While models like Llama 3.1 are pushing context lengths to 128K tokens, achieving even longer contexts for training or inference remains a significant challenge dueailing to GPU memory limitations. This talk, presented by Jinghan Yao from Ohio State University and based on collaborative work with the Microsoft DeepSpeed team, introduces the Fully Pipelined Distributed Transformer (FPDT), a novel approach designed to efficiently overcome these memory hurdles and enable training of LLMs with millions of tokens in context.

FPDT addresses the prohibitive memory consumption associated with activations during long-context training, particularly in the attention mechanism's backward pass. By integrating three key designs—GPU Trunking, Host Memory Offloading, and Double Buffering—FPDT significantly extends the feasible sequence length, pushing it far beyond the current state of the art. The presented results demonstrate FPDT's capability to train a 7-billion parameter model on 2 million token sequences, a 16x improvement over existing methods, while maintaining high Model Flops Utilization (MFU). This work is crucial for unlocking the full potential of LLMs to understand and generate highly coherent and contextually rich long-form content, with direct implications for both training and the pre-filling stage of inference.

Background

▶ Watch: Introduction: Challenges of ultra-long context LLM training (0:00)

Training large language models with extended context windows presents a formidable challenge, primarily due to the escalating demand for GPU memory. During model training, four main components contribute to GPU memory consumption: model parameters, optimizer states, gradients, and activations. While the first three are largely fixed once the model architecture and optimizer are chosen, activations are highly dynamic. Their memory footprint scales directly with the sequence length, denoted as 'S'. As 'S' increases, activations consume disproportionately more memory, quickly leading to out-of-memory (OOM) errors on GPUs.

A deeper analysis reveals that the most significant memory spikes occur during the backward pass of the attention mechanism within a transformer block. Even with optimized attention implementations like FlashAttention, the backward pass for attention can demand up to eight times the GPU memory compared to its forward counterpart. This is because, in addition to storing the query (Q), key (K), and value (V) projections and the attention output from the forward pass, the backward pass requires storing gradients for these tensors, leading to a substantial, albeit transient, memory peak. These spikes often exceed the memory used by model parameters and optimizer states, making attention the primary bottleneck for long-context training.

Existing context parallel strategies, such as Ring Attention and DeepSpeed Ulysses, attempt to distribute the sequence across multiple GPUs. While effective to a degree, they often require an impractically large number of GPUs for extremely long contexts. For instance, training a relatively small 1.2 billion parameter GPT model with a 1 million token context using DeepSpeed Ulysses alone would necessitate at least 64 GPUs to avoid OOM, highlighting the severe limitations of current approaches in scaling to truly ultra-long contexts. The problem thus lies in designing a system that can effectively manage and reduce these activation memory spikes without sacrificing computational efficiency, thereby enabling million-level sequence lengths on more modest hardware configurations.

Key Findings

▶ Watch: Backward attention pass: the primary memory bottleneck (4:00)

The core contribution of this work is the Fully Pipelined Distributed Transformer (FPDT), an innovative framework built upon the existing DeepSpeed Ulysses context parallel design. FPDT introduces three synergistic architectural designs aimed at drastically reducing GPU memory consumption and enabling ultra-long context training:

1. GPU Trunking: This technique involves slicing the input sequence into smaller, manageable "trunks." By processing these trunks iteratively, FPDT significantly reduces the peak activation memory footprint on the GPU's High Bandwidth Memory (HBM), effectively turning large, unmanageable tensors into a series of smaller, sequential operations.

1. Host Memory Offloading: To support context lengths that far exceed even the reduced HBM requirements of trunking, FPDT strategically offloads less immediately needed sequence data (specifically, past key and value trunks) to the host CPU memory. This allows the GPU to focus its precious HBM on the currently active trunks, pushing the theoretical maximum context length into the millions.

1. Double Buffer Design: Recognizing that offloading data to host memory can introduce significant latency, FPDT incorporates a double buffering mechanism. This design intelligently overlaps the data transfer (pre-fetching the next trunk from host memory) with the computation of the current trunk on the GPU. This ensures that the GPU remains busy and minimizes idle time, thereby maintaining high Model Flops Utilization (MFU) despite the reliance on slower host memory.

The combination of these designs yields remarkable results:

• FPDT enables the training of a 7-billion parameter model with a 2 million token sequence length on a single A100 80GB GPU. This represents a 16x increase in supported context length compared to the prior state of the art.
• Crucially, FPDT achieves this while maintaining a high Model Flops Utilization (MFU) of over 55%, demonstrating its efficiency.
• When combined with other advanced optimization techniques like tensor parallelism, activation checkpointing, and Zero family optimizers, FPDT can further extend context lengths to 4 million tokens for an 8-billion parameter Llama 3 model on an 8-GPU DGX node, still maintaining over 55% MFU.
• The methodology is versatile, applicable not only to model training and fine-tuning but also to the pre-filling stage of inference, where similar memory spikes can occur.
• The entire solution is implemented using standard PyTorch functions and hooks, without requiring custom kernel modifications, and its code is publicly available within the DeepSpeed repository.

Technical Deep Dive

▶ Watch: Introducing FPDT: trunking, offloading, and double buffering (4:40)

FPDT builds upon and significantly enhances the DeepSpeed Ulysses context parallel design. In Ulysses, before the attention operation, the sequence is distributed such that each GPU holds only a piece of the sequence but all the attention heads. During the attention computation, an all-to-all communication gathers the entire sequence onto each GPU, while simultaneously scattering the attention heads. This means each GPU momentarily holds the full sequence but only a subset of heads. After attention, another all-to-all operation scatters the processed data back to its original distribution. This strategy helps manage memory by distributing the sequence, but it still requires each GPU to temporarily hold the entire sequence during attention, which becomes prohibitive for very long contexts.

FPDT addresses this by introducing GPU Trunking and Host Memory Offloading. Instead of treating the entire sequence as a single unit, FPDT slices the input into multiple smaller "trunks." For an input tensor of size S (sequence length) by B (batch size) by D (hidden dimension), FPDT slices it into trunks where each trunk has a sub-sequence length (e.g., S/N, where N is the number of trunks). This allows for fine-grained control over how many tokens reside in HBM at any given time.

Consider the auto-regressive nature of LLMs: a query (Q) token only needs to attend to key (K) and value (V) tokens that precede it or are simultaneous. When processing the first trunk (Q1, K1, V1), Q1 attends to K1 and V1. The output is generated, and crucially, K1 and V1 are then cached to host memory. As subsequent trunks (e.g., QM, KM, VM for the M-th trunk) are processed, QM must attend to its own KM and VM, as well as all previously processed and offloaded K and V trunks (K1, V1, ..., K(M-1), V(M-1)).

This iterative process is similar in principle to techniques like online softmax or how FlashAttention processes attention in blocks. Each time, only a subset of the sequence (the current trunk and relevant previous trunks) is brought into HBM. A partial attention output is computed, and statistics (like maximum logits for stable softmax) are updated. This process repeats for all trunks until the final output for the current query trunk is obtained. By controlling the size of these trunks, FPDT effectively caps the peak memory consumption of activations, preventing OOM errors.

However, relying on host memory for caching introduces a new challenge: the significant latency of host-to-device and device-to-host data transfers. To mitigate this, FPDT incorporates a Double Buffer Design. If data is simply loaded, computed, then unloaded in a sequential manner, the GPU will frequently stall, waiting for data from the host, leading to low MFU. FPDT exploits the computational complexity of the attention mechanism. While element-wise operations (like FFN or ReLU) have O(N) complexity (where N is the sequence length or trunk size), attention has O(N^2) complexity. This quadratic scaling of computation means that for a sufficiently large trunk size, the time required to compute the attention for the current trunk can perfectly overlap the time required to pre-fetch the next trunk from host memory.

Profiling on an A100 GPU with a PCI Gen 4 system revealed a "sweet spot" for trunk size. When the sequence length (trunk size) reached approximately 64K tokens, the latency of pre-fetching the next trunk from host memory became roughly equal to the latency of the attention computation on the current trunk. Below this point, pre-fetching latency could not be fully hidden. Above this point, while overlap was perfect, increasing the trunk size further would simply consume more HBM without yielding additional T-flops, thus wasting precious GPU memory. This empirical finding of 64K tokens as an optimal trunk size for A100/PCI Gen 4 systems is critical for balancing memory efficiency and computational throughput.

The overall FPDT pipeline orchestrates these components using three distinct CUDA streams: one for all-to-all communication, one for computation, and one for offloading (host-to-device and device-to-host transfers). By carefully analyzing data dependencies and scheduling operations across these streams, FPDT maximizes overlap. For example, host-to-device transfers for the next trunk are initiated while the current trunk's attention is being computed, and device-to-host transfers for previous K/V trunks are performed asynchronously. The entire implementation leverages standard PyTorch building functions and hooks, avoiding the need for custom CUDA kernels, which simplifies integration and maintenance within the DeepSpeed framework.

Experimental Setup & Results

▶ Watch: Background: DeepSpeed Ulysses context parallel design (6:40)

The effectiveness of FPDT was rigorously evaluated across various model sizes and hardware configurations, demonstrating its significant advancements in long-context training.

Models and Baselines:

• Initial analysis and motivation were drawn from observations on a 1.2 billion parameter GPT model, highlighting that achieving a 1 million context length with DeepSpeed Ulysses alone required 64 GPUs.
• Memory profiling was conducted on a 40-layer, 30 billion parameter GPT model to precisely identify the memory spikes during forward and backward passes, particularly in the attention and FFN layers.
• The primary benchmarks for FPDT's performance were a 7-billion parameter model and an 8-billion parameter Llama 3 model.
• Comparisons were implicitly made against existing "state-of-the-art" context parallel solutions, including DeepSpeed Ulysses, which FPDT enhances.

Hardware:

• The main experimental platform utilized NVIDIA A100 80GB GPUs on systems equipped with PCI Gen 4 interconnects.
• For the most extensive context lengths (4 million tokens), an 8-GPU DGX node was used.
• A crucial discussion during the Q&A session highlighted adaptability to NVIDIA H100 systems, which typically feature PCI Gen 5 interconnects.

Metrics:

• Maximum Supported Sequence Length: The primary metric, demonstrating how many tokens FPDT can handle without OOM.
• Model Flops Utilization (MFU): A measure of computational efficiency, indicating how effectively the GPU's theoretical peak performance is utilized.
• Memory Footprint: Detailed profiling to show the reduction in peak activation memory.
• Latency Analysis: Profiling of pre-fetching and computation latencies to determine optimal trunk sizing.

Headline Numbers and Results:

• 16x Context Length Increase: On a single A100 80GB GPU with a 7-billion parameter model, FPDT successfully enabled training with a 2 million token sequence length. This was explicitly stated as 16 times longer than the prior state of the art.
• High MFU: Across these demanding scenarios, FPDT consistently maintained an MFU of over 55%. This is a critical achievement, as often memory-saving techniques can lead to significant drops in computational efficiency.
• 4 Million Token Context: By combining FPDT with other DeepSpeed optimizations (such as tensor parallelism, activation checkpointing, offloading checkpointing, and Zero family optimizers), the system achieved 4 million token context length for an 8-billion parameter Llama 3 model on an 8-GPU DGX node, again with MFU exceeding 55%.
• Memory Spike Reduction: Post-optimization memory profiling graphs visually demonstrated the dramatic reduction and flattening of memory spikes during the backward pass, particularly for the FFN and attention layers. By dynamically adjusting the number of trunks, FPDT allowed control over peak memory consumption, enabling a balance between activation memory and model parameter memory.
• Optimal Trunk Size: Empirical profiling on PCI Gen 4 systems indicated that a trunk size of 64K tokens provided the ideal balance, where the latency of pre-fetching the next trunk from host memory was effectively overlapped by the attention computation of the current trunk. Beyond this point, T-flops did not significantly improve, but HBM consumption increased unnecessarily.
• H100 Adaptability: The Q&A session revealed that the optimal trunk size remains 64K on H100 systems. This is because, while PCI Gen 5 doubles the throughput of PCI Gen 4, the attention computation throughput on H100 also roughly doubles compared to A100, maintaining the balance between I/O and compute. However, specific configurations (e.g., A100 with PCI Gen 5) might require re-tuning the trunk size for optimal overlap.

The experimental results definitively validate FPDT as a highly effective and efficient solution for addressing the long-standing memory bottleneck in training ultra-long context language models.

Practical Implications

▶ Watch: Detailing trunking and offloading for memory control (8:00)

The development of FPDT carries substantial practical implications for the entire AI/ML ecosystem, from model developers to infrastructure teams and deployers. Its ability to scale context windows to millions of tokens fundamentally changes what is possible with large language models.

For practitioners and model builders, FPDT directly enables the creation and training of LLMs that can ingest and reason over unprecedented amounts of information. This means models can process entire books, extensive codebases, lengthy legal documents, or years of conversational history, leading to significantly enhanced capabilities in tasks requiring deep contextual understanding, consistency, and coherence over very long spans. Use cases such as advanced document summarization, complex code generation and analysis, long-form content creation, and highly personalized chatbots become more feasible and performant. The ability to fine-tune models on such long sequences also allows for the adaptation of general-purpose LLMs to highly specialized, context-rich domains.

For infrastructure teams and deployers, FPDT offers a pathway to more efficient resource utilization. By supporting multi-million token contexts on a relatively modest number of A100 or H100 GPUs (e.g., 2M tokens on a single A100 80GB, or 4M tokens on an 8-GPU DGX node), it reduces the prohibitive GPU requirements typically associated with such scale. This translates to lower operational costs, less power consumption, and the ability to achieve higher computational density within existing data center footprints. FPDT is also applicable to the pre-filling stage of inference, which often faces similar memory challenges when preparing long prompts for generation. This extends its utility beyond just training, directly impacting the deployment and serving of long-context LLMs.

Tradeoffs and Limitations:

While FPDT offers significant advantages, it's important to consider its inherent tradeoffs. The reliance on host memory offloading introduces a dependency on the speed of the host-device interconnect (e.g., PCIe bandwidth). Although the double buffering mechanism effectively hides much of this latency, the system's performance is ultimately bounded by the interplay between host memory bandwidth and GPU compute power.

The optimal trunk size is hardware-dependent. As demonstrated in the Q&A, while 64K tokens works well for both A100 (PCI Gen 4) and H100 (PCI Gen 5) due to a balanced scaling of interconnect bandwidth and GPU compute, atypical configurations (e.g., A100 with PCI Gen 5) would require re-profiling and tuning. This implies that for maximum efficiency, infrastructure teams may need to perform specific benchmarking for their exact hardware stack.

Furthermore, while FPDT drastically reduces activation memory, it doesn't eliminate the need for significant GPU memory entirely, especially for very large models. The 4M context length was achieved with the aid of other DeepSpeed optimizations, highlighting that FPDT is a powerful component within a broader suite of memory management techniques rather than a standalone panacea for all memory issues. However, its integration into DeepSpeed, using standard PyTorch functions, makes it highly accessible and relatively easy to adopt for existing DeepSpeed users.

In summary, FPDT represents a crucial step forward in making ultra-long context LLMs a practical reality, offering both enhanced model capabilities and improved resource efficiency, albeit with careful consideration of hardware-specific tuning.

Key Takeaways

• Activation Memory Bottleneck: The primary barrier to training ultra-long context language models is the escalating GPU memory consumption of activations, particularly the significant spikes during the backward pass of the attention mechanism.
• FPDT's Core Innovations: The Fully Pipelined Distributed Transformer (FPDT) combines three synergistic designs: GPU Trunking (slicing sequences into smaller parts), Host Memory Offloading (moving less active data to CPU memory), and a Double Buffer (overlapping data transfer with computation).
• Million-Level Contexts: FPDT enables training with sequence lengths up to 2 million tokens on a single A100 80GB GPU for a 7B model (16x improvement) and 4 million tokens for an 8B Llama 3 model on an 8-GPU DGX node when combined with other optimizations.
• High Efficiency: Despite relying on slower host memory, FPDT maintains high computational efficiency, achieving over 55% Model Flops Utilization (MFU) by cleverly exploiting the O(N^2) complexity of attention to hide data transfer latencies.
• Hardware Adaptability: The optimal trunk size (e.g., 64K tokens) is determined by the balance between interconnect bandwidth (PCIe Gen 4/5) and GPU compute power (A100/H100), demonstrating the solution's robust design across different hardware generations.
• Practical & Accessible: Implemented using standard PyTorch functions and hooks, FPDT is available in the DeepSpeed repository, making it readily accessible for practitioners to build and deploy more capable, long-context LLMs for both training and inference pre-filling.

About the Speaker(s)

The primary presenter for this work was Jinghan Yao, a researcher from Ohio State University. The research presented, concerning the Fully Pipelined Distributed Transformer (FPDT), was conducted during Jinghan Yao's internship at the Microsoft DeepSpeed team. This collaborative effort also involved several co-authors from Microsoft and Ohio State University, including Sam Ade Jacobs, Masahiro Tanaka, Olatunji Ruwase, and Dhabaleswar Panda, highlighting a strong academic-industrial partnership in addressing critical challenges in large-scale AI/ML systems.