字節跳動發布分布式訓練技術：讓模型訓練像搭積木一樣靈活高效

這項由字節跳動種子實驗室主導的研究發表于2026年2月的arXiv預印本論文庫，論文編號為arXiv:2602.22437v1。有興趣深入了解的讀者可以通過該編號查詢完整論文內容。

在人工智能快速發展的今天，訓練超大規模語言模型就像建造摩天大樓一樣復雜。每當我們想要訓練一個擁有數千億參數的AI模型時，就需要成千上萬個GPU協同工作，這個過程中最關鍵的技術之一就是"全分片數據并行"技術，英文名叫FSDP。不過，現有的FSDP系統就像一套死板的積木玩具，每個積木的形狀都是固定的，無法根據具體需求進行調整。字節跳動的研究團隊發現了這個問題，并開發出了一種全新的解決方案——veScale-FSDP系統。

現有FSDP系統的問題就像用固定尺寸的積木搭建復雜模型一樣困難。當研究人員想要使用一些先進的訓練技術時，比如矩陣優化器Muon或者8位量化訓練時，傳統系統就顯得力不從心了。這些新技術需要將數據按照特定的塊狀結構進行處理，但傳統FSDP系統只能按照固定的方式切分數據，就像只能橫切或豎切的刀具，無法滿足復雜切割需求。

更糟糕的是，傳統系統在性能上也存在明顯不足。當訓練規模擴展到上萬個GPU時，系統的通信效率和內存使用都會出現問題，就像一個城市的交通系統在高峰期出現擁堵一樣。研究人員發現，現有系統在內存管理方面缺乏統一規劃，導致大量內存碎片，同時通信過程中的數據拷貝操作也造成了顯著的性能損失。

**一、革命性的靈活切分技術**

veScale-FSDP的核心創新在于引入了一種叫做"RaggedShard"的新型數據分片格式。如果把傳統的數據分片比作用固定模具制作餅干，那么RaggedShard就像是可以隨意調整形狀的橡皮泥模具，能夠根據具體需求靈活調整分片的大小和形狀。

這種靈活性的實現原理可以用拼圖游戲來理解。傳統FSDP系統就像強制要求所有拼圖塊都必須是標準的長方形，而RaggedShard允許拼圖塊有各種不同的形狀。當我們需要處理Muon優化器時，它需要完整的2D矩陣結構，RaggedShard就可以確保每個設備上的數據塊恰好包含完整的矩陣，避免了矩陣被切斷的問題。

對于8位量化訓練來說，這種靈活性更是至關重要。量化訓練需要將參數按照32×32的塊進行處理，就像制作瓷磚時需要按照固定尺寸切割一樣。傳統系統由于無法保證切分邊界與量化塊邊界對齊，往往需要額外的通信來收集完整的量化塊，這就像為了得到完整的瓷磚而需要從多個工廠收集碎片一樣低效。RaggedShard則能確保每個設備上的數據恰好按照量化塊的邊界進行切分，每個設備都能獨立完成量化操作。

更重要的是，RaggedShard還能與現有的并行化策略完美融合?，F代大模型訓練通常需要同時使用多種并行化技術，包括張量并行、專家并行等。RaggedShard的設計使得它能夠與這些技術無縫配合，就像一個萬能適配器，可以連接各種不同的設備一樣。

**二、智能規劃算法解決通信難題**

僅有靈活的數據分片還不夠，如何高效地進行通信才是關鍵。veScale-FSDP開發了一套智能規劃算法，專門解決如何將不規則形狀的數據塊高效地打包和傳輸。

這個問題的復雜性可以用快遞打包來類比。當我們需要郵寄一堆形狀各異的物品時，如何將它們裝入標準尺寸的快遞箱中，既要確保物品不被損壞，又要最大化空間利用率，同時還要保證每個快遞箱的重量平衡？這正是veScale-FSDP面臨的挑戰。

研究團隊將這個問題數學化為一個優化問題。他們需要在滿足三個關鍵約束條件的前提下，最小化通信緩沖區的大小。第一個約束是"非分片塊約束"，確保數據塊不會在通信過程中被意外切斷；第二個是"連續內存約束"，保證數據在內存中的連續性以避免碎片化的拷貝操作；第三個是"嚴格負載平衡約束"，確保每個設備的通信負載相等。

理論上，這是一個NP困難問題，意味著尋找最優解需要指數級的時間。但研究團隊發現了一個巧妙的解決方案。他們觀察到Transformer模型具有高度規律的結構特性：線性層權重占據了參數總數的絕大部分，而且分片塊的大小在不同層之間往往是一致的?；谶@個觀察，他們設計了一個多項式時間的動態規劃算法，能夠在實際應用中找到接近最優的解決方案。

算法的核心思想是案例分析。對于每個數據塊，算法分析它與分片邊界的三種可能關系：完全位于單個分片內、跨越兩個相鄰分片但不包含完整分片、完全包含至少一個分片。通過這種分類，算法能夠有效地搜索解空間，找到最佳的數據布局方案。

**三、分布式緩沖區實現零拷貝通信**

為了進一步提升性能，veScale-FSDP引入了一個叫做"分布式緩沖區"（DBuffer）的新原語。這個組件可以理解為一個智能的內存管理器，專門為分布式通信優化。

傳統FSDP系統在通信時需要進行大量的數據拷貝操作，就像搬家時需要把物品從房間搬到走廊，再搬到卡車上一樣低效。DBuffer則實現了零拷貝通信，就像在房間和卡車之間建立了一條直通管道，物品可以直接傳輸。

DBuffer的設計有四個關鍵特性。首先，它提供了全局緩沖區語義，抽象掉了N維設備拓撲的復雜性，讓開發者可以像操作單機內存一樣操作分布式內存。其次，它支持組級別的操作，能夠將多個張量的相同操作（如加法、縮放、清零）融合成一個批量操作，減少了內核啟動的開銷。第三，它通過RaggedShard的規劃算法提供了持久的地址映射，每個張量的數據指針都有固定的地址，避免了通信前后的數據拷貝。最后，它支持原地通信和計算，進一步減少了內存占用。

這種設計的效果非常顯著。在實際測試中，DBuffer使得通信開銷降低了約7%，而規劃算法的優化則帶來了34%的性能提升。更重要的是，這些優化是自動進行的，開發者無需修改任何模型代碼就能享受到性能提升。

**四、突破性的性能表現**

veScale-FSDP在各種規模的實驗中都展現出了卓越的性能。在端到端訓練性能對比中，該系統在MoE稀疏模型上比所有基線系統快11%到66%，在LLaMA-3-70B密集模型上也有5%的性能提升。這些提升主要來自于三個方面的優化：優化的通信重疊、基于DBuffer的零拷貝集合通信、以及避免填充開銷的靈活分片粒度。

內存使用方面的改善同樣令人印象深刻。veScale-FSDP在各個測試中都將峰值預留內存降低了16%到30%。這種內存節省源于確定性的批量內存管理策略：系統顯式管理流依賴關系以實現可預測的內存釋放，并且通過批量分配來減少碎片化。相比之下，傳統系統由于隱式的record_stream機制導致不確定性的內存釋放，往往會阻止緩存分配器重用緩沖區，使峰值預留內存增加20%。

擴展性測試顯示，veScale-FSDP能夠高效地擴展到一萬個GPU的規模。在弱擴展性測試中，系統保持了接近線性的擴展性，這是因為FSDP的通信成本和每GPU的計算成本都與GPU數量無關，只依賴于模型和輸入大小。在強擴展性測試中，系統在128M token的全局批處理規模下能夠線性擴展到一萬個GPU，即使在16M token的較小批處理下，從1K到8K GPU也能實現3.4倍的吞吐量提升。

**五、支持前沿優化技術的靈活性驗證**

為了驗證系統的靈活性，研究團隊特別測試了veScale-FSDP對兩種前沿優化技術的支持能力：8位Adam優化器和分布式Muon優化器。

8位Adam優化器通過將梯度統計信息量化為INT8格式來大幅減少優化器狀態的內存占用。要實現這種量化，系統需要支持32×32的塊級量化粒度。在veScale-FSDP中，開發者只需要設置參數的量化粒度，系統就會自動確保每個設備上的數據分片與量化塊邊界完美對齊。這樣，每個設備都可以獨立進行量化操作，無需任何額外的通信。相比之下，現有的FSDP系統由于無法原生地跟蹤塊邊界，要么需要侵入性的系統修改，要么需要手動實現復雜的集合通信來交換量化元數據。

Muon優化器的情況更加復雜。這種矩陣符號預調節器需要在原始2D參數矩陣形狀上進行Newton-Schulz迭代。veScale-FSDP通過RaggedShard的不均勻分片能力，讓用戶可以用標準的SPMD方式編寫分布式Muon算法。具體來說，系統通過負載平衡選擇一個根設備，然后使用DTensor的redistribute操作將完整的2D參數收集到根設備上，在那里執行Newton-Schulz更新，最后再將更新后的參數分發回原始設備。整個過程對其他設備來說Newton-Schulz更新變成了無操作，實現了清晰的SPMD語義。

實驗結果顯示，無論是8位Adam還是分布式Muon，veScale-FSDP都能在幾乎不修改代碼的情況下實現高效支持。8位Adam的損失曲線與傳統DDP實現高度一致，只是偶爾出現一些量化精度導致的小波動。分布式Muon的收斂速度明顯快于AdamW，在訓練約800億token后穩定在比AdamW低0.01的損失水平，這與之前的研究結果一致。

**六、工程實現的精巧設計**

veScale-FSDP的工程實現體現了"站在巨人肩膀上"的設計哲學。系統沒有重新發明輪子，而是巧妙地擴展了PyTorch的DTensor抽象。RaggedShard被實現為DTensor的一個新的放置策略，這使得它能夠無縫集成現有的并行化基礎設施，包括張量并行、專家并行，以及成熟的訓練工具如分布式檢查點。

系統的實現包含了7600行Python代碼，完全兼容標準的PyTorch分布式運行時和廣泛的PyTorch版本范圍。更重要的是，它作為FSDP2的即插即用后端，使用相同的PyTorch原生fully_shard API，這意味著現有的訓練代碼幾乎無需修改就能享受到性能提升。

規劃算法的實際運行開銷極小，在所有實驗中算法運行時間都少于0.3秒，這是一次性的初始化成本，在分布式訓練的整個生命周期中可以忽略不計。算法的時間復雜度為O(|T|?m log(E) log(|T|m))，其中T是張量數量，m是設備數量，E是總元素數量。

**七、深度性能分析與組件貢獻**

研究團隊進行了詳細的消融實驗來量化各個組件的貢獻。結果顯示，DBuffer和規劃算法貢獻了大部分的性能提升：禁用DBuffer會使吞吐量降低7.2%，禁用規劃算法則會導致34.6%的性能下降。

DBuffer的性能提升主要來自于消除了集合通信周圍的拷貝輸入/拷貝輸出開銷。當通信緩沖區需要拷貝時，這些操作會成為性能瓶頸。規劃算法的影響更為顯著，因為當禁用規劃時，量化塊無法保證完全包含在設備的本地分片內，系統必須回退到DTensor重分布來在每塊量化之前組裝所需的優化器狀態，這會產生大量額外的通信開銷。

至于RaggedShard本身，它不僅僅是一個優化——它是使得塊級8位Adam在不進行侵入性模型/優化器修改或手寫集合通信的情況下可用的關鍵抽象。禁用RaggedShard會使系統變得幾乎不可用，用戶必須要么仔細修改每個模型和優化器張量使得32×32塊邊界與分片邊界對齊，要么手動實現復雜的集合通信來恢復塊級語義。

**八、生產環境的實戰驗證**

veScale-FSDP不僅在研究環境中表現出色，更重要的是它已經在字節跳動的生產環境中得到了實戰驗證。系統成功支持了超過一萬個GPU的大規模訓練任務，訓練的模型參數規模達到了2.4萬億。這種規模的訓練對系統的穩定性、效率和可維護性都提出了極高的要求。

在實際部署中，研究團隊總結了三個重要的經驗教訓。首先，小規模實驗能夠準確預測大規模性能。FSDP工作負載的性能可以通過各層的計算時間和FSDP通信時間來準確估算，因為計算完全在GPU內部進行，而FSDP通信時間在GPU數量增加時基本保持不變。這使得團隊可以在64個GPU上進行性能分析，然后外推到數千個GPU的規模。

其次，在成熟抽象基礎上設計系統抽象的重要性。DTensor提供了一個強大的抽象層，已經支持廣泛的并行化技術。通過將RaggedShard設計為DTensor的擴展，veScale-FSDP能夠無縫集成現有的并行化策略，最小化工程投入的同時為更廣泛的社區做出貢獻。

第三，將模型定義與系統優化解耦的價值。模型架構的快速演進要求頻繁更新模型定義，但像Megatron-LM這樣的框架將系統級并行化優化與模型代碼緊密耦合，使得研究人員難以修改或擴展架構。veScale-FSDP通過將模型定義從系統框架中解耦出來，讓研究人員能夠專注于模型設計，同時保持在上萬個GPU上的線性擴展性。

說到底，veScale-FSDP的意義遠超一個技術系統的范疇。它代表了一種新的設計哲學：在追求極致性能的同時，不犧牲靈活性和易用性。就像一個優秀的工具應該既強大又好用一樣，veScale-FSDP讓研究人員能夠專注于AI模型本身的創新，而不用被復雜的分布式訓練細節所束縛。

這項技術的開源釋放，為整個AI社區提供了一個強大的工具。無論是學術研究機構還是產業界，都可以利用這項技術來訓練更大、更復雜的AI模型。而隨著AI模型規模的不斷增長，類似veScale-FSDP這樣的技術創新將變得越來越重要，它們正在默默地為AI技術的發展提供堅實的基礎設施支撐。

Q&A

Q1：RaggedShard是什么技術？

A：RaggedShard是veScale-FSDP系統的核心創新，它是一種新型的數據分片格式。傳統FSDP系統只能按照固定方式切分數據，而RaggedShard可以根據具體需求靈活調整分片的大小和形狀，就像可調節的模具一樣。這種靈活性讓它能夠完美支持Muon優化器和8位量化訓練等先進技術。

Q2：veScale-FSDP比傳統系統性能提升多少？

A：veScale-FSDP在不同測試中表現出顯著的性能優勢。在MoE稀疏模型上，吞吐量比其他系統高11%到66%，在密集模型上也有5%的提升。同時，內存使用降低了16%到30%。這些提升主要來自優化的通信重疊、零拷貝集合通信和智能的內存管理策略。

Q3：普通研究團隊能使用veScale-FSDP嗎？

A：完全可以。veScale-FSDP已經開源，包含7600行Python代碼，完全兼容標準PyTorch環境。它作為即插即用的模塊，使用與PyTorch FSDP2相同的API，現有訓練代碼幾乎無需修改就能享受性能提升。無論是學術研究還是產業應用，都可以直接使用這項技術。