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文件或聯絡團隊。