# veScale-FSDP: Flexible and High-Performance FSDP at Scale > [!info] Talk metadata > - **会議:** [[MLSys2026]] Day 4 (May 21 / Thu)、Industry Track Oral "LLM Training 4"（14:45 PDT 開始） > - **登壇者:** Zezhou Wang / Youjie Li（ByteDance Seed / University of Washington） > - **全著者:** Zezhou Wang\*^{1,2}, Youjie Li\*^1, Zhiqi Lin\*^1, Jiacheng Yang\*^1, Cong Xie^1, Guanyu Feng^1, Zheng Zhong^1, Ziyue Huang^1, Hongyu Zhu^1, Zhi Zhang^1, Yanghua Peng^1, Xin Liu^1（\* equal contribution） > - **所属:** ^1 ByteDance Seed、^2 University of Washington > - **対応著者:** Youjie Li、Yanghua Peng > - **システム名:** veScale-FSDP > - **関連リンク:** https://github.com/volcengine/veScale > [!abstract] 概要（論文 PDF アブストラクト・忠実日本語訳） > FSDP（Fully Sharded Data Parallelism）は、メモリ効率の良い学習パラダイムとして広く採用されている。しかし既存の FSDP 実装は柔軟性と性能のトレードオフに直面しており、構造認識型の技術（ブロック単位量子化やテンソル構造に基づくオプティマイザなど）との統合を困難にしている。我々はこの課題を、FSDP のパラメータシャーディング設計に起因する根本的問題として特定する。すなわち要素単位・行単位といった均等シャーディングは、ゼロコピー通信を可能にしつつもテンソルの構造的意味論を壊すか、構造を保つがインターリーブドコピーのオーバーヘッドを招く。この対立を解消するため、我々は veScale-FSDP を提案する。本システムは 3 つの主要コンポーネントからなる。(1) `RaggedShard`：構造を保持しつつゼロコピー通信を実現するブロック単位の不均等シャーディング抽象、(2) 構造認識型のプランニングアルゴリズム：バケット内のシャードレイアウトを最適化しパディングオーバーヘッドを最小化、(3) `DBuffer`：計画されたレイアウトを決定論的メモリ管理とゼロコピー通信で実行する分散バッファ。実験により、veScale-FSDP は Muon オプティマイザなどの新興技術とシームレスに統合しつつ、既存システムを大幅に上回り、スループットで約 5〜66% 向上、メモリ使用量で 16〜30% 削減を達成し、数万基の GPU へ効率的にスケールすることを示す。 ## 問題設定 ### FSDP のシャーディングにおけるジレンマ FSDP は ZeRO-3 を実装する分散学習手法であり、パラメータ・勾配・オプティマイザ状態をデバイス間でシャードし、必要時に AllGather で復元する（論文 Section 2）。しかし LLM の進化に伴い、ブロック単位量子化（8-bit Adam など）やテンソル構造認識型オプティマイザ（Muon など）の利用が進んでおり、既存の FSDP シャーディング方式はこれらとの共存で本質的な困難を抱える。 既存のシャーディング方式には 2 種類ある（論文 Table 1 / スライド p.5--7）: 1. **要素単位シャーディング**（DeepSpeed ZeRO, PyTorch FSDP1）: テンソルを 1 次元に平坦化して均等分割する。ゼロコピー通信は可能だが、テンソルの形状・レイアウト意味論が失われ、ブロック単位量子化の境界が崩れ、構造認識型オプティマイザのワークフローが壊れる。 2. **行単位均等シャーディング**（PyTorch FSDP2, Megatron-FSDP）: 行方向にテンソル形状を保ってシャードする。DTensor レイアウトとして表現可能だが、通信バッファとの間でインターリーブドコピーのオーバーヘッド（12〜31%）が発生し、量子化ブロック境界が崩れる場合があり、構造認識型オプティマイザで冗長な表現が生じる。 > 論文 Table 1 の比較: 要素単位はゼロコピー通信とフラグメント管理に優れるが柔軟性に欠ける。行単位はテンソル並列統合に優れるが通信コストが高い。**いずれの方式も柔軟性と性能を両立できない**。 ### ゼロコピー通信の重要性 従来の行単位シャーディングでは、パラメータごとのコピーイン・コピーアウト（per-param copy-in/copy-out）とインターリーブドアドレスによる追加コピーが必要であり、通信オーバーヘッドが 12〜31% に達する（スライド p.5）。大規模クラスタではこのオーバーヘッドがスループットに直結するため、ゼロコピー通信の実現が不可欠である。 ## 提案手法 veScale-FSDP は 3 つのコンポーネントで構成される（スライド p.9 / 論文 Section 3--5）。 ### 1. `RaggedShard`: ブロック単位不均等シャーディング `RaggedShard` は PyTorch DTensor の拡張として実装された新たなシャーディング抽象である（論文 Section 3 / スライド p.8）。従来の均等シャーディングと異なり、ユーザ定義のブロックサイズでテンソルを分割し、各デバイスが異なる数のブロックを保持することを許容する。 利点: - **汎化されたシャーディングレイアウト**: 要素単位・行単位の均等シャーディングを特殊ケースとして包含する - **ゼロコピー通信**: ブロックが連続メモリ上に配置されるため、追加コピーなしで AllGather/ReduceScatter が実行可能 - **同期不要のブロック単位量子化**: 量子化ブロック境界がシャード境界と整合するため、デバイス間の通信なしに各デバイスがローカルシャードを独立に量子化できる - **DTensor 再分配による構造認識型オプティマイザ**: `RaggedShard` の不均等シャーディング対応により、Muon オプティマイザの Newton-Schulz 反復などをクリーンな SPMD パターンで記述可能。ルートランクに完全 2D パラメータ行列を再分配し、他ランクでは no-op とする `RaggedShard` は DTensor のオプショナルな placement として実装されており、テンソル並列やエキスパート並列などの既存インフラと容易に連携する（論文 Section 3.4）。 ### 2. 構造認識型プランニングアルゴリズム `RaggedShard` の柔軟性を活かすには、通信バッファ内のテンソル配置を最適化し、パディングオーバーヘッドを最小化する必要がある（論文 Section 4 / スライド p.10--13）。 FSDP ではバケッティング（複数テンソルのグループ化）が通信効率に不可欠である。テンソルをバケットに配置する際、`RaggedShard` ではブロック粒度が異なるテンソル間でシャード境界を揃える必要があり、テンソル間にパディングが挿入される。このパディングの最小化は **NP 困難**な最適化問題として定式化される（論文 Theorem 1, multiprocessor scheduling からの帰着）。 veScale-FSDP はヒューリスティック順序付け + 動的計画法による多項式時間アルゴリズムを提案する: 1. **ヒューリスティック順序付け**: テンソルをブロックサイズで降順ソート（大きいブロックから配置） 2. **DP 配置**: 各テンソルの配置オフセットを制約条件下で最適化。単調性を利用した状態圧縮により、計算量は $O(|\mathcal{T}|^2 m \log(E) \log(|\mathcal{T}|m))$ に抑えられる（$|\mathcal{T}|$ はテンソル数、$m$ はデバイス数、$E$ は最大テンソルサイズ） 実行時間は全実験で 0.3 秒未満であり、分散学習初期化において無視できる（論文 Section 6.4）。DeepSeek-V3-671B と GPT-OSS-120B の評価では、1 倍および 16 倍行粒度でパディングオーバーヘッドは全 FSDP サイズを通じて 3% 未満に収まる（論文 Fig.11）。 ### 3. `DBuffer`: 分散バッファによるゼロコピー実行 `DBuffer` は、プランニングアルゴリズムが決定したレイアウトを実行する分散通信バッファである（論文 Section 5 / スライド p.14）。 主要な設計: - **ゼロコピー**: グループレベルオペレータでテンソルを `DBuffer` にゼロコピーで配置し、インプレース AllGather を実行して復元後もゼロコピーで各テンソルを取り出す - **決定論的メモリ管理**: PyTorch の `record_stream` メカニズム（非決定論的メモリ解放）を排し、明示的なストリーム依存関係管理とバッチ割り当てにより予測可能なメモリ解放を実現。これにより予約メモリを 16〜30% 削減 - **フラグメンテーション低減**: 通信バッファの一括割り当てとバッチ解放によりメモリフラグメンテーションを抑制 ## 実験・評価 ### 実験設定 - **ハードウェア**: H800 GPU クラスタ（最大 10K GPU）、各ノード 8 GPU、NVLink intra-node、IB inter-node（論文 Section 6.1） - **ワークロード**: 8-bit Adam オプティマイザ、Muon オプティマイザ、GPT-OSS-120B（MoE モデル）、社内 MoE モデル（最大 800B パラメータ・2.4T パラメータ）、LLaMA-3-70B（デンスモデル） - **ベースライン**: DeepSpeed ZeRO, PyTorch FSDP1, PyTorch FSDP2, Megatron-FSDP（すべて ZeRO-3 + 混合精度構成） ### エンドツーエンド性能（GPT-OSS-120B、8-bit Adam） 128〜1024 GPU（4×256）での評価（論文 Fig.5, Table 1 / スライド p.16--17）: - **スループット**: veScale-FSDP は全ベースラインに対し **5〜66% 高い**スループットを達成。256 GPU では DeepSpeed が OOM、512 GPU（2×256）では FSDP1 がエラー・FSDP2 が OOM、1024 GPU（4×256）では FSDP1 がエラー・FSDP2 が OOM - **ピークメモリ**: veScale-FSDP は **16〜30% 低い**ピーク予約メモリ。128 GPU で veScale-FSDP は約 57 GB、他は 62〜77 GB - **スケーラビリティ**: DeepSpeed は 256 GPU で OOM、FSDP2 は 512 GPU で OOM であるのに対し、veScale-FSDP は 1024 GPU まで安定動作 ### LLaMA-3-70B（デンスモデル、8-bit Adam） 128〜512 GPU での評価（論文 Fig.6）: - veScale-FSDP は Megatron-FSDP を **11〜25%** 上回るスループット - FSDP1 は一貫して OOM - 512 GPU で veScale-FSDP は Megatron-FSDP と比較して高スループットかつ低メモリ ### メモリ削減の要因 veScale-FSDP のメモリ削減は決定論的な `DBuffer` メモリ管理に起因する（論文 Section 6.1）: - DeepSpeed と FSDP1 は PyTorch の暗黙的 `record_stream` メカニズムを継承し、キャッシュアロケータがバッファを再利用できないため予約メモリが 20% 膨張する - FSDP2 の行単位シャーディングはパディング膨張バッファを必要とし、MoE 実験でピークメモリが 33% 増加する - Megatron の混合精度サポートは低精度バッファを常駐させ、LLaMA-3 実験で veScale-FSDP より 24% 多いメモリを消費 ### スケーラビリティ（社内 MoE モデル、最大 10K GPU） 800B パラメータ MoE モデル（入力 2K〜16K トークン/GPU）を 1K〜8K GPU で評価（論文 Fig.9 / スライド p.18）: - **弱スケーリング**: GPU あたりバッチサイズ固定で GPU 数を増加。全入力サイズで**ほぼ線形のスケーラビリティ**を示す。FSDP の通信コストと GPU あたりの計算コストがいずれも GPU 数に依存せず一定であるため - **強スケーリング**: グローバルバッチ（120M トークン）を固定し GPU 数を増加。1K から 8K GPU で **3.4 倍**のスループット向上（16M トークンバッチ）。エキスパート並列との組み合わせで FSDP 通信時間を削減 ### モデルスケーリング（1K GPU） GPU 数を 1K に固定し、モデルサイズを 400B から 2.4T パラメータまで拡大（論文 Fig.9d）。`DBuffer` の効率的メモリ管理により、2.4T パラメータモデルでも 1K GPU のみで性能劣化なく学習可能。MFU はモデルサイズの増大に伴いわずかに向上する（計算密度の増加による）。 ### ケーススタディ: 8-bit Adam と分散 Muon - **8-bit Adam**（論文 Section 6.3 / スライド p.20）: `RaggedShard` の `orig_param_policy` インタフェースにより、パラメータごとに 32×32 ブロック・32 行粒度の量子化設定が可能。各デバイスがローカルシャードを独立に量子化し、ブロック境界が保持されるため**同期不要**。損失曲線は DDP 8-bit Adam とほぼ一致 - **分散 Muon**（論文 Section 6.3, Algorithm 2 / スライド p.19）: `RaggedShard` の不均等シャーディングと DTensor `redistribute` により、Muon の Newton-Schulz プリコンディショナをクリーンな SPMD パターンで実装。ルートランクに完全 2D パラメータを集約し、Newton-Schulz 反復を実行後に元の placement に再分配する。256 GPU の Hopper GPU 上で **47.3% MFU** を達成。損失曲線は DDP Muon と一致し、AdamW より約 80B トークン学習後に 0.01 低い損失に安定化 ### コンポーネント寄与分析 32 GPU、GPT-OSS 型モデル、8-bit Adam での ablation（論文 Table 2）: | コンポーネント | 正規化スループット | |---|---| | 全コンポーネント結合 | 100.0% | | `DBuffer` 無効化 | 92.8%（7.2% 低下） | | プランニングアルゴリズム無効化 | 65.4%（34.6% 低下） | | `RaggedShard` 無効化 | N/A（実行不能） | `DBuffer` はゼロコピー通信による 7.2% の寄与、プランニングアルゴリズムは量子化ブロックのシャード内完結保証による 34.6% の寄与を持つ。`RaggedShard` はブロック単位 8-bit Adam を成立させる前提であり、無効化するとモデル・オプティマイザの大幅書き換えか手動コレクティブが必要となり実質的に動作不能である。 ## 教訓（Lessons Learned） ByteDance での 10K GPU 超の実運用から得られた 3 つの教訓（論文 Section 7）: 1. **小規模ワークロードで大規模性能を予測可能**: FSDP の通信コストと計算コストは GPU 数に依存しないため、約 64 GPU でのプロファイルから数千 GPU の性能を高精度に外挿できる。弱スケーリング実験がこの観測を裏付ける 2. **巨人の肩の上にシステム抽象を設計せよ**: DTensor は既に多様な並列化技術を支援する強力な抽象を提供している。`RaggedShard` を DTensor のオプショナル placement として実装することで、テンソル並列・エキスパート並列・分散チェックポイントなどの既存インフラとシームレスに協調できる。`RaggedShard` は PyTorch 公式ロードマップにも計画機能として掲載されている 3. **モデル定義とシステム最適化を分離せよ**: モデルアーキテクチャの急速な進化に対応するため、Megatron-LM のようにモデルコードとシステム最適化を密結合するのではなく、veScale-FSDP ではモデル定義をシステムフレームワークから分離し、研究者がモデル設計に集中できるようにしている ## 結論・Takeaway - veScale-FSDP は `RaggedShard`（柔軟なシャーディングレイアウト）、構造認識型プランニングアルゴリズム（通信効率最適化）、`DBuffer`（ゼロコピー・低フラグメンテーション実行）の 3 層構成により、FSDP における柔軟性と高性能を両立する - 8-bit Adam や Muon など構造認識型オプティマイザとシームレスに統合し、既存 FSDP 実装（DeepSpeed ZeRO, PyTorch FSDP1/FSDP2, Megatron-FSDP）に対して**スループット 5〜66% 向上、メモリ 16〜30% 削減**を達成 - 最大 2.4T パラメータモデル・10K Hopper GPU までスケールし、ほぼ線形のスケーラビリティを実証 - OSS 公開: https://github.com/volcengine/veScale - `RaggedShard` は PyTorch 公式ロードマップに掲載されており、エコシステムへの還元も進んでいる