StateFlow: Sequence Pipeline Parallelism for Long-Context Modeling with Linear Recurrence

2026-08-07

AI Summary

This work addresses the high memory consumption and low training throughput of recurrent and hybrid models in long-context training, which stem from their limited parallelizability. The authors propose StateFlow, a novel sequence pipeline parallelism method for linear recurrent models that partitions sequences into chunks and propagates boundary states and gradients across them, thereby substantially shortening activation lifetimes. A key innovation is a performance-aware non-uniform chunking strategy that dynamically balances computation between recurrence and Softmax attention, while overlapping low-parallelism state transitions with surrounding operations. Evaluated on a 32B-parameter model with 256K context length, StateFlow achieves a 2.22× higher training throughput and 2.45× lower memory usage compared to conventional pipeline parallelism, enabling large-scale long-context training configurations previously deemed infeasible.

Abstract

Long-context training is increasingly important for large language models, and linear attention and state space models have become popular for improving long-context efficiency. However, efficiently parallelizing long-sequence training for recurrent and hybrid models remains challenging. We present StateFlow, a sequence pipeline parallelism system for models with linear recurrence. StateFlow partitions each sequence into chunks and schedules their execution while propagating boundary states and gradients across chunks, thereby reducing activation lifetimes and improving training throughput. StateFlow further uses profile-guided nonuniform chunking to balance recurrence and softmax attention computation in hybrid models, and overlaps state transitions that expose limited parallelism with surrounding computation. Applying StateFlow to models with up to 32B parameters and 256K context length, we achieve up to \(2.22\times\) throughput improvements and \(2.45\times\) memory reduction compared to conventional pipeline parallelism, enabling otherwise infeasible configurations.

Problem

• long-context modeling
• sequence pipeline parallelism
• linear recurrence
• training efficiency
• activation memory

Innovation

• sequence pipeline parallelism
• linear recurrence
• nonuniform chunking
• long-context modeling
• state propagation