在 PyTorch 訓練指令碼中使用 SMDDP 程式庫 - Amazon SageMaker AI
對於 PyTorch DDP 或 FSDP對於 DeepSpeed 或 Megatron-DeepSpeed

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

在 PyTorch 訓練指令碼中使用 SMDDP 程式庫

從 SageMaker AI 分散式資料平行化 (SMDDP) 程式庫 v1.4.0 開始，您可以使用程式庫做為 PyTorch 分散式套件的後端選項。若要使用 SMDDP AllReduce 和 AllGather 集體操作，您只需在訓練指令碼開頭匯入 SMDDP 程式庫，並在程序群組初始化期間將 SMDDP 設定為 PyTorch 分散式模組的後端。只需一行後端規範，您就可以保持所有原生 PyTorch 分散式模組和整個訓練指令碼不變。下列程式碼片段示範如何使用 SMDDP 程式庫做為 PyTorch 型分散式訓練套件的後端：PyTorch 分散式資料平行化 (DDP)、PyTorch 全碎片資料平行化 (FSDP)、DeepSpeed 和 Megatron-DeepSpeed。

對於 PyTorch DDP 或 FSDP

初始化程序群組，如下所示。

import torch.distributed as dist
import smdistributed.dataparallel.torch.torch_smddp

dist.init_process_group(backend="smddp")
注意

(僅適用於 PyTorch DDP 任務) smddp 後端目前不支援使用 torch.distributed.new_group() API 建立子程序群組。您也無法與其他程序群組後端 (例如 NCCL 和 Gloo) 同時使用 smddp 後端。

對於 DeepSpeed 或 Megatron-DeepSpeed

初始化程序群組，如下所示。

import deepspeed
import smdistributed.dataparallel.torch.torch_smddp

deepspeed.init_distributed(dist_backend="smddp")
注意

若要在 使用 SageMaker Python SDK 透過 SMDDP 啟動分散式訓練任務 中使用 SMDDP AllGather 搭配 mpirun 型啟動器 (smdistributed 和 pytorchddp)，您也需要在訓練指令碼中設定下列環境變數。

export SMDATAPARALLEL_OPTIMIZE_SDP=true

如需撰寫 PyTorch FSDP 訓練指令碼的一般指引，請參閱 PyTorch 文件中的使用全碎片資料平行化 (FSDP) 的進階模型訓練。

如需撰寫 PyTorch DDP 訓練指令碼的一般指引，請參閱 PyTorch 文件中的分散式資料平行化入門。

完成訓練指令碼的調整後，請繼續前往使用 SageMaker Python SDK 透過 SMDDP 啟動分散式訓練任務。