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

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

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

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

注意

SageMaker AI 資料平行化程式庫 v1.5.0 及更新版本提供 PyTorch Lightning 支援。

  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 – 建立 DDPStrategy 類別的物件並將 process_group_backend 設為 "smddp"、將 accelerator 設為 "gpu"，然後將其傳遞給 Trainer 類別。

    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 – 建立 FSDPStrategy 類別的物件 (搭配自行選擇的包裝政策) 並將 process_group_backend 設為 "smddp"、將 accelerator 設為 "gpu"，然後將其傳遞給 Trainer 類別。

    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 Python SDK 透過 SMDDP 啟動分散式訓練任務 中建構 SageMaker AI PyTorch 估算器並提交訓練任務請求時，您必須提供 requirements.txt 以在 SageMaker AI PyTorch 訓練容器中安裝 pytorch-lightning 與 lightning-bolts。

# requirements.txt
pytorch-lightning
lightning-bolts

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