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

SageMaker AI 분산 데이터 병렬 처리(SMDDP) 라이브러리 v1.4.0부터이 라이브러리를 PyTorch 분산 패키지의 백엔드 옵션으로 사용할 수 있습니다. SMDDP AllReduce 및 AllGather 집합 작업을 사용하려면 훈련 스크립트 시작 시 SMDDP 라이브러리를 가져와서 프로세스 그룹 초기화 중에 SMDDP를 PyTorch 분산 모듈의 백엔드로 설정하기만 하면 됩니다. 한 줄의 백엔드 사양을 사용하면 모든 기본 PyTorch 분산 모듈과 전체 훈련 스크립트를 변경하지 않을 수 있습니다. 다음 코드 조각은 SMDDP 라이브러리를 PyTorch 분산 데이터 병렬(DDP), PyTorch 완전 샤딩 데이터 병렬(FSDP), DeepSpeed 및 Megatron-DeepSpeed 등 PyTorch 기반 분산 훈련 패키지의 백엔드로 사용하는 방법을 보여줍니다.

PyTorch DDP 또는 FSDP의 경우

다음과 같이 프로세스 그룹을 초기화합니다.

import torch.distributed as dist
import smdistributed.dataparallel.torch.torch_smddp

dist.init_process_group(backend="smddp")

DeepSpeed 또는 Megatron-DeepSpeed의 경우

다음과 같이 프로세스 그룹을 초기화합니다.

import deepspeed
import smdistributed.dataparallel.torch.torch_smddp

deepspeed.init_distributed(dist_backend="smddp")

export SMDATAPARALLEL_OPTIMIZE_SDP=true

PyTorch FSDP 훈련 스크립트 작성에 대한 일반적인 지침은 PyTorch 설명서의 FSDP(Fully Sharded Data Parallel)를 사용한 고급 모델 훈련을 참조하세요.

PyTorch DDP 훈련 스크립트 작성에 대한 일반 지침은 PyTorch 설명서의 분산 데이터 병렬 시작하기를 참조하세요.

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