# Adjoint sharding for very long context training of state space models

arXiv: https://arxiv.org/abs/2501.00692 · HTML: https://mlsanity.com/paper/2501.00692/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 40 of its 54 references that are in the map.
- First posted on arXiv 1 Jan 2025.

## Builds on (most cited first, 30 of 40)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Extending Context Window of Large Language Models via Positional Interpolation](https://mlsanity.com/paper/2306.15595/) (2023)
- [Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers](https://mlsanity.com/paper/2110.13985/) (2021)
- [YaRN: Efficient Context Window Extension of Large Language Models](https://mlsanity.com/paper/2309.00071/) (2023)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Hungry Hungry Hippos: Towards Language Modeling with State Space Models](https://mlsanity.com/paper/2212.14052/) (2022)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [Resurrecting Recurrent Neural Networks for Long Sequences](https://mlsanity.com/paper/2303.06349/) (2023)
- [Augmented Neural ODEs](https://mlsanity.com/paper/1904.01681/) (2019)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens](https://mlsanity.com/paper/2402.13753/) (2024)
- [Focal Self-attention for Local-Global Interactions in Vision Transformers](https://mlsanity.com/paper/2107.00641/) (2021)
- [Attention Mechanisms in Computer Vision: A Survey](https://mlsanity.com/paper/2111.07624/) (2021)
- [World Model on Million-Length Video And Language With Blockwise RingAttention](https://mlsanity.com/paper/2402.08268/) (2024)
- [A Survey on Distributed Machine Learning](https://mlsanity.com/paper/1912.09789/) (2019)
- [An Empirical Study of Mamba-based Language Models](https://mlsanity.com/paper/2406.07887/) (2024)
- [Focused Transformer: Contrastive Training for Context Scaling](https://mlsanity.com/paper/2307.03170/) (2023)
- [How to Train Your HiPPO: State Space Models with Generalized Orthogonal Basis Projections](https://mlsanity.com/paper/2206.12037/) (2022)

## Abstract (arXiv)

Despite very fast progress, efficiently training large language models (LLMs) in very long contexts remains challenging. Existing methods fall back to training LLMs with short contexts (a maximum of a few thousands tokens in training) and use inference time techniques when evaluating on long contexts (above 1M tokens context window at inference). As opposed to long-context-inference, training on very long context input prompts is quickly limited by GPU memory availability and by the prohibitively long training times it requires on state-of-the-art hardware. Meanwhile, many real-life applications require not only inference but also training/fine-tuning with long context on specific tasks. Such applications include, for example, augmenting the context with various sources of raw reference information for fact extraction, fact summarization, or fact reconciliation tasks. We propose adjoint sharding, a novel technique that comprises sharding gradient calculation during training to reduce memory requirements by orders of magnitude, making training on very long context computationally tractable. Adjoint sharding is based on the adjoint method and computes equivalent gradients to backpropagation. We also propose truncated adjoint sharding to speed up the algorithm while maintaining performance. We provide a distributed version, and a paralleled version of adjoint sharding to further speed up training. Empirical results show the proposed adjoint sharding algorithm reduces memory usage by up to 3X with a 1.27B parameter large language model on 1M context length training. This allows to increase the maximum context length during training or fine-tuning of a 1.27B parameter model from 35K tokens to above 100K tokens on a training infrastructure composed of five AWS P4 instances.