Adjoint Sharding for Memory-Efficient Long-Context Training of State Space Models
A technique that makes it possible to train recurrent and state-space language models on million-token contexts without running out of GPU memory. It replaces standard backpropagation with the adjoint method, expressing each parameter gradient as a sum of independent vector-Jacobian products that can be sharded across the sequence and layer dimensions and computed with roughly constant memory. Students learn how a truncated variant trades a little accuracy for speed and how distributed and parallel variants spread the computation across GPUs.
2501.00692
Adjoint sharding is a memory-reduction technique for training recurrent and state-space language models on very long context inputs. Standard backpropagation stores all intermediate activations, so i…