投稿日: Dec 20, 2023
本日、AWS は Amazon SageMaker モデル並列ライブラリ (SMP) のメジャーバージョンリリースを発表しました。このリリースでは、PyTorch 完全シャーディングデータ並列 (FSDP) API との互換性がサポートされるようになり、深層学習モデルのトレーニングを最大 20% 高速化できます。SMP により、モデルを複数のアクセラレーターとコンピューティングインスタンスに自動的に分割して分散させることで、数十億のパラメータを持つ大規模モデルのトレーニングを加速できます。数分で SMP の使用を開始し、わずか数行のコードで既存の PyTorch FSDP トレーニングスクリプトを高速化できます。
PyTorch FSDP はよく利用される分散型トレーニング手法であり、クラスター内のアクセラレーター間でモデルの重み、勾配、オプティマイザーの状態をシャーディングすることで、トレーニングのメモリ使用量を削減します。今回のリリースにより、SageMaker モデル並列ライブラリの新しい API で PyTorch FSDP トレーニングスクリプトとの互換性がサポートされるようになり、PyTorch FSDP トレーニングスクリプトがさらに高速化されます。そのため、お客様は SageMaker でトレーニングする際に既存のワークロードを簡単にアップグレードできます。わずか数行のコードで、ハイブリッドシャーディングデータ並列処理などの最先端のトレーニング技術を実現できます。これにより、モデルのシャーディングのレベルを変更して、トレーニングジョブのメモリと通信の要件を制御できます。この新しいリリースでは FSDP の機能も拡張され、SageMaker のお客様向けのテンソル並列トレーニング技術が追加されました。これにより、モデルのレイヤーをさまざまなアクセラレーターデバイスに分割して分散させることで、数千億のパラメータを持つモデルのトレーニングが可能になります。SageMaker モデル並列の使用を開始するには、AWS のドキュメントを参照してください。