SageMaker Python SDK에서 PyTorch 프레임워크 추정기 사용

SageMaker AI PyTorch 프레임워크 예측기에 distribution 인수를 추가하여 분산 훈련을 시작할 수 있습니다. SageMaker AI 분산 데이터 병렬 처리(SMDDP) 라이브러리는 PyTorch 분산 훈련을 지원합니다.

PyTorch 분산 훈련을 시작하는 데 다음 시작 관리자 옵션을 사용할 수 있습니다.

NCCL AllGather를 SMDDP AllGather로 교체하도록 선택한 경우 세 가지 옵션을 모두 사용할 수 있습니다. 사용 사례에 맞는 옵션을 하나 선택합니다.

NCCL AllReduce을 SMDDP AllReduce로 교체하기로 선택한 경우 mpirun 기반 옵션인 smdistributed 또는 pytorchddp 중 하나를 선택해야 합니다. 다음과 같이 MPI 옵션을 추가할 수도 있습니다.

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

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

다음 코드 샘플은 분산 훈련 옵션이 있는 ModelTrainer의 기본 구조를 보여줍니다.

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