veScale-FSDP: 拡張性と高パフォーマンスの Large-Scale FSDP

arXiv:2602.22437v3 Announce Type: replace-cross 摘要：フルシャーディングデータ並列（FSDP）、ゼロ余分度最適化器（ZeRO）とも呼ばれ、メモリ効率の良さやモデルコードへの介入が最小限であるため、大規模モデルのトレーニングで広く採用されています。しかし、既存の FSDP システムは、ブロック構造の計算と衝突する固定した要素ごとのまたは行ごとのシャーディング形式に依存しており、現代の構造感知トレーニング手法、例えばブロックごとの量化やシャンプー・ムーンのような非要素ごとの最適化器をサポートするのが困難です。さらに、現在の実装は数十万台の GPU にわたるスケールにおいて、効率が低下する通信およびメモリのオーバーヘッドを伴います。本研究では、RaggedShard という柔軟なシャーディング形式と構造感知計画アルゴリズムを組み合わせ、拡張性とパフォーマンスを両立させる新しい FSDP システム「veScale-FSDP」を発表します。veScale-FSDP はゼロコピー FSDP コミュニケーションを可能にし、ブロックごとの量化や非要素ごとの最適化器をネイティブにサポートし、既存の FSDP システムに比較して 5% から 66% の高出力性と 16% から 30% の低メモリ使用量を実現し、数十万台の GPU にわたって効率的にスケーリングします。

arXiv:2602.22437v3 Announce Type: replace-cross Abstract: Fully Sharded Data Parallel (FSDP), also known as Zero Redundancy Optimizer (ZeRO), is widely used for large-scale model training, because of its memory efficiency and minimal intrusion on model code. However, existing FSDP systems rely on fixed element-wise or row-wise sharding formats that conflict with block-structured computations. As a result, they struggle to support modern structure-aware training methods, including block-wise quantization and non-element-wise optimizers such as Shampoo and Muon. In addition, today's implementations incur communication and memory overheads that degrade efficiency at the scale of tens of thousands of GPUs. We introduce veScale-FSDP, a novel FSDP system that combines RaggedShard, a flexible sharding format, with a structure-aware planning algorithm to deliver both flexibility and performance. veScale-FSDP enables zero-copy FSDP communications and natively supports block-wise quantization and non-element-wise optimizers, achieving 5% to 66% higher throughput and 16% to 30% lower memory usage than existing FSDP systems, while scaling efficiently to tens of thousands of GPUs.