Cloud TPU PyTorch/XLA 사용자 가이드

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

이 가이드의 절차를 시작하기 전 Google Cloud 프로젝트 준비에 설명된 대로 TPU VM을 설정하고 ssh로 연결합니다.

Pytorch/XLA에는 두 가지 버전의 사용 가능한 TPU 소프트웨어가 있습니다.

  1. 사용 가능한 최신 이미지: tpu-vm-pt-1.11
  2. tpu-vm-pt-1.10
  3. v2-alpha

tpu-vm-pt-1.11 이미지는 PyTorch 1.11.0 및 Pytorch/XLA 1.11.0에 사전 설치되어 있습니다. tpu-vm-pt-1.10 이미지는 PyTorch 1.10.0 및 Pytorch/XLA 1.10.0에 사전 설치되어 있습니다. v2-alpha 이미지는 PyTorch 1.8.1에 사전 설치되어 있습니다.

기본 설정

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

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

상당한 크기의 자주 할당되는 모델의 경우 tcmalloc가 C/C++ 런타임 함수 malloc에 비해 성능을 개선합니다. TPU VM에 사용되는 기본 malloctcmalloc입니다. LD_PRELOAD 환경 변수를 설정 해제하여 TPU VM 소프트웨어가 표준 malloc를 사용하도록 강제할 수 있습니다.

   (vm)$ unset LD_PRELOAD

PyTorch 버전 변경

TPU VM에 사전 설치된 PyTorch 버전을 사용하지 않으려면 사용하려는 버전을 설치합니다. 예를 들어 1.9를 사용하려면 다음 안내를 따르세요.

(vm)$ sudo bash /var/scripts/docker-login.sh
(vm)$ sudo docker rm libtpu || true
(vm)$ sudo docker create --name libtpu gcr.io/cloud-tpu-v2-images/libtpu:pytorch-1.9 "/bin/bash"
(vm)$ sudo docker cp libtpu:libtpu.so /lib
(vm)$ sudo pip3 uninstall --yes torch torch_xla torchvision
(vm)$ sudo pip3 install torch==1.9.0
(vm)$ sudo pip3 install torchvision==0.10.0
(vm)$ sudo pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-1.9-cp38-cp38-linux_x86_64.whl

간단한 계산 수행

  1. TPU VM에서 Python 인터프리터를 시작합니다.

    (vm)$ python3
    
  2. 다음 PyTorch 패키지를 가져옵니다.

    import torch
    import torch_xla.core.xla_model as xm
    
  3. 다음 스크립트를 입력합니다.

    dev = xm.xla_device()
    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 alpha compute tpus list --zone=${ZONE}를 실행하여 리소스가 삭제되었는지 확인합니다.

고급 설정

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

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

이 접근 방식의 장점은 학습을 실행할 때 컴파일 캐시가 지속된다는 것입니다. 개별 프로세스로 XLA 서버를 실행하면 서버 측 로깅 정보가 /tmp/xrt_server_log에 기록됩니다.

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

TPU VM 성능 프로파일링

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

포드

PyTorch/XLA에서는 모든 TPU VM이 모델 코드 및 데이터에 액세스할 수 있어야 합니다. TPU VM 포드를 만들 때 다음 시작 스크립트를 사용합니다. 모든 TPU VM에 데이터가 다운로드됩니다.

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

    $ sudo /opt/google-cloud-sdk/bin/gcloud components update
    
  2. 다음 환경 변수를 내보냅니다.

    $ export PROJECT_ID=project-id
    $ export TPU_NAME=tpu-name
    $ export ZONE=zone
    $ export RUNTIME_VERSION=tpu-vm-pt-1.11
    
  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)$으로 시작하는 각 명령어를 TPU VM 셸에서 실행합니다.

  4. worker 0과 같은 아무 TPU 워커에 ssh로 연결합니다. 데이터/모델 다운로드가 완료되었는지 확인합니다. 다운로드 속도는 빠르지만, 파일 처리에 몇 분 정도 걸릴 수 있습니다. 다운로드 디렉터리에서 ls -al을 실행하여 파일 처리가 완료되었는지 확인할 수 있습니다. 파일은 시작 스크립트에 지정된 순서로 다운로드됩니다. 예를 들면 다음과 같습니다.

    download file1
    download file2
    download file3
    

    file3ls 명령어의 출력에 있으면 파일 처리가 완료된 것입니다.

    파일 처리가 완료되면 ssh-key를 생성하여 포드의 VM 작업자 간에 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 alpha compute tpus list --zone=${ZONE}를 실행하여 리소스가 삭제되었는지 확인합니다. 삭제하는 데 몇 분 정도 걸릴 수 있습니다.

원격 조정자가 있는 포드

TPU 유지보수 이벤트가 발생할 때 원격 조정자를 사용하여 포드가 자동 복구되도록 하는 것이 좋습니다. 원격 조정자는 TPU VM이 아닌 표준 Compute Engine VM입니다. 포드의 TPU VM에 명령을 실행합니다.

  1. 다음 환경 변수를 내보냅니다.

    $ export TPU_NAME=tpu-name
    $ export ZONE=zone
    $ export PROJECT_ID=project-id
    
  2. TPU VM 포드 슬라이스를 만듭니다.

    $ gcloud alpha compute tpus tpu-vm create tpu-name \
     --zone europe-west4-a --project tpu-prod-env-one-vm --accelerator-type v3-32 \
     --version tpu-vm-pt-1.11 --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.11 환경을 활성화하고 여기에서 학습을 실행합니다.

    (vm)$ gcloud compute config-ssh
    
    (vm)$ conda activate torch-xla-1.11
    (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