使用 SageMaker Python 中的 PyTorch 架構估算器 SDK - Amazon SageMaker

本文為英文版的機器翻譯版本，如內容有任何歧義或不一致之處，概以英文版為準。

使用 SageMaker Python 中的 PyTorch 架構估算器 SDK

您可以透過將distribution引數新增至 SageMaker 架構估算器PyTorch或 來啟動分散式訓練TensorFlow。如需更多詳細資訊，請從下列選項中選擇 SageMaker 分散式資料平行處理 （SMDDP） 程式庫支援的其中一個架構。

PyTorch

下列啟動器選項可用於啟動 PyTorch 分散式訓練。

  • pytorchddp – 此選項會執行mpirun和設定在 上執行 PyTorch 分散式訓練所需的環境變數 SageMaker。若要使用此選項，請將下列字典傳遞至 distribution 參數。

    { "pytorchddp": { "enabled": True } }

  • torch_distributed – 此選項會執行torchrun和設定在 上執行 PyTorch 分散式訓練所需的環境變數 SageMaker。若要使用此選項，請將下列字典傳遞至 distribution 參數。

    { "torch_distributed": { "enabled": True } }

  • smdistributed – 此選項也會執行smddprun，mpirun但會設定在 上執行 PyTorch 分散式訓練所需的環境變數 SageMaker。

    { "smdistributed": { "dataparallel": { "enabled": True } } }

如果您選擇將 取代NCCLAllGather為 SMDDP AllGather，您可以使用全部三個選項。選擇一個適合您使用案例的選項。

如果您選擇NCCLAllReduce以 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"
        }
    }
}

下列程式碼範例顯示具有分散式訓練選項的 PyTorch 估算器基本結構。

from sagemaker.pytorch import PyTorch

pt_estimator = PyTorch(
    base_job_name="training_job_name_prefix",
    source_dir="subdirectory-to-your-code",
    entry_point="adapted-training-script.py",
    role="SageMakerRole",
    py_version="py310",
    framework_version="2.0.1",

    # 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 data parallel library: 
    # ml.p4d.24xlarge, ml.p4de.24xlarge
    instance_type="ml.p4d.24xlarge",

    # 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_estimator.fit("s3://bucket/path/to/training/data")
注意

PyTorch Lightning 及其公用程式程式庫，例如 Lightning Bolts，不會預先安裝在 中 SageMaker PyTorch DLCs。建立下列 requirements.txt 檔案並儲存在存放訓練指令碼的來源目錄中。

# requirements.txt
pytorch-lightning
lightning-bolts

例如，tree-structured 目錄看起來應該如下所示。

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

如需指定來源目錄以放置requirements.txt檔案以及訓練指令碼和任務提交的詳細資訊，請參閱 Amazon SageMaker Python SDK 文件中的使用第三方程式庫。

啟用SMDDP集體操作和使用正確分散式訓練啟動器選項的考量

  • SMDDP AllReduce 和 目前SMDDPAllGather不可相互相容。

  • SMDDP AllReduce 根據預設， 會在使用 smdistributed或 時啟用pytorchddp，這是mpirun以 為基礎的啟動器，並且NCCLAllGather會用到 。

  • SMDDP AllGather 根據預設， 會在使用torch_distributed啟動器時啟動，並AllReduce回復為 NCCL。

  • SMDDP AllGather 也可以在將 mpirun型啟動器與額外的環境變數集搭配使用時啟用，如下所示。

    export SMDATAPARALLEL_OPTIMIZE_SDP=true
TensorFlow
重要

SMDDP 程式庫已停止支援 TensorFlow ，且 不再提供 2.11.0 TensorFlow 版之後DLCs的 。若要尋找已安裝程式SMDDP庫的先前 TensorFlow DLCs ，請參閱 TensorFlow (已廢除)。

from sagemaker.tensorflow import TensorFlow

tf_estimator = TensorFlow(
    base_job_name = "training_job_name_prefix",
    entry_point="adapted-training-script.py",
    role="SageMakerRole",
    framework_version="2.11.0",
    py_version="py38",

    # 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 data parallel library: 
    # ml.p4d.24xlarge, ml.p3dn.24xlarge, and ml.p3.16xlarge
    instance_type="ml.p3.16xlarge",

    # Training using the SageMaker data parallel distributed training strategy
    distribution={ "smdistributed": { "dataparallel": { "enabled": True } } }
)

tf_estimator.fit("s3://bucket/path/to/training/data")