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
    
  2. 다음 명령어를 실행하여 리소스가 삭제되었는지 확인합니다. TPU가 더 이상 나열되지 않았는지 확인합니다. 삭제하는 데 몇 분 정도 걸릴 수 있습니다.

    $ gcloud compute tpus tpu-vm list \
      --zone europe-west4-a