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

기계 번역으로 제공되는 번역입니다. 제공된 번역과 원본 영어의 내용이 상충하는 경우에는 영어 버전이 우선합니다.

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

PyTorch Lightning 훈련 스크립트를 가져와 SageMaker AI에서 분산형 데이터 병렬 훈련 작업을 실행하려는 경우, 훈련 스크립트를 최소한으로 변경하여 훈련 작업을 실행할 수 있습니다. 필요한 변경 사항에는 smdistributed.dataparallel 라이브러리의 PyTorch 모듈을 가져오고, SageMaker 훈련 도구 키트에서 사전 설정된 SageMaker AI 환경 변수를 수락하도록 PyTorch Lightning의 환경 변수를 설정하고, 프로세스 그룹 백엔드를 "smddp"로 설정하여 SMDDP 라이브러리를 활성화하는 작업이 포함됩니다. 자세히 알아보려면 코드 예제를 사용하여 단계를 세분화하는 다음 지침을 살펴보세요.

참고

PyTorch Lightning 지원은 SageMaker AI 데이터 병렬 라이브러리 v1.5.0 이상에서 사용할 수 있습니다.

  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의 경우 - 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
)

훈련 스크립트 조정을 완료한 후 SageMaker Python SDK를 사용하여 SMDDP로 분산 훈련 작업 시작 섹션으로 넘어갑니다.

참고

SageMaker AI PyTorch 예측기를 구문화하고 SageMaker Python SDK를 사용하여 SMDDP로 분산 훈련 작업 시작에서 훈련 작업 요청을 제출하는 경우, SageMaker AI PyTorch 훈련 컨테이너에 pytorch-lightning 및 lightning-bolts를 설치하도록 requirements.txt를 제공해야 합니다.

# requirements.txt
pytorch-lightning
lightning-bolts

훈련 스크립트 및 작업 제출과 함께 requirements.txt 파일을 배치할 소스 디렉터리를 지정하는 방법에 대한 자세한 내용은 Amazon SageMaker AI Python SDK 설명서의 서드 파티 라이브러리 사용을 참조하세요.