ParaRNN: Unlocking Parallel Training of Nonlinear RNNs for Large Language Models

TL;DR

ParaRNN tackles the scalability bottleneck of nonlinear RNNs by recasting the forward pass on a sequence of length $L$ as a nonlinear system solved in parallel via Newton's method and a specialized parallel reduction for the resulting block bi-diagonal system. It adapts GRU and LSTM cells with diagonal Jacobians to enable efficient parallelization, and provides a PyTorch+CUDA software stack with progressively fused solvers to maximize performance. Empirically, ParaRNN achieves runtime parity or superiority to Mamba and competitive perplexities against Transformer baselines at scales up to $7\mathrm{B}$ parameters, demonstrating that nonlinear RNNs can be trained effectively at large scale. This work offers a practical pathway to diversify sequence-modeling options beyond attention-based architectures and provides an open-source platform to foster further nonlinear RNN innovations for language modeling.

Abstract

Recurrent Neural Networks (RNNs) laid the foundation for sequence modeling, but their intrinsic sequential nature restricts parallel computation, creating a fundamental barrier to scaling. This has led to the dominance of parallelizable architectures like Transformers and, more recently, State Space Models (SSMs). While SSMs achieve efficient parallelization through structured linear recurrences, this linearity constraint limits their expressive power and precludes modeling complex, nonlinear sequence-wise dependencies. To address this, we present ParaRNN, a framework that breaks the sequence-parallelization barrier for nonlinear RNNs. Building on prior work, we cast the sequence of nonlinear recurrence relationships as a single system of equations, which we solve in parallel using Newton's iterations combined with custom parallel reductions. Our implementation achieves speedups of up to 665x over naive sequential application, allowing training nonlinear RNNs at unprecedented scales. To showcase this, we apply ParaRNN to adaptations of LSTM and GRU architectures, successfully training models of 7B parameters that attain perplexity comparable to similarly-sized Transformers and Mamba2 architectures. To accelerate research in efficient sequence modeling, we release the ParaRNN codebase as an open-source framework for automatic training-parallelization of nonlinear RNNs, enabling researchers and practitioners to explore new nonlinear RNN models at scale.

Figures (12)

• Figure 1: Our ParaRNN framework makes it possible to apply classical RNNs in parallel, dramatically speeding up their training, and allowing them to be used competitively for language modeling.
• Figure 2: Timing results for the parallel reduction operations \ref{['alg::pcr']} (left), and the full parallel RNN application \ref{['alg::newton']} (right), applied to sequences of varying length. Colors refer to different variants of the solvers: for diagonal Jacobians (blue, used by ParaGRU) and 2x2 block-diagonal ones (green, used by ParaLSTM). Heavier lines refer to progressively more efficient implementations: in PyTorch, in CUDA, and (for the RNN forward pass) with the fully-fused CUDA kernel. For reference, equivalent timing results for Mamba are also included (yellow). Actual speedup measurements reported in the main text refer to $L=2^9$ (highlighted).
• Figure 3: Tokens throughput at generation for ParaGRU (blue), ParaLSTM (green), Mamba (yellow), and Transformers (purple).
• Figure 4: Newton's method convergence behavior for ParaGRU (top) and ParaLSTM (bottom) cells with different input sequence lengths $L$, for freshly initialized RNN cells (left), and for the RNN cells in the last layer of our trained 400M models (right). Inputs are randomly selected from the evaluation datasets, and the residuals reported are the $\max$ on batches of size 8.
• Figure 5: Timings of different application modalities for the ParaGRU (left) and ParaLSTM (right) RNN cells, applied to sequences of varying length. Heavier lines refer to progressively more efficient applications: sequential, using parallel reduction in PyTorch (\ref{['alg::newton']}), using parallel reduction in CUDA and system assembly in PyTorch, and using a fully-fused CUDA kernel for the whole Newton routine. Equivalent timing results for Mamba are also included for reference (yellow). Same setup as for \ref{['fig::kernel_and_model_own_vs_mamba']} (right).