Baseten 工程團隊發佈了其增量權重同步系統的最新進展:對於約770GB的GLM-5.3檢查點,同步時間已從平均約36秒縮短至15秒,請求暫停時間從6秒降至0.8秒,吞吐量提升最高可達3.2倍。這一優化不僅刷新了性能記錄,還解鎖了此前因CPU內存限制而無法運行的Kimi K3強化學習rollout。
此前,增量權重同步系統需要依靠一個CPU副本作為解碼基礎。訓練器將完整檢查點和稀疏增量發佈到對象存儲,每個增量張量是父版本字節與新版本字節的XOR結果,再經zstd壓縮。rollout集羣通過Baseten Delivery Network拉取對象,重建新策略並原地更新GPU。然而,在舊設計中,每個副本必須在主機內存中保留完整的父檢查點,以便將XOR字節解碼為新值,並在每次更新後修補CPU副本。
這個CPU副本帶來了三重代價。首先是內存壓力:GLM-5.3的完整父版本約770GB,在8塊B300的節點上,它與1.1TB的主機KV緩存無法共存,緩存不得不縮減到275GB。其次是時間成本:每次更新後,將CPU副本提升為新策略需要約45秒,阻塞下一個增量的處理。第三是雙數據源風險:每次更新必須同時寫入GPU和CPU副本,否則下一個增量將基於錯誤的父版本解碼。
新設計的核心思想是:GPU上已經保存了策略N的每一個字節,只需要一種方法讀取發生變化的那幾個字節。啓動時,系統從本地磁盤加載模型,按照引擎的實際GPU佈局驗證每個服務張量,記錄每個張量的Adler-32校驗和,然後刪除本地副本。此後,GPU內存成為唯一的父版本。
當增量到達時,原生解碼器直接遍歷zstd命令,只保留非零的XOR字節:對於每個變化的元素,記錄其位置和XOR掩碼。它從不物化稠密的XOR張量,也從不讀取模型權重。在舊策略繼續服務的同時,每個GPU工作線程只收集變化位置上的舊值,應用XOR掩碼得到絕對新值,用訓練器的校驗和驗證結果,並將新值封存到引擎原生張量佈局的寫入計劃中。
校驗和驗證也經過了專門優化。訓練器發佈每個變化張量新字節的Adler-32值。樸素做法是重新哈希整個張量,效率很低。由於Adler-32中每個字節的貢獻僅取決於其位置和值,變化字節對校驗和的影響可以在常數時間內計算。系統從啓動時記錄的校驗和出發,為每個變化字節添加一次修正,再與訓練器的值比較,未變化的字節永遠不會被讀取。這一優化使整個更新過程最多快2.3倍,例如Kimi K3在30GB增量上的更新時間從約230秒降至約100秒,同時確保任何損壞下載或錯誤父版本在任何GPU寫入前被拒絕。
由於大量準備工作在暫停之前完成,暫停窗口只需執行三個步驟:關閉准入並原地暫停運行中的請求;原地寫入新值並刷新派生權重;發佈新策略版本並重新開放准入。在GLM-5.3上,暫停時間現為0.8秒,此前為6秒。45秒的CPU提升過程也已消失,下一個增量可以在落地後立即應用。
進一步優化解決了更新工作與在線服務共享GPU和CPU資源的問題。並行感知的增量解碼器根據張量並行和專家並行將每個變化分配給負責的GPU,每個GPU將變化直接映射到其打包佈局中的地址,包括swizzled FP8和FP4縮放因子。為避免與SGLang調度器的Python事件循環爭搶GIL,更新計劃在獨立的無CUDA輔助進程中運行,通過共享固定內存與GPU工作線程交換批量數據。下載和解碼也實現了流水線化,每個分片一落地就立即解碼,運行在持續的、低優先級的原生線程上,並預留16個核心用於服務。
Kimi K3的案例尤為關鍵。其服務態MXFP4佈局大小為1.7TB。如果保留CPU檢查點,每個GB300節點僅其一半權重就需要約838GB主機內存,幾乎耗盡整個服務容器的內存限制。由於GB300硬件預配置,這是一項硬約束。現在,Kimi K3可以運行在兩個各四塊GPU的GB300節點上,每個節點只解碼、驗證和暫存自己擁有的張量,絕不持有另一節點的權重。在兩個節點都成功寫入後,新策略才會發佈並恢復請求。對於8.4GB的增量,Kimi K3從同步請求到權重生效僅需75秒,暫停4.1秒。
釋放出的主機內存被重新用於KV緩存卸載。被驅逐的GPU KV頁面寫回主機內存而非丟棄,後續具有相同前綴的請求可以重新加載,而不是重新運行prefill。對於RL rollout的多輪長共享前綴場景,從CPU內存重新加載150k token的驅逐前綴需要0.48秒,而全新prefill需要4.9秒。在AgentX基準測試中,併發為16時KV緩存已完全適配GPU,兩種方案無差異;併發為48時,代理前綴超出GPU內存但適合主機內存,卸載將未命中轉為命中,緩存命中率從2%躍升至92%,每GPU輸出吞吐量提升3.2倍。
所有結果均基於單副本、訓練器形狀的固定增量測量,從訓練器的增量同步請求計到副本開始服務新策略,即以S3標記到第一個新響應計時。團隊表示,希望將此方案用於其訓練棧的用户可以查閲RL rollout文檔或聯繫團隊。