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

이 빠른 시작에서는 Cloud TPU를 만들고 PyTorch를 설치하고 Cloud TPU에서 간단한 계산을 실행하는 방법을 보여줍니다. Cloud TPU에서 모델을 학습시키는 방법을 보여주는 자세한 튜토리얼은 Cloud TPU PyTorch 튜토리얼 중 하나를 참조하세요.

시작하기 전에

이 빠른 시작을 수행하기 전에 Google Cloud Platform 계정을 만들고 Google Cloud CLI를 설치하고 gcloud 명령어를 구성해야 합니다. 자세한 내용은 계정 및 Cloud TPU 프로젝트 설정을 참조하세요.

gcloud로 Cloud TPU 만들기

기본 사용자 프로젝트, 네트워크, 영역에서 TPU VM을 생성하려면 다음을 실행합니다.

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central1-a \
   --accelerator-type=v3-8 \
   --version=tpu-ubuntu2204-base

명령어 플래그 설명

zone
Cloud TPU를 만들려는 영역입니다.
accelerator-type
가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.
version
Cloud TPU 소프트웨어 버전입니다.

TPU를 만드는 동안 기본 네트워크 및 서브네트워크를 지정하려면 추가로 --network 및 --subnetwork 플래그를 전달할 수 있습니다. 기본 네트워크를 사용하지 않으려면 --network 플래그를 전달해야 합니다. --subnetwork 플래그는 선택사항이며 사용 중인 네트워크(기본 또는 사용자 지정)에 대한 기본 서브네트워크를 지정하는 데 사용할 수 있습니다. 이러한 플래그에 대한 자세한 내용은 gcloud API 참조 페이지를 확인하세요.

Cloud TPU VM에 연결

   $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central1-a

TPU VM에 PyTorch/XLA 설치

   (vm)$ pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html

TPU 런타임 구성 설정

PyTorch/XLA 런타임이 TPU를 사용해야 합니다.

   (vm) $ export PJRT_DEVICE=TPU

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

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

    import torch
import torch_xla
import torch_xla.core.xla_model as xm
print(f'PyTorch can access {len(torch_xla.devices())} TPU cores')

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

    (vm)$ python3 tpu-count.py

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

    PyTorch can access 8 TPU cores

기본 계산 수행

  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)$ 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 계정에 청구되지 않도록 하려면 다음 단계를 수행합니다.

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

    (vm)$ exit

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

  2. Cloud TPU를 삭제합니다.

    $ gcloud compute tpus tpu-vm delete tpu-name \
  --zone=us-central1-a

이 명령어의 출력에서 TPU가 삭제되었는지 확인할 수 있습니다.

다음 단계

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