Cloud TPU PyTorch/XLA 사용자 가이드

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

이 가이드의 절차를 시작하기 전 Cloud TPU VM 사용자 가이드에 설명된 대로 TPU VM 및 ssh를 설정합니다.

Pytorch/XLA에 사용할 수 있는 TPU 소프트웨어 버전 목록은 PyTorch 지원 버전을 참조하세요.

기본 설정

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

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

   (vm)$ export TPU_NUM_DEVICES=4

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

   (vm)$ unset LD_PRELOAD

PyTorch 버전 변경

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

(tpuvm):$ cd /usr/share/
(tpuvm):$ sudo git clone -b release/1.13 --recursive https://github.com/pytorch/pytorch
(tpuvm):$ cd pytorch/
(tpuvm):$ sudo git clone -b r1.13 --recursive https://github.com/pytorch/xla.git
(tpuvm):$ cd xla/
(tpuvm):$ yes | sudo pip3 uninstall torch_xla
(tpuvm):$ yes | sudo pip3 uninstall torch
(tpuvm):$ yes | sudo pip3 uninstall torch_vision
(tpuvm):$ sudo pip3 install torch==1.13.0
(tpuvm):$ sudo pip3 install torchvision==0.14.0
(tpuvm):$ sudo pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-1.13-cp38-cp38-linux_x86_64.whl
(tpuvm):$ sudo rm -rf /usr/local/lib/python3.8/dist-packages/libtpu*
(tpuvm):$ sudo pip3 install torch_xla[tpuvm]

간단한 계산 수행

TPU VM에서 Python 인터프리터를 시작합니다.
```
(vm)$ python3
```

다음 PyTorch 패키지를 가져옵니다.

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)
```
참고: OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op 오류 메시지는 예상된 문제이며 향후 출시 버전에서 바로잡을 예정입니다.

다음 출력이 표시됩니다.
```
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 compute tpus tpu-vm delete tpu-name \
--zone=zone

삭제하는 데 몇 분 정도 걸릴 수 있습니다. gcloud 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 Pod 예시

TPU VM Pod에서 PyTorch/XLA를 실행하는 설정 정보와 예시는 PyTorch TPU VM Pod를 참조하세요.

TPU VM의 Docker

이 섹션에서는 PyTorch/XLA가 사전 설치된 TPU VM에서 Docker를 실행하는 방법을 보여줍니다.

사용 가능한 Docker 이미지

GitHub 리드미를 참조하여 사용 가능한 모든 TPU VM Docker 이미지를 찾을 수 있습니다.

TPU VM에서 Docker 이미지 실행

(tpuvm): sudo docker pull gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker run --privileged  --shm-size 16G --name tpuvm_docker -it -d  gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker exec --privileged -it tpuvm_docker /bin/bash
(pytorch) root:/#

libtpu 확인

libtpu가 설치되어 있는지 확인하려면 다음을 실행합니다.

(pytorch) root:/# ls /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/ | grep libtpu

다음과 비슷한 출력이 생성되어야 합니다.

libtpu
libtpu_nightly-0.1.dev20220518.dist-info

결과가 표시되지 않으면 다음을 사용하여 해당 libtpu를 수동으로 설치할 수 있습니다.

(pytorch) root:/# pip install torch_xla[tpuvm]

`tcmalloc` 확인하기

tcmalloc은 TPU VM에서 사용되는 기본 malloc입니다. 자세한 내용은 이 섹션을 참조하세요. 이 라이브러리는 최신 TPU VM Docker 이미지에 사전 설치되어 있어야 하지만 항상 수동으로 확인하는 것이 좋습니다. 다음 명령어를 실행하면 라이브러리가 설치되었는지 확인할 수 있습니다.

(pytorch) root:/# echo $LD_PRELOAD