PyTorch용 분산형 학습 구성

이 문서에서는 분산 PyTorch 학습 작업을 만드는 방법을 설명합니다. 분산 학습 작업을 만들 때 AI Platform Training은 분산 PyTorch 학습을 지원하는 환경 변수와 함께 노드라고도 하는 가상 머신(VM) 인스턴스의 클러스터에서 코드를 실행합니다. 이렇게 하면 대규모 데이터를 처리하도록 학습 작업을 확장하는 데 도움이 됩니다.

이 가이드에서는 PyTorch 시작에 설명된 것처럼 학습을 위해 사전 빌드된 PyTorch 컨테이너를 사용한다고 가정합니다. 분산 학습을 위해 PyTorch 코드를 수정할 때는 최소한의 변경만 필요합니다.

학습 클러스터 구조 지정

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

  • 마스터 작업자: 순위가 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는 이러한 환경 변수를 사용하여 클러스터를 초기화합니다.

다음 단계