来自 UIUC SSAIL 实验室、Anyscale 与 Snowflake 的研究团队,于 4 月 29 日通过 PyTorch 官方博客发布了 AutoSP —— 一个基于编译器的 DeepSpeed 扩展,能把标准的 transformer 训练代码自动转换为序列并行(sequence-parallel,SP)代码,用于跨多 GPU 的长上下文 LLM 训练。其卖点是:在 10 万 token 以上的上下文上做训练时,无需 SP 历来要求的那种侵入式代码改写。AutoSP 集成进 DeepSpeed 的编译器生态 DeepCompile;用户只需 import AutoSP、编译模型,SP 即自动启用。它能与现有并行策略(如 ZeRO)组合,基于编译器的实现也具备跨硬件厂商的性能可移植性。
序列并行是这里真正在解决的工程痛点。在 10 万 token 以上的上下文上,即便是 ZeRO/FSDP 也会撞上 OOM;把 token 横切到多设备上做分割(SP),就是这道难题的出口。但手写 SP 需要:对输入上下文与中间激活做分区、插入通信集合操作、把通信与计算交错重叠 —— 而且要在前向与反向两路上各做一遍。多年来,想做长上下文能力的研究者一直在按模型、按硬件目标重复这一类工作。AutoSP 把「分区 / 通信集合 / 重叠」这一套逻辑下推到编译器:你写出标准 PyTorch 风格的训练代码,编译器自动产出SP-aware 版本。团队报告称相对手写基线「运行时开销很小」—— 也就是说,自动化没有让你失去手写 SP 原本能拿到的那部分性能。
两条值得关注的模式。第一,这是「ML 系统层走向编译器化并行」这一更大趋势的延续。PyTorch 的 torch.compile、NVIDIA 的 NeMo Megatron、Google 的 Pathways、更宽泛意义上的 pjit 谱系 —— 它们都把「并行决策」推到编译器层,因为手写并行根本扛不住「跨模型架构、跨硬件代际」的 scale。AutoSP 是这条线上最新的一个例子,并且坐在了真正会被用起来的衬底上(DeepSpeed 已经被广泛采纳)。第二,「长上下文训练」如今是一块真实的市场。1M+ token 上下文的模型 —— Gemini、Claude、我们这周报道过的 Poolside Laguna XS.2 —— 已经在生产中。训练侧的瓶颈,已经从「我们能不能训这个模型」转成了「我们能不能在这么长的上下文上训这个模型」。AutoSP 正是这场转变里需要的工具。
对 builders 三件具体事情。第一,如果你在训练任何面向长上下文用例的模型 —— 基于大文档的 RAG、跨多小时会话的 agentic 工作流、文 / 图 / 音多模态训练 —— 在动手写 SP 之前,先评估 AutoSP。手写 SP 是实打实的工程时间;编译器自动版本只是一行 import。第二,SSAIL / Anyscale / Snowflake 这种合作组合,本身就是 ML 系统研究正在收敛的信号。Anyscale 出 Ray;Snowflake 出数据基础设施;UIUC 出系统层研究。未来留意来自这一联盟、把更多东西搬进 DeepSpeed 编译器栈的工作。第三,「跨硬件性能可移植」是这篇里那一条带光环的承诺。如果 AutoSP 实测开销在「不同 GPU 厂商之间」相对手写都很小,那它会被迅速采纳;但若它只在 NVIDIA Hopper 一类硬件上保持小开销,采纳就会很慢。在把训练流水线接到AutoSP 上之前,先把完整论文里的 benchmark 方法读一遍。