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

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

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

從 SageMaker 分散式資料平行處理 （SMDDP） 程式庫 v1.4.0 開始，您可以使用程式庫作為PyTorch 分散式套件 的後端選項。若要使用 SMDDPAllReduce和 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）SMDDP 使用 SageMaker Python 啟動分散式訓練任務 SDK，您也需要在訓練指令碼中設定下列環境變數。

export SMDATAPARALLEL_OPTIMIZE_SDP=true

如需撰寫 PyTorch FSDP訓練指令碼的一般指引，請參閱 文件中的 PyTorch具有完整陰影資料平行的進階模型訓練 （FSDP）。

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

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