翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

Amazon SageMaker AI で分散トレーニングを開始する

以下のページでは、Amazon SageMaker AI で分散トレーニングを始めるために必要な手順について説明します。既に分散トレーニングに慣れている場合は、以下のオプションからお好みの戦略やフレームワークに合わせて 1 つを選択して始めてください。分散型トレーニング全般について知りたい場合は、分散トレーニングの概念 を参照してください。

SageMaker AI 分散トレーニングライブラリは SageMaker トレーニング環境向けに最適化されており、分散トレーニングジョブを SageMaker AI に適応させ、トレーニングの速度とスループットを向上させるのに役立ちます。ライブラリは、データ並列とモデル並列の両方のトレーニング戦略を提供します。ソフトウェアとハードウェアのテクノロジーを組み合わせて GPU 間およびノード間の通信を改善し、SageMaker AI のトレーニング機能は、トレーニングスクリプトにわずかなコード変更を加えるだけで済む組み込みオプションによって拡張されます。 

始める前に

SageMaker Training は、単一インスタンスと複数インスタンスの分散トレーニングをサポートしているため、あらゆる規模のトレーニングを大規模に実行できます。SageMaker Python SDK の PyTorch や TensorFlow などのフレームワーク推定クラスを使用することをお勧めします。これらは、さまざまな分散型トレーニングオプションを備えたトレーニングジョブランチャーです。推定器オブジェクトを作成すると、そのオブジェクトは分散トレーニングのインフラストラクチャをセットアップし、バックエンドで CreateTrainingJob API を実行します。また、現在のセッションが実行されているリージョンを特定し、深層学習フレームワーク、分散型トレーニングフレームワーク、EFA ドライバーなどの多数のライブラリがあらかじめパッケージ化された、ビルド済みの AWS 深層学習コンテナの 1 つ取得します。FSx ファイルシステムをトレーニングインスタンスにマウントする場合は、VPC サブネットとセキュリティグループ ID を estimator に渡す必要があります。SageMaker AI で分散型トレーニングジョブを実行する前に、基本的なインフラストラクチャの設定に関する以下の一般的なガイダンスをお読みください。

アベイラビリティーゾーンとネットワークバックプレーン

複数のインスタンス (ノードとも呼ばれる) を使用する場合は、インスタンスを接続するネットワーク、インスタンスがトレーニングデータを読み取る方法、インスタンス間で情報を共有する方法を理解しておくことが重要です。例えば、分散データ並列トレーニングジョブを実行する場合、コンピューティングリソースの最適な使用とトレーニング速度の高速化を実現するには、AllReduce オペレーションを実行する計算クラスターのノード間の通信や、Amazon Simple Storage Service または Amazon FSx for Lustre 内のノードとデータストレージ間のデータ転送など、さまざまな要因が重要な役割を果たします。通信オーバーヘッドを減らすには、インスタンス、VPC サブネット、データストレージを必ず同じ AWS リージョン およびアベイラビリティーゾーンに設定してください。

より高速なネットワークと高スループットのストレージを備えた GPU インスタンス

技術的には、どのインスタンスでも分散トレーニングに使用できます。大規模言語モデル (LLM) や拡散モデルなど、より高速なノード間通信を必要とする大規模モデルをトレーニングするためにマルチノード分散トレーニングジョブを実行する必要がある場合は、SageMaker AI がサポートする EFA 対応 GPU インスタンスを推奨します。特に、SageMaker AI で最もパフォーマンスの高い分散型トレーニングジョブを実現するには、NVIDIA A100 GPU を搭載した P4d および P4de インスタンスを推奨します。これらには、高スループット、低レイテンシーのローカルインスタンスストレージ、およびより高速なノード内ネットワークも搭載されています。データストレージには、トレーニングデータセットとモデルチェックポイントを高いスループットで保存できる Amazon FSx for Lustre をお勧めします。

SageMaker AI 分散データ並列処理ライブラリ (SMDDP) を使用する

SMDDP ライブラリは、AWS ネットワークインフラストラクチャと Amazon SageMaker AI ML インスタンストポロジ向けに最適化された AllReduce および AllGather 集合通信演算の実装を通じて、ノード間の通信を改善します。SMDDP ライブラリは PyTorch ベースの分散トレーニングパッケージ (PyTorch Distributed Data Parallel (DDP)、PyTorch Fully Sharded Data Parallelism (FSDP)、DeepSpeed、Megatron-DeepSpeed) のバックエンドとして使用できます。次のコード例は、2 つの ml.p4d.24xlarge インスタンスで分散トレーニングジョブを開始する場合の PyTorch 推定器の設定方法を示しています。

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    ...,
    instance_count=2,
    instance_type="ml.p4d.24xlarge",
    # Activate distributed training with SMDDP
    distribution={ "pytorchddp": { "enabled": True } }  # mpirun, activates SMDDP AllReduce OR AllGather
    # distribution={ "torch_distributed": { "enabled": True } }  # torchrun, activates SMDDP AllGather
    # distribution={ "smdistributed": { "dataparallel": { "enabled": True } } }  # mpirun, activates SMDDP AllReduce OR AllGather
)

トレーニングスクリプトを準備し、SageMaker AI で分散データ並列トレーニングジョブを開始する方法については、「SageMaker AI 分散データ並列処理ライブラリを使用して分散トレーニングを実行する」を参照してください。

SageMaker AI モデル並列処理ライブラリ (SMP) を使用する

SageMaker AI は SMP ライブラリを提供し、シャーディングデータ並列処理、パイプライン、テンソル並列処理、オプティマイザステートシャーディングなど、さまざまな分散トレーニング技術をサポートしています。SMP ライブラリが提供するものの詳細については、SageMaker モデル並列処理ライブラリの主要機能 を参照してください。

SageMaker AI のモデル並列処理ライブラリを使用するには、SageMaker AI フレームワーク推定器の distribution パラメータを設定します。サポートされているフレームワーク推定器は PyTorch と TensorFlow です。次のコード例は、2 つの ml.p4d.24xlarge インスタンスでモデル並列処理ライブラリを使用して分散トレーニング用のフレームワーク推定器を構築する方法を示しています。

from sagemaker.framework import Framework

distribution={
    "smdistributed": {
        "modelparallel": {
            "enabled":True,
            "parameters": {
                ...   # enter parameter key-value pairs here
            }
        },
    },
    "mpi": {
        "enabled" : True,
        ...           # enter parameter key-value pairs here
    }
}

estimator = Framework(
    ...,
    instance_count=2,
    instance_type="ml.p4d.24xlarge",
    distribution=distribution
)

トレーニングスクリプトの適合、estimator クラスでの分散パラメータの設定、および分散トレーニングジョブを起動する方法については、「SageMaker AI's model parallelism library」(「SageMaker Python SDK ドキュメント」の「Distributed Training APIs」も参照) を参照してください。

オープンソースの分散型トレーニングフレームワークを使用する

SageMaker AI は、バックエンドでの mpirun と torchrun を操作する以下のオプションもサポートしています。