Configurer une tâche de formation avec un cluster hétérogène dans Amazon AI SageMaker

Pour configurer des groupes d’instances d’un cluster hétérogène pour une tâche d’entraînement, utilisez la classe sagemaker.instance_group.InstanceGroup. Vous pouvez spécifier un nom personnalisé pour chaque groupe d’instances, le type d’instance et le nombre d’instances pour chaque groupe d’instances. Pour plus d'informations, consultez sagemaker.instance_group. InstanceGroupdans la documentation du SDK SageMaker AI Python.

L'exemple de code suivant montre comment configurer deux groupes d'instances composés de deux ml.c5.18xlarge CPU-only instances nommées instance_group_1 et d'une instance ml.p3dn.24xlarge GPU nomméeinstance_group_2, comme indiqué dans le schéma suivant.

Le schéma précédent montre un exemple conceptuel de la manière dont les processus de pré-entraînement, tels que le prétraitement des données, peuvent être affectés au groupe d’instances CPU et transmettre les données prétraitées au groupe d’instances 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
)

À l'aide des objets du groupe d'instances, configurez les canaux d'entrée d'entraînement et attribuez des groupes d'instances aux canaux via l'instance_group_namesargument. L’argument instance_group_names accepte une liste de chaînes de noms de groupes d’instances.

L’exemple suivant montre comment définir deux canaux d’entrée d’entraînement et attribuer les groupes d’instances créés dans l’exemple de l’étape précédente. Vous pouvez également spécifier des chemins de compartiment Amazon S3 pour les groupes d'instances afin de traiter les données à des fins d'utilisation.

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

Pour plus d’informations sur les arguments de InputData, consultez les liens suivants.

Configurez un objet d'entraînement SageMaker AI avec l'instance_groupsargument, comme indiqué dans l'exemple de code suivant. L’argument instance_groups accepte une liste de InstanceGroup objets.

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]
)

Démarrez la tâche de formation avec les canaux d'entrée de formation configurés avec les groupes d'instances.

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