來自 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、跨多小時 sessions 的 agentic 工作流、文 / 圖 / 音多模態訓練 —— 在動手寫 SP 之前,先評估 AutoSP。手寫 SP 是實打實的工程時間;編譯器自動版本只是一行 import。第二,SSAIL / Anyscale / Snowflake 這種合作組合,本身就是 ML 系統研究正在收斂的訊號。Anyscale 出 Ray;Snowflake 出資料基礎設施;UIUC 出系統層研究。未來留意來自這一聯盟、把更多東西搬進 DeepSpeed 編譯器堆疊的工作。第三,「跨硬體效能可攜」是這篇裡那一條帶光環的承諾。如果 AutoSP 實測開銷在「不同 GPU 廠商之間」相對手寫都很小,那它會被迅速採納;但若它只在 NVIDIA Hopper 一類硬體上保持小開銷,採納就會很慢。在把訓練流水線接到AutoSP 上之前,先把完整論文裡的 benchmark 方法讀一遍。