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

2026-04-29

AI Summary

This work addresses the lack of automated support for long-context optimization in existing large language model training frameworks, which currently require users to manually integrate complex strategies at high development cost. We propose the first compiler-driven, automated optimization framework tailored for long-context training, seamlessly integrating sequence parallelism and long-context-aware activation checkpointing without any user code modifications. Our approach substantially reduces implementation complexity and enhances training scalability, enabling maximum context lengths up to 2.7× and 2.5× those of the baseline on NVIDIA and AMD hardware, respectively, while incurring negligible throughput degradation.

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.

Problem

• long-context training
• sequence parallelism
• LLM optimization
• training efficiency
• developer productivity

Innovation

• sequence parallelism
• compiler-based optimization
• long-context LLM training
• activation checkpointing
• automated model compilation