TPU Pod 슬라이스에서 PyTorch 코드 실행

PyTorch/XLA에서는 모든 TPU VM이 모델 코드 및 데이터에 액세스할 수 있어야 합니다. 시작 스크립트를 사용하여 모델 데이터를 모든 TPU VM에 배포하는 데 필요한 소프트웨어를 다운로드할 수 있습니다.

TPU VM을 Virtual Private Cloud(VPC)에 연결하는 경우 포트 8470~8479에 대한 인그레스를 허용하는 방화벽 규칙을 프로젝트에 추가해야 합니다. 방화벽 규칙 추가에 대한 자세한 내용은 방화벽 규칙 사용을 참조하세요.

환경 설정하기

Cloud Shell에서 다음 명령어를 실행하여 gcloud의 최신 버전을 실행하세요.
```
$ gcloud components update
```
gcloud를 설치해야 하는 경우 다음 명령어를 사용합니다.
```
$ sudo apt install -y google-cloud-sdk
```

몇 가지 환경 변수를 만듭니다.

$ export PROJECT_ID=project-id
$ export TPU_NAME=tpu-name
$ export ZONE=us-central2-b
$ export RUNTIME_VERSION=tpu-ubuntu2204-base
$ export ACCELERATOR_TYPE=v4-32

TPU VM 만들기

$ gcloud compute tpus tpu-vm create ${TPU_NAME} \
--zone=${ZONE} \
--project=${PROJECT_ID} \
--accelerator-type=${ACCELERATOR_TYPE} \
--version ${RUNTIME_VERSION}

학습 스크립트 구성 및 실행

프로젝트에 SSH 인증서를 추가합니다.
```
ssh-add ~/.ssh/google_compute_engine
```

모든 TPU VM 작업자에 PyTorch/XLA를 설치합니다.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --worker=all --command="
  pip install torch~=2.3.0 torch_xla[tpu]~=2.3.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"

모든 TPU VM 작업자에서 XLA를 클론합니다.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --worker=all --command="git clone -b r2.3 https://github.com/pytorch/xla.git"

모든 작업자에서 학습 스크립트를 실행합니다.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --worker=all \
  --command="PJRT_DEVICE=TPU python3 ~/xla/test/test_train_mp_imagenet.py  \
  --fake_data \
  --model=resnet50  \
  --num_epochs=1 2>&1 | tee ~/logs.txt"

학습에 5분 정도 걸립니다. 완료되면 다음과 비슷한 메시지가 표시됩니다.

Epoch 1 test end 23:49:15, Accuracy=100.00
10.164.0.11 [0] Max Accuracy: 100.00%

삭제

TPU VM 사용이 완료되었으면 다음 단계에 따라 리소스를 삭제하세요.

Compute Engine에서 연결을 해제합니다.
```
(vm)$ exit
```
다음 명령어를 실행하여 리소스가 삭제되었는지 확인합니다. TPU가 더 이상 나열되지 않았는지 확인합니다. 삭제하는 데 몇 분 정도 걸릴 수 있습니다.
```
$ gcloud compute tpus tpu-vm list \
  --zone europe-west4-a
```