本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。
在 L PyTorch ightning 训练脚本中使用 SMDDP 库
如果您想引入 PyTorchLightningsmdistributed.dataparallel
库的 PyTorch 模块,设置 L PyTorch ightning 的环境变量以接受 SageMaker 培训工具包预设的 SageMaker 环境变量,以及通过将流程组后端设置为来激活 SMDDP 库。"smddp"
要了解详情,请仔细阅读以下分别介绍了各个步骤并提供代码示例的说明。
注意
PyTorch Lightning 支持在 SageMaker 数据并行库 v1.5.0 及更高版本中提供。
-
导入
pytorch_lightning
库和smdistributed.dataparallel.torch
模块。import lightning as pl import smdistributed.dataparallel.torch.torch_smddp
-
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"])
-
对于 PyTorch DDP — 使用 for
process_group_backend
和"gpu"
for 创建 ddpStrategy 类 "smddp"
的对象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 )
对于 PyTorch FSDP — 使用 for 和 for 创建 for 的 fsdpStrategy 类对象(可选择包装策略
) accelerator
,process_group_backend
然后"gpu"
将其传递给 Trainer 类。"smddp"
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 )
调整完训练脚本后,继续到 第 2 步:使用 SageMaker Python 软件开发工具包启动分布式训练作业。
注意
当你构造 SageMaker PyTorch 估算器并在中提交训练任务请求时第 2 步:使用 SageMaker Python 软件开发工具包启动分布式训练作业,你需要在 SageMaker PyTorch 训练容器lightning-bolts
中提供requirements.txt
安装pytorch-lightning
和。
# requirements.txt pytorch-lightning lightning-bolts
有关指定存放requirements.txt
文件以及训练脚本和作业提交的源目录的更多信息,请参阅 Amaz SageMaker on Python SDK 文档中的使用第三方库