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

이 문서에서는 PyTorch 및 Cloud TPU 작업에 대해 간략히 안내합니다.

시작하기 전에

이 문서의 명령어를 실행하기 전에 Google Cloud 계정을 만들고 Google Cloud CLI를 설치하고 gcloud 명령어를 구성해야 합니다. 자세한 내용은 Cloud TPU 환경 설정을 참조하세요.

필요한 역할

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

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

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

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

`gcloud`를 사용하여 Cloud TPU 만들기

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

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 소프트웨어 버전입니다.

다음 명령어를 실행하여 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']

기본 계산 수행

현재 디렉터리에 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)

스크립트를 실행합니다.

(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')

삭제

이 페이지에서 사용한 리소스 비용이 Google Cloud 계정에 청구되지 않도록 하려면 다음 단계를 수행합니다.

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

Cloud TPU를 삭제합니다.

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

다음 명령어를 실행하여 리소스가 삭제되었는지 확인합니다. TPU가 더 이상 나열되지 않는지 확인합니다. 삭제하는 데 몇 분 정도 걸릴 수 있습니다.
```
$ gcloud compute tpus tpu-vm list \
    --zone=$ZONE
```

다음 단계

Cloud TPU VM에 대해 자세히 알아보세요.