PyTorch トレーニングスクリプトで SMDDP ライブラリを使用する - Amazon SageMaker AI
PyTorch DDP または FSDP の場合DeepSpeed または Megatron-DeepSpeed の場合

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

PyTorch トレーニングスクリプトで SMDDP ライブラリを使用する

SageMaker AI 分散データ並列処理 (SMDDP) ライブラリ v1.4.0 以降では、ライブラリを PyTorch 分散パッケージのバックエンドオプションとして使用できます。SMDDP の AllReduce および AllGather 集合演算を使用するには、トレーニングスクリプトの冒頭で SMDDP ライブラリをインポートし、プロセスグループの初期化中に SMDDP を PyTorch 分散モジュールのバックエンドとして設定するだけで済みます。1 行でバックエンドを指定するだけで、ネイティブの PyTorch 分散モジュールはすべてそのまま残すことができ、トレーニングスクリプト全体を変更する必要はありません。次のコードスニペットは、PyTorch ベースの分散トレーニングパッケージ (PyTorch Distributed Data Parallel (DDP)、PyTorch Fully Sharded Data Parallelism (FSDP)、DeepSpeed、Megatron-DeepSpeed) のバックエンドとして SMDDP ライブラリを使用する方法を示しています。

PyTorch DDP または FSDP の場合

プロセスグループを次のように初期化します。

import torch.distributed as dist
import smdistributed.dataparallel.torch.torch_smddp

dist.init_process_group(backend="smddp")
注記

(PyTorch DDP ジョブのみ) smddp バックエンドは、現時点では、torch.distributed.new_group() API によるサブプロセスグループの作成には対応していません。smddp バックエンドは、NCCL や Gloo などの他のプロセスグループバックエンドと同時に使用することはできません。

DeepSpeed または Megatron-DeepSpeed の場合

プロセスグループを次のように初期化します。

import deepspeed
import smdistributed.dataparallel.torch.torch_smddp

deepspeed.init_distributed(dist_backend="smddp")
注記

SMDDP の AllGather を SageMaker Python SDK を使用して SMDDP で分散トレーニングジョブを開始する の mpirun ベースのランチャー (smdistributed と pytorchddp) で使用するには、トレーニングスクリプトで次の環境変数を設定する必要があります。

export SMDATAPARALLEL_OPTIMIZE_SDP=true

PyTorch FSDP トレーニングスクリプトの記述に関する一般的なガイダンスについては、PyTorch ドキュメントの「Advanced Model Training with Fully Sharded Data Parallel (FSDP)」を参照してください。

PyTorch DDP トレーニングスクリプトの記述に関する一般的なガイダンスについては、PyTorch ドキュメントの「Getting started with distributed data parallel」を参照してください。

トレーニングスクリプトの調整が完了したら、SageMaker Python SDK を使用して SMDDP で分散トレーニングジョブを開始する に進みます。