Cloud TPU Pod에서 PyTorch 모델 학습

이 가이드에서는 모델을 단일 Cloud TPU(v2-8 또는 v3-8)에서 Cloud TPU Pod로 확장하는 방법을 보여줍니다. TPU Pod의 Cloud TPU 가속기는 높은 대역폭의 상호 연결로 연결되므로 학습 작업을 확장하는 데 효율적입니다.

Cloud TPU Pod 제품에 대한 자세한 내용은 Cloud TPU 제품 페이지 또는 Cloud TPU 프레젠테이션을 참조하세요.

다음 다이어그램은 분산 클러스터 설정의 개요를 제공합니다. VM의 인스턴스 그룹이 TPU pod에 연결됩니다. 8개의 TPU 코어 그룹마다 하나의 VM이 필요합니다. VM은 TPU 코어에 데이터를 공급하며 모든 학습은 TPU Pod에서 발생합니다.

이미지

목표

  • PyTorch/XLA를 사용한 교육을 위해 Compute Engine 인스턴스 그룹 및 Cloud TPU Pod 설정
  • Cloud TPU Pod에서 PyTorch/XLA 학습 실행

시작하기 전에

Cloud TPU Pod에서 분산 학습을 시작하기 전에 모델이 단일 v2-8 또는 v3-8 Cloud TPU 기기에서 학습하는지 확인합니다. 단일 기기에서 모델에 심각한 성능 문제가 발생하는 경우 권장사항문제해결 가이드를 참조하세요.

단일 TPU 기기의 학습이 성공적으로 실행되면 다음 단계에 따라 설정하고 Cloud TPU Pod에서 학습합니다.

  1. gcloud 명령어를 구성합니다.

  2. [선택사항] VM 디스크 이미지를 VM 이미지로 캡처

  3. VM 이미지에서 인스턴스 템플릿 만들기

  4. 인스턴스 템플릿에서 인스턴스 그룹 만들기

  5. Compute Engine VM에 SSH를 통해 연결

  6. VM 간 통신을 허용하도록 방화벽 규칙 확인

  7. Cloud TPU Pod 만들기

  8. Pod에서 분산 학습 실행

  9. 삭제

gcloud 명령어 구성

gcloud를 사용하여 GCP 프로젝트를 구성합니다.

프로젝트 ID의 변수를 만듭니다.

export PROJECT_ID=project-id

gcloud에서 프로젝트 ID를 기본 프로젝트로 설정합니다.

gcloud config set project ${PROJECT_ID}

gcloud로 기본 영역을 구성합니다.

gcloud config set compute/zone us-central1-a

[선택사항] VM 디스크 이미지 캡처

이미 데이터 세트와 패키지 등이 설치되어 있는 단일 TPU 학습에 사용한 VM의 디스크 이미지를 사용할 수 있습니다. 이미지를 만들기 전에 gcloud 명령어를 사용하여 VM을 중지합니다.

gcloud compute instances stop vm-name

다음으로 gcloud 명령어를 사용하여 VM 이미지를 만듭니다.

gcloud compute images create image-name  \
    --source-disk instance-name \
    --source-disk-zone us-central1-a \
    --family=torch-xla \
    --storage-location us-central1

VM 이미지에서 인스턴스 템플릿 만들기

기본 인스턴스 템플릿을 만듭니다. 인스턴스 템플릿을 만들 때 위의 단계에서 만든 VM 이미지를 사용할 수도 있고 또는 Google이 제공하는 공개 PyTorch/XLA 이미지를 사용할 수도 있습니다. 인스턴스 템플릿을 만들려면 gcloud 명령어를 사용합니다.

gcloud compute instance-templates create instance-template-name \
    --machine-type n1-standard-16 \
    --image-project=${PROJECT_ID} \
    --image=image-name \
    --scopes=https://www.googleapis.com/auth/cloud-platform

인스턴스 템플릿에서 인스턴스 그룹 만들기

gcloud compute instance-groups managed create instance-group-name \
    --size 4 \
    --template template-name \
    --zone us-central1-a

Compute Engine VM에 SSH를 통해 연결

인스턴스 그룹을 만든 후 인스턴스 그룹의 인스턴스(VM) 중 하나에 SSH로 연결합니다. 다음 명령어를 사용하여 gcloud 명령어를 그룹화하는 인스턴스의 모든 인스턴스를 나열합니다.

gcloud compute instance-groups list-instances instance-group-name

list-instances 명령어에서 나열된 인스턴스 중 하나로 SSH로 연결합니다.

gcloud compute ssh instance-name --zone=us-central1-a

인스턴스 그룹의 VM이 서로 통신할 수 있는지 확인

nmap 명령어를 사용하여 인스턴스 그룹의 VM이 서로 통신할 수 있는지 확인합니다. 연결된 VM에서 nmap 명령어를 실행합니다. 이때 instance-name을 인스턴스 그룹의 다른 VM 인스턴스 이름으로 바꿉니다.

(vm)$ nmap -Pn -p 8477 instance-name
Starting Nmap 7.40 ( https://nmap.org ) at 2019-10-02 21:35 UTC
Nmap scan report for pytorch-20190923-n4tx.c.jysohntpu.internal (10.164.0.3)
Host is up (0.00034s latency).
PORT     STATE  SERVICE
8477/tcp closed unknown

STATE 필드가 filtered로 되어 있지 않다면 방화벽 규칙이 올바르게 설정된 것입니다.

Cloud TPU Pod 만들기

gcloud compute tpus create tpu-name \
    --zone=us-central1-a \
    --network=default \
    --accelerator-type=v2-32 \
    --version=1.6

Pod에서 분산 학습 실행

  1. VM 세션 창에서 Cloud TPU 이름을 내보내고 conda 환경을 활성화합니다.

    (vm)$ export TPU_NAME=tpu-name
    (vm)$ conda activate torch-xla-1.6
    
  2. 학습 스크립트를 실행합니다.

    (torch-xla-1.6)$ python -m torch_xla.distributed.xla_dist \
          --tpu=$TPU_NAME \
          --conda-env=torch-xla-1.6 \
          --env XLA_USE_BF16=1 \
          --env ANY_OTHER=ENV_VAR \
          -- python /usr/share/torch-xla-1.6/pytorch/xla/test/test_train_mp_imagenet.py \
          --fake_data
    

위의 명령어를 실행하면 다음과 비슷한 출력이 표시됩니다(--fake_data 사용에 유의). 학습에 소요되는 시간은 v3-32 TPU Pod에서 약 30분입니다.

2020-08-06 02:38:29  [] Command to distribute: "python" "/usr/share/torch-xla-nightly/pytorch/xla/test/test_train_mp_imagenet.py" "--fake_data"
2020-08-06 02:38:29  [] Cluster configuration: {client_workers: [{10.164.0.43, n1-standard-96, europe-west4-a, my-instance-group-hm88}, {10.164.0.109, n1-standard-96, europe-west4-a, my-instance-group-n3q2}, {10.164.0.46, n1-standard-96, europe-west4-a, my-instance-group-s0xl}, {10.164.0.49, n1-standard-96, europe-west4-a, my-instance-group-zp14}], service_workers: [{10.131.144.61, 8470, v3-32, europe-west4-a, pytorch-nightly, my-tpu-slice}, {10.131.144.59, 8470, v3-32, europe-west4-a, pytorch-nightly, my-tpu-slice}, {10.131.144.58, 8470, v3-32, europe-west4-a, pytorch-nightly, my-tpu-slice}, {10.131.144.60, 8470, v3-32, europe-west4-a, pytorch-nightly, my-tpu-slice}]}
2020-08-06 02:38:31 10.164.0.43 [0]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:31 10.164.0.43 [0]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2757      0 --:--:-- --:--:-- --:--:--  3166
2020-08-06 02:38:34 10.164.0.43 [0]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:34 10.164.0.43 [0]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2623      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:37 10.164.0.46 [2]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:37 10.164.0.46 [2]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2583      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:37 10.164.0.49 [3]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:37 10.164.0.49 [3]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2530      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:37 10.164.0.109 [1]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:37 10.164.0.109 [1]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2317      0 --:--:-- --:--:-- --:--:--  2375
2020-08-06 02:38:40 10.164.0.46 [2]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:40 10.164.0.49 [3]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:40 10.164.0.46 [2]                                  Dload  Upload   Total   Spent    Left  Speed
2020-08-06 02:38:40 10.164.0.49 [3]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2748      0 --:--:-- --:--:-- --:--:--  3166
100    19  100    19    0     0   2584      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:40 10.164.0.109 [1]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:40 10.164.0.109 [1]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2495      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:43 10.164.0.49 [3]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:43 10.164.0.49 [3]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2654      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:43 10.164.0.43 [0]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:43 10.164.0.43 [0]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2784      0 --:--:-- --:--:-- --:--:--  3166
2020-08-06 02:38:43 10.164.0.46 [2]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:43 10.164.0.46 [2]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2691      0 --:--:-- --:--:-- --:--:--  3166
2020-08-06 02:38:43 10.164.0.109 [1]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:43 10.164.0.109 [1]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2589      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/14 Epoch=1 Step=0 Loss=6.87500 Rate=258.47 GlobalRate=258.47 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/15 Epoch=1 Step=0 Loss=6.87500 Rate=149.45 GlobalRate=149.45 Time=02:38:57
2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.43 [0] Epoch 1 train begin 02:38:52
2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:1/0 Epoch=1 Step=0 Loss=6.87500 Rate=25.72 GlobalRate=25.72 Time=02:38:57
2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:0/6 Epoch=1 Step=0 Loss=6.87500 Rate=89.01 GlobalRate=89.01 Time=02:38:57
2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:0/1 Epoch=1 Step=0 Loss=6.87500 Rate=64.15 GlobalRate=64.15 Time=02:38:57
2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:0/2 Epoch=1 Step=0 Loss=6.87500 Rate=93.19 GlobalRate=93.19 Time=02:38:57
2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:0/7 Epoch=1 Step=0 Loss=6.87500 Rate=58.78 GlobalRate=58.78 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] Epoch 1 train begin 02:38:56
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:1/8 Epoch=1 Step=0 Loss=6.87500 Rate=100.43 GlobalRate=100.43 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/13 Epoch=1 Step=0 Loss=6.87500 Rate=66.83 GlobalRate=66.83 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/11 Epoch=1 Step=0 Loss=6.87500 Rate=64.28 GlobalRate=64.28 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/10 Epoch=1 Step=0 Loss=6.87500 Rate=73.17 GlobalRate=73.17 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/9 Epoch=1 Step=0 Loss=6.87500 Rate=27.29 GlobalRate=27.29 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/12 Epoch=1 Step=0 Loss=6.87500 Rate=110.29 GlobalRate=110.29 Time=02:38:57
2020-08-06 02:38:57 10.164.0.46 [2] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.46 [2] | Training Device=xla:0/20 Epoch=1 Step=0 Loss=6.87500 Rate=100.85 GlobalRate=100.85 Time=02:38:57
2020-08-06 02:38:57 10.164.0.46 [2] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.46 [2] | Training Device=xla:0/22 Epoch=1 Step=0 Loss=6.87500 Rate=93.52 GlobalRate=93.52 Time=02:38:57
2020-08-06 02:38:57 10.164.0.46 [2] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.46 [2] | Training Device=xla:0/23 Epoch=1 Step=0 Loss=6.87500 Rate=165.86 GlobalRate=165.86 Time=02:38:57

삭제

이 가이드에서 사용한 리소스 비용이 Google Cloud Platform 계정에 청구되지 않도록 하려면 다음 안내를 따르세요..

  1. Compute Engine VM의 연결을 해제합니다.

    (vm)$ exit
    
  2. 인스턴스 그룹을 삭제합니다.

    gcloud compute instance-groups managed delete instance-group-name
    
  3. TPU Pod를 삭제합니다.

    gcloud compute tpus delete ${TPU_NAME} --zone=us-central1-a
    

다음 단계

PyTorch Colab 사용: