Reproduce “Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers for Long Contexts” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Danil Sivtsov, Ivan Rodkin, Gleb Kuzmin, Yuri Kuratov, Ivan Oseledets
Diagonal Batching is a scheduling technique that reorders computations in Parallel Recurrent Memory Transformers (PRMTs) to unlock parallelism, addressing their sequential processing bottleneck. Applying this method, LLaMA-1B ARMT achieves a 3.3x speedup over standard LLaMA-1B for 131,072-token sequences and a 1.8x speedup over sequential ARMT, while preserving model accuracy.
Key insights
- The inherent dependency graph of Parallel Recurrent Memory Transformers (PRMTs) contains hidden parallelism that can be exposed by reordering computations along 'diagonals' (where segment index + layer index is constant).
- Executing computations along these diagonals reduces the total sequential steps from N_segments × N_layers to N_segments + N_layers - 1, enabling efficient GPU utilization.
- Grouping operations like linear layers and layer normalization across different layers and segments into larger batched operations significantly improves hardware efficiency and FLOPS/s.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.