![]()
![]()
![]()
在大模型分布式訓練領域,DeepSpeed作為業界主流的優化框架,持續通過版本迭代為開發者提供更高效、更穩定、更兼容的訓練能力。近日,DeepSpeed正式推出v0.18.9版本,本次更新聚焦于AutoTP自動張量并行能力增強、ZeRO優化器適配、SuperOffload內存卸載優化、PyTorch兼容性修復、環境變量適配完善以及CI流程與文檔更新等多個核心方向,同時完成了AutoSP功能的合并入庫,進一步完善了序列并行相關能力,修復了多項影響訓練穩定性與環境適配的問題,為大模型訓練提供了更可靠的底層支撐。
本次版本更新共計16次提交、56個文件變更,由16位社區貢獻者共同完成,代碼層面實現3891行新增與321行刪除,覆蓋功能新增、缺陷修復、兼容性優化、文檔更新、CI流程調整等全維度內容,本文將基于官方更新日志,對DeepSpeed v0.18.9的所有變更進行全面、詳細的解析,幫助開發者快速掌握版本核心升級點,順利完成版本升級與業務適配。
一、環境適配與硬件檢測優化
DeepSpeed v0.18.9針對環境變量適配、GPU硬件檢測邏輯進行了多項優化,解決了不同硬件平臺、自定義環境路徑下的適配問題,提升了框架的環境兼容性與部署靈活性。
1. 支持$TRITON_HOME環境變量
本次更新新增對 環 境 變 量 的 適 配 , 框 架 會 自 動 識 別 該 環 境 變 量 配 置 的 路 徑 , 解 決 了 相 關 組 件 在 自 定 義 安 裝 路 徑 下 無 法 被 正 確 加 載 的 問 題 。 在 過 往 版 本 中 , 僅 默 認 讀 取 系 統 默 認 路 徑 下 的 相 關 文 件 , 當 用 戶 通 過 自 定 義 路 徑 安 裝 時 , 會 出 現 路 徑 查 找 失 敗 、 組 件 調 用 異 常 等 問 題 , 而 通 過 適 配 TRITON_HOME環境變量,允許用戶自由指定Triton的安裝目錄,大幅提升了框架在自定義部署環境下的適配能力,尤其適用于多版本Triton共存、離線部署等復雜場景。
2. ROCm GPU架構檢測移除不必要的shell=True
針對AMD ROCm平臺的GPU架構檢測邏輯,本次更新移除了代碼中不必要的shell=True參數。在Python的subprocess調用中,shell=True會帶來潛在的安全風險與執行效率問題,同時可能引發命令解析異常。DeepSpeed v0.18.9優化了ROCm GPU架構檢測的代碼實現,采用無shell模式執行硬件檢測命令,既提升了檢測邏輯的安全性,又避免了因shell環境差異導致的GPU架構識別失敗問題,保障了AMD GPU平臺上DeepSpeed的穩定運行。
3. 支持$DS_IGNORE_CUDA_DETECTION環境變量跳過本地GPU檢測
新增$DS_IGNORE_CUDA_DETECTION環境變量控制邏輯,當用戶設置該環境變量時,DeepSpeed會跳過本地CUDA GPU的檢測流程。該優化主要針對無GPU環境、遠程訓練節點、容器化部署等場景,在這些場景下,本地GPU檢測不僅無意義,還可能因檢測失敗導致框架初始化異常。通過該環境變量,開發者可以靈活控制GPU檢測行為,避免非必要的硬件檢測流程,提升框架在純CPU環境、云服務器遠程調度場景下的初始化效率。
4. 修復Triton自動調優緩存的NFS路徑判斷問題
修復了Triton自動調優緩存模塊中,is_nfs_path函數對不存在路徑的處理異常問題。在過往版本中,當判斷的NFS路徑不存在時,該函數會拋出異常,導致Triton自動調優流程中斷,影響模型訓練的性能調優。v0.18.9優化了路徑存在性校驗邏輯,當目標路徑不存在時,函數會返回合理的默認值,不會中斷調優流程,保障了Triton自動調優功能在任意路徑配置下的穩定運行,提升了大模型訓練的性能調優可靠性。
二、AutoTP自動張量并行核心功能增強
AutoTP作為DeepSpeed實現自動張量并行的核心模塊,是大模型分布式訓練的關鍵能力,本次v0.18.9版本對AutoTP進行了多項功能升級,大幅提升了其兼容性、易用性與模型適配范圍。
1. 新增AutoTP通用Checkpoint功能
本次更新為AutoTP新增通用Checkpoint(通用 checkpoint)功能,在checkpoint/constants.py中新增UNIVERSAL_CHECKPOINT_VERSION_KEY、UNIVERSAL_CHECKPOINT_VERSION_VALUE、DS_AUTOTP_UC_META等常量定義,同時在universal_checkpoint.py中實現了AutoTP通用 checkpoint 的加載與解析邏輯。
該功能支持AutoTP模式下的模型權重、并行元信息的統一存儲與恢復,新增_get_param_uc_restore_meta函數用于獲取AutoTP通用 checkpoint 的參數元數據,新增_resolve_autotp_partition函數實現AutoTP并行參數的分區解析與恢復,支持分區維度、邏輯形狀、子參數形狀、子參數大小等元信息的讀取與處理,同時支持復制型參數的特殊處理。
通用Checkpoint功能解決了AutoTP模式下模型 checkpoint 兼容性差、跨節點恢復困難的問題,支持不同并行配置下的模型權重遷移,尤其適用于大模型訓練中斷恢復、多節點并行訓練切換、模型微調等場景,大幅提升了AutoTP模式下訓練任務的可遷移性與可靠性。
2. 支持HuggingFace官方tp_plan配置
新增對HuggingFace Transformers庫中tp_plan的原生支持,AutoTP可自動識別并使用Llama、Qwen、Gemma2等模型內置的base_model_tp_plan配置,無需開發者手動設置preset_model或partition_config參數。該優化實現了DeepSpeed AutoTP與HuggingFace模型生態的無縫銜接,開發者無需額外編寫并行分區規則,直接使用HuggingFace官方模型即可開啟AutoTP并行訓練,支持colwise與rowwise兩種核心分區類型,大幅降低了AutoTP的使用門檻。
同時,官方在huggingface-tp/README.md中更新了相關文檔,詳細說明HuggingFace tp_plan的適配邏輯與使用方法,明確了AutoTP對自定義分區的支持方式,讓開發者可以快速上手基于HuggingFace模型的自動張量并行訓練。
三、PyTorch兼容性與訓練穩定性修復
DeepSpeed的核心能力依賴PyTorch生態,本次版本針對低版本PyTorch的兼容性、訓練過程中的內存與梯度更新問題進行了關鍵修復,保障了不同PyTorch版本下訓練任務的穩定運行。
1. 修復PyTorch 2.4以下版本torch.amp.custom_fwd兼容性問題
修復了torch.amp.custom_fwd接口在PyTorch 2.4以下版本的向后兼容問題。在過往版本中,DeepSpeed的自動混合精度訓練模塊調用custom_fwd接口時,會因PyTorch版本差異出現接口不兼容、訓練中斷等異常,v0.18.9通過接口適配與邏輯兼容處理,確保在PyTorch 2.4以下版本中,自動混合精度訓練功能可以正常運行,無需開發者升級PyTorch版本,提升了框架對老舊環境的適配能力。
2. SuperOffload模塊修復多分組共享CPU緩沖區梯度更新問題
修復了SuperOffload模塊中,多分組梯度更新使用共享CPU緩沖區時的數據丟失問題。SuperOffload是DeepSpeed實現內存卸載、降低GPU顯存占用的核心功能,在大模型訓練中,多分組梯度更新場景下,共享CPU緩沖區會導致梯度數據覆蓋、更新失效等問題,嚴重影響訓練收斂效果。v0.18.9優化了緩沖區管理邏輯,保障多分組梯度更新時的數據獨立性,確保梯度更新結果正確保留,提升了SuperOffload模式下大模型訓練的穩定性與收斂效果。
四、ZeRO優化與Muon優化器適配擴展
ZeRO作為DeepSpeed的核心內存優化技術,本次更新對ZeRO Stage 3進行了優化,同時擴展了Muon優化器的支持范圍,提升了分布式訓練的內存效率與優化器兼容性。
1. 擴展Muon優化器對ZeRO Stage 3的支持
新增Muon優化器在ZeRO Stage 3模式下的完整支持,解決了過往版本中Muon優化器與ZeRO Stage 3不兼容、參數分區異常、優化器狀態同步失敗等問題。ZeRO Stage 3是DeepSpeed顯存優化的最高級別,可將模型參數、梯度、優化器狀態全部分區存儲,大幅降低單卡顯存占用,而Muon優化器作為新型優化器,在大模型訓練中具備收斂速度快、穩定性強的優勢,本次適配讓開發者可以在ZeRO Stage 3模式下使用Muon優化器,兼顧顯存效率與訓練性能。
2. ZeRO相關底層邏輯優化
配合Muon優化器適配,DeepSpeed v0.18.9對ZeRO Stage 3的參數分區、狀態同步、內存管理邏輯進行了微調,確保優化器狀態與模型分區參數的正確映射,避免了顯存泄漏、狀態同步延遲等問題,進一步提升了ZeRO Stage 3模式下的訓練穩定性。
五、AutoSP功能合并入庫
本次版本最大的功能升級之一,是將AutoSP(自動序列并行)功能正式合并入DeepSpeed主分支,完成了AutoSP模塊的全量集成。AutoSP是DeepSpeed針對長序列大模型訓練推出的自動序列并行技術,基于torch.compile與FX圖優化實現序列維度的自動分片,無需手動修改模型代碼,即可實現序列并行訓練,大幅提升長上下文大模型的訓練效率。
1. AutoSP核心測試用例新增
在tests/unit/v1/compile/目錄下新增大量AutoSP相關測試用例,覆蓋SDPA節點檢測、輸入ID/標簽ID/位置ID節點識別、符號序列長度節點查找、分片偏移量生成、符號切片索引創建、張量分片節點替換等核心功能測試,測試序列長度覆蓋64、128、256等常用場景,確保AutoSP功能在不同序列長度下的穩定性。
2. AutoSP與Ulysses SP對齊驗證
新增compare_sp_loss測試函數,實現AutoSP編譯模型與Ulysses SP基準模型的損失值對齊驗證,在分布式環境下對比兩種序列并行方案的訓練損失,確保AutoSP的計算精度與Ulysses SP一致,誤差控制在RTOL=0.1、ATOL=0.01范圍內,保障自動序列并行的訓練精度。
3. AutoSP圖優化邏輯完善
新增create_gm_nodes函數實現模型FX圖捕獲,通過自定義torch.compile后端捕獲模型計算圖,確保AutoSP可以正確識別序列維度與輸入節點;新增find_sym_seq_node函數實現符號序列長度節點的查找,解決了torch 2.9 bf16模式下符號節點順序異常的問題,保障AutoSP在不同PyTorch版本與精度模式下的正常運行。
六、CI流程、文檔與版權信息更新 1. CI工作流優化
對.github/workflows/aws-torch-latest-full.yml進行調整,優化AWS L40S GPU集群的CI測試流程,新增容量 fallback 邏輯,當L40S GPU資源不足時,自動切換至8xA100節點執行測試;優化提交檢測邏輯,簡化GitHub Actions的命令執行格式,統一標點符號規范;新增GPU架構自動檢測步驟,通過Python代碼實時獲取GPU計算能力與設備數量,動態設置TORCH_CUDA_ARCH_LIST環境變量,提升CI測試的兼容性與靈活性。
同時新增測試 fallback 邏輯,當部分測試用例執行失敗時,自動切換至完整測試模式,確保測試覆蓋率,保障版本發布的穩定性。
2. 開發文檔更新
對AGENTS.md與CLAUDE.md文檔進行優化,在CI開發規范中新增pre-commit命令說明,明確要求開發者提交代碼前執行pre-commit run --files 命令,僅檢查修改文件而非全量代碼庫,提升開發效率;移除文檔中Microsoft Corporation版權聲明,統一版權規范;更新pre-commit配置說明,明確代碼格式化、靜態檢查的執行要求。
3. 項目README與新聞更新
更新README.md的Latest News板塊,新增兩項重要資訊:一是DeepSpeed團隊在ASPLOS 2026發布大模型系統構建教程;二是SuperOffload相關工作獲得ASPLOS 2026最佳論文榮譽提名,同步更新相關鏈接,提升項目社區影響力。
同時新增ASPLOS 2026最佳論文獎相關新聞條目,完善項目榮譽展示,增強社區開發者對DeepSpeed技術價值的認可。
4. 版本號更新
將version.txt文件中的版本號從0.18.8更新至0.18.9,完成正式版本的標識,方便開發者通過版本號區分迭代內容,順利進行版本升級。
七、版本更新總結與升級建議
DeepSpeed v0.18.9作為一次迭代型版本更新,聚焦于功能增強、兼容性修復、穩定性提升、生態適配四大核心方向,所有變更均圍繞大模型分布式訓練的實際痛點展開,無破壞性變更,適合所有使用DeepSpeed的開發者升級。
本次更新的核心價值體現在:
1. AutoTP能力全面升級,通用Checkpoint與HuggingFace tp_plan支持,大幅降低自動張量并行使用門檻;
2. 硬件與環境適配優化,支持自定義環境變量、多平臺GPU檢測,提升框架部署靈活性;
3. 訓練穩定性修復,解決PyTorch兼容、SuperOffload梯度更新、Triton路徑異常等關鍵問題;
4. AutoSP正式合并,長序列自動并行能力落地,適配長上下文大模型訓練;
5. ZeRO與優化器適配擴展,支持Muon優化器在ZeRO Stage 3模式下運行;
6. CI與文檔完善,提升開發協作效率與項目規范性。
升級建議:
1. 正在使用DeepSpeed v0.18.x系列版本的開發者,可直接升級至v0.18.9,無需修改業務代碼;
2. 使用HuggingFace模型進行張量并行訓練的開發者,升級后可直接使用官方tp_plan,簡化配置;
3. 長序列大模型訓練場景,可嘗試AutoSP功能,提升序列并行訓練效率;
4. AMD ROCm平臺、自定義Triton路徑、無GPU環境的開發者,升級后可解決環境適配問題;
5. 使用PyTorch 2.4以下版本的開發者,升級后可修復混合精度訓練接口兼容問題。
代碼地址:github.com/deepspeedai/DeepSpeed
DeepSpeed v0.18.9版本通過社區協同開發,完成了多項實用功能升級與缺陷修復,進一步鞏固了其在大模型分布式訓練領域的領先地位。AutoTP的生態適配、AutoSP的正式落地、ZeRO與優化器的深度優化,讓大模型訓練更高效、更簡單、更穩定;環境適配與兼容性修復,保障了框架在不同硬件、不同軟件版本下的可靠運行。
我們相信人工智能為普通人提供了一種“增強工具”,并致力于分享全方位的AI知識。在這里,您可以找到最新的AI科普文章、工具評測、提升效率的秘籍以及行業洞察。 歡迎關注“福大大架構師每日一題”,發消息可獲得面試資料,讓AI助力您的未來發展。
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.