SageMaker 分散モデル並列処理のベストプラクティス

SageMaker モデル並列処理ライブラリ v2 (SMP v2) で分散トレーニングジョブを実行する場合は、以下のガイドラインに従ってください。

分散トレーニングに適した構成をセットアップする

SMP v2 が提供する分散トレーニング手法を適用するために、何から着手すれば最適かを検討および確認するには、次のリストを参考にしてください。リストの各項目では、SageMaker モデル並列処理ライブラリ v2 の主要機能 を使用する利点と潜在的なトレードオフについて説明します。

設定のヒント

このセクションでは、グローバルバッチサイズの要件を満たしつつ、最適なスループットを得るための、最善のモデル構成を決定する際のガイドラインを紹介します。

まず、モデルのサイズに関係なく、以下のセットアップを推奨します。

1. 使用可能なインスタンスタイプのうち、最強のタイプを使用します。

2. 混合精度は常に有効にしてください。パフォーマンスの向上とメモリ使用量の削減が大幅に期待できます。float16 よりも精度の高い bfloat16 の使用を推奨します。

3. SageMaker 分散データ並列処理ライブラリを (NCCL を使用する代わりに) 適用できる場合は、有効にします。「AWS インフラストラクチャ用に最適化された SMDDP ライブラリとの互換性」を参照してください。ただし、テンソル並列処理のみのユースケース (hybrid_shard_degree = 1 および tensor_paralle_degree > 1) だけは例外です。

4. モデルのパラメータ数が 600 億を超える場合は、パラメータの遅延初期化 の使用を推奨します。モデルを問わず、パラメータの遅延初期化を行い、初期化を高速化することもできます。

5. アクティベーションチェックポイント を有効にすることを推奨します。

モデルのサイズに応じて、以下のガイダンスを参考にしてください。

1. シャーディングデータ並列処理を使用します。

   1. GPU メモリに収めるバッチサイズに応じて、適切なシャーディングデータ並列度を選択してください。通常は、ネットワーク通信のオーバーヘッドを最小限に抑えつつ、モデルが GPU メモリに収まる最小限の並列度から始めます。キャッシュフラッシュが発生したという警告が表示される場合は、シャーディング並列度を引き上げることを推奨します。

   2. ローカルバッチの最大サイズと、必要なグローバルバッチサイズ (指定されている場合) に基づいて、world_size を決定します。

   3. アクティベーションオフロードを試すことができます。これなら、シナリオによっては、シャーディング並列度を上げなくてもメモリ要件を満たすことができ、通信量を減らせます。

2. PyTorch FSDP のシャーディングデータ並列処理と SMP v2 のテンソル並列処理を同時に使用します。「テンソル並列性」を参照してください。

   1. 大規模なクラスターでトレーニングを行う場合、FSDP 単体だとグローバルバッチサイズが大きくなりすぎ、モデルの収束に問題が生じることがあります。通常、ほとんどの研究では、バッチサイズを 400 万トークン以下に抑えます。この場合、PyTorch FSDP と SMP v2 のテンソル並列処理を併用することで、バッチサイズを小さくして問題を解決できます。

      例えば、ノード数が 256 で、シーケンス長が 4096 の場合、GPU あたりのバッチサイズが 1 でも、グローバルバッチサイズは 800 万トークンになります。ただし、テンソル並列処理を使用し、並列度を 2、テンソル並列グループあたりのバッチサイズを 1 にした場合、GPU あたりのバッチサイズは 1/2 になり、トークン数は 400 万になります。

   2. コンテキスト長が 8k や 16k のように長いトレーニングでは、アクティベーションメモリの使用量が非常に高くなる場合があります。FSDP はアクティベーションをシャーディング (分割) しないため、アクティベーションが原因で GPU がメモリ不足になる可能性があります。このようなシナリオでは、PyTorch FSDP を SMP v2 のテンソル並列処理と組み合わせることで、トレーニングを効率化できます。

リファレンス設定

SageMaker モデル並列処理トレーニングチームは、次の基準値を提供しています。これらは、torch.sagemaker.transform を使用して Llama 2 モデルを SMP トランスフォーマーモデルに変換し、シーケンス長 4096 および混合精度 (FP16 または BF16) を使用して ml.p4d.24xlarge インスタンスでトレーニングした実験に基づいています。

モデル	モデルサイズ (モデルパラメータの数)	インスタンスの数	シャーディングデータ並列度	テンソル並列度	アクティベーションチェックポイント	アクティベーションオフロード	バッチサイズ
Llama 2	7B	1	8	1	TRUE	FALSE	4
	70B	32	256	1	TRUE	FALSE	2
	175B	64	128	4	TRUE	TRUE	6

前の構成から外挿して、モデル設定の GPU メモリ使用量を推定することができます。例えば、10 億パラメータモデルのシーケンス長を増やしたり、モデルのサイズを 20 億に増やす場合は、最初にバッチサイズを小さくすることができます。それでもモデルが適合しない場合は、テンソル並列性の度合いを上ることができます。

SageMaker AI コンソールと Amazon CloudWatch を使用したトレーニングジョブのモニタリングとログ記録

CPU メモリ使用率、GPU メモリ使用率、GPU 使用率などのシステムレベルのメトリクスをモニタリングするには、SageMaker AI コンソールで提供される視覚化を使用します。

1. 左のナビゲーションペインで [トレーニング] を選択します。

2. [トレーニングジョブ] を選択します。

3. メインペインで、詳細を表示するトレーニングジョブ名を選択します。

4. メインペインの [モニタリング] セクションを参照し、自動ビジュアライゼーションを確認します。

5. トレーニングジョブログを表示するには、[モニタリング] セクションの [ログの表示] を選択します。CloudWatch でトレーニングジョブの分散トレーニングジョブログにアクセスできます。マルチノード分散トレーニングを開始した場合は、[algo-n-1234567890] 形式のタグを含む複数のログストリームが表示されます。algo-1 ログストリームは、メイン (0 番目) ノードからのトレーニングログを追跡します。

詳細については、「トレーニングジョブのモニタリングと分析を行うための Amazon CloudWatch メトリクス」を参照してください。

アクセス許可

モデル並列処理を使用して SageMaker トレーニングジョブを実行するには、IAM ロールに次のような適切なアクセスが許可があることを確認してください。

• FSx for Lustre を使用するには、AmazonFSxFullAccess を追加します。

• Amazon S3 をデータチャネルとして使用するには、AmazonS3FullAccess を追加します。

• Docker を使用するには、独自のコンテナを構築して Amazon ECR にプッシュし、AmazonEC2ContainerRegistryFullAccess を追加します。

• SageMaker AI 機能のスイート全体を使用するためのフルアクセスを持つには、 を追加しますAmazonSageMakerFullAccess。