> [!abstract] 概要(abstract 日本語訳) > Fully Sharded Data Parallel(FSDP、Zero Redundancy Optimizer とも呼ばれる)は、メモリ効率とモデルコードへの最小限の侵襲性のために大規模モデル訓練で広く使われている。しかし、既存の FSDP システムは固定的な要素単位または行単位のシャーディングフォーマットに依存しており、これはブロック構造を持つ計算と衝突する。その結果、ブロック単位量子化や Shampoo・Muon のような非要素単位オプティマイザを含む、現代の構造認識学習手法をサポートするのに苦労している。さらに、今日の実装は通信・メモリのオーバーヘッドを抱えており、数万 GPU 規模での効率を低下させている。我々は veScale-FSDP という新しい FSDP システムを導入する。これは柔軟なシャーディングフォーマットである RaggedShard と、構造認識プランニングアルゴリズムを組み合わせ、柔軟性と性能の両方を提供する。veScale-FSDP はゼロコピー FSDP 通信を可能にし、ブロック単位量子化と非要素単位オプティマイザをネイティブにサポートしながら、既存 FSDP システムに対しスループット 5〜66% 向上・メモリ使用量 16〜30% 削減を達成し、数万 GPU への効率的なスケーリングを実現する。 ## 論文情報 - タイトル: veScale-FSDP: Flexible and High-Performance FSDP at Scale - 著者: Zezhou Wang*, Youjie Li*, Zhiqi Lin*, Jiacheng Yang*, Cong Xie, Guanyu Feng, Zheng Zhong, Ziyue Huang, Hongyu Zhu, Zhi Zhang, Yanghua Peng, [[Xin Liu]](*は Equal contribution) - 所属: [[ByteDance Seed]] / [[University of Washington]](第 2 著者インターン当時の所属) - 媒体・発表年: 第 9 回 MLSys Conference(Industry Track)、Bellevue, WA, USA, 2026 - コード: RaggedShard 実装が OSS 化されている(https://github.com/volcengine/veScale) - 発表スライド: MLSys 2026 Industry Track、Slide ID 3860(94 ページ、段階的ビルドアニメーション形式) ## 概要 veScale-FSDP は、PyTorch ネイティブの `fully_shard` API を維持したまま FSDP2 のバックエンドを再設計した FSDP システムである。新しいシャーディングフォーマット RaggedShard により任意のブロック粒度でのシャーディングを可能にし、構造認識プランニングアルゴリズムで通信バッファのレイアウトを最適化し、Distributed Buffer(DBuffer)というゼロコピー通信プリミティブでこれらを支える。ByteDance Seed の本番環境で 10K GPU 級のワークロードに展開されている。 ## 問題設定 - **入力**: PyTorch `fully_shard` API で並列化される任意の大規模モデル(密モデル・MoE モデル)と、構造化されたテンソルブロックを要求するオプティマイザ(Shampoo、Muon)や量子化手法(ブロック単位 INT8/FP8)。 - **前提条件**: 既存 FSDP システムは DTensor の `Shard(dim)` プレースメントを使い、要素単位(DeepSpeed・FSDP1)または行単位(FSDP2・Megatron-FSDP)にテンソルを均等分割する。この均等分割はブロック境界と噛み合わず、非要素単位計算・ブロック単位量子化を阻害する。 - **必要なデータ**: 訓練対象モデルのパラメータ構造(行列形状・ブロックサイズ)と、デバイス数・ネットワークトポロジ(NVLink・NCCL コレクティブ設定)。 ## 提案手法 ### アーキテクチャ veScale-FSDP は 3 層構成を取る(Figure 3)。 - **RaggedShard DTensor**: 複雑な大規模モデル・構造認識オプティマイザを PyTorch ネイティブ API の下で表現する新しい DTensor プレースメント。 - **Planning Algorithm**: RaggedShard DTensor 群を効率的な通信のためにグループ化・レイアウトする多項式時間アルゴリズム。 - **Distributed Buffer(DBuffer)**: 上記のレイアウトを実際のゼロコピー通信プリミティブへマッピングする基盤。 **Figure 3: veScale-FSDP のシステム概要** ![[_attachments/48_veScale_FSDP_Flexible_and_H/fig03-vescale-fsdp-overview.png]] (Figure 3. veScale-FSDP overview. モデル・オプティマイザが PyTorch-Native API を介して veScale-FSDP に接続し、RaggedShard DTensor → Planning Algorithm → Distributed Buffer(DBuffer)の 3 層を通じて 10K デバイスへスケールする構成を示す。) ### RaggedShard フォーマット(§4、柔軟性) 既存のシャーディングフォーマットは 2 種類に大別される: 構造を無視して任意に分割する Element-wise Shard(DeepSpeed・FSDP1 の基盤)と、テンソル次元に沿って均等分割する Row-wise (Even) Shard(FSDP2 の基盤)。前者は非要素単位演算・コレクティブ・量子化のいずれとも相性が悪く、後者は非要素単位演算とコレクティブは扱えるがブロック単位量子化とは整合しない。 RaggedShard は JaggedTensor/NestedTensor(PyTorch/TensorFlow の既存プリミティブ)に着想を得て、**任意のシャーディング粒度**(要素・行・任意次元のブロック)と**任意のシャーディング分布**(各デバイスに配置するブロック数)を DTensor の追加プレースメントとして提供する。粒度を 1 行に設定すると Row-wise RaggedShard(デバイスごとに行数が異なりうる)、粒度をカスタムブロック形状にすると Block-wise RaggedShard(1 テンソルを複数ブロックに分割し、デバイスごとに異なる個数のブロックを配置)になる。Block-wise RaggedShard はブロック境界とシャード境界を完全に一致させられるため、ブロック単位量子化を通信なしで実現できる。 **Figure 4: シャーディングフォーマットの柔軟性比較** ![[_attachments/48_veScale_FSDP_Flexible_and_H/fig04-sharding-format-comparison.png]] (Figure 4. Flexibility comparison of different sharding formats. Element-wise Shard・Row-wise Shard・Row-wise RaggedShard・Block-wise RaggedShard の 4 種類を、非要素単位計算・Redistribute(通信)・ブロック単位量子化・汎用性の 4 軸で比較し、Block-wise RaggedShard のみが全軸を満たすことを示す。) RaggedShard は既存の DTensor プレースメント(Replicate・Partial)と直交し、Shard と組み合わさる際は専用の `StridedRaggedShard` プレースメントで並べ替え・ストライドのメタデータを扱う。これにより Tensor Parallelism(`Shard(0)`/`Shard(1)`)・Expert Parallelism(`Shard(0)`)といった既存の 2D 並列スキームとシームレスに合成できる。 **Figure 6 の準備として**、まず既存の Element-wise Shard がなぜ非要素単位演算に破綻するか(Figure 1)を確認する。 **Figure 1: DTensor による柔軟な通信・計算** ![[_attachments/48_veScale_FSDP_Flexible_and_H/fig01-dtensor-communication.png]] (Figure 1. Distributed Tensor (DTensor) for flexible communication and computation. 行シャード(`Shard(0)`)の DTensor を `redistribute` で複製(`Replicate`)に変換し、列シャード(`Shard(1)`)の DTensor と行列積(`matmul`)することで、列シャードの DTensor を得る例。青色部分が各デバイスに実体化されているローカルテンソルを示す。) ### プランニングアルゴリズム(§5、性能) RaggedShard DTensor を素朴に通信バッファへ連結すると 3 種の非効率が生じる(Figure 6(a)): ①ブロックがデバイス境界を跨いで分断される(Sharded block)、②コレクティブのアライメント要件のためにテンソル内部にパディングが挿入され連続性が壊れる(Non-contiguous tensor memory)、③テンソルサイズ・ブロックサイズ・パディングサイズの違いによりデバイスごとのバッファサイズが不均衡になる(Imbalanced load)。 veScale-FSDP はまずテンソルを並べ替え(Permute Tensors)、次にテンソル内部ではなくテンソル間にパディングを挿入する(Pad Between Tensors)という 2 段階アプローチでこれを解決する(Figure 6(b))。 **Figure 6: グループ化 RaggedShard DTensor の通信レイアウト** ![[_attachments/48_veScale_FSDP_Flexible_and_H/fig06-grouped-communication-layout.png]] (Figure 6. Communication layouts for grouped RaggedShard DTensors. (a) 素朴な連結戦略はブロックを分断・非連続メモリ・不均衡負荷を招き非効率な通信になる。(b) 構造認識プランナーはテンソルの並べ替え→テンソル間パディングにより、非シャードブロックの保持・連続メモリ・均衡負荷を同時に達成し効率的な通信を実現する。) この問題は形式的に、$m$ デバイスに分割された RaggedShard DTensor 集合 $T$ に対し、各デバイスあたりの通信バッファサイズ $S$ を最小化する最適化問題として定式化される。制約は「非シャードブロック(各テンソルの区間 $[\ell_t, r_t)$ が連続かつシャードサイズ内に収まる)」「連続テンソルメモリ(テンソル同士が重ならない)」「均衡負荷($kS$ がテンソル境界を跨ぐならブロック粒度 $g_t$ の倍数)」の 3 つである。この問題は古典的な Partition 問題(Garey & Johnson, 1975)に帰着でき、NP-hard であることが証明される。 ILP ソルバーは数十分オーダーの実行時間がかかり実運用では非現実的なため、veScale-FSDP はトランスフォーマーモデルの構造的規則性(線形重みが大部分を占め、シャーディングブロックが層間で一貫する)を利用した多項式時間の動的計画法(DP)ヒューリスティックを提案する(Algorithm 1)。テンソルの並び順は既定順序・ブロックサイズ順・テンソル形状順の 3 通りを検討したが、実験上は既定順序で最適または準最適な解が得られたためこれを採用する。DP は「各テンソルが単一シャードに収まる」「隣接シャードに跨るが 1 シャードを完全には含まない」「1 シャード以上を完全に含む」の 3 ケース解析に基づき、時間計算量 $O(|T|^2 m \log(E) \log(|T|m))$ で最小シャードサイズ $S^\star$ を求める。アルゴリズム自体の実行時間は全実験で 0.3 秒未満であり、訓練初期化の一回限りのコストとして無視できる(§6.4)。 ### Distributed Buffer(DBuffer、§5) DBuffer は DTensor に着想を得て、N 次元デバイストポロジ上のグローバルバッファセマンティクスを提供する新しいプリミティブである。3 つの性能上の工夫を持つ: 1. グループレベル演算子: テンソルごとに個別の CUDA カーネル(add・scale・zero・copy)を発行する代わりに、同一カーネルをテンソル群にわたって融合し通信をブロックするフラグメント化計算を削減する。 2. ゼロコピーアクセス: RaggedShard のプランニングアルゴリズムを活用し、各テンソルのデータポインタへの永続的なアドレスマッピングを提供することで通信前後のコピーを不要にする。 3. インプレース通信・計算: 通信・計算をバッファ上で直接実行する。 **Figure 7: Distributed Buffer(DBuffer)による高性能通信** ![[_attachments/48_veScale_FSDP_Flexible_and_H/fig07-dbuffer.png]] (Figure 7. Distributed Buffer (DBuffer) for high-performance communication. パラメータ AllGather 用の 2D DBuffer(上段)と、`(Partial, Partial)` から `(Replicate, Shard)` へ再分配し ReduceScatter + AllReduce による 2D 勾配リダクションを実装する DBuffer(下段)の 2 例を示す。ゼロコピーでテンソルとバッファ間を往復する。) 対して既存の FSDP2 は per-parameter の `Shard(0)` DTensor 設計のため、AllGather 後に interleaved アドレスからパラメータをコピーアウトし、ReduceScatter 前にコピーインする必要がある(Figure 2)。この Copy-In/Copy-Out オーバーヘッドは訓練イテレーションの最大 14% を消費しうる。 **Figure 2: FSDP2 設計におけるコピーオーバーヘッド** ![[_attachments/48_veScale_FSDP_Flexible_and_H/fig02-fsdp2-copy-overhead.png]] (Figure 2. The copy overhead in the FSDP2 design. 各パラメータは `Shard(0)` DTensor として均等シャードされる。AllGather 前に各デバイスの通信入力バッファへ全パラメータのローカルシャードをコピーインし、AllGather 後は出力バッファ内で各パラメータが interleaved アドレスに配置されるため、計算に必要な連続アドレスへコピーアウトしなければならない。ReduceScatter パスは逆方向で interleaved Copy-In オーバーヘッドを負う。Table 1(GPT-OSS-120B、64 H800)では `Shard(0)` で AllGather 43.71ms に対し Copy-Out 5.22ms、ReduceScatter 94.24ms に対し Copy-In 12.37ms が計測されている。) ## 新規性 - **既存研究との違い**: DeepSpeed ZeRO はフラグメント化した AllGather と非効率なメモリ管理に苦しむ。PyTorch FSDP1 は AllGather の非効率性を一部解消するが遅い ReduceScatter を残す。FSDP2 は per-parameter の均等シャーディングでメモリ管理を改善したが、上記の interleaved コピーオーバーヘッドを新たに導入した。Megatron-FSDP はゼロコピー設計だが、DTensor の `Shard(0)` セマンティクスと互換に保つために追加パディングが必要で、MoE モデルで最大 33% のバッファ膨張(padding inflation)を招く。 - **veScale-FSDP の解決策**: 固定シャーディング粒度を捨て、RaggedShard による任意粒度と構造認識プランニングによるパディング最小化を同時に達成することで、上記のトレードオフ(柔軟性 vs 性能)を両立させる。 - **抽象化の設計思想**: RaggedShard は DTensor の上に構築されたオプションのプレースメントとして実装されており、既存の Tensor Parallelism・Expert Parallelism・分散チェックポイント基盤(PyTorch Distributed Checkpoint)とそのまま連携できる(§7 Lesson-2)。この設計は PyTorch チームの公式ロードマップ(2026 H1)にも将来機能として掲載されている。 ## 実験設定 - **ハードウェア**: §6.1・§6.4・§6.5 は NVIDIA H800 クラスタ(1 ノードあたり 8×H800、979 BF16 TFLOPS、80GB HBM、400GB/s NVLink)。§6.2・§6.3 は NVIDIA Hopper クラスタ。 - **実装**: Python 7.6K 行(LoC)。FSDP2 のバックエンドを透過的に置き換え、PyTorch ネイティブの `fully_shard` API を維持するプラグアンドプレイモジュール。 - **ベースライン**: DeepSpeed ZeRO v0.17.6、PyTorch 2.7.1 FSDP1、PyTorch 2.7.1 FSDP2(`fully_shard`)、Megatron-FSDP。すべて ZeRO-3 + 混合精度(FP32 マスター重み・BF16 前向き後退)で統一。 - **モデル**: Llama-3-70B(密モデル)、GPT-OSS-120B(MoE)、内部 MoE モデル(Internal-Model-160B、および §6.2 では最大 2.4T パラメータ)。 - **評価指標**: 正規化スループット(1024 GPU 構成を基準に正規化した tokens/s)、ピーク GPU あたりメモリ(GB)、パディングオーバーヘッド率(%)、Model FLOPS Utilization(MFU、%)。 ## 実験結果 ### エンドツーエンド性能(§6.1、Figure 8) **Figure 8: FSDP 訓練性能(スループット・メモリ)** ![[_attachments/48_veScale_FSDP_Flexible_and_H/fig08-throughput-memory-results.png]] (Figure 8. FSDP training performance. 上段: 正規化された集約スループット(tokens/s)。下段: ピーク GPU あたりメモリ(GB)。128/256 GPU での FSDP(ZeRO-3)と、2 way/4 way レプリケーションの HSDP(2*256, 4*256 GPU)を掃引。Internal-Model-160B・GPT-OSS-120B・Llama-3-70B の 3 モデルで、veScale-FSDP(赤)が一貫して最高スループット・最低メモリを示す。GPT-OSS-120B では FSDP2 が 256 GPU 超で OOM、Megatron-FSDP が ERROR になる区間がある。) - MoE モデルでは veScale-FSDP がベースラインに対し 11〜66% 高いスループットを示す。Llama-3-70B ではベースラインより 5% 高速。 - メモリはベンチマーク全体でピーク予約メモリを 16〜30% 削減。DBuffer の決定論的・バッチ化されたメモリ管理(明示的なストリーム依存管理と一括アロケーション)による。 - FSDP2 は GPT-OSS-120B を 128 デバイスでは訓練できるが 256 デバイスで OOM する。128 エキスパートを 256 デバイスに分散すると per-parameter シャーディングのパディングにより AllGather バッファが実質倍増するためである。 - Megatron-FSDP は固定 `Shard(0)` 粒度のため MoE モデルで 33% のバッファパディング膨張が生じ、混合精度バッファの低精度保持により Llama-3 実験でメモリを veScale-FSDP より 24% 多く消費する。 ### スケーラビリティ・合成性(§6.2、Figure 9) **Figure 9: veScale-FSDP のスケーラビリティ** ![[_attachments/48_veScale_FSDP_Flexible_and_H/fig09-scalability.png]] (Figure 9. Scalability of veScale-FSDP. (a) 8K GPU までの弱スケーリング。(b) 8K GPU までの強スケーリング。(c) より大規模での強スケーリング(正規化スループット・正規化 GPU 数)。(d) 1K GPU 上でのモデルスケーリング(2.4T パラメータまで)。) - **弱スケーリング**: 800B パラメータの内部 MoE モデルを 1K〜8K GPU で、GPU あたり入力サイズ 2K〜16K トークン固定で訓練し、ほぼ線形にスケール。 - **強スケーリング**: グローバルバッチを 16M〜128M トークンに固定し、120M トークンのグローバルバッチで 10K GPU まで線形にスケールしつつ、1K→8K GPU で 3.4 倍のスループット向上を達成。GPU 数が非常に多くなると GPU あたりトークン数が減り FSDP 通信(AllGather/ReduceScatter)が支配的になるため、クロスノード Expert Parallelism を併用して通信時間を削減する(トークン交換のための計算コスト増とのトレードオフ)。 - **モデルスケーリング**: GPU 数 1K・入力 8K トークン/GPU 固定で 400B〜2.4T パラメータへスケール。DBuffer の効率的なメモリ管理により、モデルサイズが大きいほど計算密度が上がり MFU がわずかに向上する。 - **合成性**: Expert Parallelism との組み合わせにより、10K の Hopper GPU で最大 2.4T パラメータの内部モデルまで訓練可能。 ### 構造認識オプティマイザ・量子化ケーススタディ(§6.3) - **8-bit Adam**: 勾配統計にブロック単位 INT8 量子化を適用しオプティマイザ状態メモリを削減する手法。veScale-FSDP は `orig_param_policy` インターフェースでパラメータごとの量子化粒度を設定可能にし、32×32 ブロック・32 行粒度の実験では各デバイスが通信なしでローカルシャードを量子化できる(RaggedShard がブロック境界を保存するため)。DDP 8-bit Adam との損失曲線はほぼ一致(Figure 10a、本文中で言及、画像は未埋め込み)。 - **Distributed Muon**: Muon の行列符号プレコンディショナ(Newton–Schulz)は元の 2D 形状を保持した全体行列を要求する。RaggedShard の不均等シャーディング対応により、ルートランクへの `redistribute` → Newton–Schulz 更新(他ランクは no-op)→ 更新の再分配、という SPMD な実装が可能になる(Algorithm 2)。256 Hopper GPU で 47.3% MFU を達成(通信・計算のオーバーラップと `torch.compile` の活用による)。 ### プランニング品質と性能内訳(§6.4・§6.5) - **パディングオーバーヘッド(Figure 11、本文中で言及、画像は未埋め込み)**: DeepSeek-V3-671B・GPT-OSS-120B で行粒度 1×/16×/128× を掃引すると、1×・16× 粒度ではパディングオーバーヘッドが全 FSDP サイズで 3% 未満に収まる。128 行(DeepSeek 式 128×128 タイリング相当)では GPT-OSS が最大 18% までのステップ状スパイクを示す(全エキスパートを単一パラメータテンソルに融合しているため per-expert パディングの余地が少ない一方、DeepSeek-V3 はエキスパートを個別パラメータとして扱うためパディングの融通が利く)。 - **コンポーネントアブレーション(Table 2、GPT-OSS 系モデル・8-bit Adam・32 GPU)**: DBuffer を無効化するとスループットが 92.8% に低下(7.2% 減、コピーによるオーバーヘッド)。プランニングアルゴリズムを無効化すると 65.4% に低下(34.6% 減、量子化ブロックがシャード境界を跨ぎ DTensor の再分配にフォールバックするため)。RaggedShard を無効化すると実質的に動作不能(N/A)——32×32 ブロック境界とシャード境界を手動で一致させるか、per-block メタデータ交換を手書きする必要があるため。 ## 考察 - RaggedShard は「最適化」ではなく「抽象化」である点が Table 2 のアブレーションから読み取れる: DBuffer・Planning Algorithm の無効化はスループット低下(それぞれ 7.2%・34.6% 減)にとどまるが、RaggedShard の無効化はブロック単位量子化そのものを侵襲的な変更なしには実現不可能にする。 - パディングオーバーヘッドの分析(§6.4)は「モデル構造(全エキスパート融合 vs エキスパート個別化)」がプランニングアルゴリズムの有効性に影響することを示しており、モデル設計者がハードウェア co-design の観点から隠れ層サイズを小さな合成因子の倍数に選ぶといった実践的指針を導いている。 - Lesson-1(§7)は「小規模ワークロードの性能特性から大規模性能を予測できる」という実践的知見であり、GPU あたり計算時間と FSDP 通信時間がいずれも GPU 数に依存しないという弱スケーリングの原理に基づく。ただしこの外挿はネットワークトポロジ・コレクティブアルゴリズム/プロトコル・十分な帯域飽和を仮定する。 ## 強み / 弱点・課題 **強み** - RaggedShard という単一の抽象化で、柔軟性(構造認識学習のサポート)と性能(ゼロコピー通信)を同時に達成している。 - NP-hard な最適化問題を実用的な多項式時間 DP ヒューリスティックに帰着させ、実行時間 0.3 秒未満というオーバーヘッドの小ささを実証。 - PyTorch ネイティブ API・DTensor 抽象化の上に構築されているため、既存の Tensor/Expert Parallelism・分散チェックポイントとの合成性が高く、可搬性がある(内部インフラに非依存)。 - ByteDance Seed の本番環境で 10K GPU 超の実運用にすでに投入されている。 **弱点・課題** - ヒューリスティックのテンソル並び順は既定順序を採用しており、トランスフォーマー以外のアーキテクチャでの最適性は「代替順序をプラグイン可能」と述べるにとどまり、非トランスフォーマーでの実証データは示されていない。 - GPT-OSS-120B のようにエキスパートを単一パラメータテンソルへ融合するモデル構造では、大きな粒度(128 行)でパディングオーバーヘッドが最大 18% までスパイクする(§6.4)。実践的な回避策(FSDP グループサイズの抑制、HSDP との併用)は述べられているが、根本解決ではない。 - 8-bit Adam・Muon のケーススタディは損失曲線と MFU の実証にとどまり、これら構造認識オプティマイザ自体の収束性・精度への影響についての深い分析は本論文のスコープ外。