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

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

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

始める前に

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

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

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

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

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

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

SMDDP ライブラリは、 AWS ネットワークインフラストラクチャAllReduceと Amazon SageMaker AI ML インスタンストポロジに最適化された および 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 のモデル並列処理ライブラリ」を参照してください (SageMaker Python SDK ドキュメントの「分散トレーニング APIs」も参照してください）。

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

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

• SageMaker AI で PyTorch DistributedDataParallel (DDP) をmpirunバックエンドで使用するには、 distribution={"pytorchddp": {"enabled": True}}を PyTorch 推定器に追加します。詳細については、SageMaker Python SDK ドキュメントのPyTorch Distributed Training」と「SageMaker AI PyTorch Estimator」の引distribution数も参照してください。 SageMaker PyTorch SageMaker

  注記

  このオプションは、PyTorch 1.12.0 以降で使用できます。

  from sagemaker.pytorch import PyTorch
  
  estimator = PyTorch(
      ...,
      instance_count=2,
      instance_type="ml.p4d.24xlarge",
      distribution={"pytorchddp": {"enabled": True}}  # runs mpirun in the backend
  )

• SageMaker AI は、P3 や P4 などの GPU ベースの Amazon EC2 インスタンスでの分散トレーニング用の PyTorch torchrunランチャーと、 AWS Trainium デバイスを搭載した Trn1 をサポートしています。

  SageMaker AI で PyTorch DistributedDataParallel (DDP) をtorchrunバックエンドで使用するには、 distribution={"torch_distributed": {"enabled": True}}を PyTorch 推定器に追加します。

  注記

  このオプションは、PyTorch 1.13.0 以降で使用できます。

  次のコードスニペットは、SageMaker AI PyTorch 推定器を構築して、 torch_distributed 分散オプションを使用して 2 つのml.p4d.24xlargeインスタンスで分散トレーニングを実行する例を示しています。

  from sagemaker.pytorch import PyTorch
  
  estimator = PyTorch(
      ...,
      instance_count=2,
      instance_type="ml.p4d.24xlarge",
      distribution={"torch_distributed": {"enabled": True}}   # runs torchrun in the backend
  )

  詳細については、SageMaker Python SDK ドキュメントの「分散 PyTorch トレーニング」と「SageMaker AI PyTorch 推定器の引distribution数」を参照してください。 SageMaker PyTorch SageMaker

  Trn1 での分散トレーニングに関する注意事項

  Trn1 インスタンスは最大 16 台の Trainium デバイスで構成され、各 Trainium デバイスは 2 つの NeuronCores で構成されます。 AWS Trainium デバイスの仕様については、 AWS Neuron ドキュメントの Trainium アーキテクチャを参照してください。

  Trainium 搭載のインスタンスでトレーニングするには、Trn1 インスタンスコード を SageMaker AI PyTorch 推定器クラスの instance_type引数に文字列ml.trn1.*で指定するだけで済みます。使用可能な Trn1 インスタンスタイプを見つけるには、「AWS Neuron ドキュメント」の「AWS Trn1 Architecture」を参照してください。

  注記

  Amazon EC2 Trn1 インスタンスでの SageMaker トレーニングは現在、 v1.11.0 以降の AWS Deep Learning Containers for PyTorch Neuron の PyTorch フレームワークでのみ使用できます。PyTorch Neuron のサポートされているすべてのバージョンのリストについては、AWS Deep Learning Containers GitHub リポジトリの「Neuron Containers」を参照してください。

  SageMaker Python SDK を使用して Trn1 インスタンスでトレーニングジョブを起動すると、SageMaker AI は AWS Deep Learning Containers が提供する Neuron Containers から適切なコンテナを自動的に取得して実行します。Neuron Containers には、トレーニングジョブを SageMaker トレーニングプラットフォームと Amazon EC2 Trn1 インスタンスに簡単に適応させることができるように、トレーニング環境設定と依存関係があらかじめパッケージ化されています。

  注記

  SageMaker AI を使用して Trn1 インスタンスで PyTorch トレーニングジョブを実行するには、xlaトレーニングスクリプトを変更してバックエンドでプロセスグループを初期化し、PyTorch/XLA を使用する必要があります。XLA 導入プロセスをサポートするために、 AWS Neuron SDK は XLA を使用して PyTorch オペレーションを Trainium 指示に変換する PyTorch Neuron を提供します。トレーニングスクリプトを変更する方法については、「AWS Neuron ドキュメント」の「Developer Guide for Training with PyTorch Neuron (torch-neuronx)」を参照してください。

  詳細については、SageMaker Python SDK ドキュメントのTrn1 インスタンスでの PyTorch Neuron による分散トレーニング」と「SageMaker AI PyTorch 推定器の distribution引数」を参照してください。 SageMaker PyTorch SageMaker

• SageMaker AI で MPI を使用するには、 distribution={"mpi": {"enabled": True}}を推定器に追加します。MPI ディストリビューションオプションは、MXNet、PyTorch、TensorFlow の各フレームワークで使用できます。

• SageMaker AI でパラメータサーバーを使用するには、 distribution={"parameter_server": {"enabled": True}}を推定器に追加します。パラメータサーバーオプションは、MXNet、PyTorch、TensorFlow の各フレームワークで使用できます。

  ヒント

  フレームワークごとの MPI およびパラメータサーバーオプションの使用方法の詳細については、「SageMaker Python SDK ドキュメント」への次のリンクを参照してください。

  • MXNet 分散トレーニングと SageMaker AI MXNet 推定器の distribution引数

  • PyTorch 分散トレーニングと SageMaker AI PyTorch 推定器の引distribution数

  • TensorFlow 分散トレーニングと SageMaker AI TensorFlow 推定器の distribution引数。