Adjoint sharding for very long context training of state space models

While our method generally applies to all recurrent models, we illustrate the idea using state-space models ($\mathrm{SSM}$s), which have shown performances at least on par with transformers at small to medium scale [14]. Given an input token sequence $\{\mathbf{x}_{t}\}_{t=1}^{T}$, the $\mathrm{SSM}$s first calculate the corresponding matrices $\mathbf{A}^{t}$, $\mathbf{B}^{t}$, and $\mathbf{C}^{t}$ to evolve the dynamics as follows:

The $\mathrm{SSM}$s evolve a latent dynamics $\mathbf{h}^{t}$, whose initial condition $\mathbf{h}^{0}$ is often assumed to be zero. With $\mathbf{h}^{0}$ and $\mathbf{A}^{t},\,\mathbf{B}^{t}$ defined, the dynamics evolves as:

The matrices $\mathbf{C}^{t}$ then maps the latent dynamics $\mathbf{h}^{t}$ back to token space as $\mathbf{y}^{t}=\mathbf{C}^{t}\mathbf{h}^{t}$, with $\mathbf{y}^{t}$ being the predicted token at $t$. For a sequence of $T$ tokens, we denote:

In the most general case, we have $\mathbf{H}\in\mathbb{R}^{T\times N},\mathbf{A}\in\mathbb{R}^{T\times N\times N},\mathbf{B}\in\mathbb{R}^{T\times N\times P},\mathbf{C}\in\mathbb{R}^{T\times P\times N},\mathbf{X}\in\mathbb{R}^{T\times P},\mathbf{Y}\in\mathbb{R}^{T\times P}$, where $N$ is the hidden state dimension, and $P$ is the input/output dimension. We evolve the dynamics for $t=1,\dots,T$, and assume that $\mathbf{h}^{0}$ is a fixed and predefined constant.

The input to an $\mathrm{SSM}$ is $\mathbf{X}$ and $\mathbf{h}^{0}$, and the output is $\mathbf{Y}$. We define $\mathrm{SSM}(\cdot)$ as performing the following five steps:

1. 1.

   $\begin{aligned} \{\mathbf{A}^{t}\}_{t=1}^{T}=\{\boldsymbol{\mathcal{A}}(\mathbf{x}^{t})\}_{t=1}^{T},\end{aligned}$

2. 2.

   $\begin{aligned} \{\mathbf{B}^{t}\}_{t=1}^{T}=\{\boldsymbol{\mathcal{B}}(\mathbf{x}^{t})\}_{t=1}^{T},\end{aligned}$

3. 3.

   $\begin{aligned} \{\mathbf{C}^{t}\}_{t=1}^{T}=\{\boldsymbol{\mathcal{C}}(\mathbf{x}^{t})\}_{t=1}^{T},\end{aligned}$

4. 4.

   $\begin{aligned} \{\mathbf{h}^{t}\}_{t=1}^{T}=\{\mathbf{A}^{t}\mathbf{h}^{t-1}+\mathbf{B}^{t}\mathbf{x}^{t}\}_{t=1}^{T};\end{aligned}$

5. 5.

   $\begin{aligned} \{\mathbf{y}^{t}\}_{t=1}^{T}=\{\mathbf{C}^{t}\mathbf{h}^{t}\}_{t=1}^{T}.\end{aligned}$

The input to the five steps is $\mathbf{X}$, and the output is $\mathbf{Y}$. We can then write $\mathrm{SSM}(\mathbf{X})=\mathbf{Y}$. SSMs decrease the quadratic computational complexity with sequence length on transformers to linear and decrease the large inference-time memory requirements from the key-value cache. SSM-based models at a small to medium scale have shown performances on par with or better than transformer-based models. For instance, [51, 1] shows that SSM-based mixture-of-experts (MOE) model outperforms baseline transformer-based MOE model on model sizes as big as 2400M parameters. [62] performed an extensive empirical study and found that while SSMs outperform transformers on various tasks, they underperform on tasks which require strong copying, in-context learning, or long-context reasoning abilities. [62] also experimented with a SSM-transformer hybrid model, which outperforms transformers and is up to eight times faster when generating tokens at inference time. [37] trained a 52B parameter model and further affirmed the hybrid models performances.

In practice, we have $K$ $\mathrm{SSM}$s stacked together, and we have a large language head (LLH) $\Omega\in\mathbb{R}^{\mathbb{T}\times P}$, where $\mathbb{T}$ is the number of all possible tokens. To predict a token, we have $\mathbf{o}^{t}=\Omega\hat{\mathbf{y}}_{K}^{t}$. Define $(\mathbf{y}_{K}^{1},\dots,\mathbf{y}_{K}^{T})=\mathbf{Y}_{K}$, a ResNet computes $\mathbf{Y}_{K}$ as follows:

where $\hat{\mathbf{Y}}_{k}=(\hat{\mathbf{y}}_{k}^{1},\dots,\hat{\mathbf{y}}_{k}^{T})=(\mathrm{Norm}(\mathbf{y}_{k}^{1}),\dots,\mathrm{Norm}(\mathbf{y}_{k}^{T}))$ and $\mathrm{SSM}_{k}(\hat{\mathbf{Y}}_{k-1})=\tilde{\mathbf{Y}}_{k}$. Therefore, for a latent state at time $t$ we have $\mathbf{y}_{K}^{t}=\mathbf{y}_{0}^{t}+\sum_{k=1}^{K}\tilde{\mathbf{y}}_{k}^{t}$.

ResNet has been the foundation of numerous modern networks, including the transformers, diffusion models, segmentation models, SSMs, and more [29, 26, 34, 48]. ResNet’s residual structure allows for a separation between gradients of each layer by applying differentiation on summations.

The adjoint method is concerned with optimizing $\mathbf{y}(\mathbf{h}(\boldsymbol{\theta}),\boldsymbol{\theta})$ with respect to $\boldsymbol{\theta}$, where $\mathbf{h}(\boldsymbol{\theta})\in\mathbb{R}^{P}$ is the solution to $\mathbf{f}(\mathbf{h}(\boldsymbol{\theta}),\boldsymbol{\theta})=0$ [8]. To employ gradient based algorithms like the stochastic gradient descent (SGD) or the Adam, we compute the derivative of $\mathbf{y}$ regarding $\boldsymbol{\theta}\in\mathbb{R}^{|\boldsymbol{\theta}|}$:

After computing adjoint states $\{\boldsymbol{\lambda}^{i}\}_{i=0}^{t}$, the computation of the elements of $\boldsymbol{\lambda}^{i}(\partial\mathbf{f}(i,\mathbf{h}^{i-1},\boldsymbol{\theta})/\partial\boldsymbol{\theta})$ are independent, allowing parallelism. This computation is a vector-Jacobian product ($\mathrm{vjp}$), with $\boldsymbol{\lambda}^{i}$ as the vector and $\partial\mathbf{f}(i,\mathbf{h}^{i-1},\boldsymbol{\theta})/\partial\boldsymbol{\theta}$ as the Jacobian. $\mathrm{vjp}$s can be evaluated with the reverse-mode automatic differentiation and initializing the reverse phase with $\boldsymbol{\lambda}^{i}$ [3]. As each $\mathrm{vjp}$ only requires saving their corresponding computation graph, and can be disposed after the computation, we can compute $\mathrm{vjp}$s in parallel on modern GPUs. We will discuss this in more details in subsection 4.5. Adjoint sharding aims to use the adjoint method to replace backpropagation, which solves:

The backpropagation requires a sequential accumulation of the gradients, computing from the outmost layer inwards, therefore needs to save the computation graph for computations at all time $t$’s and creates memory bottlenecks.

Large scale neural networks are usually trained with the autograd framework [4, 47]. However, this framework suffers from a high memory cost when used with networks of recurrent nature [4]. Although activation checkpointing has been developed, which discards part of the intermediate values and recomputes them later on the fly, the memory cost is still high [30]. We employ the adjoint method for recurrence relations to further reduce the memory cost, and more importantly, to break the temporal dependencies of activations and parallelize their computations.

The proof of proposition 2 is in section A.1. The gradient for parameters of $\boldsymbol{\mathcal{A}}$, and $\boldsymbol{\mathcal{B}}$ are each separated into $\{\mathrm{vjp}_{\boldsymbol{\mathcal{A}}^{i}}(\frac{\mathrm{d}l^{t}}{\mathrm{d}\mathbf{y}^{t}}\boldsymbol{\lambda}^{t,i}\otimes\mathbf{h}^{i-1})\}_{i=1}^{t}$, $\{\mathrm{vjp}_{\boldsymbol{\mathcal{B}}^{i}}(\frac{\mathrm{d}l^{t}}{\mathrm{d}\mathbf{y}^{t}}\boldsymbol{\lambda}^{t,i}\otimes\hat{\mathbf{x}}^{i}\}_{i=1}^{t}$, and the gradient for parameters of $\boldsymbol{\mathcal{C}}$ only depend on inputs at time $t$. After computing the adjoint states, these $\mathrm{vjp}$ computations are separate from each other on both the network and the temporal level.

We provide the proof to proposition 3 in section A.2. Define $\boldsymbol{\Lambda}_{k}^{t}=\{\boldsymbol{\lambda}^{t,\tau}_{k}\}_{\tau=1}^{t}$, proposition 3 shows that the gradients of each network’s parameters computed with each token only correlate through the adjoint states $\{\boldsymbol{\Lambda}_{k}^{t}\}_{k,t=1,1}^{K,T}$. The adjoint states can be easily computed after a forward pass. The adjoint states can also be computed on the fly in the gradient computation phase, as it only depends on $\mathbf{C}_{k}^{t}$ and $\mathbf{A}_{k}^{t}$ and has no dependencies on the network Jacobians regarding the network parameters. The adjoint sharding method breaks down the backpropagation computation both layer-wise and token-wise into foundational $\mathrm{vjp}$ computations that do not have any dependencies on each other.

We show a schematic of the computations to $\mathrm{d}l^{t}/\mathrm{d}\boldsymbol{\theta}_{\boldsymbol{\mathcal{A}}_{k}}$, $\mathrm{d}l^{t}/\mathrm{d}\boldsymbol{\theta}_{\boldsymbol{\mathcal{B}}_{k}}$, and $\mathrm{d}l^{t}/\mathrm{d}\boldsymbol{\theta}_{\boldsymbol{\mathcal{C}}_{k}}$ in Figure 5 and a schematic for computing the adjoint states in Figure 4.

One limitation of adjoint sharding is that the number of $\mathrm{vjp}$s performed increases polynomially regarding the number of tokens $T$. In particular, adjoint sharding computes the $\mathrm{vjp}$ for $\boldsymbol{\mathcal{A}_{k}}$ and $\boldsymbol{\mathcal{B}_{k}}$ $(1+T)T/2$ times, and for $\boldsymbol{\mathcal{C}_{k}}$ $T$ times. When training large networks with many layers and long context length $T$, applying adjoint sharding becomes computationally expensive. We propose truncated adjoint sharding, with which we argue that we can get similar results by computing a linearly growing number of $\mathrm{vjp}$s, and empirically showcase its performance.

Attention mechanisms have suffered from the $\mathcal{O}(T^{2})$ complexities arising from the self-attention structure [60]. To enable training with longer context lengths, global-local attention has been proposed, where we divide the contexts into sections, and compute the attention between sections rather than tokens [67]. [57] proposed truncated backpropagation through time (T-BPTT) to avoid gradient explosion/vanishing when training with long contexts by only counting a fixed number of state transitions. Here, inspired by global-local attention and T-BPTT, instead of computing the full gradient given in Equation 11, we propose to train the $\mathrm{SSM}$s to depend on up to $\bar{T}$ states:

As shown in Equation 7 above, we perform the same computations for $t=1,\dots,\bar{T}$ as before, and only perform the $\mathrm{vjp}$s back to the last $\bar{T}$ states for $t>\bar{T}$. With truncated adjoint sharding, we perform $\bar{T}T+\bar{T}(\bar{T}-1)/2$ $\mathrm{vjp}$s, which grows linearly. We show the number of $\mathrm{vjp}$s performed with and without truncated adjoint sharding in Figure 6. When $\bar{T}=2000$, truncated adjoint sharding reduces $64\%$ of the $\mathrm{vjp}$s when training with a context length of $10\mathrm{K}$.

The essence of the truncated adjoint sharding method is that we only explicitly count gradients related to the last $\bar{T}$ states. As each state depends on its prior state, states still implicitly depend on all their prior states. We leave investigation of $\bar{T}$’s impact on performances for future works.

We now discuss how to distribute the storage and compute of the adjoint sharding method, assuming that we have $\Upsilon$ GPUs. Given the networks $\{\mathcal{A}_{k},\mathcal{B}_{k},\mathcal{C}_{k}\}_{k=1}^{K}$, initial tokens $\{\hat{\mathbf{y}}_{0}^{t}\}_{t=1}^{T}=\{\mathrm{Norm}(\mathbf{x}^{t})\}_{t=1}^{T}$, and initial conditions $\{\mathbf{h}^{0}_{k}\}_{k=1}^{K}$ (usually set to $\mathbf{0}$), we can call algorithm 1 to get all necessary vectors for computing the gradient with adjoint sharding.

As shown in algorithm 3, to compute the $\mathrm{vjp}$s’ for token index $t$ and ResNet index $k$, we only need $t,k,\mathrm{d}l(\mathbf{o}^{t})/\mathrm{d}\mathbf{y}_{K}^{t},\{\mathbf{h}_{k}^{i}\}_{i=0}^{t},\mathbf{C}_{k}^{t},\{\hat{\mathbf{y}}_{k-1}^{i}\}_{i=1}^{t},\{\mathbf{A}_{k}^{i}\}_{i=2}^{t}$. To compute all the gradients for layer $k$, we only need $\mathbf{A}$, $\mathbf{h}$, and $\mathbf{C}$ from the $k$-th layer, and $\hat{\mathbf{y}}$ from the $k-1$-th layer. Therefore, we can divide the $K$ layers into $\Upsilon$ pieces, as shown in the appendix A.4.

As the computations are fully independent and we compute the gradients using only data on local devices, we additionally distribute the model and the gradients, as shown in Table 6, where $\boldsymbol{\theta}_{k}$ represents the parameters of $\mathcal{A}_{k}$, $\mathcal{B}_{k}$, and $\mathcal{C}_{k}$, and $\mathrm{Gradient}_{k}$ represents the optimizer states for $\boldsymbol{\theta}_{k}$.

The complete training streamline is then as shown in algorithm 4. We fully distribute the activations, computations, gradients, and optimization states across $\Upsilon$ devices. While the forward evaluation pass results across different devices, as shown in algorithm 1, the computation of gradients is parallel across the $\Upsilon$ devices. This will speed up the training as the gradient computation takes most of the computation budget. We will also get a memory per GPU close to $\mathrm{Mem}/\Upsilon$, with $\mathrm{Mem}$ being the memory cost if we only have a single GPU. If we have $\Upsilon>K$ devices, we can further speed up the forward evaluation by first evaluating $\mathcal{A}$, $\mathcal{B}$, $\mathcal{C}$ in parallel, and then sequentially add them together on the distributed devices.

Adjoint sharding converts the sequential process of backpropagation gradient computation into individual independent $\mathrm{vjp}$s, allowing for parallel computation. We analyze the time and memory cost of $\mathrm{vjp}_{\boldsymbol{\mathcal{A}}^{i}_{k}}((\mathrm{d}l^{t}/\mathrm{d}\mathbf{y}_{K}^{t})\boldsymbol{\lambda}^{t,i}_{k}\otimes\mathbf{h}^{i-1}_{k})$, $\mathrm{vjp}_{\boldsymbol{\mathcal{B}}^{i}_{k}}((\mathrm{d}l^{t}/\mathrm{d}\mathbf{y}_{K}^{t})\boldsymbol{\lambda}^{t,i}_{k}\otimes\hat{\mathbf{y}}^{i}_{k-1})$, and $\mathrm{vjp}_{\boldsymbol{\mathcal{C}}^{t}_{k}}((\mathrm{d}l^{t}/\mathrm{d}\mathbf{y}_{K}^{t})\otimes\mathbf{h}^{t}_{k})$.

$\mathrm{vjp}$ has a similar time complexity as a forward pass, and a memory complexity of $\mathrm{bs}(|\boldsymbol{\theta}|+\mathbb{O})+|\boldsymbol{\theta}|$, where $\mathrm{bs}$ is the batch size, $\mathbb{O}$ is the number of elements in the network output, and $|\boldsymbol{\theta}|$ is the number of parameters [42]. We provide the memory and FLOPs required to compute the $\mathrm{vjp}$s in Table 1 [43].

We analyze training with a dataset containing contexts of lengths $T$, with $\Upsilon$ NVIDIA H100 GPUs, and performing computations in FP16. We use a selective diagonal SSM with $K$ layers, and each $\boldsymbol{\mathcal{A}}_{k}$, $\boldsymbol{\mathcal{B}}_{k}$, and $\boldsymbol{\mathcal{C}}_{k}$ network is a single-layer multi-layer perceptron (MLP).

For each data point $\{\mathbf{x}^{t}\}_{t=1}^{T}$, we store $\{\mathbf{A}^{t}_{k},\mathbf{C}^{t}_{k},\mathbf{h}^{t}_{k},\mathbf{y}^{t}_{k}\}_{(t,k)=(1,1)}^{(T,K)}$ and $\{\mathrm{d}l(\mathbf{o}^{t})/\mathrm{d}\mathbf{y}_{K}^{t}\}_{t=1}^{T}$, which is $TK(2N+P)+TP$ FP16 numbers. We also save $\boldsymbol{\theta}_{\boldsymbol{\mathcal{A}}}$, $\boldsymbol{\theta}_{\boldsymbol{\mathcal{B}}}$, and $\boldsymbol{\theta}_{\boldsymbol{\mathcal{C}}}$, each taking $PN+N$ FP16 numbers. We need to store $T(2NK+PK+P)+3N(P+1)$ FP16 numbers before computing the $\mathrm{vjp}$.

As computing all adjoint state sequences takes up to $N(2P+1)(1+T)T/2$ FLOPs, it takes $NP(1+T)/T$ FLOPs on average for each adjoint state. For $T$ large enough, $(1+T)/T\approx 1$, and we approximate the average FLOPs for each adjoint state with $NP$. Each $\mathrm{vjp}$ then takes $\mathrm{bs}(7NP+3N)$ FLOPs of computation.

When computing with a selective diagonal SSM with $P=128$, $N=225$, and $\mathrm{bs}=8$, while storing and performing computations in FP16, computing $\mathrm{vjp}_{\boldsymbol{\mathcal{A}}}$, $\mathrm{vjp}_{\boldsymbol{\mathcal{B}}}$, and $\mathrm{vjp}_{\boldsymbol{\mathcal{C}}}$ each takes around $0.6\mathrm{MB}$ memory and $1798144$ FLOPs. The capacity of a modern GPU is mostly characterized by FLOPs/sec, which measures the computation speed; GPU memory bandwidth, which is the rate at which a GPU can move data between its memory and processing cores; GPU Memory, which is the amount of data a GPU can hold; and number of Multi-Instance GPU (MIG) instances, which is the number of fully isolated GPU instances with its own high-bandwidth memory, cache, and compute cores a GPU can host.

An NVIDIA H100 Tensor Core GPU has a GPU memory bandwidth $3.35\mathrm{TB/s}$ and performs $1,979$ tera FP16 FLOPS per second. Therefore, the memory bandwidth allows computing $(3.35\mathrm{TB/s})/0.6\mathrm{MB}=5.58\times 10\mathrm{E}6$ batches of $\mathrm{vjp}$s per second, and the computing speed allows computing $(1979\mathrm{tera/s})/1798144=3.76\times 1.1\mathrm{E}9$ batches of $\mathrm{vjp}$s per second. At the same time, since the H100 GPU has $80\mathrm{GB}$ memory, it can hold up to $80\mathrm{GB}/(0.6\mathrm{MB}/\mathrm{vjp})=133$ batches of $\mathrm{vjp}$s at the same time if we do not consider any memory overhead. As each H100 GPU can hold up to $7$ instances in parallel, we perform the adjoint sharding algorithm with $7\Upsilon$ instances, offering as much as a 56x speedup on one AWS P4 instance (8 H100 GPUs). Such speedup cannot be achieved for backpropagation because of its sequential nature.