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

如果您想要攜帶 PyTorch Lightning 訓練指令碼，並在 中執行分散式資料平行訓練任務 SageMaker，則可以在訓練指令碼中以最小的變更執行訓練任務。必要的變更包括：匯入smdistributed.dataparallel程式庫的 PyTorch 模組、設定 PyTorch Lightning 的環境變數，以接受 SageMaker 訓練工具組預設 SageMaker 的環境變數，以及將程序群組後端設定為 來啟用SMDDP程式庫"smddp"。若要進一步了解，請逐步執行以程式碼範例分解步驟的下列指示。

注意

PyTorch Lightning 支援可在 SageMaker 資料平行程式庫 1.5.0 版及更新版本中使用。

PyTorch Lightning == v2.1.0 和 PyTorch == 2.0.1

1. 匯入 pytorch_lightning 程式庫和 smdistributed.dataparallel.torch 模組。

   import lightning as pl
   import smdistributed.dataparallel.torch.torch_smddp

2. 現化 LightningEnvironment。

   from lightning.fabric.plugins.environments.lightning import LightningEnvironment
   
   env = LightningEnvironment()
   env.world_size = lambda: int(os.environ["WORLD_SIZE"])
   env.global_rank = lambda: int(os.environ["RANK"])

3. 適用於 PyTorch DDP – 使用"smddp"適用於 process_group_backend和"gpu"適用於 的 建立 DDPStrategy 類別的物件accelerator，並將其傳遞給訓練人員類別。

   import lightning as pl
   from lightning.pytorch.strategies import DDPStrategy
   
   ddp = DDPStrategy(
       cluster_environment=env, 
       process_group_backend="smddp", 
       accelerator="gpu"
   )
   
   trainer = pl.Trainer(
       max_epochs=200, 
       strategy=ddp, 
       devices=num_gpus, 
       num_nodes=num_nodes
   )

   適用於 PyTorch FSDP – 使用"smddp"適用於 的 process_group_backend和 "gpu" 建立 FSDPStrategy 類別的物件 （搭配包裝政策選擇）accelerator，並將其傳遞給訓練人員類別。

   import lightning as pl
   from lightning.pytorch.strategies import FSDPStrategy
   
   from functools import partial
   from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
   
   policy = partial(
       size_based_auto_wrap_policy, 
       min_num_params=10000
   )
   
   fsdp = FSDPStrategy(
       auto_wrap_policy=policy,
       process_group_backend="smddp", 
       cluster_environment=env
   )
   
   trainer = pl.Trainer(
       max_epochs=200, 
       strategy=fsdp, 
       devices=num_gpus, 
       num_nodes=num_nodes
   )

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

注意

當您在 中建構 SageMaker PyTorch 估算器並提交訓練任務請求時SMDDP 使用 SageMaker Python 啟動分散式訓練任務 SDK，您需要提供 requirements.txt才能lightning-bolts在 SageMaker PyTorch 訓練容器中安裝 pytorch-lightning和 。

# requirements.txt
pytorch-lightning
lightning-bolts

如需指定來源目錄以放置requirements.txt檔案以及訓練指令碼和任務提交的詳細資訊，請參閱 Amazon SageMaker Python SDK 文件 中的使用第三方程式庫。