在人工智能模型訓練領域，超大規模語言模型的訓練如同搭建精密復雜的科技大廈，需要成千上萬個GPU協同工作。其中，“全分片數據并行”技術（FSDP）是關鍵支撐，但傳統FSDP系統存在明顯局限，如同固定尺寸的積木難以應對復雜搭建需求。字節跳動研究團隊針對這一問題，開發出veScale - FSDP系統，為AI訓練帶來新突破。
傳統FSDP系統問題諸多。當研究人員嘗試運用矩陣優化器Muon或8位量化訓練等先進技術時，傳統系統因只能按固定方式切分數據，無法滿足新技術對數據塊狀結構的特定處理要求，就像刀具只能橫切或豎切，難以完成復雜切割。而且，在性能方面，當訓練規模擴展到上萬個GPU時，系統通信效率和內存使用問題凸顯，內存管理缺乏統一規劃導致大量內存碎片，通信中的數據拷貝操作也造成顯著性能損失。
veScale - FSDP系統的核心創新在于引入“RaggedShard”新型數據分片格式。與傳統固定模具式的數據分片不同，RaggedShard如同可隨意調整形狀的橡皮泥模具，能根據具體需求靈活調整分片大小和形狀。以拼圖游戲類比，傳統系統要求拼圖塊都是標準長方形，而RaggedShard允許拼圖塊有各種形狀。處理Muon優化器時，它能確保每個設備上的數據塊包含完整矩陣，避免矩陣被切斷；對于8位量化訓練，它能保證每個設備上的數據按量化塊邊界切分，使每個設備可獨立完成量化操作，還能與張量并行、專家并行等現有并行化策略完美融合。
僅有靈活的數據分片還不夠，高效通信是關鍵。veScale - FSDP開發了智能規劃算法，解決不規則形狀數據塊的高效打包和傳輸問題，類似快遞打包中如何將形狀各異的物品裝入標準快遞箱，既要保證物品完好、空間利用率高，又要確保每個快遞箱重量平衡。研究團隊將此問題數學化為優化問題，在滿足“非分片塊約束”“連續內存約束”“嚴格負載平衡約束”三個關鍵條件下，最小化通信緩沖區大小。基于Transformer模型的結構特性，他們設計了多項式時間的動態規劃算法，通過案例分析數據塊與分片邊界的三種可能關系，找到最佳數據布局方案。
為進一步提升性能，veScale - FSDP引入“分布式緩沖區”（DBuffer）新原語。傳統FSDP系統通信時需大量數據拷貝操作，如同搬家時物品需多次搬運，效率低下。DBuffer實現零拷貝通信，像在房間和卡車間建立直通管道，物品可直接傳輸。它具有全局緩沖區語義、支持組級別操作、提供持久地址映射、支持原地通信和計算等關鍵特性，在實際測試中，使通信開銷降低約7%，規劃算法優化帶來34%的性能提升，且這些優化自動進行，開發者無需修改模型代碼。
veScale - FSDP在各種規模實驗中表現卓越。在端到端訓練性能對比中，該系統在MoE稀疏模型上比所有基線系統快11%到66%，在LLaMA - 3 - 70B密集模型上有5%的性能提升，這些提升源于優化的通信重疊、基于DBuffer的零拷貝集合通信以及避免填充開銷的靈活分片粒度。內存使用方面，它將峰值預留內存降低16%到30%，得益于確定性的批量內存管理策略。擴展性測試顯示，它能高效擴展到一萬個GPU規模，在弱擴展性測試中保持接近線性擴展性，在強擴展性測試中也有良好表現。
為驗證系統靈活性，研究團隊測試了veScale - FSDP對8位Adam優化器和分布式Muon優化器的支持能力。8位Adam優化器需支持32×32的塊級量化粒度，veScale - FSDP中開發者只需設置參數量化粒度，系統就能自動確保數據分片與量化塊邊界完美對齊，每個設備可獨立量化操作。Muon優化器需在原始2D參數矩陣形狀上進行Newton - Schulz迭代，veScale - FSDP通過RaggedShard的不均勻分片能力，讓用戶用標準SPMD方式編寫分布式Muon算法。實驗結果顯示，該系統能高效支持這兩種前沿優化技術，且性能表現良好。
veScale - FSDP的工程實現巧妙擴展了PyTorch的DTensor抽象，RaggedShard作為DTensor的新放置策略，無縫集成現有并行化基礎設施。系統包含7600行Python代碼，完全兼容標準PyTorch分布式運行時和廣泛版本范圍，作為FSDP2的即插即用后端，使用相同API，現有訓練代碼幾乎無需修改就能享受性能提升。規劃算法實際運行開銷極小，時間復雜度為O(|T|?m log(E) log(|T|m))。
研究團隊通過消融實驗量化各組件貢獻，結果顯示DBuffer和規劃算法貢獻大部分性能提升，RaggedShard是支持塊級8位Adam的關鍵抽象。veScale - FSDP已在字節跳動生產環境得到實戰驗證，成功支持超一萬個GPU、2.4萬億參數規模的大規模訓練任務。在實際部署中，研究團隊總結出小規模實驗可準確預測大規模性能、在成熟抽象基礎上設計系統抽象、將模型定義與系統優化解耦等重要經驗。
veScale - FSDP已開源，為整個AI社區提供了強大工具。無論是學術研究機構還是產業界，都能利用它訓練更大、更復雜的AI模型，滿足AI模型規模不斷增長的需求。
