PyTorch Lightning 훈련 스크립트에서 SMDDP 라이브러리 사용

pytorch_lightning 라이브러리 및 smdistributed.dataparallel.torch 모듈을 가져옵니다.

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

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"])

PyTorch DDP의 경우 - process_group_backend용 "smddp" 및 accelerator용 "gpu"를 사용하여 DDPStrategy 클래스의 객체를 생성하고 이를 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의 경우 - process_group_backend용 "smddp" 및 accelerator용 "gpu"를 사용하여 FSDPStrategy 클래스의 객체(래핑 정책 선택)를 생성하고 이를 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
)