在 PyTorch Lightning 训练脚本中使用 SMDDP 库 - 亚马逊 SageMaker AI
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异，请参阅 中国的 Amazon Web Services 服务入门 (PDF)。

本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。

在 PyTorch Lightning 训练脚本中使用 SMDDP 库

如果您想自带 PyTorch Lightning 训练脚本并在 SageMaker AI 中运行分布式数据并行训练作业，则只需对训练脚本进行极少更改即可运行训练作业。所需更改包括以下内容：导入 smdistributed.dataparallel 库的 PyTorch 模块，为 PyTorch Lightning 设置环境变量以接受 SageMaker 训练工具包预设的 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. For PyTorch DDP：创建一个 DDPStrategy 类对象，其中 "smddp" 为 process_group_backend，"gpu" 为 accelerator，并将其传递给 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
)

    For PyTorch FSDP：创建一个 FSDPStrategy 类的对象（使用选择的包装策略），其中 "smddp" 为 process_group_backend，"gpu" 为 accelerator，并将其传递给 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 PyTorch 估算器并提交训练作业请求时，您需要在 SageMaker AI PyTorch 训练容器中提供 requirements.txt 以安装 pytorch-lightning 和 lightning-bolts。

# requirements.txt
pytorch-lightning
lightning-bolts

有关指定放置 requirements.txt 文件以及训练脚本和作业提交的源目录的更多信息，请参阅 Amazon SageMaker AI Python SDK 文档中的使用第三方库。