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

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

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

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

PyTorch

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

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

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

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

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

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

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

如果您選擇將 NCCL 取代AllGather為 SMDDP AllGather，您可以使用這三個選項。選擇一個符合您使用案例的選項。

如果您選擇以 AllReduce SMDDP 取代 NCCLAllReduce，您應該選擇其中一個 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 AI 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 AI 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 AI Python SDK 文件》中的使用第三方程式庫。

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

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

  • 使用 smdistributed或 時，預設AllReduce會啟用 SMDDPpytorchddp，這是mpirun以 為基礎的啟動器，並使用 NCCLAllGather。

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

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

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

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

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