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
)

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

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

from sagemaker.train.configs import InputData
from sagemaker.core.shapes import Channel, DataSource, S3DataSource

training_input_channel_1 = Channel(
    channel_name='training',
    data_source=DataSource(
        s3_data_source=S3DataSource(
            s3_data_type='S3Prefix',
            s3_uri='s3://your-training-data-storage/folder1',
            s3_data_distribution_type='FullyReplicated', # Available Options: FullyReplicated | ShardedByS3Key
            instance_group_names=["instance_group_1"],
        )
    ),
    input_mode='File', # Available Options: File | Pipe | FastFile
)

training_input_channel_2 = Channel(
    channel_name='dummy-input-channel',
    data_source=DataSource(
        s3_data_source=S3DataSource(
            s3_data_type='S3Prefix',
            s3_uri='s3://your-training-data-storage/folder2',
            s3_data_distribution_type='FullyReplicated',
            instance_group_names=["instance_group_2"],
        )
    ),
    input_mode='File',
)

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

次のコード例に示すように、 instance_groups引数を使用して SageMaker AI トレーニングオブジェクトを設定します。instance_groups 引数には InstanceGroup オブジェクトのリストを使用できます。

PyTorch

from sagemaker.train import ModelTrainer
from sagemaker.train.configs import SourceCode
from sagemaker.core import image_uris

training_image = image_uris.retrieve(
    framework="pytorch", region=region, version='x.y.z',    # 1.10.0 or later
    py_version='pyxy', instance_type='ml.p3dn.24xlarge',
    image_scope="training"
)

model_trainer = ModelTrainer(
    ...
    source_code=SourceCode(entry_script='my-training-script.py'),
    training_image=training_image,
    job_name='my-training-job-with-heterogeneous-cluster',
    instance_groups=[instance_group_1, instance_group_2]
)

TensorFlow

from sagemaker.train import ModelTrainer
from sagemaker.train.configs import SourceCode
from sagemaker.core import image_uris

training_image = image_uris.retrieve(
    framework="tensorflow", region=region, version='x.y.z', # 2.6.0 or later
    py_version='pyxy', instance_type='ml.p3dn.24xlarge',
    image_scope="training"
)

model_trainer = ModelTrainer(
    ...
    source_code=SourceCode(entry_script='my-training-script.py'),
    training_image=training_image,
    job_name='my-training-job-with-heterogeneous-cluster',
    instance_groups=[instance_group_1, instance_group_2]
)

インスタンスグループで設定されたトレーニング入力チャネルを使用してトレーニングジョブを開始します。

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