Usa gli stimatori del PyTorch framework nell'SDK di Python SageMaker

Puoi avviare un training distribuito aggiungendo l'distributionargomento allo stimatore del PyTorch framework SageMaker AI. La libreria SageMaker AI Distributed Data Parallelism (SMDDP) supporta la formazione distribuita. PyTorch

Le seguenti opzioni di avvio sono disponibili per avviare la formazione PyTorch distribuita.

Se si sceglie di sostituire AllGather NCCL con AllGather SMDDP, è possibile utilizzare tutte le tre opzioni. Scegli un’opzione adatta al tuo caso d’uso.

Se hai scelto di sostituire AllReduce NCCL con AllReduce SMDDP, devi scegliere una delle opzioni basate su mpirun: smdistributed o pytorchddp. Puoi anche aggiungere ulteriori opzioni MPI come segue.

{ 
    "pytorchddp": {
        "enabled": True, 
        "custom_mpi_options": "-verbose -x NCCL_DEBUG=VERSION"
    }
}

{ 
    "smdistributed": { 
        "dataparallel": {
            "enabled": True, 
            "custom_mpi_options": "-verbose -x NCCL_DEBUG=VERSION"
        }
    }
}

Il seguente esempio di codice mostra la struttura di base di un ModelTrainer con opzioni di formazione distribuite.

from sagemaker.train import ModelTrainer
from sagemaker.train.configs import SourceCode, Compute, InputData
from sagemaker.core import image_uris

# Retrieve the training image for the desired PyTorch version
training_image = image_uris.retrieve(
    framework="pytorch",
    region="us-west-2",
    version="2.0.1",
    py_version="py310",
    instance_type="ml.p4d.24xlarge",
    image_scope="training"
)

source_code = SourceCode(
    source_dir="subdirectory-to-your-code",
    entry_script="adapted-training-script.py"
)

compute = Compute(
    # For running a multi-node distributed training job, specify a value greater than 1
    # Example: 2,3,4,..8
    instance_count=2,

    # Instance types supported by the SageMaker AI data parallel library: 
    # ml.p4d.24xlarge, ml.p4de.24xlarge
    instance_type="ml.p4d.24xlarge"
)

pt_model_trainer = ModelTrainer(
    training_image=training_image,
    base_job_name="training_job_name_prefix",
    source_code=source_code,
    role="SageMakerRole",
    compute=compute,

    # Activate distributed training with SMDDP
    distribution={ "pytorchddp": { "enabled": True } }  # mpirun, activates SMDDP AllReduce OR AllGather
    # distribution={ "torch_distributed": { "enabled": True } }  # torchrun, activates SMDDP AllGather
    # distribution={ "smdistributed": { "dataparallel": { "enabled": True } } }  # mpirun, activates SMDDP AllReduce OR AllGather
)

pt_model_trainer.train(input_data_config=[
    InputData(channel_name="training", data_source="s3://bucket/path/to/training/data")
])

# requirements.txt
pytorch-lightning
lightning-bolts

├── pytorch_training_launcher_jupyter_notebook.ipynb
└── sub-folder-for-your-code
    ├──  adapted-training-script.py
    └──  requirements.txt