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

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

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

환경 설정하기

1. Cloud Shell에서 다음 명령어를 실행하여 gcloud의 최신 버전을 실행하세요.

   $ gcloud components update

   gcloud를 설치해야 하는 경우 다음 명령어를 사용합니다.

   $ sudo apt install -y google-cloud-sdk

2. 다음과 같이 몇 가지 환경 변수를 만듭니다.

   $ 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}

학습 스크립트 구성 및 실행

1. 프로젝트에 SSH 인증서를 추가합니다.

   ssh-add ~/.ssh/google_compute_engine

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

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

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

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

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

   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 사용이 완료되었으면 다음 단계에 따라 리소스를 삭제하세요.

1. Compute Engine 인스턴스에서 연결을 해제합니다.

   (vm)$ exit

   프롬프트가 username@projectname으로 바뀌면 Cloud Shell에 있는 것입니다.

2. Cloud TPU 및 Compute Engine 리소스를 삭제합니다.

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

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

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