Cloud TPU PyTorch/XLA 사용자 가이드

PyTorch/XLA를 사용하여 ML 워크로드 실행

이 가이드의 절차를 시작하기 전에 GCP 프로젝트 준비의 설명대로 TPU VM과 ssh를 설정합니다. 그러면 이 가이드의 명령어를 실행하는 데 필요한 리소스가 설정됩니다.

PyTorch 1.8.1 및 PyTorch/XLA 1.8.1은 TPU VM에 사전 설치되어 있습니다.

기본 설정

XRT TPU 기기 구성을 설정합니다.

   (vm)$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"

상당한 크기의 자주 할당되는 모델의 경우 tcmalloc은 기본 malloc에 비해 성능을 크게 개선하므로 TPU VM에 사용되는 기본 malloctcmalloc입니다. 그러나 워크로드에 따라(예를 들어 임베딩 테이블에 대한 대규모 할당이 있는 DLRM) tcmalloc이 느려질 수 있으며, 이 경우 대신 기본 malloc을 사용하여 다음 변수를 설정 해제할 수 있습니다.

   (vm)$ unset LD_PRELOAD

간단한 계산 실행

Python 3 인터프리터 시작:

(vm)$ python3
   Python 3.6.9 (default, Jan 26 2021, 15:33:00)
   [GCC 8.4.0] on linux
   Type "help", "copyright", "credits" or "license" for more information.

Pytorch 3 인터프리터에서 다음 PyTorch 패키지 가져오기:

import torch
import torch_xla.core.xla_model as xm

다음 계산 수행:

dev = xm.xla_device()
이 명령의 출력은 다음과 유사합니다.
2021-04-01 23:20:23.268115: E   55362 tensorflow/core/framework/op_kernel.cc:1693] OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op: TPURoundRobin
2021-04-01 23:20:23.269345: E   55362 tensorflow/core/framework/op_kernel.cc:1693] OpKernel ('op: "TpuHandleToProtoKey" device_type: "CPU"') for unknown op: TpuHandleToProtoKey
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

명령어에서 다음과 같은 출력이 표시됩니다.

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

단일 기기 TPU에서 Resnet 실행

이제 PyTorch / XLA 코드를 원하는 대로 실행할 수 있습니다. 예를 들어 다음을 사용하여 가짜 데이터로 ResNet을 실행할 수 있습니다.

   (vm)$ git clone --recursive https://github.com/pytorch/xla.git
   (vm)$ python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1

ResNet 샘플은 1세대 동안 학습하며 약 7분 정도 걸립니다. 그러면 다음과 비슷한 출력이 표시됩니다.

Epoch 1 test end 20:57:52, Accuracy=100.00 Max Accuracy: 100.00%

ResNet 학습이 끝나면 TPU VM을 삭제하세요.

   (vm)$ exit
   $ gcloud alpha compute tpus tpu-vm delete tpu-name \
   --zone=zone

gcloud list를 실행하여 리소스가 삭제되었는지 확인합니다. 삭제하는 데 몇 분 정도 걸릴 수 있습니다. 다음과 같은 응답이 나타나면 인스턴스가 성공적으로 삭제되었다는 의미입니다.

   $ gcloud alpha compute tpus list --zone=zone
   

   Listed 0 items.
   

고급 설정

위 예시(간단한 계산 및 ResNet50)에서 PyTorch/XLA 프로그램은 동일한 프로세스로 로컬 XRT 서버를 시작합니다. 별도의 프로세스에서 XRT 로컬 서비스를 시작할 수도 있습니다.

(vm)$ python3 -m torch_xla.core.xrt_run_server --port 51011 --restart

이 방식의 장점은 학습 전체에서 컴파일 캐시가 지속된다는 점입니다. 이 경우 /tmp/xrt_server_log에서 서버 측 로깅을 찾을 수 있습니다.

(vm)$ ls /tmp/xrt_server_log/
server_20210401-031010.log

TPU VM 성능 프로파일링

TPU VM에서 모델을 프로파일링하는 방법에 대한 자세한 내용은 PyTorch XLA 성능 프로파일링을 참조하세요.

pod

PyTorch/XLA를 사용하려면 모든 TPU VM에서 모델 코드 및 데이터에 액세스할 수 있어야 합니다. TPU VM pod를 만들 때 다음 시작 스크립트를 사용하면 쉽게 시작할 수 있습니다. 모든 TPU VM에서 데이터를 다운로드합니다.

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

    $ sudo /opt/google-cloud-sdk/bin/gcloud components update
    
  2. Pod TPU VM을 만드는 데 필요한 다음 환경 변수를 내보냅니다.

    $ export PROJECT_ID=project-id
    $ export TPU_NAME=tpu-name
    $ export ZONE=zone
    $ export RUNTIME_VERSION=v2-alpha
    
  3. TPU VM 만들기

    $ gcloud alpha compute tpus tpu-vm create ${TPU_NAME} \
    --zone ${ZONE} --project ${PROJECT_ID} --accelerator-type v3-32 \
    --version ${RUNTIME_VERSION} --metadata startup-script='#! /bin/bash
    cd /usr/share/
    git clone --recursive https://github.com/pytorch/pytorch
    cd pytorch/
    git clone --recursive https://github.com/pytorch/xla.git
    EOF'
    

    안내를 계속 따르면서 (vm)$으로 시작하는 각 명령어를 VM 세션 창에서 실행합니다.

  4. ssh를 사용해 TPU 워커(예: 작업자 0)에 연결하고 데이터/모델 다운로드가 완료되었는지 확인한 다음(여기에는...걸립니다), pod의 작업자 간 ssh 연결을 위한 SSH 키를 생성한 후 학습을 시작합니다.

    (vm)$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone ${ZONE} --project ${PROJECT_ID}
    

    프로젝트 ssh 메타데이터를 업데이트합니다.

    (vm)$ gcloud compute config-ssh
    
    (vm)$ export TPU_NAME=tpu-name
    
    (vm)$ python3 -m torch_xla.distributed.xla_dist \
       --tpu=${TPU_NAME} -- python3 /usr/share/pytorch/xla/test/test_train_mp_imagenet.py \
       --fake_data --model=resnet50 --num_epochs=1
    

    학습에는 약 3분이 소요됩니다. 완료되면 다음과 비슷한 메시지가 표시됩니다.

    Epoch 1 test end 23:49:15, Accuracy=100.00
    10.164.0.11 [0] Max Accuracy: 100.00%
    

    각 작업자의 /tmp/xrt_server_log 아래에서 서버 측 로깅을 찾을 수 있습니다.

       (vm)$ ls /tmp/xrt_server_log/
    
    server_20210401-031010.log
    

    서버가 비정상 상태로 유지되는 경우 XRT_SERVER를 다시 시작하려면 xla_dist를 실행할 때 --restart-tpuvm-pod-server를 전달하면 됩니다. LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4 같은 환경 변수를 비롯한 새로운 XRT 서버 설정은 서버를 다시 시작할 때까지 선택되지 않습니다.

  5. ResNet 학습이 끝나면 TPU VM을 삭제하세요.

    (vm)$ exit
    
    $ gcloud alpha compute tpus tpu-vm delete tpu-name \
    --zone=zone
    
  6. gcloud list를 실행하여 리소스가 삭제되었는지 확인합니다. 삭제하는 데 몇 분 정도 걸릴 수 있습니다. 다음과 같은 응답이 나타나면 인스턴스가 성공적으로 삭제되었다는 의미입니다.

    $ gcloud alpha compute tpus list --zone=zone
    
    Listed 0 items.
    

원격 조정자가 있는 pod

pod 학습을 시작하는 데 사용할 원격 조정자를 설정하는 것을 권장합니다. 원격 조정자는 TPU VM이 아닌 표준 Compute Engine VM입니다. 원격 조정자가 TPU VM Pod에 명령어를 실행합니다. 원격 조정자를 사용하면 TPU 유지보수 이벤트가 발생할 때 학습이 중단되어 TPU VM Pod에 대한 액세스 권한을 잃을 수 있지만 원격 조정자가 TPU VM을 자동 복구하도록 설정할 수 있다는 장점이 있습니다.

  1. TPU VM을 만드는 데 필요한 다음 환경 변수를 내보냅니다.

    $ export TPU_NAME=tpu-name
    $ export ZONE=zone
    $ export PROJECT_ID=project-id
    
  2. TPU VM Pod 슬라이스 만들기

    gcloud 명령어를 사용하여 TPU Pod 슬라이스를 만듭니다. 예를 들어 v2-32 Pod 슬라이스를 만들려면 다음 명령어를 사용합니다.

    $ gcloud alpha compute tpus tpu-vm create tpu-name \
    --zone europe-west4-a --project tpu-prod-env-one-vm --accelerator-type v3-32 \
    --version v2-alpha --metadata startup-script='#! /bin/bash
    cd /usr/share/
    git clone --recursive https://github.com/pytorch/pytorch
    cd pytorch/
    git clone --recursive https://github.com/pytorch/xla.git
    EOF'
    
  3. 원격 조정자 VM을 만드는 데 필요한 다음 환경 변수를 내보냅니다.

    $ export VM_NAME=vm-name
    $ export ZONE=zone
    
  4. 다음을 실행하여 원격 조정자 VM을 만듭니다.

    (vm)$ gcloud compute --project=project-id instances create vm-name \
      --zone=zone  \
      --machine-type=n1-standard-1  \
      --image-family=torch-xla \
      --image-project=ml-images  \
      --boot-disk-size=200GB \
      --scopes=https://www.googleapis.com/auth/cloud-platform
    

    gcloud compute 명령어 실행이 끝나면 셸 프롬프트가 username@projectname에서 username@vm-name으로 변경되었는지 확인합니다. 변경되었다면 원격 조정자 VM에 로그인되었다는 의미입니다.

  5. ssh를 사용하여 원격 조정자 인스턴스에 연결합니다.

    (vm)$ gcloud compute ssh vm-name --zone=zone
    
  6. torch-xla-1.8.1 환경을 활성화하고 여기에서 학습을 실행합니다.

    (vm)$ gcloud compute config-ssh
    
    (vm)$ conda activate torch-xla-1.8.1
    (vm)$ python3 -m torch_xla.distributed.xla_dist --tpu=tpu-name --restart-tpuvm-pod --env LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4 -- python3 /usr/share/pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
    

    학습을 실행하는 데 약 3분이 소요되며 다음과 비슷한 메시지가 생성됩니다.

      Epoch 1 test end 23:19:53, Accuracy=100.00
      Max Accuracy: 100.00%
      

  7. ResNet 학습이 끝나면 TPU VM을 종료하고 원격 조정자 VM과 TPU VM을 삭제합니다.

    (vm)$ exit
    
    $ gcloud compute instances delete vm-name  \
      --zone=zone
    
    $ gcloud alpha compute tpus tpu-vm delete tpu-name \
      --zone zone