Amazon SageMaker AI 分散式訓練入門

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

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

下頁提供在 Amazon SageMaker AI 中開始分散式訓練所需的步驟資訊。如果您已經熟悉分散式訓練，請選擇下列其中一個符合您偏好策略或架構的選項以開始使用。如果您想要了解一般的分散式訓練，請參閱分散式訓練概念。

SageMaker AI 分散式訓練程式庫已針對 SageMaker 訓練環境進行最佳化，有助於將分散式訓練任務調整為 SageMaker AI，並提高訓練速度和輸送量。這些資料庫提供資料平行和模型平行訓練策略。它們結合軟體和硬體技術來改善 GPU 和節點間通訊，並使用內建選項來擴展 SageMaker AI 的訓練功能，這些選項需要對訓練指令碼進行最少的程式碼變更。 

開始之前

SageMaker 訓練支援單一執行個體和多個執行個體的分散式訓練，因此您可以大規模執行任何規模的訓練。我們建議您在 SageMaker Python SDK 使用架構估算器類別，例如 PyTorch 以及 TensorFlow，這是具有各種分散式訓練選項的訓練工作啟動器。當您建立估算器物件時，物件會設定分散式訓練基礎設施、在後端執行 CreateTrainingJob API、尋找目前工作階段執行所在的區域，以及提取預先建置的 AWS 深度學習容器，其中包含數個程式庫，包括深度學習架構、分散式訓練架構和 EFA 驅動程式。如果想要將 FSx 檔案系統掛載至訓練執行個體，則需要將 VPC 子網路及安全群組 ID 傳遞給估算器。在 SageMaker AI 中執行分散式訓練任務之前，請閱讀下列基本基礎設施設定的一般指引。

可用性區域與網路後擋板

使用多個執行個體 （也稱為節點) 時，請務必了解連接執行個體的網路、它們如何讀取訓練資料，以及它們如何在自己之間共用資訊。例如，當您執行分散式資料平行訓練工作時，許多因素 (例如用於執行 AllReduce 作業的運算叢集節點之間的通訊) 以及 Amazon Simple Storage Service 或 Amazon FSx for Lustre 中的節點及資料儲存之間的資料傳輸，都扮演著至關重要的角色，以達到最佳化運算資源和更快的訓練速度。若要降低通訊開銷，請確定您在相同的 和可用區域中設定執行個體、VPC 子網路 AWS 區域 和資料儲存。

具有更快網路及高輸送量儲存的 GPU 執行個體

您可以在技術上使用任何執行個體進行分散式訓練。對於需要執行多節點分散式訓練任務來訓練大型模型的情況，例如大型語言模型 (LLMs) 和需要更快速節點間整流的擴散模型，我們建議使用 SageMaker AI 支援的啟用 EFA 的 GPU 執行個體。特別是，為了在 SageMaker AI 中實現效能最佳的分散式訓練任務，我們建議使用配備 NVIDIA A100 GPUs 的 P4d 和 P4de 執行個體。這些還配備了高輸送量、低延遲的本機執行個體儲存以及更快速的節點內部網路。對於資料儲存，我們建議使用 Amazon FSx for Lustre 提供高輸送量來存放訓練資料集及模型檢查點。

使用 SageMaker AI 分散式資料平行處理 (SMDDP) 程式庫

SMDDP 程式庫透過針對 AWS 網路基礎設施AllReduce和 Amazon SageMaker AI ML 執行個體拓撲最佳化的 和AllGather集體通訊操作實作，改善節點之間的通訊。您可以使用 SMDDP 程式庫作為 PyTorch 型分散式訓練套件的後端：PyTorch 分散式資料平行處理 (DDP)、PyTorch 全碎片資料平行處理 (FSDP)、DeepSpeed 和 Megatron-DeepSpeed。下列程式碼範例示範如何設定PyTorch估算器，以在兩個ml.p4d.24xlarge執行個體上啟動分散式訓練任務。

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    ...,
    instance_count=2,
    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
)

若要了解如何準備訓練指令碼，並在 SageMaker AI 上啟動分散式資料平行訓練任務，請參閱 使用 SageMaker AI 分散式資料平行處理程式庫執行分散式訓練。

使用 SageMaker AI 模型平行處理程式庫 (SMP)

SageMaker AI 提供 SMP 程式庫並支援各種分散式訓練技術，例如碎片資料平行處理、管道、張量平行處理、最佳化器狀態分割等。若要進一步了解 SMP 程式庫提供的功能，請參閱SageMaker 模型平行處理程式庫的核心功能。

若要使用 SageMaker AI 的模型平行處理程式庫，請設定 SageMaker AI 架構估算器的 distribution 參數。支援的架構估算器是 PyTorch 和 TensorFlow。下列程式碼範例示範如何透過在兩個 ml.p4d.24xlarge 執行個體使用模型平行程式庫，建立用於分散式訓練的架構估算器。

from sagemaker.framework import Framework

distribution={
    "smdistributed": {
        "modelparallel": {
            "enabled":True,
            "parameters": {
                ...   # enter parameter key-value pairs here
            }
        },
    },
    "mpi": {
        "enabled" : True,
        ...           # enter parameter key-value pairs here
    }
}

estimator = Framework(
    ...,
    instance_count=2,
    instance_type="ml.p4d.24xlarge",
    distribution=distribution
)

若要了解如何調整訓練指令碼、在 estimator類別中設定分佈參數，以及啟動分散式訓練任務，請參閱 SageMaker AI 的模型平行處理程式庫 （另請參閱 SageMaker Python SDK 文件中的分散式訓練 APIs)。

使用開放原始碼分散式訓練架構

SageMaker AI 也支援下列選項，以在後端操作 mpirun和 torchrun 。