AutoSP: Unlocking Long-Context LLM Training Via Compiler-Based Sequence Parallelism

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.

• The paper introduces AutoSP, a compiler-based framework that automates sequence parallelism and activation checkpointing to extend trainable sequence lengths up to 5.6× over baseline methods.
• The methodology automates tensor shape analysis and communication collective injection, significantly reducing activation memory usage on both NVIDIA and AMD hardware.
• Empirical results demonstrate that AutoSP maintains 97% of native speed with minimal throughput loss, making long-context LLM training more practical and portable.

Compiler-Based Sequence Parallelism for Long-Context LLM Training: AutoSP

Motivation and Problem Formulation

Recent advancements in LLMs have demonstrated substantial utility in handling tasks with long input contexts, including document understanding, multi-step reasoning, and multi-turn dialogue generation. These applications impose extreme activation memory demands due to input sequences spanning tens to hundreds of thousands of tokens. Traditional distributed training libraries, such as those utilizing ZeRO-3/FSDP, Tensor, and Pipeline parallelism, focus on increasing model parameter counts, but do not directly address the memory bottleneck incurred by long-context training. Sequence Parallelism (SP) has emerged as a pivotal technique to scale input sequence lengths by distributing the sequence dimension of activations across devices, leveraging aggregated GPU memory. However, existing SP implementations are manually engineered, tightly coupled with specific frameworks, require considerable code refactoring, and lack easy composability and portability across hardware platforms.

AutoSP proposes a compiler-based framework for scalable long-context LLM training, automating both sequence parallelism and activation checkpointing (AC) within the PyTorch-2.0 compiler stack. This approach minimizes manual engineering effort, ensures correctness, and enables rapid adoption across architectures and backends. AutoSP integrates targeted optimizations as compiler passes that significantly enhance the trainability of LLMs at negligible cost to throughput, thus unlocking practical training at sequence lengths previously unattainable due to memory constraints.

Figure 1: AutoSP architecture overview illustrating the compiler-based automation of sequence parallelism and long-context activation checkpointing across PyTorch IRs.

AutoSP System Architecture

AutoSP consists of two principal compiler passes:

1. Automated Sequence Parallelism (SP) Pass: This pass operates primarily at the Torch-IR level, where the abstraction closely resembles user-defined neural network layers. The SP-pass automatically analyzes tensor shapes, injects communication collectives (e.g., all-to-all), appropriately resizes intermediate buffers, and recomputes manual indexing (such as causal masks for attention layers). By working at Torch-IR, AutoSP avoids the complexity of finer-grained operators at Aten-IR, preserves semantic correctness, and simplifies the conjugate gradient handling for backward computations.
2. Sequence-Parallel Aware Activation Checkpointing (SAC) Pass: AC is essential for reducing memory consumption during training. PyTorch-2.0’s automated AC is conservative, restricting rematerialization for compute-intensive operators to avoid performance penalties; however, in long-context regimes, linear-projection and MLP matmuls account for a sublinear fraction of the total compute, thus enabling efficient rematerialization. AutoSP relaxes these conservative constraints, permitting the targeted recomputation of such operators, significantly reducing memory usage without adversely impacting iteration time.

Implementation Challenges and Technical Solutions

AutoSP addresses several key challenges:

• Recovering Model Parameters: Compiler passes must infer sequence, batch, and hidden dimensions, which are not explicit in IRs. AutoSP utilizes program analysis of graph inputs and traverses to attention layers to deduce tensor shapes.
• Resizing Intermediate Buffers: Proper placement and resizing of tokens and position id buffers are necessary, differing across layers. AutoSP curates sets of relevant operators and manages buffer transformations automatically.
• Communication Collectives Injection: Operators that require full input context (e.g., self-attention) necessitate cross-device communication. AutoSP inserts all-to-all collectives at layer boundaries deterministically.
• Composing SP with AC: Naive composition can result in extraneous communications during backward passes. AutoSP strategically removes constraints on recomputation for specific operator types, exploiting FLOP distribution, as sequence length increases.

Empirical Evaluation

AutoSP’s efficacy is validated across NVIDIA and AMD hardware (including GH200-96GB, A100-80GB, MI250-64GB) with PyTorch-2.7, comparing against baselines such as ZeRO-3/FSDP, DeepSpeed-Ulysses, RingAttention, and their compiler-optimized variants.

Figure 2: Maximum sequence length trainable prior to OOM for different model sizes, demonstrating AutoSP's substantial increases in trainability across all evaluated configurations.

Trainability Results

• Sequence Length Scaling: AutoSP enables training with up to 2.7$\times$ longer input contexts (NVIDIA) and 2.5$\times$ (AMD) over competitive hand-written SP baselines, with negligible throughput degradation.
• Baseline Comparison: Against ZeRO-3, AutoSP achieves 5$\times$ to 5.6$\times$ longer sequences for 3B to 8B models, and 2.5$\times$ for 13B models. Compared to DeepSpeed-Ulysses and RingAttention, gains range from 1.6$\times$ to 3$\times$.
• Activation Memory Reduction: Attention and MLP operator memory usage decreased by 13.03$\times$ and 2.22$\times$, respectively, due to aggressive rematerialization via SAC-pass. Memory gains are especially pronounced for models with larger hidden dimensions.

Performance and Portability

• Runtime Efficiency: AutoSP maintains 97% of hand-written SP implementation speed, matching iteration times while vastly increasing trainability. Overheads for rematerialization are isolated chiefly in backward passes (1.14$\times$ cost), with forward pass times unaffected.
• Hardware Portability: AutoSP demonstrates performance-portability, supporting both NVIDIA and AMD platforms seamlessly with consistent gains in sequence length and minimal speed penalty.

Ablation Studies

• Impact of Individual Passes: The incremental trainability gain of the SAC-pass over the SP-pass alone is 1.66$\times$0 with only a 7% decrease in runtime performance. This empirically validates the theoretical expectation that rematerializing mat-muls becomes increasingly cheap in long-context regimes due to the FLOP distribution.

Related Work and Distinction

Existing parallel strategies (ZeRO-3/FSDP, Tensor, Pipeline, Expert-parallelism) focus on scaling parameter dimensions but not sequence length. Automated optimizations (Deep-Compile, GSPMD, TVM, Mirage, AITemplate) typically target inference or parameter sharding, neglecting sequence-based parallelism. Previous AC techniques either rely on search-based optimization or static policies, requiring manual configuration and invasive code rewrites. AutoSP distinguishes itself as the first compiler-based, PyTorch-native solution automating SP and sequence-aware AC for long-context LLM training, removing the need for manual intervention and enabling full portability.

Implications and Future Directions

Practically, AutoSP’s compiler-based abstractions drastically lower the barrier of entry for long-context LLM training, democratizing access to high performance for document processing, long-form dialogue, and reasoning tasks. Theoretically, the results reveal that memory bottlenecks in long-context LLMs can be alleviated not solely by improved hardware, but also by sophisticated compiler-level optimizations targeting the unique structure of the compute graph.

Future research directions include:

• Extending compiler passes for other forms of parallelism (e.g., expert-parallelism, sparse MoEs) and integrating with broader compiler stacks for generalized distributed training.
• Dynamic profiling-guided optimization to further reduce memory usage and iteration time through precise operator-level rematerialization.
• Adapting AutoSP's principles for complex multi-modal architectures and progenitor LLMs with vastly increased sequence lengths.

Conclusion

AutoSP establishes a new paradigm for scalable LLM training with extreme sequence lengths by automating sequence parallelism and activation checkpointing within the PyTorch compiler stack. The system achieves significant trainability improvements, robust performance-portability, and ease of adoption. Compiler-driven automation, as demonstrated by AutoSP, is a practical and theoretically sound foundation for future developments in long-context model training (2604.27089).