在 PyTorch Lightning 訓練指令碼中使用 SMDDP 程式庫 - Amazon SageMaker AI

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

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

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

注意

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

  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和 建立 FSDPStrategy 類別的物件 "gpu" （使用包裝政策選擇）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
)

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

注意

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

# requirements.txt
pytorch-lightning
lightning-bolts

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