veScale-FSDP: Flexible and High-Performance FSDP at Scale

既存の FSDP システムが抱える構造化トレーニング手法や非要素単位のオプティマイザへの対応不足、および通信・メモリ効率の課題を解決するため、柔軟な「RaggedShard」形式と構造認識型計画アルゴリズムを採用した veScale-FSDP を提案し、これにより既存システムと比較してスループットを最大 66% 向上させ、メモリ使用量を最大 30% 削減しながら数万 GPU 規模での効率的なスケーリングを実現しました。

要約

この論文は、**「超巨大な AI（大規模言語モデル）を、何万枚もの GPU（高性能な計算チップ）を使って、より速く、より安く、より柔軟に訓練するための新しい仕組み」**を紹介しています。

タイトルは「veScale-FSDP」ですが、これを日常の言葉で説明しましょう。

🧩 従来の問題：「レゴブロック」の詰め込み方

AI を訓練するには、何兆もの「パラメータ（知識の断片）」を何千枚もの GPU に分け持って計算させる必要があります。これを**「FSDP（フルシャードデータ並列）」**と呼びます。

これまでのシステム（Deepspeed や PyTorch の標準機能など）は、このパラメータを**「均等な大きさの箱」**に詰めていました。

• 問題点 1：形が合わない。
  最新の AI は、計算の仕方が「ブロック単位」だったり、特定の行列（2 次元の表）全体が必要だったりします。でも、従来のシステムは「均等な箱」しか用意できないので、**「あ、このブロックは 2 つの箱にまたがっちゃった！」「この計算には全部のデータが必要なのに、バラバラになっちゃった！」**という事態が起きます。
  • アナロジー: 大きなピザを均等に切り分けるのは簡単ですが、ピザの上に「大きな具材（ブロック）」が乗っていて、それが「2 人の境界線」をまたいでいると、誰がその具材を処理すればいいか分からなくなります。その結果、具材をやり取りするために余計な手間（通信）がかかり、遅くなります。
• 問題点 2：無駄な隙間。
  均等に割ろうとすると、最後の端っこが小さくなりすぎます。それを埋めるために「ダミーのデータ（パディング）」を無理やり入れないと計算できません。
  • アナロジー: 100 人の人が 10 台のバスに乗る時、10 人ずつ均等に座らせます。でも、もし「10 人組の家族」がいて、その家族が 1 台のバスに収まらず、2 台に分かれて座らざるを得ない場合、バスには空席が生まれたり、家族を運ぶために余計なバスが必要になったりします。これが「メモリ（記憶容量）の無駄」や「通信の遅延」になります。

✨ 解決策：veScale-FSDP の「3 つの魔法」

この論文の著者（ByteDance のチーム）は、この問題を解決するために、**「veScale-FSDP」**という新しいシステムを作りました。

1. 「RaggedShard（ラギッド・シャード）」：形に合わせた自由な箱

従来の「均等な箱」ではなく、**「必要な形に合わせて自由に大きさを調整できる箱」**を使います。

• アナロジー: 従来のシステムが「同じサイズの段ボール箱」しかなかったのに対し、veScale は**「ピザの具材の形にぴったり合う、カスタムメイドの箱」**を用意します。
• 効果: 最新の AI が使う「ブロック単位の計算」や「特殊な最適化アルゴリズム」を、システム側が邪魔することなく、そのまま実行できます。開発者は AI のコードを無理やり書き換える必要がなくなります。

2. 「計画アルゴリズム」：最高の積み方

「自由な箱」を使うと、逆に「どう並べれば一番効率的か？」という問題が生まれます。これを解決するのが、**「計画アルゴリズム」**です。

• アナロジー: 荷物をトラックに積む時、ただ適当に積むと隙間ができてしまいます。veScale は**「賢いロボット」**が、どの箱をどの位置に置けば、トラック（GPU）の容量を最大限に使い、かつ積み下ろし（通信）の時間が最小になるかを瞬時に計算します。
• 効果: 無駄な隙間（パディング）を極限まで減らし、通信速度を最大化します。

3. 「Distributed Buffer（DBuffer）」：ゼロコピーの高速道路

データを GPU から GPU へ移動させる際、従来のシステムは「一度コピーして、また移動させる」という無駄な作業をしていました。veScale は**「直接指差しで渡す（ゼロコピー）」**仕組みを作りました。

• アナロジー: 従来の方法は「手紙をコピーして、相手に渡す」でしたが、veScale は**「手紙そのものを、相手の手元にあるファイル箱から直接取り出して見せる」**ようなものです。コピー作業が不要なので、スピードが劇的に上がります。

🚀 結果：どれくらいすごい？

この新しいシステムを使うと、以下のような劇的な改善が実現しました。

• 速度アップ: 既存のシステムより5%〜66% 高速に訓練できます。
• メモリ節約: GPU のメモリ使用量を16%〜30% 削減できます。
  • アナロジー: 100 台のバスで 1 万人を運んでいたのが、同じ 100 台のバスでより快適に、あるいは 70 台のバスで同じ人数を運べるようになったイメージです。
• 大規模対応: 何万枚もの GPU を使った超巨大な AI 訓練でも、効率よく動きます。
• 柔軟性: 「8 ビット Adam」や「Muon」といった、最新の高度な AI 技術も、特別な改造なしで使えます。

💡 まとめ

この論文は、**「AI 訓練という巨大なパズルを、従来の『均等な箱』という硬いルールで無理やり解こうとするのではなく、『形に合わせた自由な箱』と『賢い積み方』で解くことで、圧倒的に速く、安く、大きくできる」**という画期的なアプローチを示しています。

これにより、より高性能な AI が、より少ないコストで、より早く作れるようになることが期待されています。

技術概要

veScale-FSDP: 大規模スケールにおける柔軟性と高性能を両立する FSDP の技術的概要

本論文は、ByteDance Seed によって提案されたveScale-FSDPについて詳述しています。これは、大規模言語モデル（LLM）の学習において広く用いられている「Fully Sharded Data Parallel (FSDP)」を再設計したシステムです。既存の FSDP システムが抱える構造化された学習手法への対応不足や、通信・メモリ効率の限界を克服し、柔軟性と高性能を両立させることを目指しています。

以下に、問題定義、手法、主要な貢献、評価結果、および意義について詳細をまとめます。

---

1. 背景と課題 (Problem)

大規模モデルの学習には、数千枚の GPU にモデル状態とオプティマイザ状態を分散させる FSDP（ZeRO）が不可欠です。しかし、既存の FSDP システム（DeepSpeed ZeRO, PyTorch FSDP1/2, Megatron-FSDP など）には以下の重大な課題がありました。

• 構造化学習手法への非対応:
  • 最先端モデル（Gemini, Kimi K2, DeepSeek-V3 など）では、非要素単位のオプティマイザ（Shampoo, Muon）やブロック単位量子化（Block-wise Quantization）が採用されています。
  • これらの手法は、テンソルを特定の「ブロック単位」で処理する必要があります。
  • 既存の FSDP は「要素単位（Element-wise）」または「行単位（Row-wise）」の均等なシャリングしかサポートしておらず、ブロック境界とシャリング境界が一致しません。その結果、モデル開発者はコードを侵入的に修正するか、システム側で複雑なパディングや通信ロジックを実装せざるを得ませんでした。
• 通信とメモリ効率の限界:
  • 既存の実装では、通信バッファの未整列や、パラメータのコピーオーバーヘッド（FSDP2 におけるインターリーブされたコピー）、パディングによるメモリ浪費が発生しています。
  • これにより、数万枚の GPU へのスケーリングが困難であり、メモリ不足（OOM）やスループット低下を招いています。

2. 提案手法 (Methodology)

veScale-FSDP は、柔軟なシャリング形式、構造認識型の計画アルゴリズム、高性能なプリミティブを組み合わせて、上記の課題を解決します。

2.1 RaggedShard（不規則なシャリング形式）

既存の均等なシャリングに代わる、新しいシャリング形式RaggedShardを導入しました。

• 特徴: 任意の粒度（ブロック単位など）と、デバイス間の任意の分散配置を表現できます。
• 柔軟性: 非要素単位のオプティマイザやブロック単位量子化において、シャリング境界がブロック境界と完全に一致するように設定可能です。
• 互換性: PyTorch の DTensor プリミティブに統合されており、既存のテンソルパラレル（TP）やエキスパートパラレル（EP）と組み合わせて使用できます。

2.2 構造認識型計画アルゴリズム (Structure-Aware Planning)

RaggedShard テンソルを通信バッファに配置する際、パディングを最小化し通信効率を最大化する計画アルゴリズムを開発しました。

• 最適化問題: テンソルを通信バッファに配置する問題は、パディング最小化と負荷分散の観点から NP 困難問題として定式化されます。
• ヒューリスティック解法: 大規模スケールでの実用性を考慮し、多項式時間の動的計画法（DP）とヒューリスティック（テンソルの並べ替え戦略）を用いて、高品質な解を高速に導出します。
• 効果: 通信バッファ内のブロックが連続的に配置され、不要なパディングや断片化を抑制します。

2.3 Distributed Buffer (DBuffer)

通信と計算の効率化を実現する新しいプリミティブDBufferを導入しました。

• ゼロコピー: 計画アルゴリズムにより、テンソルデータへのポインタを直接マッピングし、通信前後のコピー（Copy-in/Copy-out）を排除します。
• バッチ割り当て: メモリ割り当てをバッチ処理することで、メモリフラグメンテーションを削減し、メモリ使用量を最適化します。
• インプレース操作: 通信と計算をインプレースで行い、オーバーヘッドを最小化します。

3. 主要な貢献 (Key Contributions)

1. RaggedShard の導入: 任意のシャリング粒度と配置を可能にする新しい DTensor プレイスメント形式。これにより、ブロック単位量子化や非要素単位オプティマイザを侵入的変更なしでサポート可能になりました。
2. NP 困難な配置問題への実用的解法: 大規模クラスターでも実行可能な、高速かつ高品質な通信バッファ配置計画アルゴリズムの提案。
3. DBuffer による高性能化: ゼロコピー通信とバッチメモリ管理による、通信オーバーヘッドとメモリ使用量の大幅な削減。
4. 実生産環境での検証: 1 万枚以上の GPU クラスターでのスケーリング実績と、Muon オプティマイザや 8-bit Adam などの最先端手法へのネイティブ対応。

4. 評価結果 (Results)

大規模モデル（LLaMA-3-70B, GPT-OSS-120B, 内部 MoE モデルなど）を用いた評価により、以下の成果が確認されました。

• スループットの向上:
  • 既存の FSDP システム（DeepSpeed, FSDP1, FSDP2, Megatron-FSDP）と比較して、5%〜66% のスループット向上を達成しました。
  • 特に MoE モデルや非要素単位オプティマイザを使用するケースで顕著な改善が見られました。
• メモリ使用量の削減:
  • ピークメモリ使用量が16%〜30% 削減されました。
  • DBuffer によるメモリ管理の最適化と、パディングの削減が寄与しています。
• スケーラビリティ:
  • 1 万枚の GPUまで効率的にスケーリング可能であることを実証しました。
  • 弱スケーリング（入力サイズ固定）、強スケーリング（バッチサイズ固定）、モデルサイズスケーリング（400B〜2.4T パラメータ）のすべてにおいて、高い効率を維持しています。
• 柔軟性の検証:
  • 8-bit Adam: RaggedShard によりブロック単位量子化をネイティブにサポートし、追加の通信オーバーヘッドなしで実装可能になりました。
  • Muon オプティマイザ: 非要素単位の行列演算を、複雑な通信ロジックなしで分散実行可能にしました。

5. 意義と結論 (Significance)

veScale-FSDP は、大規模モデル学習の次の段階における重要な基盤技術です。

• モデル開発とシステム最適化の分離: 研究者がモデルアーキテクチャやオプティマイザに集中できる環境を提供し、システム側の最適化（FSDP）を透明化しました。これにより、最先端の学習手法の導入が容易になりました。
• PyTorch エコシステムとの統合: PyTorch の DTensor 上に構築されており、既存の分散チェックポイントや並列化戦略とシームレスに連携します。
• 産業応用への寄与: 1 万枚以上の GPU を活用する大規模クラスターにおいて、メモリ制約と通信ボトルネックを解消し、より効率的で高品質なモデル学習を可能にしました。

本システムは、ByteDance 内部で実生産環境で battle-tested されており、オープンソース（veScale）として公開されています。これは、大規模 AI 学習における FSDP の新たな標準となり得る技術です。