[DL] PyTorch: TPU사용하기

https://github.com/pytorch/xla

GitHub - pytorch/xla: Enabling PyTorch on XLA Devices (e.g. Google TPU)

 다음 문서도 참고할 것.

https://github.com/pytorch/xla/blob/master/API_GUIDE.md

 

TPU로 PyTorch를 사용하기 위해서는 PyTorch / XLA 를 이용한다.

---

PyTorch / XLA

PyTorch로 구현된 model 및 PyTorch code 들을

• cpu와 gpu 를 넘어서서
• XLA (Accelerated Linear Albebra) API를 지원하는 H/W 에서도
• 동작하도록 해줌.

 

XLA를 지원하는 대표적인 device는 바로 구글의 TPU(Tensor Processing Unit)임.

 

PyTorch는 기본적으로 CPU와 GPU를 지원하지만, XLA를 통해 TPU와 같은 추가 하드웨어 가속기에서도 효율적으로 동작할 수 있음.

---

PyTorch XLA의 주요 역할:

1. 하드웨어 추상화:
   1. PyTorch XLA는 PyTorch와 TPU 사이의 중간 계층으로 작동
   2. PyTorch 코드가 TPU 상에서 실행될 수 있게 해 줌.
   3. 개발자는 TPU의 복잡한 세부 사항에 신경을 쓰지 않고, PyTorch 코드 작성에 집중하며 XLA가 나머지를 처리해주는 방식.
2. 성능 최적화:
   1. XLA는 복잡한 tensor 연산을 최적화하여 TPU와 같은 XLA를 지원하는 하드웨어에서 효과적으로 실행시켜줌.
   2. 이는 메모리 접근 최적화, 병렬 처리 최적화 등을 통해 이루어짐.
3. 확장성:
   1. TPU는 높은 병렬 처리 능력을 가지고 있고 특히 GPU 대비 보다 큰 메모리 지원을 함.
   2. XLA를 통해 이러한  TPU의 장점을 효과적으로 활용할 수 있음.
   3. 이는 큰 규모의 머신 러닝 모델을 훈련하거나, 대량의 데이터를 처리할 때 특히 유용.
4. 이식성:
   1. PyTorch XLA를 사용하여 XLA를 지원하는 시스템에서 기존의 PyTorch 코드를 최소한의 변경으로실행할 수 있음.

---

Colab의 TPU사용하기

https://gist.github.com/dsaint31x/8290dab796f1dc2382ef7245abaa32d7

dl_torch_tpu_xla.ipynb