修改 PyTorch Lightning 脚本

导入 pytorch_lightning 库和 smdistributed.dataparallel.torch 模块。

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

设置 LightningEnvironment 类对象的世界大小和秩。在 SageMaker 中启动训练作业时，SageMaker 训练工具包会设置环境变量 "RANK"、"LOCAL_RANK" 和 "WORLD_SIZE"。这些环境变量分别代表进程的全局秩、其局部秩和世界大小。使用这些 SageMaker 环境变量来配置 LightningEnvironment。

import os
from pytorch_lightning.plugins.environments.lightning_environment \
  import LightningEnvironment

env = LightningEnvironment()
env.world_size = lambda: int(os.environ["WORLD_SIZE"])
env.global_rank = lambda: int(os.environ["RANK"])

使用 PyTorch Lightning DDPStrategy 模块设置分布式训练策略，创建 PyTorch Lightning Trainer 对象，然后对其进行调整以使用 SageMaker 数据并行库。

（推荐）适用于 PyTorch Lightning v1.6.0 及更高版本

创建 DDPStrategy 类的对象（以下代码示例中的 ddp），然后为 process_group_backend 参数指定 "smddp"。配置 PyTorch Lightning Trainer 对象时，请使用 SageMaker 环境变量指定 GPU 集群的规模和用于设置分布式训练策略的 ddp 对象。

from pytorch_lightning.strategies import DDPStrategy

ddp = DDPStrategy(
  cluster_environment=env, 
  process_group_backend="smddp", 
  accelerator="gpu"
)
  
world_size = int(os.environ["WORLD_SIZE"])
num_gpus = int(os.environ["SM_NUM_GPUS"])
num_nodes = int(world_size/num_gpus)

trainer = pl.Trainer(
  devices=num_gpus, 
  num_nodes=num_nodes,
  max_epochs=10,
  strategy=ddp
)

（可选）适用于 PyTorch Lightning v1.5.10

如果您使用已弃用的功能 DDPPlugin，请设置分布式策略，如以下代码示例所示。

from pytorch_lightning.plugins.training_type.ddp import DDPPlugin

os.environ["PL_TORCH_DISTRIBUTED_BACKEND"] = "smddp"

ddp = DDPPlugin(
  parallel_devices=[torch.device("cuda", d) for d in range(num_gpus)],
  cluster_environment=env
)
  
world_size = int(os.environ["WORLD_SIZE"])
num_gpus = int(os.environ["SM_NUM_GPUS"])
num_nodes = int(world_size/num_gpus)

trainer = pl.Trainer(
  gpus=num_gpus, 
  num_nodes=num_nodes, 
  max_epochs=10, 
  strategy=ddp
)

运行 trainer.fit 以启动 PyTorch 模型的训练作业。以下代码示例显示了通过 PyTorch Lightning Trainer 的拟合方法和 PyTorch Lightning MNIST 数据模块包装的 PyTorch 模型对象。

from pl_bolts.datamodules import MNISTDataModule

trainer.fit(model, datamodule=MNISTDataModule(batch_size=32))