在 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 建立子程序群組。您也無法將smddp後端與其他程序群組後端同時使用，例如 NCCL和 Gloo。

適用於 DeepSpeed 或 Megatron-DeepSpeed

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

import deepspeed
import smdistributed.dataparallel.torch.torch_smddp

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

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

export SMDATAPARALLEL_OPTIMIZE_SDP=true

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

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

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