PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel

要約

大規模なモデルには、幅広いドメインにわたって優れたパフォーマンスを提供する可能性があることが広く知られています。
機械学習システム研究の分野では目覚ましい進歩があり、大規模なモデルの開発と探索が可能になっているにもかかわらず、そのような能力は依然として一部の先進ユーザーや業界リーダーに限定されており、その結果、より広範なコミュニティにとって暗黙の技術的障壁となっています。
これらのテクノロジーにアクセスして活用することができます。
このペーパーでは、大規模モデルのトレーニングのための業界グレードのソリューションとして、PyTorch Fully Sharded Data Parallel (FSDP) を紹介します。
FSDP は、Tensor 実装、ディスパッチャ システム、CUDA メモリ キャッシュ アロケータなどのいくつかの主要な PyTorch コア コンポーネントと緊密に共同設計されており、非侵入的なユーザー エクスペリエンスと高いトレーニング効率を提供します。
さらに、FSDP には、さまざまなハードウェア構成全体でリソースの使用率を最適化するためのさまざまな技術と設定がネイティブに組み込まれています。
実験結果は、FSDP が分散データ並列と同等のパフォーマンスを達成しながら、TFLOPS の点でほぼ線形のスケーラビリティを備えた大幅に大規模なモデルのサポートを提供できることを示しています。

要約(オリジナル)

It is widely acknowledged that large models have the potential to deliver superior performance across a broad range of domains. Despite the remarkable progress made in the field of machine learning systems research, which has enabled the development and exploration of large models, such abilities remain confined to a small group of advanced users and industry leaders, resulting in an implicit technical barrier for the wider community to access and leverage these technologies. In this paper, we introduce PyTorch Fully Sharded Data Parallel (FSDP) as an industry-grade solution for large model training. FSDP has been closely co-designed with several key PyTorch core components including Tensor implementation, dispatcher system, and CUDA memory caching allocator, to provide non-intrusive user experiences and high training efficiency. Additionally, FSDP natively incorporates a range of techniques and settings to optimize resource utilization across a variety of hardware configurations. The experimental results demonstrate that FSDP is capable of achieving comparable performance to Distributed Data Parallel while providing support for significantly larger models with near-linear scalability in terms of TFLOPS.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google