PyTorch를 사용하여 Cloud TPU VM에서 계산 실행

gcloud를 사용하여 Cloud TPU 만들기

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

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-east5-a
   export ACCELERATOR_TYPE=v5litepod-8
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite

   환경 변수 설명

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

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

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

Cloud TPU VM에 연결

다음 명령어를 사용하여 SSH를 통해 TPU VM에 연결합니다.

$ gcloud compute tpus tpu-vm ssh $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

SSH를 사용해 TPU VM에 연결하지 못하는 경우, 해당 TPU VM에 외부 IP 주소가 없기 때문일 수 있습니다. 외부 IP 주소가 없이 TPU VM에 액세스하려면 공개 IP 주소 없이 TPU VM에 연결의 안내를 따르세요.

TPU VM에 PyTorch/XLA 설치

$ (vm) sudo apt-get update
$ (vm) sudo apt-get install libopenblas-dev -y
$ (vm) pip install numpy
$ (vm) pip install torch torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

PyTorch가 TPU에 액세스할 수 있는지 확인

다음 명령어를 사용하여 PyTorch가 TPU에 액세스할 수 있는지 확인합니다.

$ (vm) PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

이 명령어의 출력은 다음과 같이 표시되어야 합니다.

['xla:0', 'xla:1', 'xla:2', 'xla:3', 'xla:4', 'xla:5', 'xla:6', 'xla:7']

기본 계산 수행

1. 현재 디렉터리에 tpu-test.py이라는 파일을 만들고 이 파일에 다음 스크립트를 복사하여 붙여넣습니다.

   import torch
   import torch_xla.core.xla_model as xm
   
   dev = xm.xla_device()
   t1 = torch.randn(3,3,device=dev)
   t2 = torch.randn(3,3,device=dev)
   print(t1 + t2)
   

2. 스크립트를 실행합니다.

   (vm)$ PJRT_DEVICE=TPU python3 tpu-test.py

   스크립트의 출력에 계산 결과가 표시됩니다.

   tensor([[-0.2121,  1.5589, -0.6951],
           [-0.7886, -0.2022,  0.9242],
           [ 0.8555, -1.8698,  1.4333]], device='xla:1')