PaperlayerView exact revision on arXiv

Source-linked AI summary

Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking

Ravi Ghadia, Maksim Abraham, Sergei Vorobyov, Max Ryabinin

arXiv:2602.21196v2cs.LGcs.DC

TL;DR

Long-context Transformer training is limited by activation memory, and existing context-parallel approaches retain sequence-length-dependent memory bottlenecks. UPipe chunks attention at the head level to improve memory efficiency while preserving performance. It supports substantially longer contexts, including 5M tokens on a single 8×H100 node, with comparable throughput.

  • Problem

    Context parallelism scales long-context training across accelerators, but activation memory per device still grows linearly with sequence length and becomes a bottleneck beyond 2M tokens.

  • Method

    UPipe executes attention in multiple stages by processing heads in chunks, improving memory reuse while remaining compatible with existing attention algorithms and GQA.

  • Results

    5 million tokens are supported on a single H100 node and 8M on two nodes, exceeding prior supported context lengths by 25% and 33%, respectively, with comparable throughput.

  • Takeaways & Limitations

    UPipe provides a practical, composable path to extending context length while reducing attention memory and preserving training speed.

  • Takeaways & Limitations

    The memory analysis assumes batch size 1 and mixed-precision tensors, with some intermediate tensors requiring 4 bytes.

Abstract

from arXiv · show

Efficiently processing long sequences with Transformer models usually requires splitting the computations across accelerators via context parallelism. The dominant approaches in this family of methods, such as Ring Attention or DeepSpeed Ulysses, enable scaling over the context dimension but do not focus on memory efficiency, which limits the sequence lengths they can support. More advanced techniques, such as Fully Pipelined Distributed Transformer or activation offloading, can further extend the possible context length at the cost of training throughput. In this paper, we present UPipe, a simple yet effective context parallelism technique that performs fine-grained chunking at the attention head level. This technique significantly reduces the activation memory usage of self-attention, breaking the activation memory barrier and unlocking much longer context lengths. Our approach reduces intermediate tensor memory usage in the attention layer by as much as 87.5$\%$ for 32B Transformers, while matching previous context parallelism techniques in terms of training speed. UPipe can support the context length of 5M tokens when training Llama3-8B on a single 8$\times$H100 node, improving upon prior methods by over 25$\%$.

1. Introduction

Long-context training is constrained by activation memory even with distributed context parallelism. UPipe addresses this bottleneck by chunking attention heads, enabling longer contexts while retaining comparable throughput.

  • At sequence lengths above 2M, activation memory per device remains a bottleneck despite context parallelism sharding computation across accelerators.
  • 5 million tokens on one H100 node and 8M on two nodes exceed prior maximum supported context lengths by 25% and 33%, respectively.UPipe’s throughput remains comparable to other context-parallel techniques.
  • UPipe processes attention heads in smaller chunks, improving memory reuse through staged attention execution.The method is designed around the observation that a subset of heads can saturate the GPU during long-context training.
  • UPipe is a plug-and-play context-parallel replacement that uses the same attention kernels as non-distributed training.
  • 5 million tokens are supported on a single 8×H100 node, with negligible performance differences from other context-parallel methods.

2. Background

Long-context Transformer training combines sequence-length-dependent memory demands with attention communication and activation overheads. The paper analyzes these costs and motivates UPipe’s headwise chunking as a way to reduce attention intermediates while preserving efficiency.

  • 2.1. Multi-head Attention: Standard self-attention requires access to the entire sequence, making activation memory a central obstacle for long-context training.A reported 14B Diffusion Transformer reaches 8 TB of activation memory at 1M tokens.
  • 2.1. Multi-head Attention: Grouped-Query Attention shares key/value heads across query-head groups, reducing key/value memory by a factor of G.
  • 2.2. Context Parallelism: Ring Attention shards context across devices with O(C) communication calls per attention operation, whereas DeepSpeed-Ulysses uses a single all-to-all collective.
  • 2.3. Activation memory for a Transformer model: The forward pass comprises embedding, attention, feed-forward, and cross-entropy phases with memory costs that scale with sequence length.The analysis assumes batch size 1 and mixed-precision parameters, activations, and gradients, while some intermediates use 4 bytes.
  • 2.4. Mitigating memory overheads: Cross-entropy is identified as the first memory obstacle, while tiled computation, activation offloading, and UPipe target different intermediate-memory sources.UPipe is proposed specifically for the attention layer after tiled loss and feed-forward computation and layerwise checkpointing with CPU offloading.
  • 2.4. Mitigating memory overheads: UPipe consumes ν times less intermediate QKV and all-to-all activation memory than Ulysses with activation offloading.FPDT can use arbitrary chunk sizes and lower memory usage, but the passage reports performance degradation.

3. Untied Ulysses

DeepSpeed-Ulysses redistributes sequence shards across devices for attention but retains a memory bottleneck from QKV and communication buffers. UPipe addresses this by executing attention in headwise chunks, reusing buffers across stages and reducing peak activation memory.

  • 3.1. DeepSpeed-Ulysses: DeepSpeed-Ulysses shards sequences across devices, then uses all-to-all operations to redistribute QKV tensors so each device accesses the full sequence for assigned heads.Attention outputs are reshared to the original sequence-sharded layout with a second all-to-all operation.
  • 3.2. Memory Usage: The peak memory bottleneck in DeepSpeed-Ulysses comes from storing QKV tensors and similarly sized all-to-all buffers for all heads.These intermediate tensors scale proportionally with the number of attention heads.
  • 3.3. Untied Ulysses: UPipe: UPipe executes attention across H/U stages, processing U heads per stage and reusing intermediate buffers between stages.The design performs input and output all-to-all communication for each head chunk rather than all heads simultaneously.
  • 3.4. Memory Savings: With U = C, UPipe’s peak attention activation memory becomes independent of the number of heads.The resulting maximum memory usage is 12 · S · dhead bytes.
  • 3.4. Memory Savings: 87.5% lower intermediate activation memory: UPipe requires 12 · S · dhead bytes versus 96 · S · dhead bytes for DS-Ulysses on Qwen3-32B with H = 64 and C = 8.The comparison concerns self-attention intermediate activations on a single 8×H100 node.

4. Implementation

UPipe is integrated as a drop-in implementation alongside unified sequence-parallel baselines, using shared framework components for fair comparison. Its GQA schedule reorders query processing so communicated key/value tensors can be reused, reducing communication when G > 1.

  • 4. Implementation: The training framework and Unified Sequence Parallelism baselines are kept consistent to support fair comparisons across context-parallel techniques.UPipe and USP use the same training framework.
  • 4. Implementation: UPipe is integrated into TorchTitan through a drop-in replacement of existing optimized layers.The implementation uses the same code structure as USP and can be extended to a hybrid Ulysses-plus-Ring setup.
  • 4.1. Grouped-Query Attention Scheduling: UPipe retains GQA compatibility by rearranging query tensors so key/value tensors communicated in a previous stage can be reused.This preserves the memory advantage of GQA while changing the order in which heads are processed.
  • 4.1. Grouped-Query Attention Scheduling: For every G stages, the schedule communicates keys, values, and queries in the first stage, then only queries in subsequent stages.This avoids redundant key/value communication for grouped queries.
  • 4.1. Grouped-Query Attention Scheduling: O((3 + G −1) · H/C · G · (C −1)) communication is always less than naive processing when G > 1.The reduction follows from reusing communicated key/value heads across each group of stages.

5. Experiments

The experiments compare UPipe with context-parallel and offloading-based baselines across models, hardware configurations, and sequence lengths. UPipe reduces memory usage while maintaining competitive throughput and extending supported context lengths.

  • Single-node training: At sequence lengths ≥2M, UPipe matches Ulysses throughput while using much lower GPU memory and incurs only slight short-sequence overhead from additional stage launches.The launch overhead is amortized at longer contexts.
  • Single-node training: 5M tokens is supported for Llama 3-8B on a single 8×H100 node, exceeding FPDT’s 4M-token limit by 25%.UPipe supports longer sequences with negligible performance differences compared with other methods.
  • Multi-node training: 8M tokens is supported for Llama 3-8B on 16 H100 GPUs, improving maximum context length over USP-Hybrid by 33% while retaining comparable throughput.Figure 5 compares peak GPU memory and normalized throughput between UPipe and USP-Hybrid from 512K to 6M tokens.
  • Multi-node training: For Qwen3-32B on 16×H100s, UPipe outperforms all other methods at sequence lengths ≥2M and supports 4M tokens, twice Ulysses’ 2M limit.UPipe also delivers 8.3% better performance than FPDT in this setup.
  • Ablation on the head chunk size: Increasing head-chunk size U raises memory usage but lowers runtime; setting U = C maximizes memory savings at the cost of slight kernel-launch overhead.The overhead is amortized at longer context lengths.

6. Reinvesting Memory Savings

The paper examines how UPipe’s freed GPU memory can be reinvested to improve runtime or increase batch size. Strided offloading and selective checkpointing improve performance while retaining lower memory usage than Ulysses.

  • Memory reinvestment: UPipe’s lower memory usage creates headroom for performance-oriented activation-management strategies.The section explores strided activation offloading and selective activation checkpointing.
  • Strided Activation Offloading: At 128K context, strided activation offloading significantly improves UPipe performance by skipping CPU offloading every seventh layer.The experiment uses a single 8×H100 node with a total of 3M tokens.
  • Selective Activation Checkpointing: Selective Activation Checkpointing further improves UPipe performance while keeping memory usage below DeepSpeed-Ulysses with or without SAC.SAC retains attention output activations to avoid costly backward-pass recomputation.
  • Batch size ablation: At 128K context, UPipe outperforms Ulysses at the largest tested batch size while using less memory across all batch sizes.Ulysses slows at batch size 10 because memory pressure triggers frequent CUDA allocation retries.

7. Conclusion

The paper concludes that UPipe reduces attention activation memory through head-level chunking, enabling longer context training without sacrificing throughput. It presents the method as a simple and composable building block for future long-context systems.

  • Contribution: UPipe reduces attention activation memory by chunking attention at the head level.The method executes attention in multiple stages that process head chunks.
  • Results: On Llama 3-8B, UPipe reaches 5M tokens on one 8×H100 node and 8M tokens on 16 H100 GPUs while keeping throughput comparable to common baselines.The single-node result is 25% beyond FPDT.
  • Results: For larger models, UPipe reduces intermediate attention-tensor memory by up to 87.5%.The conclusion links this reduction to avoiding allocation retries that can degrade training performance.
  • Implications: UPipe is characterized as simple, composable with existing methods, and capable of substantial memory reductions while preserving training speed.The authors identify it as a potential building block for future systems requiring larger sequence lengths.

Impact Statement

The work targets memory-efficient training of Transformer models on long-context inputs. It extends the maximum context length that fits on given hardware and is intended for broad long-context settings rather than a specific use case.

  • Impact: The work improves memory efficiency when training Transformer-based models on long-context inputs.Its stated scope is broad applicability across long-context Transformer settings.
  • Impact: The method extends the maximum context length that can fit on a given hardware configuration without targeting a particular use case.The statement frames the contribution as broadly applicable across long-context settings.

Supplementary Material

The supplementary material derives per-stage peak activation-memory accounting for Ulysses and UPipe, including forward and backward attention stages. UPipe reuses chunk slots and processes heads in stages, while backward execution requires a separate full-sized gradient accumulator.

  • Forward memory accounting: DeepSpeed-Ulysses forward peaks at γ + 2 units during sequential Q, K, V all-to-all communication.The peak comprises the saved input, projected Q/K/V, and one in-flight destination buffer.
  • Forward memory accounting: UPipe divides attention into ν = H/U stages, each processing U heads, and pre-allocates a full-sized final output buffer.During each stage, a freed output slot can be reused as the destination buffer for the head-distributed query.
  • Backward memory accounting: UPipe backward keeps X, final out, and δfinal out live while consuming one output slot per stage through chunked all-to-alls.The freed slot is reused for the corresponding intermediate buffer, but cannot store the full-sized δX accumulator.
  • Backward memory accounting: The input all-to-all peak accounts for X, δX, remaining output slots, and per-group gradient accumulators after the first stage begins.At ν = 1, the expression reduces to γ + 2, matching Ulysses with activation offloading.
  • Backward memory accounting: The FlashAttention backward peak is the maximum of stage-0 and stage-1 formulas, with stage 1 dominating for ν ≥ 3.At ν = 1, the value matches Ulysses with activation offloading; at ν = 2, the two expressions are equal.
  • Backward memory accounting: The output all-to-all is reported relative to stage 0, while the pre-backward-attention peak remains constant at 3 units because δX is not yet allocated.For ν ≥ 2, a higher per-stage peak occurs when δX is live, but it is not the global peak.

B. Memory usage comparison

UPipe provides the best memory efficiency among the compared context-parallel methods while matching Ulysses throughput, whereas FPDT has lower allocated memory but incurs CPU overhead and fails beyond 4M tokens.

  • Memory usage comparison: FPDT execution fails at sequence lengths greater than 4M tokens.Table 8 marks these cases as out of memory or failed execution.
  • Memory usage comparison: FPDT exhibits the best memory usage but performs poorer because of CPU overhead.Its lower allocated memory does not translate into the strongest overall execution behavior.
  • Runtime comparison: UPipe and DeepSpeed-Ulysses are compared by FlashAttention forward time, backward time, all-to-all communication, and total training-step runtime.UPipe is slower at shorter sequence lengths because of multiple kernel launches, while this disadvantage is amortized at longer lengths.

D. Training Loss Convergence

The convergence experiment compares UPipe and DeepSpeed-Ulysses training loss for Llama 3-8B on C4 over 1000 steps at 128K-token context length. UPipe uses the same attention computation but changes head-processing order, which may cause minor numerical differences.

  • Training Loss Convergence: Table 10 compares the training-loss trajectories of UPipe and Ulysses during the convergence experiment.The comparison is specifically designed to verify whether UPipe affects training-loss convergence.
  • Training Loss Convergence: UPipe performs the exact same attention computation as Ulysses while changing the order in which attention heads are processed.Because floating-point operations are non-associative, this scheduling difference could produce minor numerical differences.
  • Training Loss Convergence: The experiment trains Llama 3-8B on C4 for 1000 steps with sequence length 128K, batch size 1, and learning rate 10^-5.Training loss is reported every 100 steps.

E. Scope and Limitations

UPipe modifies only multi-head self-attention, so the paper presents it as applicable across Transformer-based models and modalities. Its improved memory efficiency can also support larger batch sizes and Selective Activation Checkpointing.

  • Scope: UPipe applies to LLMs, Vision Transformers, and Diffusion Transformers regardless of modality because it modifies only multi-head self-attention.The scope follows from restricting the architectural change to the self-attention component.
  • Limitations: UPipe’s better memory efficiency enables larger batch sizes and Selective Activation Checkpointing despite lower throughput at shorter contexts.The passage identifies shorter-context throughput as the principal trade-off in this scope discussion.
  • Limitations: SM-free all-to-all communication could further improve throughput by overlapping communication with attention computation.This is presented as a possible improvement rather than an evaluated result.
Loading 2602.21196v2…