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

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

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

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

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

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

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

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

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

**二、智能規(guī)劃算法解決通信難題**

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

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

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

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

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

**三、分布式緩沖區(qū)實現(xiàn)零拷貝通信**

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

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

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

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

**四、突破性的性能表現(xiàn)**

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

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

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

**五、支持前沿優(yōu)化技術的靈活性驗證**

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

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

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

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

**六、工程實現(xiàn)的精巧設計**

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

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

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

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

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

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

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

**八、生產環(huán)境的實戰(zhàn)驗證**

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

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

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

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

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

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

Q&A

Q1：RaggedShard是什么技術？

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

Q2：veScale-FSDP比傳統(tǒng)系統(tǒng)性能提升多少？

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

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

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