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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Gupta, Ahan, Wang, Zhihao, Dani, Neel, Tanaka, Masahiro, Ruwase, Olatunji, Zhang, Minjia
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866913073297620992
author Gupta, Ahan
Wang, Zhihao
Dani, Neel
Tanaka, Masahiro
Ruwase, Olatunji
Zhang, Minjia
author_facet Gupta, Ahan
Wang, Zhihao
Dani, Neel
Tanaka, Masahiro
Ruwase, Olatunji
Zhang, Minjia
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2604_27089
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle AutoSP: Unlocking Long-Context LLM Training Via Compiler-Based Sequence Parallelism
Gupta, Ahan
Wang, Zhihao
Dani, Neel
Tanaka, Masahiro
Ruwase, Olatunji
Zhang, Minjia
Machine Learning
Distributed, Parallel, and Cluster Computing
Performance
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.
title AutoSP: Unlocking Long-Context LLM Training Via Compiler-Based Sequence Parallelism
topic Machine Learning
Distributed, Parallel, and Cluster Computing
Performance
url https://arxiv.org/abs/2604.27089