Activation sharding for scalable training of large models
Authors: Xingzi Xu
Amir Tavanaei
Kavosh Asadi
Karim Bouyarmane
Venue: Accepted by TMLR ()

Abstract: Despite fast progress, efficiently training large language models (LLMs) in extremely long contexts remains challenging. 
Existing methods fall back to training LLMs with short contexts (up to a few thousand tokens) and use inference time techniques when evaluating on very long contexts (above 1M tokens). 
Training on very long contexts is limited by GPU memory availability and the prohibitively long training times it requires on state-of-the-art hardware. 
Meanwhile, many real-life applications require training/fine-tuning with long context on specific tasks. Such applications include, for example, augmenting the context with various sources of raw reference information for extraction, summarization, or fact reconciliation tasks. 
We propose adjoint sharding, a novel technique that comprises sharding gradient calculation during training to reduce memory requirements by orders of magnitude, making training on very long contexts computationally tractable. At the core of our adjoint sharding algorithm lies the adjoint method, which efficiently computes gradients that are provably equivalent to the gradients computed using standard backpropagation. 
We also propose truncated adjoint sharding to accelerate the algorithm while maintaining performance. 
We provide a distributed and a parallel-computing version of adjoint sharding to speed up training and to show that adjoint sharding is compatible with these standard memory-reduction techniques. 
Empirical results show the proposed adjoint sharding algorithm reduces memory usage by up to 3$\times$ on a large language model with 1.27B parameters on 1M context length training. This reduction in memory usage allows increasing the maximum context length of training a 1.27B parameter model from 35K tokens to above 100K tokens on a training infrastructure composed of five AWS P4 instances.

Published in Transactions on Machine Learning Research (10/2025)
Adjoint sharding for very long context training of state space models
Xingzi Xu, Amir Tavanaei, Kavosh Asadi, Karim Bouyarmane
Amazon
ataraaa atavanae, kavasadi, bou.kariiaamazzn.coooo
Reviewed on OpenReview:
Abstract
Despite fast progress, efficiently training large language models (LLMs) in extremely long contexts remains challenging. Existing methods fall back to training LLMs with short contexts (up to a few thousand tokens) and use inference time techniques when evaluating on very long contexts (above 1M tokens). Training on very long contexts is limited by GPU memory availability and the prohibitively long training times it requires on state-of-the-ar hardware. Mcanwhile, many real-life applications require training/ ffinettuning with long context specific tasks. Such applications include, for example, augmenting the context with various sources of raw reference information for extraction, summarization, or fact reconciliation tasks. We propose adjoint sharding, a novel technique that comprises sharding gradient calculation during training to reduce memory requirements by orders of magnitude, making training on very long contexts computationally tractable. At the core of our adjoint sharding algorithm lies the adjoint method, which efficiently computes gradients that are provably equivalent to the gradients computed using standard backpropagation. We also propose truncated adjoint sharding to accelerate the algorithm while maintaining performance. We provide distributed and paralle--ompuutng version of adjoint sharding to speed up training and to show that adjoint sharding is compatible with these standard memory-reduction techniques. Empirical results show the proposed adjoint sharding algorithm reduces memory usage by up to 3x on a large language model with 1.27B parameters on 1M context length training. This reduction in memory usage allows increasing the maximum context length of training 1.27B parameter model from 35K tokens to above 100K tokens training infrastructure composed of five AWS P4 instances.
1 Introduction
Foundation models are a new paradigm in artificial intelligence research focused on building large, general- purpose models that adapt to different tasks OpenAl et al. (2024); Meta et al. (2024); Cai et al. (2024); Pióro et al. (2024). Extensive training on large datasets equips foundation models with broad capabilities. We then fine-tune the foundation models on smaller datasets for specific applications. Foundation models commonly employ the transformer architecture Vaswani et al. (2023). Despite the immense success, training hansformer-based models requires memory growing quadratically with the context length L, limiting their applications on long context tasks Li et al. 2024). Rescarchers developed varions techniques to conquer this problem, ranging from inference time context window expansion (Ding et al.) 2024bla). 10-aware algorithms
Dao et al. 2022 Dao 2023 Shah et al. 2024), and various linearly scaling language model architectures Gu & Dao 2024b Dao & Gu 2024b Pong et al.) 2023a Beltagy et al.) 2020). On another note, distributed
learning enables training large models with many GPUs, and efficient training methods like activation
checkpointing, model/gradient sharding, and mixed-precision computing have further reduced the memory

---

Published in Transactions on Machine Learning Research (10/2025)
80000
GPU memory (/GB)
60000
60
adj-shard/32M
adj-shard/63M
40
adj-shard/127M
40000
20
adj-shard/225M
adj-shard/1.227
0-
10a
10a
backprep/32M
20000
backprop/63M
backprop/127W
backprop/225M
0
backprop/1.278
101
102
10a
104
10
10
Training context length
Figure 1: Compared to backpropagation (red lines), adjoint sharding (blue lines) significantly reduces memory requirements at training. Showing memory cost to train 32M, 63M, 127M, 225M, and 1.27B parameter State Space Model (SSM) with batch size 2 and Adam optimizer one GPU.
requirement of training large model Verbracken et al. (2020): Zhao (2023): Rajbhandari al. (020b));
Mieikevicius et al.] (2018); Herrmann et al. (2019). However, current methodologies entirely based backpropagation and compute the gradient as whole, inevitably requiring memory growing rapidly with model size and context length (Damadi et al. 2023). Current sharding methods ignore the activations and only consider the model weights and optimizer states, constituting only a minor amount of the total memory cost (Sohoni et al. 2022). Activation checkpointing is among the limited techniques that consider activation values. Activation checkpointing offloads necessary intermediate states to the CPU and recomputes them on the fly, trading compute time for memory reduction (Sohoni al.) 2022 Rajbhandari substantial time required for offloading to the CPU hinders the effectiveness of activation checkpointing. We propose adjoint sharding to disassemble the gradient computation of residual and/or recurrent- casre models to achieve orders of magnilude lower memory usage during training.
al. 2020a).. The
Adjoint method The adjoint sharding method roots in the adjoint method for recurrent models CCoo ct al.) 2002; Johnson 2007). Given an optimization problem of a parametric recurrent forward process, the adjoint method is concerned with the computation of the gradients regarding the process's parameters. Backpropagation saves intermediate states calculate gradients, whereas the adjoint method relies on a backward adjoint process to compute gradients. The adjoint method is constant- optimization
technique for dynamical systems Chen al. 2019 Xu et al. (2022). In this paper, are only concerned with the adjoint method for recurrent relations.
Vector-Jacobian product Adjoint sharding disassembles the gradient computation of a large language model (LLM) into independent vector-Jacobian product (VJP) computations. By multiplying the Jacobian on the left with a vector, becomes unnecessary compute the expensive Jacobian. Modern VJPss are as fast as forward function call of the model, and can be thousands of times faster than Jacobian computations Balestriero & Baraniuk (2021). We speed up adjoint sharding by employing the VJPs.
Truncated adjoint sharding Sharding the gradient computation allows us to prioritize the essential gradients and disregard the rest, resulting in faster computation. We term this novel method truncated adjoint sharding, and empirically showcase its performance.
Distributed and parallel computation addition, we have developed distributed multi-GPU variant of adjoint sharding to improve the scalability of LLM training further. We also analyze the memory cost of parallel computation of adjoint sharding, opening up directions for massive speedups.
2

---

Published in Transactions on Machine Learning Research (10/2025)
State-space models and residual networks Residual networks (ResNets) are commonly applied neural network structure. We illustrate adjoint sharding assuming a ResNet structure (He et al. 2015). State-space models (Mamba) have achieved performances on par with attention models while possessing a linear scaling regarding the context length L, a polynomial speedup compared to the L scaling of transformers Vaswani al. (2023) : Gu & Dao (2024a).
2 Related works
Linear LLMs (De et al. 2024 Beck et al.) 2024 Peng ct al. 2023a proposed LLM architectures with linear inference time complexity. We form each linear LLM by stacking K' residual layers together, where each layer has a recurrent relation. However, their temporal relation- ships are nonlinear, which limits the application of adjoint sharding to disassemble the gradients into independent. voctor-Jacobian products.
t
k
Backpropagation through time Applying the adjoint method for recurrent models leads to backpropagation through time (BPTT) [Werbos] 1990). BPTT is a training algorithm developed for recurrent neural networks (RNNs). RNN models suffer from the exploding and vanishing gradi- ent because of the 11 0f(x², h wn))/h'-1 term Pascanu et al. 20133. SSMs provide remedies with careful parameterization of the recurrent dynamics inspired by
classical SSM theory (Fu et al.]
2023 Gu et al. 2021
Figure 2: Adjoint sharding disassembles large models' gradient computations along the sequence dimension and the layer dimension k. When evaluating the gradient at time 4, we perform vector-Jacobian products along the adjoint dimen- sion i for every layer index k.
2022
Gupta et al. 2023 Orvicto et al.
2023 Kaul 2020
Linear temporal relations allow ellicient eval ualions of the model, while preserving universal approximation capabil- ities (Wang& & Xue 2023). By a similar token, truncated adjoint sharding can be seen as a more general version of the truncated backpropagation through time (Jaeger)
2005
Tallec & Ollivier 2017).
Neural ordinary differential equations The adjoint method also applies to the optimization of continuous systems, especially the ordinary differential equations (ODEs) (Chen et al. 2019 Dupont et al. 2019). Optimizing neural ODEs with autograd requires backpropagating through numerical solvers along every step, using an unrealistic amount of memory. The adjoint method does not backpropagate through the operations of the solver and uses a constant amount of memory. However, applying the adjoint method for continuous systems requires solving costly ODE initial value problem with dimensionality equal to the number of parameters.
Low memory training methods Researchers proposed various low memory training techniques to train big models in very long contexts. ZERO provides data- and model-parallel training while retaining low communication volume, while eliminating memory redundancies (Rajbhandari et al., 2020b). PyTorch FSDP streamlines model, gradient, and data parallelization (Zhao et al. 2023). Activation checkpointing discards intermediate values during the forward step, and recomputes on the fly during the training phase (Sohoni et al. 2 CPU offloading scales large model training by offloading data and computations to the CPU, trading computing time for memory reduction (Ren et al. 2021). Ring attention leverages the blockwise computation of self-attention and feedforward to distribute long sequences across multiple devices while fully overlapping the communication ion of key-value blocks with blockwise attention computations, enabling very-long context training of attention-basee methods (Liu et al. 2023 2024). Mini-Sequence Transformers segment sequences but still require full gradient computation within each segment (Luo et al. 2024). Cut Your Losses optimizes vocabulary-level computations, which is orthogonal to our sequence-level memory reduction approach (Wijmans et al. 2025). StreamBP targets long sequences through different. memory-time trade-off
3

---

Published in Transactions on Machine Learning Research (10/2025)
strategy (Luo et al. 2025). Unlike these approaches that focus segmentation, vocabulary optimization, memory-time trade-offs, adjoint sharding fundamentally changes the gradient computation mechanism itself by decomposing gradients into independent vector-Jacobian products that can be computed in parallel for recurrent architectures. The proposed adjoint sharding distributes state-space model computations across multiple devices and multiple multi-GPU instances (MIG) to enable very-long context. training of stale-space models.
Context length extension methods The Existing context length extension method separates into two
classes. The first type is fine-tuning free methods, including Positional Interpolation (PI) [Chen et al. 2023
the NTKAware Scale ROPE (NTK) (users) 2023). ; and StreamingI.LM (Xiao et al.) 2024 The second type is
fine-tuning methods, including LongCha Chat (Li et al. 2023), LongAlpaca (Chen et al. 2024). YaRN (Peng
et al. 2023b), and LongLlama (Chen et al. 2024). Additional methods like activation beacon tune a network separate from the LLM (Zhang et al. 2024). As shown Figure 3 fine-tuning methods achieve better performances than finettnnin-free free methods at lengths where we perform fine-tuning. However, fino-tuning methods sulfer from high computational cost and require potentially intraciable amount of GPU memory during fine-tuning.
15.0
PG19
PI
10.0
NTK Streaming
YaRN-128K
CodeParrot Proof- Pole-
7.50
LongChat-32K
5.00
LongAlpaca-16n
LongLlama
2.50
4.00
2.00
10
10
Inference context length
Figure 3: Lines in red and in blue are fine-tuning free and fine-tuning methods. Finc-tuning methods achieve better performances than the imo-tuning-free method but often suffer from out-of-memort issucs (Chen et al.
2023
users 2023 Xiao et al.] 2024;
Li et. al. 2023 Chen et. al. 2024 Peng et. al. 2023b Zhang el. al. 2024
Tworkowski et 2023). Lower values are better across all three tasks.
3
Background
We first give a concise introduction to the state-space models, the residual networks, and the adjoint method.
3.1
State-space models
While our method generally applies all recurrent models, we illustrate the idea using state-space models (SSMs), which have shown performances at least. par with transformers at small to medium scale (Dao &
Gu 2024a). Given input token sequence {x, 1T the SSMs first caleulate the corresponding matrices A
B and C' to evolve the dynamics as follows:
A'= A(x': B' B(x²); C C(x').
The SSMs evolve latent dynamics h', whose initial condition hD is often assumed to be zero. With h" and A', B' defined, the dynamics evolves as:
h A'h-11 + -'x'
4

---

Published in Transactions on Machine Learning Research (10/2025)
The matrices C' then maps the latent dynamics h' back to token space as y' C'h', with y' being the predicted token at L. For sequence of T tokens, we denote:
A (A A A²), B (B) B B²),
C (C C² C²)) =hh², h, h²)
X
,....,"'),, x"), Y = (y y
In the most general case; have H E TERX A E RTRNNN B RTTXXP, C RTRXXNN X X Y RTXP, where N is the hidden state dimension, and P is the input/ /oupput dimension. We evolve the dynamics for t = I,...,T, T, and assume that h is fixed and predefined constant. The input to an SSM is X and h" and the output is Y. We define SSM(.) as performing the following five steps:
(A (A(x')E=
2. (B') (B('
(C")= (C(x
(b')== (A'h--1 +B'x
(y =
The input to the five steps is X, and the output is Y. We can then write SSM(X) Y. SSMs reduce the quadratic computational complexity with sequence length on transformers to linear, and the inference-time memory requirements from the key-value cache. SSM-based models at small to medium scales have shown performances on par with or better than transformer-based models. For instance, (Pióro al. 2024 Auth al. 2024 shows that the SSM-based mixture-of-experts (MOE) model outperforms the baseline transformer- based MOE model on model sizes as big as 2400M parameters. WWlelffe al.) 2024) performed an extensive empirical study and found that while SSMs outperform transformers on various tasks, they underperform tasks that require stellar copying, in-context learning, long-context reasoning abilities. (Waleffe et al. 2024] also experimented with an SSM-transformer hybrid model, which outperforms transformers and is up to eight times faster when generating tokens at inference time. [Lieber et al., 2024 trained a 52B parameter model and further affirmed the hybrid model's performance.
3.2 Residual Networks
In practice, we have K SSMs stacked together, and we have a large language head (LLH) 82 E Rxx where T is the number all possible tokens. To predict a token, we have o' = S2 Define (yk...., yx) = Yx, a ResNet computes YK as follows:
(yk .kk = -YK-I+ (YSM((KK)))
-Yo + SSSSSM +... SSM (YK-1)
K Yo SSM&(Yx-1) Yo + Σ Y
K
A=1
where
Y, (S),,.., 97) = (Norm(y))..... Norm (yZ)),
and SSM (--1) Yg-- Therefore, for latent state at time we have y'x = yb
Σk y/.
k=
ResNet has been the foundation of numerous modern networks, including the transformers, diffusion models, segmentation models, SSMs, and more (He al.] 2016 Guo et al., 2022 Kirillov et al. 2023 Peebles & Xie 2023). ResNet's residual structure allows for separation between gradients of each layer by applying diflerentiation summations.

---

Published in Transactions on Machine Learning Research (10/2025)
3.3 Adjoint method
The adjoint method is concerned with optimizing y(h(0), 0) with respect 0, where h(0) E EP is the solution to f(h(0), 0) 0 (Cao et al. 2002). To employ gradient-based algorithms like the stochastic gradient descent (SGD) or the Adam, we compute the derivative of y regarding E R
dy
8y
dy 0h
+
(1)
d0
00
0h 00
with d being the total derivative, and 0 being the partial derivative. The adjoint method converts computing dy/d0 to solving an adjoint equation. In our case, we need the adjoint method for recurrence relations, where is given by = y" = y(h'(0), 0), and h is given by
h" = =b(0), he = [(1,, h)-
(2)
We have
df hh-l_h0
@t h
@t h 1-1
0h'
(3)
d0
00
Oht-1
00
Proposition 3.1. (Cao et al. respect to 0 is given as:
2002
When the states h are defined as Equation the gradient of y with
dy
Oy* + + + + A 8f
00
= 0y"/Oh',
(4)
A Xi 10"6i,h'""0000h""""
Equicolently, we have
[+1
X = (0y' /Oh')
II 10,
Johnson 2007).
After computing adjoint states (X') the computation of the elements of (if h 1ati:hh-1//000 00//00))))))) indepen- dent, allowing parallelism. This computation is vector-Jacobian product (vip), with X as the vector and ii ( as the Jacobian. vjps can be evaluated with the reverse-mode automatic differentiation and initializing the reverse phase with A Baydin et al. (2018a). As each vjp only requires saving its corresponding computation graph, and can be disposed of after the computation, we can compute vjps in parallel on modern GPUs. We will discuss this further in |subsection 4.5 Adjoint sharding aimn to use the adjoint method replace backpropagation, which solves:
dy 0y Øy*
[t h 0)
@t h 1--h1000 .0)
Tf 1,h 1 -,0)
Of 1,h²-2,0)
00 + 0h
de
ae
Oh
00
0h--2
@f(( 2, 3-3
}}}
00
The backpropagation requires sequential accumulation of the gradients, computing from the outermost layer inwards, therefore needs to save the computation graph for computations at all times and creates memory bottlenecks.
4 Adjoint sharding
We now introduce the adjoint sharding technique. We first illustrate the method assuming only one layer of SSM, and generalize to K layers.

---

Published in Transactions on Machine Learning Research (10/2025)
4.1 Adjoint sharding for one SSM
Large-scale neural networks usually train with the autograd framework (Baydin et al., 2018b Paszke al, 2019). However, this framework suffers from high memory cost when used with networks of recurrent nature Baydin et al. 2018b). Although activation checkpointing proves strong tool, which discards part of the intermediate values and recomputes them later on the fly, the memory cost. is still high (Herrmann et al., 2019). We employ the adjoint method for recurrence relations to reduce the memory cost further, and more importantly, to break the temporal dependencies of aclivations and parallelize their computations.
Define = (f.. Bz fe) as A's, B's, and C's parameters, for loss 1' = I(y"), in the context of a single-layer SSM, we prove
Proposition 4.1. The gradient dl /do is given as
dl' d0
Σ
A
de A
+ Σ
vip
dl A (x4 dy
@ vipe dy h')
di
(5)
where the adjoint state A = CI A viPNer (v) = v. Neto with being Net 's parameters and being the index of Input, @ is the vector outer product, and @ is vector concatenation.
The proof of proposition 4.1 is in section A.1 The gradient for parameters of A, and B are each separated 8 and the gradient for parameters of C only depend on inputs at time t. After computing the adjoint states, these vjp computations are separated on both the network and the temporal level.
dr X into A dyr A @ h ji= )' {vip ay
di
x
X
2
X
x
C:
A
A
A
A
Figure 4: The adjoint states are computed soquentially backwards.
4.2
Adjoint sharding for multiple SSMs
We now generalize the results from subsection 4.1 to the general of K SSMs concatenated. As introduced in subsection 3.2 the outputs of each SSM layer added to the results of the last layer and normalized before feeding into the next layer. Define the loss over all token predictions L =I= 1", using the residual structure we have
dL
dl dy
d0
dy
d0
T
dl d(yt
>K v1)
dyk
d0
1=1
T
dl'
K
dy
dy
d0
K
k=1
Combining with proposition 4.1
we have

---

Published in Transactions on Machine Learning Research (10/2025)
Proposition 4.2. The gradient of the total loss L with respect to the SSM parameters 0 is given
dL
K
dl dy'k
h )
vIPA:
d0
dlc
viB dy'k
X 6 sk-1)
(6)
K
di
+
vJPe ( dy'k
@ b/))
kml
where the imput to viPe: dy @ h)),, the k-th SSM and the 8K-1= Norm SSM (-2))') (the normalized output sequence of the (k-1) 3-th SSM). The adjoint state at layer l defined as X C (I A (+11)).
vi
Lay' k
t.t @ 11 and vips; ( dl ti & 8k--) are computed with (ayx
dl(o') dy'x
hg
AA
vjp
al(o')
x²
dy
al(o") dy
12
vips;
Aip di(o')
b
As
x @ ki)
dy's
al(('))
x²
x
dl(o') dd 4
dy
82
viPag
al(o''
x 9% dy
dl(o')) d0s,
dl(a')
h 1
A
vi
A 1)
dy
dl(o')
x²
2
dy
di(o') dy
vipay
x sbh,)
dl(a')
Cs
dBc,
Figure 5: Computation schematic of dl' /d0 A. di' dl'/dBB, and d'' /d0c..
We provide the proof proposition 4.2 section A.2 Define A = {Ab ( proposition 4.2 shows that the gradients of each network's parameters computed with each token only correlate through the adjoint states (Af K,T We can efficiently compute the adjoint states after forward pass. We can also compute the adjoint states on the fly the gradient computation phase, as it only depends on C and A and has no dependencics on the network Jacobians regarding the network parameters. The adjoint sharding method breaks down the backpropagation computation from layer-wise and token-wiss into foundational vjp computations that do not have any dependencies each other.
(A,
8

---

Published in Transactions on Machine Learning Research (10/2025)
We show schematic of the computations to di' /d0 di' /d0B.: and di' /d0c, in Figure 5 and schematic for computing the adjoint states in Figure 4
4.3 Truncated adjoint sharding
One limitation of adjoint sharding that the number of vjps performed increases polynomially with the number of tokens T. In particular, adjoint sharding computes the vip for As and Bic (1 + T)T/22 times, and for Ck 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 vips, and empirically showcase its performance.
Attention mechanisms have suffered from the O(T2) complexities arising from the self-attention structure (Vaswani et al., 2023). : 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 (Yang et al. 2021). (Tallec & Ollivier 2017) proposed truncated backpropagation through time (T-BPTT) to avoid gradient explosion/vanishing when training with very long contexts by only counting a fixed number of state transitions. Here, inspired by global-local attention and T-BPTT, instead of computing we propose to train the SSMs to depend on up to T states:
the full gradient given in Equation 11.
dL
K
vipe
df
h[/)
ΣΣΣ
I
VIPA
dl'
X
&
h
d0
dy'k
dy'
du
K
di dy
Vi dy
h )
viPB:
k @
1, s)-z)
(7)
[=T+1 k= i=1+1-T
Ri i=1
T K
di
vjPs
dy
t= t 1-
As shown in Equation above, we perform the same computations for / = 1,..., T as before, and only perform the vips back to the last T states for t T. With truncated adjoint sharding. we perform TT + - 1)/2 vips, which grows lincarly with T. We show the number of vips performed with and without truncated adjoint sharding in Figure 6 When T 2000, truncated adjoint sharding reduces 64% of the vjps when training with a context length of 10K. The essence of the truncated adjoint sharding method is that we only explicitly count gradients related to the last T states. As each state depends on its prior state, states still implicitly depend on all their prior states. The choice of truncation parameter involves a trade-off between computational efficiency and gradient fidelity. In practice, we find choosing T = vT provides good balance between memory savings and gradient fidelity, as it keeps computation time in check while preserving sufficient gradient information. The optimal choice depends on the specific model architecture and sequence characteristics. We leave the investigation of T's impact on performances for future work.
We provide the algorithm for evaluating adjoint states for token index t and ResNet index k with truncated adjoint sharding T in algorithm | and the algorithm for evaluating the vjps for token index t and ResNet. index with truncated adjoint sharding T in algorithm 2
4.4 Distributed training
We now discuss how to distribute the storage and compute of the adjoint sharding method, assuming that we have Y GPUs. Given the networks (Ak, BE, ) k initial tokens (+6)T= = (Norm(x')) and initial conditions { (usually set to 0), we can call a algorithm 3 to get all necessary vectors for computing the gradient with adjoint sharding. As shown in algorithm 2 to compute the vjps' for token index and ResNet index k we only need t, kt /l(oo)/dyk. ( C₁2 1 (A To compute all the gradients for layer k we only need A, h, and C from the k-th layer, and y from the k 1-th layer. Therefore, we can divide the K layers into Y pieces, as shown in the appendix R. As the computations are fully independent.
9

---

Published in Transactions on Machine Learning Research (10/2025)
Algorithm 1 Evaluating adjoint states for token index and ResNet index & with truncated adjoint sharding
T
1: Inputs: 1, Ki T, C {A 1
2: Initialize adjoint state A = =
Compute: intermediate values:
4: Ç² =AAA A 142-1
AGA k ... A
A
AGA II.
5: Compute: adjoint states
Ao (AL+III A
C.5² 1-42-T
,..., (
.....
7: Return: A
Algorithm 2 Evaluating the vip's for token index and ResNet index & with truncated adjoint sharding T
1: Inputs: k; T, al(oo)
(hi) Cg (yk-1)) (At) (
Call alg. to compute {1 11=+11-
3: Compute: allo') & h
{ [djo
A
@ hi
1 -1=+1-77 a 8 9
dy
Compute:
sinet ark
w.()). Eistr vipAg
(illo))
dilo') X 8 88 X 6 h Eisti- vip ark
ayj
5: Return: vipeo ( dicol 8
& bbt, Eimn T vipA dila' A @ hi ", Ei vipng
( x & 881-1) is
Algorithm Forward step evaluation mode on distributed system
1: Inputs: (86)T=1, [b]] {Ax, BR, 1 Ω
2: On devices e = 1,..., Y, in parallel do
3: for SSM model index = (v - (K/K//) + I,..., ((K/TT) do 1, .............................. do
for Time step index
Compute: A An($,_1) B Bx( C. CC(9% Aqh. B C.
Compute: y yk-1 yt.
7:
Compute: 5t = Norm
8: end for
9: end for
10: Store:
Tu(K//T) (B
(9% T (L
{A {AD] (L)))))22 (L
on device e.
11: Pass: (y ( V {8' (oo(K/ (1 to device v +
12: for Time step index 1. = 1,..., do
Compute: for = Sykk VT (!(o')), d((o') 1T
13:
14: end for
15: Store: dillo')
dyi Jt=1 on all T devices.
and we compute the gradients using only data on local devices, we additionally distribute the model and the gradients, as shown in Tabbl where Og represents the parameters of Ar, BR, and Cs, and Gradients represents the optimizer states for Ox. The complete training streamline is provided in algorithm 4 We distribute the activations, computations, gradients, and optimization states across T devices. While the forward evaluation pass results across different devices, as shown in algorithm g the computation of gradients is parallel across the Y devices. This parallelization 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 Mem/T, with Mem being the memory cost we only have a single GPU. If we have Y > K devices, we can further speed up the forward evaluation by first evaluating A, B, C in parallel, and then adding them together the distributed devices.
10

---

Published in Transactions on Machine Learning Research (10/2025)
Algorithm Evaluating dL with truncated adjoint sharding T on Y devices
de
1: Inputs: (yolT=1 (halk=1> {Ax, BR, C Q, T,
Call alg. 3 for {AA, Cj hy. (C) (CA)) (1.1)* dy'
(T.K) dil('') 3m and saved on each GPU device.
3: On each device v, in parallel do
4: Initialize gradient dA d0
for Time step index t 1...., T, layer index k (v L)(K//T) t....., ((K/IT) do
6: Call alg. 2 for == vipat 1 aie @ bi)) E.s vip - x @ hi 1) Eim vipa (a A @s .)
7: Compute: =
8: end for
for Time stop index
P I...., T, layer index & = (o 1)(K//T) I...., o do
10: Call alg. 2 for
vinet (
dir'" axk
0 hg), Ei=tni VipA: d'lo') x 60 ") Eimur ving
|llاo xir so s/-1) avic
dx
11:
Compute: =
12: end for
13: Return:
de
4.5 Parallel computing
Adjoint sharding converts the sequential process of backpropagation gradient computation into individ- ual independent vips, allowing for parallel computation. We analyze the time and memory cost of ) @ hf")), viPs; [(dl"/dyke ) & 8k-1), and viPe; ((di' /dy'k) @ h2)) vip has a simi- lar time complexity as a forward pass, and memory complexity of bs( + 0) + 10), where bs is the batch size, O is the number of clements in the network output, and [0] is the number of paramcters Novak et al. 2022
AiP ((dl'/dyk))) (.
We provide the memory and FLOPs required to compute the vjps in Table INVIDIA 2024).
We analyze training with a dataset containing contexts of lengths T, with T NVIDIA H100 GPUs, and performing computations in FP16. We use selective diagonal SSM with K layers, and each As, Ba, and Ce network is a single-layer multi-layer perceptron (MLP). For each data point {x" Jt=1: {x()T we store (T,K) {A C hi y (1.)((111) and (d((o')/dyk)?+ which is TK(2N + +P) TP FP16 numbers. We also save A 0g, and Oc, each taking PN + N FP16 numbers. We need to store T(2NK + PK + P) + 1)) FP16 numbers before computing the vip. As computing all adjoint state sequences takes up to +
viPA
viPs
viPe
Memory
bs(N² + 10Al") 10xl
bs(NP + 0g(() + 0g
bs(N 00ll)) + 0cl
Unstructured SSM
FLOPs
bs(N"(2P +1) 1))
bs(NP(2P +1)
bs(NP x (2P 1))
Memory
bs(N + -00xl")+100
bs(N + (0ss"))++++ [0g]
bs(N + +(cl")) (1 10c
Diagonal SSM
FLOPs
bs(N(2P+1)
bs(N(2P+1)
ha( (e(N(P++1)
Memory
bs(1 be(1+10Al")+ [0_a])++
10al
be( + (0==+)++
bs( (e(N- (0cl') + 10c
Scalar SSM
FL.PPs
bs(2P 1PP--1
+1)
be(N(2P+1)
Table 1: Memory and FLOPs required to compute the vjps. elements of the biggest parameter vector of A, B, and C.
|0,l", 10gl", and (0c represents the number of
FLOPss takes NP + T)/T FLOPS on average for each adjoint state. For T large enough, 11+T) (T+T)/T>s 1, we approximate the average FLOPs for each adjoint state with NP. Each vip then takes bs(7N 3N) FLOPs of computation.
When computing with selective diagonal SSM with P 128, N = 225, and bs 8, while storing and performing computations in FP16, computing vjP.A- viPs: and vjpe cach takes around 0.6MB memory and 1798144 FLOPs. We characterize the capacity of a modern GPU with 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 Iastinenstance GPU (MIG) instances, which the number of fully isolated GPU instances with its own high- bandwidth memory, cache, and compute cores GPU can host. An NVIDIA H100 Tensor Core GPU has
11

---

Published in Transactions on Machine Learning Research (10/2025)
GPU memory bandwidth 3.35TB/s and performs 1,979 tera. FP16 FLOPS per second. Therefore, the memory bandwidth allows computing speed allows computing (1979tera/s)/ 17981444 = 3.76 x 1.199 batches of vips per second. At the same time, since the H100 GPU has S0GB memory, it can hold up to 80GB/ (0. 6MB/vjj = 133 batches of vjps 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 7Y instances, offering as much as 56x speedup on one AWS P4 instance (8 H100 GPUs). We can not achieve such speedup for backpropagation because of its sequential nature.
= =5.58 x 10F6 batches of vjps per second, and the computing
3.0
Limitation The adjoint sharding method provides an alternative method of computing gradients backpropagation. While we ana- lytically proved that the gradients computed from adjoint sharding are equal to those from backpropagation, adjoint sharding suffers from a polynomial time complexity regarding the training context length when computing equiv- alent gradients. We provided the truncated adjoint sharding as a linear time complexity alternative, and leave the analysis of its conver- gence and further improvements on it for future work. We also provided distributed and paral- lel computing algorithm for performing adjoint sharding. However, the overhead of naïve im- plementations such an algorithm with multi- threading or multiprocessing overweights the speedups when the training context length is small. We leave the efficient implementation of the parallel algorithm on a CUDA kernel for future work. Our work focuses specifically on State Space Models (SSMs) like Mamba, which have fundamentally different computational structures than transformers. The adjoint sharding method we developed leverages the recurrent nature of SSMs, where the sequential dependencies allow for efficient VJP decomposition. Transformers, with their attention mechanisms and different architectures, would require a separate analysis to determine how to apply adjoint methods. Unlike SSMs' Markovian recurrent structure that enables clean temporal decomposition, transformers' self-attention creates all-to-all dependencies that break the sequential dependency structure the adjoint method relies on. This architectural constraint represents an important direction for future work, as extending memory-effiieen training tochniques to Transformer architectures would significantly broaden the impact of this approach.
Training time per epoch (/day)
2.5
2.0
Truncated adjoint sharding
Adjoint sharding
1.5
1e
Backpropagation
16:06
8--05
1.0
Se-05
36:0
0.5
De+00
0.0
100
10
102
103
104
105
106
Training context length
Figure Training time (/day) per epoch comparison for adjoint sharding, truncated adjoint sharding, and backprop- agation with different context lengths. Assumed 100-layer SSM-ResNet model, 280x acceleration for adjoint sharding from parallel computing (achievable with five Amazon P4 instances), and T from 15 to 2500.
Conclusion We introduced adjoint sharding, distributed and parallel computing algorithm, to facilitate the training of LLMs on very long contexts. Unlike the sequential backpropagation, the adjoint sharding computes gradients of each LLM layer against each token independently through vector-Jacobian product, allowing for parallel computation. We propose truncated adjoint sharding to focus on essential gradients avoid the limitation of vjps increasing polynomially regarding context length. Our approach provides an algorithmic alternative to existing memory reduction techniques such as activation checkpointing, model sharding, and vocabulary optimization methods. While these techniques address different aspects of the memory bottleneck, adjoint sharding specifically targets the gradient computation mechanism in recurrent architectures. The method is currently applicable to State Space Models, where the scquential recurrent structure enables clean VJP decompasition. We analyzed the memory and FLOP cost of each computation block in adjoint sharding and proposed a method to accelerate through parallel computing. Empirical results suggest orders of magnitude of memory reduction in training while maintaining the same training results as backpropagation. By enabling training on contexts previously intractable due to memory constraints, adjoint sharding opens new possibilities for fine-tuning SSM-based models on very long sequence tasks.
12

---

Published in Transactions on Machine Learning Research (10/2025)
References
Quentin Anthony, Yury Tokpanov, Paolo Glorioso, and Beren Millidge. Blackmamba: Mixture of experts for state-space models, 2024. URL
Randall Balcstricro and Richard Baraniuk. Fast jacobian-voctor product for deep networks, 2021. URL
/taps://arxivvoor//bb////2 00219
Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: survey, 2018a. URL https://arxiv
05767
Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic diffcrentiation in machine learning: survey, 2018b. URL
05767
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Olcksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory, 2024. URL http://////rdda.ooraa////a////// 04511
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer, 2020. URL
wtw.://///a.b/0......5. 05150
Zheng Cai, Maosong Cao, Haojiong Chen, Kai Chen, Keyu Chen, Xin Chen, al. Internim2 tochnical report, 2024. URL httts://arxiv...t/.////////// www.o//g//a//......../a/ 17297
Yang Cao, Shengtai Li, and Linda Petzold. Adjoint sensitivity analysis for different ial-algebraic equations: algorithms and software. Journal of Computational and Applied Mathematics, (49())171-191, 2002. ISSN 0377-0427. doi:
URL Scientific and Engineering Computations for the 21st Century Me thodologies and Applications Proceedings of the 15th Toyota Conference.
com/ http://siaanoo/rrrrcc///////1160
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations, 2019. URL http:://aaraa..orr/bbb//86.
Shouyuan Chen, Sherman Wong, Liangiian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation, 2023. URL
wwwp:/vv..rg/bb///00000 15595
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models, 2024. URL htt.htttts///ara..org///aa 2209 12307
Saeed Damadi, Golnaz Moharrer, and Mostafa Cham. The backpropagation algorithm for math student, 2023. URL http:://aaxxivorg//aaaaa 19901.99977
Tri Dao. Flashattention-22: Faster attention with better parallelism and work partitioning, 2023. URL
Tri Dao and Albert Gu. Transformers are ssiss: Generalized models and efficient algorithms through structured state space duality, 2024a. URL https://arxiv.org/bbb wttp:///..rrraabbb2..... 20000
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, 2024b. URL https://arriv..or/raa///e/ 21060
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness, 2022. URL htttp://arxxv.orr/abb///222 14135
Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Murarru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando De Freitas, and Caglar Gulcchre. Griffin: Mixing gated linear recurrences with local attention for efficient language models, 2024. URL https
7/axiiv.co
perxiv..rr/aaae orw.orgg./.........o...oo....... attaab//..... 19427
13

---

Published in Transactions on Machine Learning Research (10/2025)
Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: 2024a.
Extending Ilm context window beyond 2 million tokens. orXiv preprint ar arxi:22402. 19753,
Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: Extending Ilm context window beyond 2 million tokens, 2024b. URL https: www.iv.org/abs//// 2402. 13753
Emilien Dupont, Arnaud Doucet, and Yee Whye Teh. Augmented neural odes, 2019. URL https://arxiv org/ http//bb//99..1666
Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models, 2023. URL https://arxiv.org/eaae 14052
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2024a. URL https://arxiv, http:://axxiv..rr/abbsss
00752
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2024b. URL https://arxiv.org/ http//aab//2120072// 00752
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state-space layers, 2021. URL https
//axxiv.org/abs/21111 13985
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Ré. How to train your hippo: State space models with generalized orthogonal basis projections, 2022. URL. stttps://arxis org/ 2206
12037
Meng-Hao Guo, Tian-Xing Xu, Jiang-Jiang Liu, Zheng-Nii Ning Liu, Peng-Tao Jiang, Tai-Jiang Mu, Song-Hai Zhang, Ralph R Martin, Ming- Ming Cheng, and Shi-Min Hu Attention mechanisms in computer vision: A survey. Compulational visual media, 8(3):331-333 11-8, 2022.
Ankit Gupta, Harsh Mehta, and Jonathan Berant. Simplifying and understanding state space models with diagonal linear rnns, 2023. URL https::/aarsn org/ htrg/abs/22122/ /2tp/221..... 00768
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015. URL
1/1//1112.00 03385
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Julien Herrmann, Olivier Beaumont, Lionel Eyraud-Dubois, Julien Hermann, Alexis Joly, and Alena Shilova. Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory, 2019. URL
13214
Herbert Jaeger. A tutorial on training recurrent neural networks covering bppt, rtrl, ekf and the echo state network approach semantic scholar. In National Research Center for Information Technology, 2002, 2005. URL wtw.r/ccpp..emmaatccccccccoaaa Sttp://sssaattccccoooooo..ooooop org/CCrrps : 192593367
Steven Johnson. Adjoint methods and sensitivity analysis for recurrence, 01 2007.
Shiva Kaul. Linear dynamical systems as a core computational primitive. H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Newral Information Processing Systems, volume 33, pp. 16808 16820. Curran Associates, Inc., 2020. URL
files/
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tote Xiao, Spencer Whitehead, Alexander Berg, Wan- Yen Lo, al. Segment anything. In Proceedings of the IRFE/CVF International Conference Computer Vision, pp. 4015-40266 2023.
14

---

Published in Transactions on Machine Learning Research (10/2025)
Dacheng Li Rulin Shao*, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph E. Gonzalez, lon Stoica, Xuezhe Ma, and Hao Zhang. How long can open-source llms truly promise on context, length?, June 2023. URL
https://lmsys: http:://mssya.orrr/l/////////// 0Tt////////0000000-999-99999 29-1ogcchat
Tianle Li, Ge Zhang, Quy Duc Do, Xiang Yue, and Wenhu Chen. Long-context 1lms struggle with long in-context learning, 2024. URL httpss//aariv.orr//aa///044 atg/abs/24044.2060
Opher Licher, Barak Lenz, Hofit. Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, and Yoav Shoham. Jamba: A hybrid transformer-mamba language model, 2024. URL
https://arxiv. wwwpaawwww...g/bb///......./ /00
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context, 2023. URL
01889
Hao Liu, Wilson Yan, Matci Zaharia, and Picter Abbool. World model on million-length vidco and language with blockwise ringattention, 2024. URL stttpa//aaxxi..orr/bbb////
08268
Chong Luo, Jiawei Zhao, Zhuoming Chen, Beidi Chen, and Anima Anandkumar. Mini-sequence transformer: Optimizing intermediate memory for long sequences training, 2024. URL https://arrxiv.or//abbbb
15892
Qijun Luo, Mengqi Li, Lei Zhao, and Xiao Li. Streambp: Memory-effiiiet exact backpropagation for long sequence training of llms, 2025. URL
Meta et al. The llaaa 3 herd of models, 2024. URL. https::/aari...rr//abs...s
21783
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training. 2018. URL
03740
Roman Novak, Jascha Sohl-Dicksteinn, and Samuel Schoenholz. Fast finite width neural tangent kernel, 2022. URL https://arriiv.or/aab///// 08720
NVIDIA. Matrix multiplication htaplippllcatio  bccccnonnnnt user's guide, 2024. URL https::/docss hatp:::iia/ccc..i.iiia.aa. nwidia. com/
html
OpenAl et al. Gpt-4 technical report, 2024. URL
Antonio Orvieto, Samuel Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences, 2023. URL https://arxiv.org.
abs/ 2///2333300664
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks, 2013. URL. sttps://aarxi..or//abbe argg.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library, 2019. URL https://arxiv. wwwpi/vv.rg/bbbb///eeeeeee http//111...01111
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195 4205, 2023.
Bo Peng, Eric Alcaide, Quentin Authony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Mattco Grella, Kranthi Kiran GV, Xuzheng Hc, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomicj Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Bolun Wang, Johan S. Wind, Stanislaw Wozniak,
15

---

Published in Transactions on Machine Learning Research (10/2025)
Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu. Rwkv: Reinventing rmns for the transformer era, 2023a.
abs//305..2. 13048
Bowen Peng, Jeftrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models, 2023b. URL htttps:/aaxxvv.rrg/ http:///bb/2339...0g
Maciej Pióro, Kamil Ciebiera, Krystian Król, Jan Ludziejewski, Michal Krutul, Jakub Krajewski, Skymon Antoniak, Piotr Milos, Marek Cygan, and Sebastian Jaszczur. Moe-mamba: Efficient selective state space models with mixture of experts, 2024. URL http:////aaa.....o///////////0/0
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong Hc. Zero: Memory optimizations toward training trillion parameter models, 2020a. URL attpp::/aaxxiw.org/bbb////
02054
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models, 2020b. URL.
02004
Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. Zero-offload: Democratizing billion-scale model training, 2021. URL
https://arxiv. org. httgg/abs/210100644)) 2101.06844
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradcep Ramani, and Tri Dao. Flashattention-B: Fast and accurate attention with asyuchrony and low-precision, 2024. URL https://arxizs org/abs/ /bs//
2407.08600
Nimit S. Sohoni, Christopher R. Aberger, Megan Leszczynski, Jian Zhang, and Christopher Ré. Low-memory neural network training: A technical report, 2022. URL
Corentin Tallec and Yann Ollivier. Unbiasing truncated backpropagation through time, 2017. URL https:
Szymon Tworkowski, Konrad Staniszewski, Mikolaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Milos.
Focused transformer: Contrastive training for context scaling, 2023. URL http:://aaxxx..r///////333/
03170
Reddit users. Ntk-aware scaled rope, 2023. URL http://www.ddditt..o//////aaaala 141z7 haaee///tttaaare_aaaadededde scaled, rope allows llama models to to__avee.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023. URL
Joost Verbraeken, Matthijs Wolting, Jonathan Katzy, Jeroen Kloppenburg, Tim Verbelen, and Jan S. Rellermeyer. A survey distributed machine learning. ACM Computing Surveys, 53(2):1 11-33, March 2020. ISSN 1557-7341. doi: 10.1145/3377454. URL http://dx.....o http//xxxxooi...///// org/10..11. 1145/3377454
Roger Waleffe, Wonmin Bycon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh, Jared Casper, Jan Kautz, Mohammad Shoeybi, and Bryan Catanzaro. An empirical study of mamba-baas based language models, 2024. URL https:///axiv: org/ //s//bb//200 240060000767777777
Shida Wang and Beichen Xue. State-space models with layer-wise noulincarity are universal approximators with exponential decaying memory, 2023. URL lttpp:///axiiv... wwtperiivorrg/bssa..e wtgg////////......./ggg/ 13414
P.J. Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE, 78(10): 1550-1560, 1990. doi: 10.1109/5.5833
Erik Wijmans, Brody Huval, Alexander Hertzberg, Vladlen Koltun, and Philipp Krähenbühl. Cut your losses in large-vocabulary language models, 2025. URL https: http://wwwvvv..r/////1/1.1...0/.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks, 2024. URL httpps://aazivvor//aaa///
17453.
16

---

Published in Transactions on Machine Learning Research (10/2025)
Xingzi Xu, Ali Hasan, Khalil Elkhalil, Jie Ding, and Vahid Tarokh. Characteristic neural ordinary differential equations, 2022. URL https://arxiv.ont org/abs/2111
Jianwei Yang, Chunyuan Li, Pengchuan Zhang, Xiyang Dai, Bin Xiao, Lu Yuan, and Jianfeng Gao. Focal self-altention for local-global interactions in vision transformers, 2021. URL
2107.00641
Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, and Zhicheng Dou. Soaring from 4k to 400k: Extending Ilm's context with activation beacon, 2024. URI. https://arxiv.. org/
03462
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chil Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Out, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. Pytorch fsdp: Experiences scaling fully sharded data parallel, 2023. URL https: :ttps://aree //txxv/rxvvva///a///// http://axx...og/ba//////// 11277
A Proof
A.1 Proof for proposition
Proof. Define 0y/Oh' expression for y from subsection 3.2 proposition 3.1 states that
y oh' 10h hig and 05/00 yo. 0h'/00 lig. by plugging in the
dy
Yn
bbg + h hb .... + h h²g 1-1 h/
+ +50
do
i=1
i=1
In the context of SSM, we have:
h -'=A'h111 B'x' h A h/ = Agh'-i + BBx , C' = Co, = -igh..
(8)
Plugging these relations, we get:
dy
6-1
t-2
C
II
A
h2
(I
A ++] >hg 0 + A'hh
h/
+ yg.
(9)
d0
Define the adjoint state ALt = -C'[[ A 141-1) we have
dy
= An + 2 + .... X 16 hi
4- + Ah + +
d0
Therefore, we have
du
dl' d(ý' + +*)
d0
dy d0 dl' dy dy' de dl
dyr ( hbg
2 + 1 + A +
Plug everything, we have
dl
dl'
dy [ (AAh + BBx x (A + +BXx)
+ - (B'x"" Cyh²
d0
dic
+ di dy
(Aghh-1+ 1- (igh-1++Bxx))
Cyp'
dyr
vip dlt dyr
dlt A &/ -hh--1 + vip
+ vipe dl
dy Aht @ xx))
A
dy & h²)
17

---

Published in Transactions on Machine Learning Research (10/2025)
where define vi) = NNo (Input'), with 0 being NN's parameters and being the index of Input. Now, as Aj dy Ar & h viPs di'
t 2 x²); and vipe dy & h'Jare separate, we have dy
dl
di Ai
dl vip
dl vipe ( 2 hh),,
d0
Σ
X hh- dy
+
Σ
s x08)) dy
dy
(10)
where @ is vector concatenation.
A.2 Proof for proposition 4.2
Proof. First, using the structure of ResNet, we have
dL
dl
dy
d0
dy'k
d0
T dl
d(yb
Σ v1)
dy'k
d0
T
dlc
K
dy
dy'k k=1
d0
K
dl'
dy d0
dyk
from proposition 4.1 have proven that for single SSM model, have
dl'
dl v'p
dl dy A
di vipe.
dy A
@
viPs:
xx)
dy @ h²),
d0
so for the ResNet model, we have
dL
T
du dy
d0
dy
d0
T
( dit
dl'
di X @ x/) e evipe dyk @ hi))
viPA:
X b
+ Σ vips; dy'k
dy'k
dl
viPet dy
@ h/))
dlc
Ao. b
viP dy
(11)
di
vjPB
A 3 xx61)
dy
K
dl'
viPc; dy'k
2 h[))
k=
di
viP
Xoa @ bi
dy
dl dy'k
viPs:
X @ ())))
18

---

Published in Transactions on Machine Learning Research (10/2025)
where the imput to viPe; dy
dr'
are computed with di ti h) and vips; dy X @ x 2 hi), viP.: (ayk
di'
the k-th SSM and the & = y = RMSNorm(y' + R (Vx=))))))) (the normalized output sequence of the (k-1)-th SSM), and the adjoint state X = C
A.3 Proof of concept for VJP computation
As proof of concept of why (dl' /dy") C'yh can computed with vip, we present an explicit and simple example. We have y = 92), h [h1, h2, ha], 0 O. We then have
dl
In
RRIP
dy
C
OØ
C
12
13
PxNx00
C1
C
C
[hi
h =
h12
E RR-
h3
With each C. = |0Ct,/001..... 00,. E R0 We have
dl dy
Coh Cha lu 11 + Chil C + h
+ Ci Cluba
+ Castuuhs
= Unr hn Loi h2 ly ha hi hyg hg)
[Cfi ch Clac c22 chal
1 va
sum
2
hh
h2
ha)
C
where is vector dot product, @ is vector outer product, is wlement-wise product, and sum means summing all elements in a matrix.
B Distributed tensors' locations
We provide the specific location for each tensors in distributed training:
Table 2: Tensors stored on each GPU, part 1.
GPU index
di(o')/dyf,
h/
==
t 1, ,,...,7 T
,,....,II T: k = I.... K//T
v=
t ,.....7 T
1,....,
T; =K//K + 1,..., ((K//))
== T - 1
t =
T
t=
",.,,,T;k==37 (2)//()////) + 1, 1..... (T-1)(K//T) T, - (1)(K/)()))))))) K,..,,
== T
t T
t=
19

---

Published in Transactions on Machine Learning Research (10/2025)
Table 3: Tensors stored on each GPU, part 2.
GPU index
==1
{ 1 ,,.., T; k K//T.
==2
1 1 ,...;
T; k K//T +1,,+..... (,k..,((K//))
...
T - 1
= ,.....T T
v == T
1 == I,....
T; k = (T ()(K/)))++1....... K
Table 4: Tensors stored on cach GPU, part 3.
GPU index
1 1, ,.... T: k 0,... K//Y -
V= 2
////TT....,(////7 - t = 1, ...... T
== T - 1
t == I....,
0= T
t = 1,....,,7 - - (T)]=(T))
Table 5: Tensors stored on each GPU, part 4.
GPU index
A
==1
t == 2, ...
T: k = 1,...K//TT
V=2 2
1 1= 2,
T: 1: = K//T ((K/))
== T --1 1
( 1= 2,.., k;
..,,..,Tkk == (T - + + -f.
== T
t t= 2, 2....,T ,, -T -1)(K//T + +1. 1....,,
Table 6: Tensors stored on each GPU, part 5.
GPU index
02
Gradienty
V=1
k: I,... K//T
k K...K//T
V= 2
k = K//T + (4...,2(K//T))
k = +
0 T -
k= (T (2)(K/// +1...., - (1-1)(k//)
k = (T (2)((K///))+++)
+ +,...., .. (1)1()m/
v == T
k= (T- -1)(K/)/))+ 1,..., K
k: = (T 1)(k//))+1.. K
20


Reviews:
[{"review_id": "Igo5N4VxEZ", "reviewer": "Reviewer_kc5n", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The paper addresses an important and practical problem — the memory cost of training long-context LLMs, which has strong applicability.\n2. The proposed method is simple yet effective, and the design is reasonably motivated.\n3. The paper provides theoretical derivation of the method’s application to ResNet and Mamba, demonstrating its generalizability beyond just Transformers.\n\nWeaknesses:\n\n1. The paper lacks discussion and comparison with several key baselines that also aim to reduce peak GPU memory usage, such as [1-3].\n2. The analysis and experiments are limited to ResNet and Mamba, with no evaluation or discussion on Transformer-based architectures, which are more commonly used in LLM training.\n\n[1] Mini-Sequence Transformers: Optimizing Intermediate Memory for Long Sequences Training, NeurIPS 2024\n\n[2] Cut Your Losses in Large-Vocabulary Language Models, ICLR 2025.\n\n[3] StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces Adjoint Sharding, an efficient training method for extremely long-context scenarios. Specifically, it reduces peak GPU memory usage by decomposing gradient computation into independent vector–Jacobian products (VJPs). The authors provide both theoretical derivation and empirical validation on ResNet and Mamba, demonstrating up to a 3× reduction in memory cost when training 1.2B-parameter LLMs with 1M context tokens.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Could you discuss and compare your method with the following necessary baselines [1-3], which also aim to reduce peak GPU memory usage in long-context LLMs training?\n\n2. Could you add an analysis of your method within the Transformer architecture, which is widely used in LLMs?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "x8qxdrLr5J", "reviewer": "Reviewer_ZhrN", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- **Novelty**: Adjoint Sharding is a compelling alternative to backpropagation for recurrent models, offering provably equivalent gradients while improving memory efficiency.\n- **Scalability**: The technique supports distributed training and shows significant memory and compute savings, enabling long-context training on modest hardware.\n- **Theoretical Rigor**: The paper provides formal derivations for both the base and truncated versions of adjoint sharding, grounded in residual and state-space models.\n- **Practical Implementation**: Includes a detailed discussion of memory/FLOPs cost and implementation guidelines, with empirical validation on large-scale models.\n\n### Weaknesses\n- **Computational Overhead**: Full adjoint sharding still scales polynomially in context length, limiting its utility without truncation.\n- **Empirical Scope**: The experiments focus primarily on memory savings and feasibility; additional benchmarks on task performance (e.g., in-context learning, retrieval, QA) would strengthen claims.\n- **Truncation Parameter**: The choice of truncation length \\( \\bar{T} \\) is left unvalidated; its effect on convergence and performance is not analyzed.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces **Adjoint Sharding**, a novel memory-efficient technique for training large language models (LLMs) on **very long context sequences** using **state space models (SSMs)**. The method decomposes the gradient computation into **independent vector-Jacobian products (VJPs)** using the adjoint method, enabling parallel computation and significantly reducing memory usage. To address the quadratic scaling of standard adjoint methods, the authors propose **Truncated Adjoint Sharding**, which limits the number of VJPs and reduces complexity to linear in sequence length. The paper also presents **distributed and parallel implementations**, showing that the method scales effectively across GPUs. Empirical results demonstrate up to **3× memory reduction** and enable training with context lengths over 100K tokens for a 1.27B parameter model.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. **Truncation Analysis**: Provide experiments or ablation studies showing how different values of \\( \\bar{T} \\) in truncated adjoint sharding affect training efficiency and model accuracy.\n2. **Overhead Quantification**: Include wall-clock training time comparisons between backpropagation, adjoint sharding, and truncated adjoint sharding on the same hardware.\n3. **Clarity Improvements**: Clarify the assumptions made in the distributed computing setup and provide pseudocode or code references to aid reproducibility.", "broader_impact_concerns": "The proposed method facilitates more memory-efficient training of LLMs on long sequences, potentially democratizing access to long-context training for smaller research labs. However, by enabling longer context lengths, it may also exacerbate concerns around **data leakage**, **privacy**, and **computational inequality**, particularly if used to train large-scale models on sensitive or proprietary data.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "g3fTuQWkS1", "reviewer": "Reviewer_zJCe", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strenghts:\n- Nice application of the adjoint method from control theory to gradient computation (by breaking gradient computation into independent vector-Jacobian products that can be parallelized);\n- Theoretical analysis shows memory requirements that scale much better than standard backpropagation;\n- The methodology shows up to 3× memory reduction on 1.27B parameter models, plus it enables training on 100K+ tokens vs 35K tokens on the same hardware infrastructure\n\n\nWeaknesses:\n- Lack of comparison of actual training convergence or final model quality;\n- Missing comparisons with other memory reduction techniques (gradient checkpointing, CPU offloading, etc.);\n- No evaluation on downstream tasks to verify model quality is preserved;\n- Without truncation, VJP count grows polynomially with sequence length $T: \\mathcal{O}(T^2)$, and the truncated version may lose important gradient information - no analysis of this trade-off.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes adjoint sharding, a training technique that reduces GPU memory requirements for training large language models on very long contexts (>1M tokens). Instead of standard backpropagation, the method uses the adjoint method from control theory to break gradient computation into independent vector-Jacobian products that can be computed in parallel.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I'd suggest to:\n- Include convergence analysis and training time comparisons;\n- Evaluate on standard benchmarks to verify model quality;\n- Possibly compare against existing memory reduction techniques;\n- Analyze the impact of truncation parameter $\\bar{T}$ on gradient quality and convergence.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)