PyTorch용 분산형 학습 구성

이 문서에서는 분산형 PyTorch 학습 작업을 만드는 방법을 설명합니다. 분산형 학습 작업을 만들면 AI Platform Training은 분산형 PyTorch 학습을 지원하는 환경 변수를 사용하여 노드라고도 하는 가상 머신(VM) 인스턴스의 클러스터에서 코드를 실행합니다. 이렇게 하면 학습 작업이 확장되어 많은 양의 데이터를 처리할 수 있습니다.

이 가이드에서는 PyTorch 시작하기에 설명된 대로 사전 빌드된 PyTorch 컨테이너를 학습에 사용한다고 가정합니다. 분산형 학습을 위해 PyTorch 코드를 조정하려면 최소한의 변경만 필요합니다.

학습 클러스터 구조 지정

분산형 PyTorch 학습의 경우, 1개의 마스터 작업자 노드와 1개 이상의 작업자 노드를 사용하도록 작업을 구성합니다. 이러한 역할은 다음과 같이 동작합니다.

  • 마스터 작업자: 순위 0을 사용하는 VM 이 노드는 클러스터의 노드 간 연결을 설정합니다.
  • 작업자: 클러스터의 나머지 노드입니다. 각 노드는 학습 애플리케이션 코드에서 지정된 대로 학습의 일부를 수행합니다.

학습 클러스터의 마스터 작업자 노드와 워커 노드를 지정하는 방법은 머신 유형 또는 확장 등급 지정을 참조하세요.

컨테이너 이미지 지정

학습 작업을 만들 때 trainingInput.masterConfig.imageUri 필드에서 사용할 마스터 작업자의 Docker 컨테이너 이미지를 지정하고 trainingInput.workerConfig.imageUri 필드에서 사용할 각 작업자의 Docker 컨테이너 이미지를 지정합니다. 사전 빌드된 PyTorch 컨테이너 목록을 참조하세요.

gcloud ai-platform jobs submit training 명령어를 사용하여 학습 작업을 만드는 경우 --master-image-uri--worker-image-uri 플래그로 이러한 필드를 지정할 수 있습니다.

하지만 trainingInput.workerConfig.imageUri 필드를 지정하지 않으면 기본값은 trainingInput.masterConfig.imageUri입니다. 대부분의 경우 모든 노드에 동일한 사전 빌드된 PyTorch 컨테이너를 사용하는 것이 좋습니다.

학습 코드 업데이트

학습 애플리케이션에서 다음 코드를 추가하여 학습 클러스터를 초기화합니다.

import torch

torch.distributed.init_process_group(
    backend='BACKEND',
    init_method='env://'
)

BACKEND를 다음 섹션에 설명된 지원되는 분산형 학습 백엔드 중 하나로 바꿉니다. init_method='env://' 키워드 인수는 PyTorch가 환경 변수를 사용하여 클러스터의 통신을 초기화하도록 지시합니다. 자세한 내용은 이 가이드의 환경 변수 섹션을 참조하세요.

또한 torch.nn.parallel.DistributedDataParallel 클래스를 사용하도록 학습 코드를 업데이트합니다. 예를 들어 코드에 model라는 PyTorch 모듈을 만들었다면 다음 줄을 추가합니다.

model = torch.nn.parallel.DistributedDataParallel(model)

분산형 학습 구성에 대한 자세한 내용은 PyTorch 문서의 분산형 학습 가이드를 참조하세요.

분산형 학습 백엔드

AI Platform Training은 분산형 PyTorch 학습을 위해 다음 백엔드를 지원합니다.

  • gloo: CPU 학습 작업에 권장
  • nccl: GPU 학습 작업에 권장

백엔드 간 차이점을 읽어보세요.

환경 변수

분산형 PyTorch 학습 작업을 만들 때 AI Platform Training은 각 노드에 다음과 같은 환경 변수를 설정합니다.

  • WORLD_SIZE: 클러스터의 총 노드 수입니다. 이 변수는 모든 노드에서 동일한 값을 가집니다.
  • RANK: 각 노드의 고유 식별자입니다. 마스터 작업자에서는 0로 설정됩니다. 각 작업자에서는 1부터 WORLD_SIZE - 1까지의 다른 값으로 설정됩니다.
  • MASTER_ADDR: 마스터 워커 노드의 호스트 이름입니다. 이 변수는 모든 노드에서 동일한 값을 가집니다.
  • MASTER_PORT: 마스터 워커 노드가 통신하는 포트입니다. 이 변수는 모든 노드에서 동일한 값을 가집니다.

PyTorch는 이러한 환경 변수를 사용하여 클러스터를 초기화합니다.

다음 단계