PyTorch AutoSP：基于编译器的工具自动将训练代码转换为序列并行，支持100k+词元上下文

2026年4月29日，PyTorch团队与UIUC SSAIL实验室、Anyscale和Snowflake的研究人员共同发布了AutoSP——一个集成在DeepSpeed/DeepCompile生态系统中的基于编译器的工具，可自动将标准单GPU Transformer训练代码转换为序列并行形式。目标：无需手动编写分布式代码，即可训练具有超长上下文（100k+词元）的大语言模型。

AutoSP解决的问题

训练长上下文大语言模型需要将序列分布到多个GPU（序列并行，SP），因为Transformer的激活值随上下文长度二次方扩展，很容易导致显存溢出错误。现有解决方案（RingFlashAttention、DeepSpeed-Ulysses）需要手动重写训练代码——词元分区、通信集合操作、复杂的注意力掩码处理。

AutoSP将这一切自动化：用户编写标准Transformer代码，编译器将其转换为SP感知的形式。

如何启用？

在DeepSpeed配置中添加三行配置：

config = {
    "compile": {
        "deepcompile": True,
        "passes": ["autosp"]
    },
    "sequence_parallel_size": 4,
    "zero_optimization": {"stage": 1},  # AutoSP与ZeRO 0/1兼容
    ...
}

加上用于准备输入的工具函数prepare_auto_sp_inputs()。底层策略采用DeepSpeed-Ulysses架构——在NVLink/fat-tree网络上，通信开销随GPU数量增长保持恒定，可扩展到注意力头数量（例如7-8B模型中的32个注意力头）。

序列感知激活检查点（SAC）

AutoSP还引入了SAC——针对长上下文训练优化的自定义检查点策略。与保守的PyTorch 2.0最大流最小割方案不同，SAC利用长上下文特有的FLOP动态——在反向传播中释放计算代价低廉的算子的中间激活并重新计算。权衡： 略微降低吞吐量，但使更长的上下文成为可能。

实验结果

在8块A100-80GB SXM节点（PyTorch 2.7，CUDA 12.8）上使用不同规模的Llama 3.1模型测试：

• 在相同资源下，最大可训练序列长度显著提升
• 与手写RingFlashAttention和DeepSpeed-Ulysses基线相比，运行时开销极小

端到端示例（包括Llama 3.1 8B）可在github.com/deepspeedai/DeepSpeedExamples/tree/master/benchmarks/autosp查阅。

局限性

目前AutoSP要求单一可编译单元（整个Transformer必须作为一个块编译），且不允许模型内部存在图中断。团队将图中断容错能力列为下一步工作。