Amazon SageMaker AI の異種クラスターで分散トレーニングを実行する

SageMaker AI ModelTrainer クラスの distribution引数を使用して、分散トレーニングを実行する特定のインスタンスグループを割り当てることができます。例えば、次の 2 つのインスタンスグループがあり、片方でマルチ GPU トレーニングを実行するとします。

from sagemaker.instance_group import InstanceGroup

instance_group_1 = InstanceGroup("instance_group_1", "ml.c5.18xlarge", 1)
instance_group_2 = InstanceGroup("instance_group_2", "ml.p3dn.24xlarge", 2)

分散トレーニング設定はインスタンスグループの 1 つに設定することができます。例えば、以下のコード例では、2 つの ml.p3dn.24xlarge インスタンスを持つ training_group_2 を分散トレーニング設定に割り当てる方法を示しています。

MPI 使用時

PyTorch

from sagemaker.train import ModelTrainer
from sagemaker.train.distributed import Torchrun

# Note: In v3, Torchrun does not support scoping to a specific instance group.
# It applies to all instances in the training job. Use instance_groups with
# Channel/S3DataSource to control which group receives training data.
model_trainer = ModelTrainer(
    ...
    instance_groups=[instance_group_1, instance_group_2],
    distributed=Torchrun()
)

TensorFlow

from sagemaker.train import ModelTrainer
from sagemaker.train.distributed import Torchrun

# Note: In v3, Torchrun does not support scoping to a specific instance group.
# It applies to all instances in the training job. Use instance_groups with
# Channel/S3DataSource to control which group receives training data.
model_trainer = ModelTrainer(
    ...
    instance_groups=[instance_group_1, instance_group_2],
    distributed=Torchrun()
)

SageMaker AI データ並列ライブラリ使用時

PyTorch

from sagemaker.train import ModelTrainer
from sagemaker.train.distributed import Torchrun

model_trainer = ModelTrainer(
    ...
    instance_groups=[instance_group_1, instance_group_2],
    distributed=Torchrun()
)

TensorFlow

from sagemaker.train import ModelTrainer
from sagemaker.train.distributed import Torchrun

model_trainer = ModelTrainer(
    ...
    instance_groups=[instance_group_1, instance_group_2],
    distributed=Torchrun()
)

SageMaker AI データ並列ライブラリの詳細については、「SageMaker AI Data Parallel Training」を参照してください。

SageMaker AI モデル並列ライブラリ使用時

PyTorch

from sagemaker.train import ModelTrainer
from sagemaker.train.distributed import Torchrun

model_trainer = ModelTrainer(
    ...
    instance_groups=[instance_group_1, instance_group_2],
    distributed=Torchrun()
)

TensorFlow

from sagemaker.train import ModelTrainer
from sagemaker.train.distributed import Torchrun

model_trainer = ModelTrainer(
    ...
    instance_groups=[instance_group_1, instance_group_2],
    distributed=Torchrun()
)

SageMaker AI モデル並列ライブラリの詳細については、「SageMaker AI Model Parallel Training」を参照してください。