Amazon SageMaker AI で異種クラスターを使用したトレーニングジョブを設定する

異種クラスターのインスタンスグループをトレーニングジョブ用に設定するには、sagemaker.instance_group.InstanceGroup クラスを使用します。各インスタンスグループのカスタム名、インスタンスタイプ、インスタンス数を指定できます。詳細については、SageMaker AI Python SDK ドキュメントの「sagemaker.instance_group.InstanceGroup」を参照してください。

以下のコード例は、次の図のように instance_group_1 という名前の 2 つの ml.c5.18xlarge CPU 専用インスタンスと、instance_group_2 という名前の 1 つの ml.p3dn.24xlarge GPU インスタンスで構成される 2 つのインスタンスグループを設定する方法を示しています。

前の図は、データ前処理などの事前トレーニングプロセスを CPU インスタンスグループに割り当て、前処理されたデータを GPU インスタンスグループにストリーミングする方法の概念的な例を示しています。

from sagemaker.instance_group import InstanceGroup

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

インスタンスグループオブジェクトを使用してトレーニング入力チャンネルを設定し、「sagemaker.inputs.TrainingInput」クラスの instance_group_names 引数を使用してインスタンスグループをチャンネルに割り当てます。instance_group_names 引数には、インスタンスグループ名の文字列のリストを使用できます。

以下の例では、2 つのトレーニング入力チャンネルを設定し、前のステップの例で作成したインスタンスグループを割り当てる方法を示します。インスタンスグループの s3_data 引数に Amazon S3 バケットパスを指定することで、利用目的に応じたデータを処理を行うこともできます。

from sagemaker.inputs import TrainingInput

training_input_channel_1 = TrainingInput(
    s3_data_type='S3Prefix', # Available Options: S3Prefix | ManifestFile | AugmentedManifestFile
    s3_data='s3://your-training-data-storage/folder1',
    distribution='FullyReplicated', # Available Options: FullyReplicated | ShardedByS3Key 
    input_mode='File', # Available Options: File | Pipe | FastFile
    instance_groups=["instance_group_1"]
)

training_input_channel_2 = TrainingInput(
    s3_data_type='S3Prefix',
    s3_data='s3://your-training-data-storage/folder2',
    distribution='FullyReplicated',
    input_mode='File',
    instance_groups=["instance_group_2"]
)

TrainingInput の引数の詳細については、以下のリンクを参照してください。

以下のコード例に示すように、instance_groups 引数を使用で SageMaker AI 推定器を設定します。instance_groups 引数には InstanceGroup オブジェクトのリストを使用できます。

インスタンスグループで設定したトレーニング入力チャネルで estimator.fit メソッドを設定し、トレーニングジョブを開始します。

estimator.fit(
    inputs={
        'training': training_input_channel_1, 
        'dummy-input-channel': training_input_channel_2
    }
)