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

SageMaker AI 推定器クラスの 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 を分散トレーニング設定に割り当てる方法を示しています。

注記

現在、ディストリビューション設定に指定できるのは、異種クラスターの 1 つのインスタンスグループだけです。

MPI 使用時

PyTorch

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    ...
    instance_groups=[instance_group_1, instance_group_2],
    distribution={
        "mpi": {
            "enabled": True, "processes_per_host": 8
        },
        "instance_groups": [instance_group_2]
    }
)

TensorFlow

from sagemaker.tensorflow import TensorFlow

estimator = TensorFlow(
    ...
    instance_groups=[instance_group_1, instance_group_2],
    distribution={
        "mpi": {
            "enabled": True, "processes_per_host": 8
        },
        "instance_groups": [instance_group_2]
    }
)

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

PyTorch

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    ...
    instance_groups=[instance_group_1, instance_group_2],
    distribution={
        "smdistributed": {
            "dataparallel": {
                "enabled": True
            }
        }, 
        "instance_groups": [instance_group_2]
    }
)

TensorFlow

from sagemaker.tensorflow import TensorFlow

estimator = TensorFlow(
    ...
    instance_groups=[instance_group_1, instance_group_2],
    distribution={
        "smdistributed": {
            "dataparallel": {
                "enabled": True
            }
        }, 
        "instance_groups": [instance_group_2]
    }
)

注記

SageMaker AI データ並列ライブラリを使用する場合は、インスタンスグループがライブラリでサポートされているインスタンスタイプで構成されていることを確認してください。

SageMaker AI データ並列ライブラリの詳細については、SageMaker AI データ並列トレーニング」を参照してください。

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

PyTorch

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    ...
    instance_groups=[instance_group_1, instance_group_2],
    distribution={
        "smdistributed": {
            "modelparallel": {
                "enabled":True,
                "parameters": {
                    ...   # SageMaker AI model parallel parameters
                } 
            }
        }, 
        "instance_groups": [instance_group_2]
    }
)

TensorFlow

from sagemaker.tensorflow import TensorFlow

estimator = TensorFlow(
    ...
    instance_groups=[instance_group_1, instance_group_2],
    distribution={
        "smdistributed": {
            "modelparallel": {
                "enabled":True,
                "parameters": {
                    ...   # SageMaker AI model parallel parameters
                } 
            }
        }, 
        "instance_groups": [instance_group_2]
    }
)

SageMaker AI モデル並列ライブラリの詳細については、SageMaker AI モデル並列トレーニング」を参照してください。