AutoSP: Unlocking Long-Context LLM Training Via Compiler-Based Sequence Parallelism
Authors: Ahan Gupta
Zhihao Wang
Neel Dani
Masahiro Tanaka
Olatunji Ruwase
Minjia Zhang
Venue: ICLR 2026 Poster (2026)
Keywords: Long context training
Sequence Parallelism
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: An automated approach for lifting Sequence Parallelism, and other targeted memory optimisations for long-context training, into the compiler.

Abstract: Large-language-models (LLMs) demonstrate enormous utility in long-context tasks which require processing prompts that consist of tens to hundreds of thousands of tokens. However, existing LLM training libraries do not provide easy to use abstractions to optimize for long-context training, instead focusing on optimizations for models with large parameter counts through ZeRO-3/FSDP, Tensor and Pipeline parallelism. This forces users to rewrite LLM training libraries to incorporate compositions of various complex long-context optimizations, such as sequence-parallelism, to training pipelines; a process that requires in-depth expertise, reducing developer productivity. To tackle these challenges, we introduce AutoSP: the first automated solution to automatically optimize LLM training for longer-contexts. AutoSP compiles models and applies a targeted set of optimizations: automated sequence parallelism, and long-context aware activation-checkpointing, to drastically enhance LLM trainability at negligible cost to throughput. Our evaluation demonstrates AutoSP's capability on both NVIDIA and AMD hardware, increasing training contexts by upto 2.7$\times$ and 2.5$\times$ respectively over competitive hand-written baseline at negligible cost to runtime performance.

Published as a conference paper at ICLR 2026
AUTOSP: UNLOCKING LONG-CONTEXT LLM TRAIN- ING VIA COMPILER-BASED SEQUENCE PARALLELISM
Ahan Gupta', Zhihao Wang", Neel Dani SSAIL Lab, University of Illinois Srbana-Champar Champaign, IL, USA {ag82, zhihaow6, neeld2)eillinoh
Masahiro Tanaka
Anyscale
San Francisco, CA, USA mt attankka@ansccaee. com
Olatunji Ruwase
Minjia Zhang
Snowflake San Mateo, CA, USA
SSAIL Lab, University of Illinois Urbana-Champaign Champaign, IL. USA Sinjiaz@1limme edd
rurj...ruas@anoowlaae com
ABSTRACT
laggellanguagemmodel (LLMs) demonstrate enormous utility in long-context tasks which require processing prompts that consist of tens to hundreds of thousands tokens. However, existing LLM training libraries do not provide casy tions to optimize for long-context training, instead focusing on optimizations for models with large parameter counts through ZeRO-3/FSDD Tensor and Pipeline parallelism. This forces users to rewrite LLM training libraries to incorporate compositions of various complex long-context optimizations, such as sequence- parallelism, to training pipelines; a process that requires in-depth expertise, re- ducing developer productivity. To tackle these challenges, we introduce AutoSP: the first automated solution to automatically optimize LLM training for longer- contexts. AutoSP compiles models and applies a targeted set of optimizations: automated sequence parallelism, and long-context aware activation-checkponntii to drastically enhance LLM trainability negligible cost to throughput. Our eval- uation demonstrates AutoSP's s capability on both NVIDIA and AMD hardware, increasing training contexts by upto 2.7 x and 2.5x respectively over competitive hand-written baseline at negligible cost to runtime performance.
abstrac-
INTRODUCTION
Large Language Models (LLMs) increasingly being trained with long-context data for scenarios such as document understanding (Han et al., 2025; Zhu et al., 2025; Landeghem et al., 2023), multi-step reasoning (Cobbe et al., 2021; Feng al., 2020), and extended multi-turn dialogue generation (Touvron et al., 2023; Bai et al., 2024). These use cases often contain input sequences ranging from tens hundreds of thousands of tokens, creating massive activation memory demands and pushing the memory and system limits of GPU clusters.
To circumvent out-of- mfmmoo- errors, researchers have explored Sequence Parallelism (SP), a key enabler for cong-context training. State-of-the SP strategies such as SeepSpeed-Ulysseee (Jacobs et al., 2024) and RingAttention (Liu et al., 2023) distribute the sequence dimension of activations across devices and allow the training engine to leverage aggregated GPU memory to train longer contexts with increasing device counts.
Despite effectively enabling long- long-context training, existing SP are implemented in eager mode and tightly coupled to specialized training frameworks such as DeepSpeed (Rajbhandari et al., 2020) and Megatron-LM (Shoeybi et al., 2020). Integrating SP to new training pipelines typically requires invasive code refactoring, which makes it difficult to apply across diverse model architectures and hardware platforms. Developers must manually insert communication collectives (e.g., al alllall) between operators that require the full input sequence (such as attention), manage activation layouts
"qqual contribution.

---

Published as a conference paper at ICLR 2026
across devices, and ensure correctness in both forward and backward passes. These manual efforts reduce scientists' productivity and limit portability.
To improve productivity, researchers have begun to lift several complex distributed training strategies such as ZeRO-3/FSDP (Rajbhandari et al., 2020) into SoTA deep-learning compilers, c.g. Py PyTorch- 2.0 (Ansel et al., 2024). Examples include: SimpleFSDP (Zhang et al., 2024) & DeepCompile (Tanaka et al., 2025), which implement ZeRO-3/FSDP as series of compiler passes in the PyTorch-2.0 ecosystem. However, each of these techniques focuses on how increase model parameter counts, uncovering different ways to shard model parameters rather than explicitly optimize for long-context training. While these efforts successfully lift data and model parallelism into compiler abstractions, they do not address parallelism strategies tailored to long-context training. This raises the question: can SP also be lifted into a deep-learning compilation stack to enable automated sequence parallelism?
In this work, we focus on lifting SP into PyTorch-2.0's compiler ecosystem as compiler-based, PyTorch-native implementa-
1 Register passes
2 preg_passes, (1' d'aauuossp sp sp_ac']]
3 Initialize
dist training #
4 Addist.ini (SP L___GGROUU_SS
tion. However, achieving this introduces sev- eral challenges. First, PyTorch-2.0's compi- lation pipeline includes multiple intermediate-
5 Compile model.
6model compile()
7 model training
representations (IRs) such Torch-IR, IR, Aten-IR, IR, and Inductor-IR, to name a few. Each IR oper- ates different abstraction level and encodes the input program at varying levels of granular- ity. This makes identifying an appropriate IR to conduct program-analssie so as to recover the necessary information to apply semantically- preserving rewrites that transform single-Gie GPU code into distributed ssquence-paralles execu- tion, major challenge. A tion will make program program-analysis (to uncover important attributes about the input model) increasingly non-trivial whilst coarse-grained abstraction will make semantic-rewrite, (to insert communication collectives, transform buffer sizes and recompute manually indexed tensors) infeasible. Second, inferring sequence-dependers tensor shapes for resizing intermediate buffers is non-trivial to do within a compiler. (1) The lowering process inserts intermediate data movement operators (such as transpositions), frequently changing the sequential axis of buffers. (2) SP strategies require resizing the sequential axis of only certain buffers (such as token and position id buffers), leaving other buffers (such as attention masks) untouched. Consequently, disambiguating which token's 's sequential axis requires resizing is challenging. Third, lifting SP into the PyTorch-2... 2-00 compilation stack has consequences other optimization passes that Py 2oroc-22.0 natively supports, notably: activation- checkpointing (AC). AC also enables memory savings for training by discarding activations in the forward pass and rematerializing them in the backwards pass to compute gradients. However, naively rematerializing activations together with SP triggers extraneous communication in the backwards pass, adversarially impacting runtime performance.
8 for batch in data_set:
9 output = model :, SP )
sntput.backwart optimizer... step (modem)
11
Listing I: AutoSP provides an easy to use interface to apply long-context optimizations to training pipelines. manages activation layouts between devices, inserts communication collectives and ensures correctness of the forward and backward training pass.
abstrac-
To tackle these challenges, we introduce AutoSP, the first compiler-based, Py Torchh-native imple- mentation of sequence parallelism. AutoSP introduces two key components: (1) sequence-parallo transformation pass that automatically inserts communication collectives and reshapes activations, and (2) a sequence-aware activation checkpointing pass that exploits the compute-memory char- acteristics of long-context training. With just few lines of code (as shown in 1), scientists can compile standard PyTorch models into distributed long-context training pipelines that scale input lengths without manual engineering. Our evaluation demonstrates that AutoSP significantly improves trainability at negligible loss to training speed on diverse hardware backends (NVIDIA and AMD GPUs), enabling training with up to 2.7x longer input context lengths compared Hand-written SP implementations such as HeepSpeed-Ulsseeee and RingAttention.
2 BACKGROUND
Sequence parallel training. Sequence parallelism (SP) is key enabler for long-context train- ing. These strategies enable scaling input sequence lengths with increasing GPU resources by
2

---

Published as a conference paper at ICLR 2026
sharding input tensors and activations across the sequence dimension. Communication collec- tives are inserted in the forward and backward pass as necessary to correctly shuffle tokens to the desired device. A popular SP strategy, and the focus of this work, is DeepSpeed-Ulyee Ulysses (Ulysses) (Jacobs et al., 2024). We illustrate how Ulysses operates with 3-SP groups in Fig. 1. First, tokens are sharded across the sequence di- mension, with different devices operating on dif- ferent parts of the input sequence. Next, linear projections form multiple O/K/V heads. Since linear projections operate pointwise across the sequential dimension, each device directly oper- ates on its own tokens with no additional com- munication. However, since attention requires the entire input-context, an all-to-all reshuffles tokens, re-sharding the activations across the head-dimension. Each device locally computes attention its respective head(s) after which another all-to-all reshuffles the tokens back to their original input sizes, re-sharding across the sequence dimension.
Attention Layer
Heads
GPUD
Linpar
Atlention
Projection
GPU
Linear
Projection
Projection
Linear
Allenton
Figure 1: DeepSpeed-Ulysses with 3-SP groups. alltoall operators toggle the layout of ac- tivations at attention-layer boundaries. Linear- projections operate on the partial sequence length, while attention-layers operate subset of the heads.
PyTorch-2.0 compiler. PyTorch-2.0 (Ansel et al., 2024) a just-in-time deep-learning compiler that targets train- ing workloads. It comprises of two components: dynamo and inductor each with a series of compiler passes that progressively lowers and optimizes code. An example code-snippet of a neural-network and its progressive low- ering through different IRs, is shown in Fig. 2.
Doa griined Operators
def aodal (a, b):
F. F.iinoar b)
Torch-IR
def mt b_b.):
torch.C. annn.nnneee 1n.1rre 12_o_, 1_b_)
Dynamo. The input to dynamo a model comprising of Py PyTorch & python operators. Dynamo then executes the function and records frace, represented as its interme- diate representation: computation graph. Each node in the computation -rappp comprises of Torch-IR statements, which loosely correspond to statements in the original in- put program. Next, AOTAutograd lowers each Torch-IR statement to Aten-IR statements, which consist of finer- grained operators. Aten-IR statements do not consist of higher-level abstractions such linear or attention layers, con-
return L_a_
Aten-IR
Fine Operatons
del nodel 1h, 1.
be dww.ddffaaaa danl O1 o cae comeene aten am, detsuiti
un torch, pernute:
zaturn torch. ops aten.
add. add.defaulte
but instead consist of (batch) matrix-multiplictioon volution, and data-mon movement operators. Each Torch-IR statement is accordingly lowered to its corresponding set of (multiple) Aten-IR statement(s), forming an FX-graph. For example, in Fig. 2, we observe that the linea op- crator in Torch-IR is lowered to two operators in Aten-IR: permute and mm. At this stage, variety of compiler- passes to optimize the FX-graph are applied, notably automated acitvation-cheecponning (AC). The AC compiler- pass (Chillee) is responsible for selecting which tensors rematerialize in the backwards pass without incurring performance penalties. It reduces this problem to a network- flow construction whose min-cut determines the tensors rematerialize. We describe it in detail in Section 3.2.
Figure A sample neural network com- piled using PyTorch-2.0. We illustrate the lowering to Torch-IR and Aten-IR that occur within Dynamo and the extra operators inserted during the lowering process.
Finally, inductor consumes the output Aten-IR FX-graph and lowers it to custom define-by- run IR, subsequently code-generating necessary kernels specialized to the backend microarchitecture.
Inductor:
3 AUTOSP
AutoSP lifts SP parallelism compiler-patt into the PyTorch-2.0 compiler stack optimize long-context training. Fig. 3 is an overview of how AutoSP's compiler-pant -papssss-passs interoperate with the PyTorch-220 compilation stack to optimize L.M training code. Section 3.1 describes how enable
3

---

Published as a conference paper at ICLR 2026
$ 3.1 SP-Pass
SAC-PaSs
Source
Input Model
W
W
W
W
Imeort forel
linear
linear
linear
matmul
matmul
Zlam teonsformeno Hoeiiiny
Lower Torch-IR
82g
Inseric
a2a
Interted
8.2
Lower to
Other Korward
gradients
Aten-IR
dk
graph
sasateed, Ale ste jersellon
Tensor 203 eatrarted
sdpa
W
Boseline
dw
Irperied
Min
Sink
Transfarmatiion
Analysis
linear
*-+-Aus Delation
Figure 3: An overview of AutoSP. AutoSP enables an automated approach to scale input context lengths for long-conte context training through a targeted set of compiler optimization passes: an automated lequence-paralleele pass (Section 3. 1), and a SP-aware long-context AC-pass: SAC (Section 3.2).
automated sequence parallelism as compiler pass, and Section 3.2 describes our long-context aware activation checkpointing strategy.
3.1 AUTOMATED SEQUENCE PARALLELISM
Challenges. To implement automated sequence parallelism, we must select (a single) PyTorch IR(s) to analyze and transform accordingly. However, recovering the necessary model information through program analysis and subsequently transforming the model through mnmantically-preservii rewrites is non-trivial for three reasons. (1) Each inserted communication collective must have token-buc bakfers instantiated to a particular size parameterized by the
batch-size and sequence length. However, these parameters not explicitly represented in any IR. (2) Existing intermediate buffers need to resized to different shapes depending on the SP group-size as well as their placement within the neural network. For example, in DeepSpeed- -Uysses,, buffers within the attention layer should be resized to operate on the full-sequence but a subset of all the heads, whilst buffers within MLP layers operate the full model-dimension but partial sequence-length. (3) Certain operations require manual indexing for correctness, e.g. indexing the causal mask to appropriately apply attention matrices, and need to be automatically recomputed.
To tackle these challenges, AutoSP's SP-pass accordingly analyzes computational structures to extract the pertinent information required for transforming single-GPU code to distributed sequence-pane parallel code. We next describe how it operates.
3.1.1 ANALYSIS AND TRANSFORMATIONS
Our pass operates in two stages. First, analyze tensor sizes to gather information about the batch, sequence and hidden dimension that parameterize the model. Next, we transform the IR by: (I) inserting communication collectives at appropriate places in the network with appropriately sized buffers to store their output, (2) adjusting the sizes of existing buffers within the graph to account for sequence sharding. (3) Recompute any manually indexed tensors appropriately. All our analysis and transformations operate on Torch-IR.
Why AutoSP analyzes and transforms Torch-IR? Our SP-pass operates Torch-IR the analysis and transformations are significantly more challenging to accomplish on Aten-IR for three reasons. (1) Torch-IR more closely resembles the neural-network programmed by the user (see Section with operators such as linear and attention-layers, making it casier to identify which parts of the graph belong which layer for appropriate tensor resizing. the other hand, equivalent operators in Aten-IR are represented as a series of finer- fner-grained operators, such as mat-muls and permutations, making it challenging to reason about which operators belong to linear- projer-pojestions and attention- operators respectively. (2) The lowering process from Torch-IR to AIen-IR inserts various data- layout transformations such as reshapes and permutes, obscuring information as which dimension corresponds the sequence, batch, and hidden sizes of tensor, making it challenging to appropriately resize the correct tensor dimension. (3) Torch-IR operates on only the forward wwward-pass resulting in our transformations operating simpler computation-grappee merely requiring each new added node
4

---

Published as a conference paper at ICLR 2026
to have registered conjugate gradient operator. On the other hand, Aten-IR operates on both the forward and backwards pass and requires more complex transformations to the computation-gras
Program analysis to uncover training param- eters. To correctly instantiate token-buffeet for communication collectives, the correct batch,
1 b-batch, s=seg, h-heads, d=dim
2 dee transform (mod, h,),
3
for node in mod, mod.nodes:
sequence and model-dimensions need to be ex- tracted. Fortunately, we can analyze the input
4
ws World Size
5
part_seq s/WS
6
Resize buffers. if node node.name name in RESIZE_
nodes of the entire computation graph to extract
7
the necessary information. Since the imput to the computation- graph is the preprocessing, it guaranteed to resemble par-
8
it node. in ATTN_OPS:
9
resize_attn (node, [b, s,wws, d]) else:
after
10
11
resize_othere (node,
ticular shape depending on the problem domain. 12
For example, in natural-language tasks, the data
Recalculate manual indexing
14
if node. in INDEX_OPS; recalc_index (node, node.args) Insert comm. collectives if name is ATTN
will be a [batch, seq_length] -sized ten-
15
sor. Next, to acquire the model-dimension we 16
traverse the graph until encounter atten-
17
tion operator and inspect its output ND tensor
but putjproj_ = Ibatch, sec, heads/s dim]
19
whose last two dimensions are: [num. heads, 20
Insert_befort (node,, all_to_all, buf_proj)
head_di m] -sized. The product of the outer two dimensions the model-dimension.
21
22
elif node name is ATTN_OPS [-1]
23
buf_attn [batch, seq/WS, heads, dim insert_afted (node, all_to_al buf_attn)
Program transtormation. After acquiring the
24
batch, sequence and model-dimensions, we have
25
26
the necessary information to transform the com- putation graph from singleGGUU to distributed sequence prrenceeparallel code. Listing 2 illustrates how we transform the existing computation-graphh mod, comprising of Torch-IR statements. We traverse through the graph, and for each node:
AutoSP's Transformation Pass: converting single- code to sequence-parallel multi-GPU code. The pass resizes buffers, recomputes manually indexed tensors, and inserts communication collectives.
Listing 2:
(I) Check if it belongs to the RESIZE BUFS set, and accordingly resize its buffers depending (2) Check if it belongs in the (3) Check if it is the first/last attention-op
on its placement in the attention or linear- linear-projection/ projection/MLP-laeertt INDEX_OPS set, and accordingly resize its tensor indexing. and accordingly instantiate communication buffers and insert the necessary a alltoall before/after the attention-layer. We manually curate the RESIZE BUFS, INDEX OPS, and ATTN CPS sets by analyzing dynamo FX-graphs of compiled hand-written transformer implementations.
3.2 SEQUENCE- AWARE ACTIVATION CHECKPOINTING
Challenges. In addition to the SP-pass, activation checkpointing (AC) is an important memory reducing optimization that enables longer context training. torch. compi le provides an automated AC-pass (Chillee) that operates on Aten-IR to compose with arbitrary neural networks. However, naively composing its AC-pass with AutoSP's s SP-pass within the compiler-stack leads to sub-optimal performance for long-context training. We briefly explain how PyTorch-2.0's automated AC-pass functions, and why it is insufficient.
Py Tyrch-2..0's AC-s PyTorch's AC-pass operates on Aten-IR, within Dynamo. Its primary function is to reduce memory consumption of model training without incurring performance penalties. The optimization reduces the problem to a network-flow construction whose min-cut determines the tensors to rematerialize. We give an example code-snippet and its equivalent network- construction in Fig. 4. The input to the optimization is an FX-graph comprising of Aten-IR statements. First, a joint-graph) of the forward and backwards graph is constructed. Next, the source node is connected to all the input tensors of the graph, and all the nodes reachable from the incoming gradients are connected to the sink. Then, capacities, representing costs, are assigned to each node. node's capacity is determined by a heuristic function comprising of various characteristics such as the output activation memory produced. Finally, the problem is converted from flow on nodes to a flow on edges. Each edge's 's capacity is set to inf and each node is split into two: an .in and .out node. Incoming edges to each node are connected to its respective .in node and outgoing edges are connected to its respective out node. An edge with the original node's capacity connects in to out. A min-cut this graph will cut only finite-capacity edges from node's 's. in out
5

---

Published as a conference paper at ICLR 2026
(highlighted by green circles in Fig. 4); only nodes on this cut stored. Intuitively, the min-cut identifies the smallest cost set of activations to preserve to compute the necessary dependencies the backwards pass. PyTorch-2.0 additionally enforees constraints which nodes can be rematerialized, resulting in its conservative nature which we explain next.
Why PyTorch-2.0's AC solution is insufficient. PyTorch22.0's AC-pass is effective identifying without compromising run-
ded modde 2. Loathh opa uben aa dww. R.
which activations time performance in polynomial time. However, it makes the conservative decision to disallow remateri- alization of many seemingly compute- intensive opera- tors such as: mat-muls, and scaled mat-muls to name few, based purely on operator type, without con- sidering runtime cost in long-context settings. This is enforced by connecting each compute-intensive node's i to the source with infinite capacity (in- dicated by the additional dotted red line in Fig. 4), enforcing that node, or some downstream value of belong to the min-cut. However, in long-context training, we observe that certain compute-intensive operators take a small fraction of overall compute and can accordingly be rematerialized without incur- ring performance penalties. AutoSP's AC strategy exploits this. We explain these observations next.
return torch wpp.. aree mut dedarii
Source
min-cut
Addttionl Bige
Pytorsh
Auso store for bod
metmul
gradients
mul
Las cocoda
Oroph
Sink
output
Figure 4: A comparison of PyTorch-2.0 vs. AutoSP's AC-passes on sample code- snippet. The red line is the additional edge from source to node in PyTorch-2.0's AC- pass enforce rematerialization of com- pute heavy operators. AutoSP's AC-pass, in removing this constraint, reduces memory- consumption at negligible cost to throughput.
Observations. We analyze the structure of com- pute operations in modern LLMs to identify regions that can be appropriately rematerialized. For trans- former with: batch-size b, sequence-length number of heads h, head dimension d, and MLP hidden dimension drras we have that:
2bhs à Attention FLOPS
@bhsd2 Linear-projection FLOOPS
4bhsdren d MLP FLOPS
When training on long-contexts, we have that s>>d, h, dern, which results in the following fraction of FLOPs linear-projection and MLP layers take over all the compute operations:
Bbhsd' + 4bhsdarad
+ oc
(I)
2bhs 8bhsd 4bhsdrrad
Indicating that the fraction of compute operations for linear-projection and MLP matrix- multiplications decreases as a function of input-sequence length. This observation underpins AutoSP's automated AC strategy.
AutoSP's AC strategy. AutoSP exploits equation 1, building upon PyTorch-2's automated AC strategy. However, instead of conservatively banning every compute-intensive operator, we permit configurations where (batch) matrix-multillications and other compute-intensie operators outside of the attention layer are rematerialized. We achieve this by iterating over the joint-graph and removing any additional edges from the source to compute-heavy operators, resulting in only inputs to the graph (tensors a and in Fig. connecting the source. We then dispatch this mutated foint-graph to PyTorch-2.0's AC strategy. This change enables traning on significantly longer context lengths at negligible cost to training throughput.
EVALUATION
We evaluate AutoSP with comprehensive set of experiments. We demonstrate its effectiveness in enhancing trainibility of various models and sizes in Section 4.1, and detailed breakdowns of the impact of each component in Section 4.2.
Setup. We evaluate AutoSP and all the baselines on NVIDIA GH200-96GB & A100-80GB and AMD MI250-64GB hardware. All experiments use PyTorch-2.7 with CUDA 12.8 (on NVIDIA

---

Published as a conference paper at ICLR 2026
GPUs), and ROCm 6.4 (on AMD GPUs). To implement AutoSP, lift the CrapSpeed-Ulsseeee SP scheme into PyTorch-2.0's compilation stack and integrate all our compiler optimizations into the DeepSpeed project, due to its popularity in training large scale LLMs.
Baselines. We compare AutoSP to both compiler-optimized distributed training solutions and hand-optimized SP solutions demonstrate the memory and compute efficiency of our approach. Specifically, we compare against ZeRO-3 (FSDP) Rajbhandari et al. (2020) optimized through torch. compile
PyTorch-2.0, and hand-written DeepSpeed-Ulysee (Jacobs et al., 2024) & RingAttention implementations when compiled under PyTorch-2.09 inductor backend. We use which
DeepSpeed SeSSeee-UUyssss'' code in the original DeepSpeed repository' and RingFlashAttentinn are both highly optimized implementations of SP strategies. We evaluate all techniques on range of model sizes: I.lama-322 1B & 3B, Llama-3.1 8B, and Ilama-2 13B, covering models with either Grouped-Query-Atrer Query-Attention (GQA) or Full-Attention.
4.1 MAIN RESULTS
In this section, we evaluate how AutoSP impacts model trainability: the maximum trainable sequence length prior encountering OOM issues.
Trainability. For different techniques, we measure the maximum sequence length trainable prior OOM on 8 NVIDIA A100-80GB in Fig. For all models, we use ZeRO-1, setting the SP group-size to 2 and the DP group-size to 4 with the exception of the ZeRO-3 augmented with torch. compilipiiioooooooc iic 5.6x and 2.5x longer input sequences for the 3B, 8B and 13B models, respectively. The trainability gains come from AutoSP's compiler-based SP-pass, an optimization that targets long-context training unlike the ZeRO-3 baseline, which instead targets models with large parameter counts. Moreover, AutoSP achieves significant trainability gains compared to both inductor compiled DS-Ulysses and RingAttention implementations. Compared to DS- DS-UUlssess, AutoSP enables training longer input sequences by upto 2.14x, 3 x and 1.88 for the 3B, 8B and 13B models, respectively. Compared to RingAttention, AutoSP enables training on upto 2.14x, 3x and 1.6x for the 3B, 8B and 13B models, respectively. The additional gains over hand-written SP implementations come from the SP-aware AC-pass that exploits equation rematerializing compute-heavy operators (such mat-muls) for low runtime costs but large memory gains. The trainability gains are especially pronounced for the 8B model compared to 3B as many the campute-heavy operators, such as linear- proearrpojections and MLPs, produce activations parameterized by the model dilden-dimeesion. These need to be stored to compute gradients in the backwards pass and result in larger models producing significantly more activation memory due to compute-heavy operators. AutoSP, in rematerializing these large tensors, alleviates memory issues at negligible runtime cost. The trainability gains are less pronounced for 13B model optimizer states begin to consume a substantial portion of memory (~50%).
baseline. Compared ZeRO-3 (FSDP), AutoSP enables training on upto 5x,
75K
205
ZeRO-3
DS-Ulysses AutoSP
199
NingArtention DS Ulysses AutoSP
RingAttn
Seq. Length
195
60K
45
45K
40K
35K
29
26
31
Time
20K-
15K
8K
15K
6K 9K
10 9
000 X
OOM
X
3B
8B
13B
24K
90K
Figure 5: Maximum sequence length prior to OOM across various model sizes. AutoSP in- creases the trainability of all model sizes.
Figure 6: End-to-End training iteration times (over 10 iterations) for Llama-3.1 8B on GPUs. AutoSP maintains similar per-training it- eration time whilst significantly increasing train- ability.

---

Published as a conference paper at ICLR 2026
Runtime Performance. For different techniques, identify the impact on training iteration times at vari- sequence lengths in Fig. 6. We focus Llama- 3.1 8B model, using ZeRO-1 with a fixed DP-size of 2 and SP-size of 4, and evaluate 8 GPUs, measur- ing time taken for 10 training iterations. We toggle the sequence-aware AC pass for AutoSP only if
DS DS-Ulysses
AutoSP
150K
Length
100K
92K
81K
Max Seq.
50K
50K
OOM issues occur. Compared to the ZeRO-3 base- line compiled with torch. compile (). AutoSP reduces per-iteration times by 5x whilst increasing trainability by an order of magnitude. ZeRO-3 intro- duces more communication collectives over 7eRO-1 with SP, resulting in slower per-iteration times. More- over, compared to RingAttention and DS-Ulysses, AutoSP increases trainability by 3.75x with negligi- observe
50K
20K
(NVIDIA GH200)
(AMD MI250)
Comparing the max sequence length prior to OOM across different hard- ware. AutoSP enables training on longer se- quences on NVIDIA superchips and AMD hardware.
Figure 7:
ble cost muntime-performancee Finally, that DS-Ulysses is faster than RingAttention due to RingAttention's s p-step (where p the SP size) com- munication latency which exchanges key/values in ring-like pattern across SP-groups. Comparitively, DS- DS-Ulysses introduces single all-to-all to re- shard token-buffers. In modern clusters, intra-node all-to-alls are fast due to pairwise communication links between devices.
4.2 ANALYSIS
In this section, additional studies to demon- strate the effectiveness of AutoSP, as well as to as- certain the impact of cach optimization different components LLM training. We compare AutoSP with DS- DS-Ulysses only (rather than Ring/tttentio)) given it enables training on similarly sized context- lengths whilst being slightly faster.
DS-Ulysses AutoSP
Llama-3.2
1.5
6.0
Time (s)
2.0
Trainability on diverse hardware. We run an additional trainibility study on different NVIDIA
0.0
0.0
3
4
3
and AMD hardware. identifying the maximum se- quence length before encountering an OOM Fig. 7
Llama-3.2
3B
6.0
to demonstrate AutoSP's portability. We focus on smaller, 1B and models, running on either GH200-96GB (NVIDIA) GPUs, or 2 AMD MI250- 64GB GPUs. On NVIDIA hardware, AutoSP en- ables training 1.58x (2.70x) longer sequence lengths on the IB (3B) models respectively. On AMD hardware, AutoSP enables training 2x (2.5x) longer input sequences on the 1B (3B) models, re- spectively. AutoSP consistently delivers significant trainability gains across diverse hardware, and model sizes through its targeted long-context optimizations.
Time (s)
4.0
2.0
2.0
1.0
OOM
OOM
Seq. Length (10k)
Seq. Length (10k)
(NVIDIA)
(AMD)
Figure 8: Average execution time of vari- ous Llama 3.2 model sizes at different se- quence lengths NVIDIA and AMD hard- ware. AutoSP matches the performance of hand-written baselines and supports longer sequence training.
Runtime performance on diverse hardware. For different techniques, we measure the per-iteration end-to-end training time (averaged across 100 train- ing iterations) on NVIDIA (GH200-96GB) and AMD MM250-64GB) hardware in Fig. 8, demonstrating the performance-portabiiin of our approach. We trigger the SP-aware AC-pass only to avoid OOM issues for AutoSP. On sequence lengths that all techniques can train on, AutoSP has the following speedups: 0.97 (1B) and 0.98x (3B) compared to the inductor baseline on NVIDIA hardware, respectively. On AMD hardware, AutoSP has the following speedups: 0.97x (1B) and 0.87 (3B) compared to the inductor baseline, respectively. We note two observations. Despite being a general and cereformanne--rrtaaaaa compiler pass, AutoSP achieves 97% of DS-Ulysses' hand-written baseline
8

---

Published as a conference paper at ICLR 2026
whilst providing upto 2.7x trainability gain. Without AutoSP's targeted optimizations, training long-contexts quickly becomes infeasible.
Breakdown analysis. We breakdown the impact of AutoSP's optimizations in Fig. on a NVIDIA GH200-96GB. We breakdown the activation memory produced by the attention and MLP operators as well the per-iteration runtime of the forward and back- ward passes when training a Llama-3.2 1B model using sequence length of 40k. Overall, AutoSP reduces memory consumption of the attention and MLP operators by 13.03x and 2.22x, respectively. The marked impact on MLP operators arises due to the presence of many mat-muls, which AutoSP's AC- pass rematerializes to alleviate memory consumption at low runtime costs. On the other hand, AutoSP incurs 1.14x cost runtime performance cost for the backward pass, whilst having similar forward pass times result of the extra rematerialized operators, which are recomputed in the backwards pass only.
25
Attention
Fwd Pass
Act. Mem. (GB)
20
MLP Op
1.25
Bwd Pass
1.00
15
15.9
0.75
0.90
0.79
0.50
5
0.25
5.53
1.22
0.28
0.30
2.48
0
0.00
DS-Ulysses AutoSP
DS-Ulysses AutoSP
Figure Breakdown of Attention and MLP operator memory consumption and forward iteration times. AutoSP reduces the activation memory of Attention and MLP operators with a marginal performance difference.
A Ablation study. Finally, we toggle different optimizations on/off to demonstrate their impact trainability and speed in Table 1. The top row indicates the maximum trainable input context length before an OOM, and the bottom row indicates the average training iteration time at fixed 40k input sequence length. Overall, AutoSP's optimizations result in a 1.58x trainability gain whilst achieving 89% of CrepSpeed-Ulysseeee
per-iteration training time. Moreover, the incremental trainability gain of the AC-pass over the SP-pass is 1.66x with a mere 7% decrease in runtime performance.. This significant trainability gain with little runtime-performance cost due to equation 1. longer contexts, the attention-operatoro FL.OPs dominate runtime, enabling traditionally compute-heavy operators to be rematerialized with marginal performance costs and large memory gains. Note that the baseline is highly hand-optimized SP implementation with several optimizations such as cramuunccationcoompaaaa overlapping using streams. Nevertheless, the SP-pass achieves 97% of the baseline's 's performance a general compiler-pass.
Table I:
Various optimizations in AutoSP against a baseline on Llama-3.1 1B. AutoSP supports training on significantly longer sequence lengths at minimal cost to performance.
Method
Max Tokens
Speed (s)
DS-Ulysses
81,000
1.06
SP-Pass
77,000 128,000
1.09
AutoSP
SP & AC-Pass
1.19
5 RELATED WORK
Parallel training strategies. ZeRO-3/FSDP (Rajbhandari al., 2020), Tensor Shoeybi al. (2020) & Pipeline (Huang et al.. 2019) parallelism are training strategies that target models with large parameter counts, reducing per-device memory consumption of optimizer, model, activation and gradient states. Expert-parallelism (Lepikhin et al., 2020; DeepSeek-AI et al., 2025; Dai et al., 2024) targets large sparse mixtures--ffexxeer which contain many intermediate expert MIPss Though effective in enhancing the trainability of large-language model training. these parallel strategies do not explicitly target long-context training and are insufficient to scale input context lengths.
Automated optimizations. Deep-compile (Tanaka et al., 2025) provides automated approach to implement SERO-3/FSDD using profile-guidee optimization. Though effective, FSDP does not explicitly target long-com oon-context training. Seneral-Single-- Proggleerrrram---- Program-Multiple Data (GSPMD) (Xu et al., 2021), 1), is an automated parallelization strategy in XLA guided through user annotations, requiring some human effort. Lastly, deep-learning compilers such as: TVM (Chen et al., 201 1018) Mirage (Wu

---

Published as a conference paper at ICLR 2026
al., 2025), and AlTemplate (Meta, 2022), focus schedule rewrites for inference workloads only and do not consider Inter-GPL parallelism, a key optimization for long cong-contett training.
Activation checkpointing. Various works propose AC techniques (Jain al., 2020; Kirisame et al., 2021; Chen et al., 2016). They primarily consist of two approaches. (1) Search- based optimization (e.g. via innger--neeuurgggramnngg models). (2) Static-pold Static-policies (e.g. checkpoint chunks of VN layers' activations), which may result in extraneous communication calls in the backward pass in the SP-setting. Comparatively, our SP-aware AC-pass exploits observations of compute & memory properties of L.LM training at long-contexts alleviate memory consumption at negligible cost to runtime performance. Additionally, AutoSP's sequence-aware AC pass is implemented in PyTorch-2.0's compiler stack, requiring no human intervention. In contrast, existing AC strategies often require invasive code-rewrites.
which may scale up to today's L.M sizes (billion parameter
CONCLUSION
In this paper, we present AutoSP, the first compiler- based, PyTorch-native solution for training lrge-languagge-model at long-contexts. Through combination of automated sequence-parallelisn (SP) and a sequence-aware AC strategy, AutoSP achieves significant sequence length extensions at negligible to training throughput. Our results demonstrate that compiler-driven, PyTorch-native automation provides a practical and portable foundation for long-context model training.
REPRODUCIBILITY STATEMENT
We have made several efforts to ensure the reproducibility of our work. We have documented the critical components of our implementation in Section 3 and our evaluation in Section 4 additionally documents our setup with a description of our hardware and software versions. We have ensured that each result is consistent with good benchmarking practices, including taking the average multiple runs. Moreover, all our code and benchmarks will be made publicly available to further enhance reproducibility.
ACKNOWLEDGMENTS
We sincerely appreciate the anonymous reviewers. Their insightful feedback helps significantly improve the quality of the paper. This research was supported by the National Science Foundation (NSF) under Grant No. 2441601. The work utilized the Delta and DeltaAl system at the National Center for Supercomputing Applications (NCSA) and Jetstream2 at Indiana University through allocation CIS240055 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) program, which is supported by National Science Foundation grants #2138259, #2138286, #2138307, #2137603. and #2138296. The Delta advanced computing resource collaborative effort between the University of Illinois Urbana-Champar and NCSA, supported by the NSF (award OAC 2005572) and the State of Illinois. UIUC SSAIL Lab supported by research funding and gift from IBM, Google, Amazon, and AMD.
REFERENCES
Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia. Will Constable, Alban Desmaison, Zachary Vito, Elias Ellison, Will Feng, Jiong Gong, Michael Gschwind, Brian Hirsh, Sherlock Huang, Kshiteej Kalambarkar, Laurent Kirsch, Michael Lazos, Mario Lezcano, Yanbo Liang, Jason Liang, Yinghai Lu, C. K. Luk, Bert Maher, Yunjie Pan, Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk, Shunting Zhang, Michael Suo, Phil Tillet, Xu Zhao, Eikan Wang, Keren Zhou, Richard Zou, Xiaodong Wang, Ajit Mathews, William Wen, Gregory Chanan, Peng Wu, and Soumith Chintala. Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, ASPLOS '24, pp. 929-9499999 997, New York, NY, USA, 2024.
10

---

Published as a conference paper at ICLR 2026
Association for Computing Machinery, ISBN 9798400703850. doi: 10.1145/362006553446666 URL https:://///.. /toii/ooooo////oo//.oo wwtg/og.oo....oooooooooo 1345/3111153333 36045556666665666 3640366.
Ge Bai, Jie Liu, Xingyuan Bu, Yancheng He, Jiaheng Liu, Zhanhui Zhou, Zhuoran Lin, Wenbo Su, Tiezheng Ge, Bo Zheng, and Wanli Ouyang. 1i-bbench-111: A fine-grained benchmark for evaluating large language models in multi-turn dialogues. In Proceedings of the 62nd Annual Meering of the Association for Computational Linguistics ( (Volume I: Long Papers), pp. 7421-7454, Association for Computational Linguisties, 2024. doi: 10.18653/v1/2022
aww.o2...looon..4.4 401. URL http:// doi. org/10..//..oooo...o.ooo. 18653866//808 /01//22222222/ acl acl-long....................mm 401.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost, 2016. URL https: //arii//aiivv.a/ wwwpxiiioorggaaooe org/abs/ 1644. 06174.
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang. Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. TVM: An automated End-to-End optimizing compiler for deep learning. In 13th USENIX Symposium Operating Systems Design and Implementation (OSDI 18), pp. 578-594, Carlsbad, CA, October 2018. USENIX Association. ISBN 978- 1-9911939-33----. URL https:/ www. usenix. org/ cenfeerenee/sdiii wtw.ele///eeeeeeeeeeeeeeeeeeeeee at atiin/ccoheeeeeeeeee /hen...
Chillee. Min-cut optimal recomputation (i.e. activation checkpointing) with aotautograd.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reicchrr Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. URL https: //aaxiv.orrg/ abs/2110. 14168.
Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experrs language models, 2024. URL https: /arxivvvv. org/ abs 2201.4011100 06066.
DeepSeek-AI, Aixin Liu, Bei Feng. Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Han Bao, Hanwei Xu, Haocheng Wang, Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Li, Hui Qu, J. L. Cai, Jian Liang. Jianzhong Guo, Jiaqi Ni, Jiashi Li, Jiawei Wang, Jin Chen, Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li, Junxiao Song, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang. Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao, Litong Wang. Liyue Zhang, Meng Li, Miaojun Wang, Mingchuan Zhang. Minghua Zhang, Minghui Tang, Mingming Li, Ning Tian, Panpan Huang, Peiyi Wang, Peng Zhang, Qiancheng Wang, Qihao Zhu, Qinyu Chen, Qiushi Du, R. J. Chen, R. Jin, Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, Runxin Xu, Ruoyu Zhang, Ruyi Chen, S. S. Li, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shaoqing Wu, Shengfeng Ye, Shengfeng Ye, Shirong Ma. Shiyu Wang. Shuang Zhou, Shuiping Yu, Shunfeng Zhou, Shuting Pan, T. Wang, Tao Yun, Tian Pei, Tianyu Sun, W. L. Xiao, Wangding Zeng. Wanjia Zhao, Wei An, Wen Liu, Wenfeng Liang, Wenjun Gao, Wenqin Yu, Wentao Zhang. X. Q. Li, Xiangyue Jin, Xianzu Wang, Xiao Bi. Xiaodong Liu, Xiaohan Wang, Xiaojin Shen, Xiaokang Chen, Xiaokang Zhang, Xiaosha Chen, Xiaotao Nie, Xiaowen Sun, Xiaoxiang Wang. Xin Cheng, Xin Liu, Xin Xie, Xingchao Liu, Xingkai Yu, Xinnan Song, Xinxia Shan, Xinyi Zhou, Xinyu Yang, Xinyuan Li, Xuecheng Su, Xuheng Lin, Y. K. Li, Y. Q. Wang, Y. X. Wei, Y. X. Zhu Yang Zhang. Yanhong Xu, Yanhong Xu, Yanping Huang, Yao Li, Yao Zhao, Yaofeng Sun. Yaohui Li, Yaohui Wang. Yi Yu, Yi Zheng, Yichao Zhang, Yifan Shi, Yiliang Xiong, Ying He, Ying Tang, Yishi Piao, Yisong Wang, Yixuan Tan, Yiyang Ma, Yiyuan Liu, Yonggiang Guo, Yu Wu, Yuan Ou, Yuchen Zhu, Yuduan Wang, Yue Gong, Yuheng Zou, Yujia He, Yukun Zha, Yunfan Xiong, Yunxian Ma. Yuting Yan, Yuxiang Luo, Yuxiang You, Yuxuan Liu, Yuyang Zhou, Z. F. Wu, Z. Z. Ren, Zehui Ren. Zhangli Sha, Zhe Fu, Zhean Xu, Zhen Huang. Zhen Zhang, Zhenda Xie, Zhengyan Zhang, Zhewen Hao, Zhibin Gou, Zhicheng Ma, Zhigang Yan, Zhihong Shao, Zhipeng Xu, Zhiyu Wu. Zhongyu Zhang, Zhuoshu Li, Zihui Gu, Zijia Zhu, Zijun Liu. Zilin Li. Ziwei Xie, Ziyang Song, Ziyi Gao, and Zizheng Pan. Deepseek-v3 technical report, 2025. URL https:ttp//......... / rxxivvvve. .rgg abs 2212 19437.
11

---

Published as a conference paper at ICLR 2026
Yanlin Feng, Xinyue Chen, Bill Yuchen Lin, Peifeng Wang, Jun Yan, and Xiang Ren. Scalable
multi-hop hop relational reasoning for knowledge-awan aware question answering, 2020. URL https:
//aax//aav.//////////eee/aaeeeva wrg. abs /0t..//000000055050
Siwei Han, Peng Xia, Ruiyi Zhang. Tong Sun, Yun Li, Hongtu Zhu, and Huaxiu Yao. Mdocagent: A multi-modal multi-agent framework for document understanding, 2025. URL https::tp///o /tatps://arxivv org/abs/ 2503. 13964.
Yanping Huang. Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2019. URL https: /ttpaaaxiiv./..////ii//eaaeaaaav orq/ /bs/////////// /1811. 06965.
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Reza Yazdani Aminadabi, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. System optimizations for enabling training of extreme long sequence transformer models. In Proceedings the 43rd ACM Symposium on Principles of Distributed Computing. PODC '24, pp. 121-130, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400706684. doi: 10.1145/36621583 158. 156.6668668...66 3313336621555 58. 1...3366288 3662806.
URL https::////.......pp...pp...tt.t //////oo.... wrg/10......
Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Kurt Keutzer, lon Stoica, and Joseph E. Gonzalez. Checkmate: Breaking the memory wall with optimal tensor rematerialization, 2020.
wwt.://rgggb/b////o/o//oo/ss awss/s19.... 02653.
Marisa Kirisame, Steven Lyubomirsky. Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, and Zachary Tatlock. Dynamic tensor rematerialization, 2021. URL https: // ///rxii..orre hbtpg/abs22000 09616.
Jordy Van Landeghem, Rubén Tito, Lukasz Borchmann, Michal Pictruszka, Pawet Józiak, Rafal Powalski, Dawid Jurkiewicz, Mickaël Coustaty, Bertrand Ackaert, Ernest Valveny, Matthew Blaschko, Sien Moens, and Tomasz Stanislawek. Document understanding dataset and evaluation (dude), 2023. URL https: : / arxia. org abs 22005/22335..22 08455.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding, 2020. URL https: tettps://arxiv/ www.aiii.orrg//.oe..aaoooooaoolt abs 2006 1 16668.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near- infinite context, 2023. URL UtLhttps: //tp:v//aaiiii.a org. wrtp:oo///aabaaagaee 23100 01889.
Meta. AlTemplate: Python framework which renders neural network into high performance CUDA/HIP C++ code., 2022. URL httpp://gt thub. htt.bb.ooofaccoc com/faceebokk www.bbboooknuubttoo..... ncubbtor/
ATTemplate.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models, 2020. URL https: /ttp/aaavvv//.e//////eee/e/.eav/ 02054.
/ttp//iiw...aaaaaaaa/aa wwt./b/gg/aaaaeeeooooo..oo.eee/ 19900.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catan- zaro. Megatron-Im: Training multi-billion parameter language models using model parallelism, 2020. URL https:ttt... //aaaiiivv.... wwg. abs 909. 08053.
Masahiro Tanaka, Du Li, Umesh Chand, Ali Zafar, Haiying Shen, and Olatunji Ruwase. Deepcompile: compiler-driven approach optimizing distributed deep learning training, 2025. URL https: // //arxivv.eo. / abs / 2504. 09983.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cris- tian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux. Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihayloy, Pushkar Mishra,
12

---

Published as a conference paper at ICLR 2026
Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, lliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023. URL https: / arxiv. wrg/ abs 2307.00000 09288.
Mengdi Wu, Xinhao Cheng, Shengyu Liu, Chuman Shi, Jianan Ji, Kit Ao, Praveen Velliengiri, Xupeng Miao, Oded Padon, and Zhihao Jia. Mirage: A multi-level superoptimizer tensor programs. In 19th USENIX Symposium on Operating Systems Design and Implementation (OSDI 25), Boston, MA, July 2025. USENIX Association.
Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake Hechtman, Yanping Huang, Rahul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, Marcello Maggioni, Ruoming Pang. Noam Shazeer, Shibo Wang, Tao Wang, Yonghui Wu, and Zhifeng Chen. Gspmd: General and scalable paralleliza- tion for ml computation graphs, 2021. URL https: h/wxixxxvvvoo./// org abs /215 04663.
Ruisi Zhang, Tianyu Liu, Will Feng, Andrew Gu, Sanket Purandare, Wanchao Liang, and Francisco Massa. Simplefsdp: Simpler fully sharded data parallel with torch.compile, 2024. URL https: //arxiv. wwwpiiivwwrr//////e/ooooooo /211 00284.
Zhaoqing Zhu, Chuwei Luo, Zirui Shao, Feiyu Gao, Hangdi Xing, Qi Zheng, and Ji Zhang. A simple yet effective layout token in large language models for document understanding, 2025, URL fittps://aret http//awvvv.rg/.e//...// htgg/bb///at 2503. 8434.
13


Reviews:
[{"review_id": "Nfo906jFLH", "reviewer": "Reviewer_K3Y6", "summary": "The paper introduces AutoSP, a compiler-based system designed to optimize the training of large language models (LLMs) for long-context scenarios. AutoSP automates sequence parallelism and activation checkpointing specifically tailored for long-contexts, aiming to enhance trainability without sacrificing runtime performance. The system is evaluated on both NVIDIA and AMD hardware, showing significant improvements in trainable sequence lengths compared to existing methods.", "questions": "see weakness above.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "The paper presents a compiler-based solution for optimizing LLM training in long-context scenarios, which maybe easy to use.", "weaknesses": "1. Lack of novelty. Automatic parallelism has been thoroughly studied and there are lots of work about the automatic or dynamic sequence parallelism including selective activation checkpointing. This article integrates these elements into the compiler, which more like an engineering project.\n 2.The paper's baseline comparisons may not fully represent the current state-of-the-art techniques, particularly in terms of hand-optimized implementations. \n3.Lots of errors. Such as line 328 , “Grouped-Query-Attention (GQA) or Full-Attention” GQA and Full-attention are completely compatible.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "EqpiCaKKpr", "reviewer": "Reviewer_SrRm", "summary": "The paper introduces two compiler passes to achieve sequence parallel. It inserts all-to-all collectives and transform single GPU model into sequence parallelsized models. It also consider activation checkpointing that selectively rematerializes activations by finding best memory and perf trade offs. Similar to existing work like simpleFSDP and deepcompile, it capture the graph single GPU model and transform the IR with collectives and IR passes. It mains simple UX and single-GPU style authoring.", "questions": "Explain why it's non-trivial to come up with IR passes according to open source eager SP implementation like DeepSpeed-Ulysses\n\nFor maximum sequence length, analyze memory snapshot for each baseline and show more insights into memory usage: when did the peak happen, % of memory on model/opt state, activation, and intermidate tensors", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "fair originality: this is indeed the 1st paper I saw to use compiler pass to achieve sequence parallel. DeepSpeed-Ulysses and RingAttention implement SP at the framework level that affects the single gpu authoring. This paper lifts SP into the compiler layer. Joint optimization with activation checkpointing is non-trivial and essential to make good memory perf trade offs. \n\nenough quality: it touched the reasoning behind choosing which layer of IRs, and how collectives are inserted and optimized for better perf.   Evaluation is done on both NVIDIA and AMD GPUs, including eager and compiler baselines\n\ndescent clarify: The main technical sections (Sections 3.1–3.2) may challenge readers unfamiliar with PyTorch internals. Fortunately figures and intros effectively lowed the bar to keep up with the content.\n\nhighly significant: it's critical to showcase how to use IR passes to achieve SP. Long-context LLM training is a critical challenge. It's espectially benifitial to minimal developer friction. This could potentially influence future designs of pytorch compiler solution for parallelsims, like simpleFSDP", "weaknesses": "originiality mainly comes from using IR to implement SP. SP itself was previously implemented in open source library like DeepSpeed-Ulysses. There are also similar work that moves FSDP into IR passes. This paper is more like an extension of the idea to more parallelsims. \n\ntechnical depth: it remains questionable how non-trivial it is to come up with IR passes to achieve SP, considering we have open source implementation in eager mode. The non-trivial evaluation should be done for people with enough understanding of pytorch 2 compiler stack. But I agree joint optimization with activation checkpoint is non-trivial", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "YSb0leqJCJ", "reviewer": "Reviewer_AhGd", "summary": "This paper introduces AutoSP, a novel compiler-based system for PyTorch-2.0 that automates optimizations for training Large Language Models (LLMs) with very long contexts.", "questions": "1.  Why was DeepSpeed Ulysses selected as the baseline? Does it exclusively support eager-mode, $O(N^2)$ attention? If so, it is an inappropriate baseline. Standard $O(N^2)$ attention is outdated; modern implementations using Flash Attention or PyTorch SDPA can easily train 16K contexts on a single 80GB GPU, especially with DeepSpeed ZeRO-3 enabled. A valid comparison would require using faster context-parallel (CP) implementations, such as Megatron-LM's CP or Ring Flash Attention [1].\n\n2.  Using `torch.compile()` for activation checkpointing is uncommon in practice. Standard LLM implementations (e.g., in Hugging Face) wrap each layer with `torch.utils.checkpoint`. This standard, LLM-specific approach should be discussed. The official PyTorch API is not complex:\n\n    ```python\n    torch.utils.checkpoint.checkpoint(layer_function, *inputs, use_reentrant=False)\n    ```\n    Please clarify the efficiency difference between your activation checkpointing method and this standard `torch.utils.checkpoint` implementation.\n\n[1] https://github.com/zhuzilin/ring-flash-attention", "limitations": "", "rating": 2, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "The article's writing is commendable, particularly for its very clear and accessible explanations of technical details.", "weaknesses": "1.  I am skeptical of the core claim that context parallelism is difficult to implement. An API-based approach, inspired by Flash Attention (e.g., `flash_attn_func(query, key, value)`), seems more suitable than integrating it to `torch.compile()`. While also a one-line modification, this API provides users with greater transparency and explicit control, rather than obscuring the underlying logic.\n\n2.  The comparison between the proposed compile-stage activation checkpointing and traditional layer-wise checkpointing is vague. The paper primarily demonstrates advantages over a standard `torch.compile()` baseline but fails to clearly articulate its benefits over conventional layer-wise gradient checkpointing. This is a significant omission, as the latter is a common, effective, and simple-to-use practice in LLM training.\n\n3.  Most critically, the ideas presented are neither groundbreaking nor particularly novel. They appear to be targeted optimizations for implementing a specific function more efficiently within an existing general purpose framework, rather than a substantive innovation. As such, this contribution seems more appropriate for a pull request to the PyTorch repository than a publication at ICLR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "SP6uxjDtO8", "reviewer": "Reviewer_4FNs", "summary": "The compilation of a PyTorch 2.0 model goes through multiple stages. Dynamo, the stage responsible for running the model for the first time and recording the computation graph, can be split into two passes: Torch-IR (a higher-level graph that still resembles model layers) and Aten-IR (a lower-level graph of primitive tensor ops, e.g. matrix multiplies, permutes, convolutions, data moves). The authors present a PyTorch compiler patch integrated into the DeepSeek library that implements two optimizations:\n\n1. Takes the implementation of sequence parallelism from DeepSpeed-Ulysses into the Torch-IR pass, so that the attention layers can be automatically distributed across GPUs for parallel processing of long context sequences. The compilerized version works with arbitrary PyTorch models with minimal code changes.\n\n2. Adds a layer of optimization to the Aten-IR pass, optimizing activation checkpointing (AC). They show that PyTorch’s stock AC is too conservative for long-sequence training because it forbids rematerializing classically compute-heavy ops (e.g., matmuls, convs). They show that as sequence length grows, MLP/linear matmuls constitute a vanishing fraction of total FLOPs, so these “heavy” ops can be rematerialized cheaply. By removing this restriction, they enable a longer context to fit into the memory.\n\nThey conduct an extensive empirical study, demonstrating that their implementation \n- is compatible with different types of hardware (NVIDEA, AMD), \n- allows long sequences to fit into memory when compared to ZeRO-3 and hand-written DS-Ulysses, \n- Scales well with SP group size", "questions": "- Comparing AutoSP to other context-length extension strategies (e.g., HuggingFace Accelerate with Context Parallelism) would make the results more trustworthy", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The greatest value of this work is that it automates the implementation of SP for practitioners, requiring minimal code changes. This has the potential to accelerate research and development of language model architectures with a wide impact. The authors present an original contribution of an optimized activation checkpointing strategy and provide empirical results that demonstrate the effectiveness of the proposed changes. The paper is clearly written and provides sufficient documentation on the empirical experiments.", "weaknesses": "- The paper lacks low-level technical details about how AutoSP manipulates Torch-IR in complex models beyond toy examples. The authors should make sure this is addressed during the release of the source code.\n- The paper overemphasizes the benefits of AutoSP without detailing potential downsides, such as the scenarios where the overhead of recomputation in AC becomes substantial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)