Introducing AutoSP

AutoSP is a DeepCompile-based solution that automatically applies Sequence Parallelism (SP) to reduce OOM issues in long-context training. Previously, developers had to manually insert token splitting, collective communication, and forward/backward overlap into DeepSpeed or HuggingFace code, but AutoSP converts this automatically.

The usage is simple.

• Tag input tokens, attention mask, and position ids with prepare_auto_sp_inputs.
• Turn on compile.deepcompile in the DeepSpeed config and add autosp to passes.
• Set sequence_parallel_size together with ZeRO stage 1, then compile with model.compile(dynamic=True).

The core transformations are DeepSpeed-Ulysses-based SP and Sequence-aware AC (SAC) for long contexts. Ulysses has the advantage that communication overhead stays constant as the number of GPUs increases on NVLink or fat-tree topologies, but the SP size is limited to the number of attention heads in the model, capping at 32 for 7-8B-class models. SAC places checkpoints less conservatively than PyTorch 2.0's automatic min-cut AC, enabling training on longer contexts at the cost of a slight reduction in throughput.

In benchmarks, multiple Llama 3.1 models were run on 8 A100-80GB SXM nodes, comparing AutoSP against RingFlashAttention, DeepSpeed-Ulysses, and ZeRO-3 in an environment with PyTorch 2.7 and CUDA 12.8. The results show that AutoSP can train longer sequences with the same resources while runtime cost increase remains modest.

The constraints are also clear.

• The entire transformer must be bundled into a single compilable artifact.
• AutoSP does not work if there is a graph break.