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

이 문서의 명령어를 실행하기 전 계정 및 Cloud TPU 프로젝트 설정의 안내를 따르도록 유의하세요.

단일 TPU VM에서 PyTorch 코드를 실행한 후에는 TPU 슬라이스에서 실행하여 코드를 수직 확장할 수 있습니다. TPU 슬라이스는 전용 고속 네트워크 연결을 통해 서로 연결된 여러 TPU 보드입니다. 이 문서에서는 TPU 슬라이스에서 PyTorch 코드를 실행하는 방법을 소개합니다.

필요한 역할

TPU를 만들고 SSH를 사용하여 연결하는 데 필요한 권한을 얻으려면 관리자에게 프로젝트에 대한 다음 IAM 역할을 부여해 달라고 요청하세요.

TPU 관리자(roles/tpu.admin)
서비스 계정 사용자(roles/iam.serviceAccountUser)
Compute 뷰어(roles/compute.viewer)

역할 부여 방법에 대한 자세한 내용은 프로젝트, 폴더, 조직에 대한 액세스 관리를 참조하세요.

커스텀 역할 또는 다른 사전 정의된 역할을 통해 필요한 권한을 얻을 수도 있습니다.

Cloud TPU 슬라이스 만들기

명령어를 더 쉽게 사용할 수 있도록 몇 가지 환경 변수를 정의합니다.

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=europe-west4-b
export ACCELERATOR_TYPE=v5p-32
export RUNTIME_VERSION=v2-alpha-tpuv5

환경 변수 설명

변수	설명
`PROJECT_ID`	Google Cloud 프로젝트 ID입니다. 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다.
`TPU_NAME`	TPU 이름입니다.
`ZONE`	TPU VM을 만들 영역입니다. 지원되는 영역에 대한 자세한 내용은 TPU 리전 및 영역을 참조하세요.
`ACCELERATOR_TYPE`	액셀러레이터 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 액셀러레이터 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.
`RUNTIME_VERSION`	Cloud TPU 소프트웨어 버전입니다.

다음 명령어를 실행하여 TPU VM을 만듭니다.

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

슬라이스에 PyTorch/XLA 설치

TPU 슬라이스를 만든 후 TPU 슬라이스에서 모든 호스트에 PyTorch를 설치해야 합니다. --worker=all 및 --command 파라미터를 사용하여 gcloud compute tpus tpu-vm ssh 명령어로 이 작업을 수행할 수 있습니다.

SSH 연결 오류로 인해 다음 명령어가 실패하면 TPU VM에 외부 IP 주소가 없기 때문일 수 있습니다. 외부 IP 주소를 사용하지 않고 TPU VM에 액세스하려면 공개 IP 주소 없이 TPU VM에 연결의 안내를 따르세요.

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

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --worker=all \
    --command="pip install torch~=2.5.0 torch_xla[tpu]~=2.5.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 https://github.com/pytorch/xla.git"

TPU 슬라이스에서 학습 스크립트 실행

모든 작업자에서 학습 스크립트를 실행합니다. 학습 스크립트는 단일 프로그램 다중 데이터(SPMD) 샤딩 전략을 사용합니다. SPMD에 대한 자세한 내용은 PyTorch/XLA SPMD U사용자 가이드를 참조하세요.

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

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

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

삭제

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

Cloud TPU 인스턴스에서 아직 연결을 해제하지 않았으면 연결을 해제합니다.
```
(vm)$ exit
```
프롬프트가 username@projectname으로 바뀌면 Cloud Shell에 있는 것입니다.

Cloud TPU 리소스를 삭제합니다.

$ gcloud compute tpus tpu-vm delete  \
    --zone=${ZONE}

gcloud compute tpus tpu-vm list를 실행하여 리소스가 삭제되었는지 확인합니다. 삭제하는 데 몇 분 정도 걸릴 수 있습니다. 다음 명령어의 출력에 이 튜토리얼에서 만든 리소스가 포함되어서는 안 됩니다.
```
$ gcloud compute tpus tpu-vm list --zone=${ZONE}
```