Configure um trabalho de treinamento com um cluster heterogêneo no Amazon AI SageMaker

Para configurar grupos de instâncias de um cluster heterogêneo para um trabalho de treinamento, use a classe sagemaker.instance_group.InstanceGroup. Você pode especificar um nome personalizado para cada grupo de instâncias, o tipo de instância e o número de instâncias para cada grupo de instâncias. Para obter mais informações, consulte sagemaker.instance_group. InstanceGroupna documentação do SDK do SageMaker AI Python.

O exemplo de código a seguir mostra como configurar dois grupos de instâncias que consistem em duas ml.c5.18xlarge CPU-only instâncias nomeadas instance_group_1 e uma instância de ml.p3dn.24xlarge GPU nomeadainstance_group_2, conforme mostrado no diagrama a seguir.

O diagrama anterior mostra um exemplo conceitual de como os processos de pré-treinamento, como o pré-processamento de dados, podem ser atribuídos ao grupo de instâncias da CPU e transmitir os dados pré-processados para o grupo de instâncias da 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
)

Usando os objetos do grupo de instâncias, configure canais de entrada de treinamento e atribua grupos de instâncias aos canais por meio do instance_group_names argumento. O argumento instance_group_names aceita uma lista de strings de nomes de grupos de instâncias.

O exemplo a seguir mostra como definir dois canais de entrada de treinamento e atribuir os grupos de instâncias criados no exemplo da etapa anterior. Você também pode especificar caminhos de bucket do Amazon S3 para que os grupos de instâncias processem dados para fins de uso.

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',
)

Para obter mais informações sobre os argumentos de InputData, consulte os seguintes links:

Configure um objeto de treinamento de SageMaker IA com o instance_groups argumento, conforme mostrado no exemplo de código a seguir. O argumento instance_groups aceita uma lista de objetos 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]
)

Comece o trabalho de treinamento com os canais de entrada de treinamento configurados com os grupos de instâncias.

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