Cloud TPU Pod에서 PyTorch 모델 학습

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

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

다음 다이어그램은 분산 클러스터 설정의 개요를 제공합니다. VM의 인스턴스 그룹이 TPU pod에 연결됩니다. TPU 코어 8개 그룹마다 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에서 학습합니다.

gcloud 명령어 구성
[선택사항] VM 디스크 이미지를 VM 이미지로 캡처
VM 이미지에서 인스턴스 템플릿 만들기
인스턴스 템플릿에서 인스턴스 그룹 만들기
Compute Engine VM에 SSH를 통해 연결
VM 간 통신을 허용하도록 방화벽 규칙 확인
Cloud TPU Pod 만들기
Pod에서 분산 학습 실행
삭제

`gcloud` 명령어 구성

gcloud로 Google Cloud 프로젝트를 구성합니다.

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

export PROJECT_ID=project-id

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

gcloud config set project ${PROJECT_ID}

새 Cloud Shell VM에서 이 명령어를 처음 실행하면 Authorize Cloud Shell 페이지가 표시됩니다. 페이지 하단에 있는 Authorize를 클릭하여 gcloud에서 사용자 인증 정보로 API를 호출하도록 허용합니다.

gcloud를 사용하여 기본 영역을 구성합니다.

gcloud config set compute/zone europe-west4-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 europe-west4-a \
    --family=torch-xla \
    --storage-location europe-west4

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 europe-west4-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=europe-west4-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=europe-west4-a \
    --network=default \
    --accelerator-type=v2-32 \
    --version=pytorch-1.13

Pod에서 분산 학습 실행

VM 세션 창에서 Cloud TPU 이름을 내보내고 conda 환경을 활성화합니다.
```
(vm)$ export TPU_NAME=tpu-name
```
```
(vm)$ conda activate torch-xla-1.13
```

학습 스크립트를 실행합니다.

(torch-xla-1.13)$ python -m torch_xla.distributed.xla_dist \
      --tpu=$TPU_NAME \
      --conda-env=torch-xla-1.13 \
      --env XLA_USE_BF16=1 \
      --env ANY_OTHER=ENV_VAR \
      -- python /usr/share/torch-xla-1.13/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 계정에 청구되지 않도록 하려면 리소스가 포함된 프로젝트를 삭제하거나 프로젝트를 유지하고 개별 리소스를 삭제하세요.

Compute Engine VM의 연결을 해제합니다.
```
(vm)$ exit
```

인스턴스 그룹을 삭제합니다.

gcloud compute instance-groups managed delete instance-group-name

TPU Pod를 삭제합니다.

gcloud compute tpus delete ${TPU_NAME} --zone=europe-west4-a

인스턴스 그룹 템플릿을 삭제합니다.

gcloud compute instance-templates delete instance-template-name

[선택사항] VM 디스크 이미지를 삭제합니다.
```
gcloud compute images delete image-name
```

다음 단계

PyTorch Colab 사용: