PyTorch를 TPU에서 사용해보기 (2)

---

지난 포스팅에서는 PyTorch의 MNIST 기본 예제를 Colab을 통해 Single-core TPU에서 실행하는 과정을 정리해보았다. 이번 포스팅에서는 마찬가지로 colab에서 Multi-core TPUs를 실행하는 과정과 함께 추가적으로 동일한 코드에 대해서 Single-core TPU와 Multi-core TPUs 그리고 GPU의 성능 결과를 비교해보려 한다.
colab에서 무료로 사용할 수 있는 TPU에는 8개 core가 있다고 한다. 8개 core가 있다는 게 device가 물리적으로 8개 분리되어 있는 것을 의미하는 것인지 일반적인 CPU처럼 하나의 물리적인 device 안에 코어가 8개 있는 것인지 정확하게 이해가 되지는 않는다. GPU를 여러 개 사용할 때 파이썬에서 cuda device를 체크해보면 cuda:1, cuda:2, …와 같이 물리적인 GPU device의 개수에 따라 표시가 되는 것을 확인할 수 있었다. 그리고 마찬가지로 Colab 상에서 TPU device를 체크해보면 xla:1, xla:2, … xla:8 처럼 8개가 잡히는 것을 확인할 수 있었다. 이러한 결과를 통해 8 코어 TPU라는 것이 물리적으로 device가 8개 나눠진게 아닌가 하는 생각이 든다. 물론 이러한 생각은 컴퓨터 공학을 전공하지 않은 비전공자 입장에서 생각한 부분이라 하드웨어 구조 등을 이해하지 못하고 작성하는 내용이다.

>> xm.get_xla_supported_devices()
['xla:1', 'xla:2', 'xla:3', 'xla:4', 'xla:5', 'xla:6', 'xla:7', 'xla:8']

---

1. Colab에서 Multi-core TPUs 사용하기

지난 포스팅과 마찬가지로 PyTorch의 MNIST 기본 예제에서 Multi-core TPUs 사용을 위해 필요한 수정 사항만을 작성하도록 한다.

1.1. PyTorch용 XLA Library

import torch_xla
import torch_xla.core.xla_model as xm

# for Multi-core TPUs
import torch_xla.distributed.xla_multiprocessing as xmp
import torch_xla.distributed.parallel_loader as pl

TPU에서 PyTorch를 사용하기 위해선 지난 포스팅(Colab에서 Single-core TPU 사용하기)에서 했던 것처럼 PyToch용 XLA 라이브러리를 import해줘야한다. 그리고 Multi-core TPUs를 사용하기 위해선 xla_multiprocessing와 parallel_loader 라이브러리를 추가적으로 import 해줘야한다. xla_multiprocessing은 여러 개의 core에서 각각 처리할 수 있도록 multiprocessing 하는 것을 도와주는 라이브러리인 것 같고, parallel_loader는 multiprocessing시에 각 프로세서에 데이터를 분배해주는 역할을 하는 것으로 생각된다. 자세한 내용은 [PyTorch용-XLA-라이브러리]를 참고하면 될 것 같다.

TPU를 XLA device라고 하는 것 같다.

1.2. map_fn 정의해주기

def map_fn(index, flags):
  ~
  device = xm.xla_device()
  ~
  sampler = torch.utils.data.distributed.DistributedSampler(dataset, ~)
  dataloader = torch.utils.data.DataLoader(dataset, sampler=sampler, ~)
  ~
  for epoch in range(max_epochs):
    ~
    parallel_loader = pl.ParallelLoader(dataloader, [device]).per_device_loader(device)
    for iter, data in enumerate(parallel_loader):
    ~
    optimizer.zero_grad()
    loss.backward()
    xm.optimizer_step(optimizer)
    ~
  ~

각각의 XLA device에서 실행할 코드를 정의해준다. 쉽게 생각해보면 실제 학습 및 테스트 코드를 함수화 한 것을 의미한다. 따라서, 위의 코드와 같이 해당 함수 안에는 매 epoch, iteration마다 진행되는 학습 코드가 들어있다. 함수의 인자 중 index는 현재 코드를 실행하고 있는 프로세스의 index를 의미한다. flags는 사용자가 함수에 전달하고 싶은 값을 넣어줄 수 있는 변수이다.
PyTorch용 XLA 라이브러리에서는 Single-core TPU와 Multi-core TPUs 코드의 큰 차이점을 다음의 세 가지로 설명하고 있다.

• Parallelloader
• xm.optimizer_step(optimizer)
• xmp.spawn

Parallelloader은 위의 코드에서 볼 수 있듯이, 기존 dataloader를 각 XLA device에 분산 시켜주는 역할을 한다. 그리고 xm.optimizer_step의 경우 ParallelLoader에서 자동적으로 XLA barrier를 생성해주기 때문에 Single-core TPU에사 사용했던 barrier=True 옵션을 제외한 xm.optimizer_step(optimizer)로 사용하면 된다. barrier 옵션을 따로 지정해주지 않으면 default는 False이다. xmp.spawn는 다음 장에서 설명하도록 한다.
나는 PyTorch용 XLA 라이브러리에서 제공하는 기본적인 3가지 수정 사항 외에 DistributedSampler를 추가적으로 사용하였다. 다양한 코드를 돌려본 것은 아니지만 적어도 MNIST 기본 예제에서 만큼은 DistributedSampler를 추가하는 것이 시간적으로 더 빨리 학습되는 것을 확인하였다.

Python에서 메소드? 함수? (참고): Python에서 메소드는 객체에 속해있는 함수를 의미한다. 따라서, 큰 의미에서 메소드는 결국 함수를 의미한다.

1.3. 각 XLA device를 실행하는 process 만들기

flags= {}
xmp.spawn(map_fn, args=(flags,), nprocs=8, start_method='fork')

마지막으로 앞서 map_fn 함수를 각 XLA device에서 실행하도록 xmp.spawn 메소드를 실행시켜주면 된다. 이때 nprocs는 device/process 개수를 의미하므로 TPU에 있는 8개 코어를 모두 사용하기 위해선 8을 입력해주도록 하자. 그리고 start_method 옵션의 경우 Colab에서는 fork 밖에 지원하지 않는다고 한다. fork 외에 다른 선택 사항이 어떤게 있는지 PyTorch용 XLA 라이브러리 사이트에도 나와있지 않고, 실행하는데 별다른 문제가 없으니 start_method는 fork를 쓰는 것으로 하자.

2. Colab에서 Single/Multi-core XLA device, GPU 성능 비교하기

이전 포스팅과 본 포스팅의 최종 목적은 딥러닝 모델 학습시에 TPU와 GPU가 시간적으로 얼마나 학습 시간 차이가 있는지 확인해보기 위함이었다. 이제 Single/Multi-core TPU의 사용 방법을 익혔으니 성능을 비교해보도록 하자. 본 포스팅에서 Device의 성능 비교를 위해 PyTorch의 MNIST 기본 예제를 사용하였고, 해당 코드를 각 device에 맞게 수정하여 사용하였다. 성능 비교에 사용된 전체 코드는 https://github.com/aithlab/colab-test/에서 확인할 수 있다.
MNIST 데이터에 대해 간단한 Convolution Network를 사용하여 10 Epoch을 학습한 결과를 비교하였다. 본 포스팅에서는 테스트 데이터에 대한 정확도의 성능 차이는 비교하지 않는다. 실험을 하다보니 이상한 점을 발견하였다. 왜 이런 현상이 발생하는지 아직 확인을 하지 못했지만, Colab에서 실행하는 ipynb 파일을 어디서 만들었는지에 따라 성능 차이가 있다.

• Colab 메모장에서 제일 아래에 있는 PyTorch용 예제 3개 중 아무거나 들어가서 이미 작성된 코드를 전부 지우고 내가 실행하고자 하는 코드를 붙여 넣은 파일
• Colab 사이트나 Google Drive에서 만든 ipynb 파일에 내가 실행하고자 하는 코드를 붙여 넣은 파일

위의 두 방법에 따라 Multi-core TPUs를 사용하는 MNIST 예제에서 거의 4배 정도 시간 차이가 났다. 코드는 완전히 똑같은 코드인데 파일 자체가 어디서 생성되었는지에 따라 시간 차이가 나는 것이 이해가 가지는 않지만 아무래도 첫 번째 방법의 경우 공식적인 예제로 만들어진 파일이다보니 조금 더 좋은 TPU에 연결될 수 있도록 만들어진 것이 아닌가 추측하고 있다.
성능 비교 표를 간단히 하기 위해 첫 번째 방법으로 생성된 파일을 예제 파일 두 번째 방법으로 생성된 파일을 Colab 파일이라고 부르겠다.

파일 종류	Multi TPUs (s)	Single TPU (s)	GPU (s)
예제 파일	32.868	81.551	54.401
Colab 파일	158.917	127.238	131.210

파일 종류	Multi TPUs (s/iter)	Single TPU (s/iter)	GPU (s/iter)
예제 파일	$0.019(\pm 0.043)$	$0.011(\pm 0.033)$	$0.006(\pm 0.004)$
Colab 파일	$0.132(\pm 0.151)$	$0.014(\pm 0.031)$	$0.011(\pm 0.005)$

파일 종류	Multi TPUs (s/epoch)	Single TPU (s/epoch)	GPU (s/epoch)
예제 파일	$2.362(\pm 1.287)$	$8.155(\pm 1.796)$	$5.440(\pm 0.074)$
Colab 파일	$14.9(\pm 0.516)$	$12.724(\pm 1.629)$	$13.121(\pm 0.168)$

* Colab의 경우 매번 같은 device가 할당 되는 것이 아니기 때문에 위의 성능 비교는 조금씩 차이가 있을 수 있다.

동일한 코드를 Colab에서 동일하게 실행시켰을 때, 각 device의 속도는 매번 다르다. 예를 들어, Colab 파일에 해당하는 결과는 때로는 Single TPU가 GPU보다 빠를 때도 있고 반대일 때도 있다. 하지만 예제 파일의 경우 대체적으로 Multi-core TPUs > GPU >= Single-core TPU 순서(빠른 속도 순서. 즉, Multi-core TPUs가 가장 빠르다)로 결과를 보인다.

위의 결과를 보면 TPU, GPU 모두에서 Colab 메모장에 있는 PyTorch용 예제로부터 생성된 파일에 MNIST 코드를 옮겨서 돌린 결과가 Colab 사이트나 Google Drive에서 생성된 파일에 MNIST 코드를 옮겨서 돌린 결과보다 빠른 속도를 보이고 있다. 원인은 아직 정확히 파악하지는 못했지만, 위의 실험 결과로부터 앞으로 TPU, GPU를 Colab에서 돌릴 때는 예제로 부터 파일을 생성해서 돌리는 것이 좋을 것 같다. 여러번 실험을 해본 결과, Colab 메모장에 있는 PyTorch용 예제로 부터 생성된 파일을 로컬이나 Google Drive 등에 저장해놓고 새로운 코드를 돌릴때 마다 해당 파일을 복사해서 코드만 바꾸어 Colab에서 돌리면 일반 Colab에서 생성된 파일 보다 더 빠른 속도를 보이는 것을 확인하였다.

추가적으로 관찰한 사항은 위의 MNIST 예제에서 TPU의 경우 Epoch이 지날수록 속도가 점점 빨라진다는 것이다. 아직 정확한 이유는 모르지만 https://stackoverflow.com/a/52599151를 참고하여 공부를 해봐야겠다.

---

참고자료

• Python 함수 vs. 메소드: https://yusulism.tistory.com/11