TPU란 무엇인가?

 

개요

TPU(Tensor Processing Unit)은 Google에서 만든 ASICs(Application-Specific Integrated Circuit)이다. machine learning workload를 가속화하기 위해 사용된다. 보통 수천 개의 TPU는 함께 묶여 TPU Pod이라는 특별한 network를 구성한다. 예컨대, TPU v4에서 single v4 pod는 4096개의 TPU chips를 포함하고 있다. HBM 메모리를 활용하여 memory bandwidth가 큰 덕분에 batch size와 model이 큰 경우에도 효율적인 학습을 가능하게 한다. 

 

Figure 1: TPU pod(사진 출처: https://techcrunch.com/2019/05/07/googles-newest-cloud-tpu-pods-feature-over-1000-tpus/)

 

 

TPU는 2016년 Google I/O에서 처음 소개되었으며, Tensorflow를 위해 디자인되었다. 버전 업그레이드도 계속 이루어지고 있다. 첫 버전인 v1은 28nm 공정으로 제작되었으며, 8GiB DDR3를 사용하여 memory bandwidth가 34GB/s에 불과했지만, 매번 발표되는 TPU는 이전 성능의 두 배 이상씩 향상되어 2021년 5월에 발표한 v4는 7nm 공정에 1.2TB/s의 bandwidth를 가지고 있다. 

 

이에 따라 Google 외의 회사에서도 AI accelerator를 만들고 있다. 

 

CPU, GPU와 비교

CPU는 적은 수의 연산을 flexible하게 처리하기 위한 고성능의 ALU를 가지고 있다. 이는 연산을 빠르게 처리할 수 있다는 장점을 가지고 있지만, 큰 데이터를 참조해야 하는 neural network의 학습 과정에서는 비효율적이다. 이는 CPU가 register에 있는 데이터를 참조하여 계산하기 때문인데, register에 데이터를 불러오려면 memory에서 참조해야 하기 때문이다. 즉, memory의 bandwidth가 bottleneck이 된다. 이를 Von Neumann bottleneck이라고 한다. 

 

이에 반해 GPU는 낮은 성능의 ALU를 수천 개 포함하고 있어 간단한 연산을 한꺼번에 많이 처리할 수 있다. 원래 graphic processing을 하기 위해 만들어진 unit이니만큼, 간단한 matrix arithmetic에 최적화되어 있다. 이는 neural network에서 효율적인 사용을 가능하게 한다. 그러나 GPU의 경우에도 operand를 read/write하기 위해서 register 또는 shared memory에 access해야 하고 이 과정에서 bottleneck이 발생한다는 문제가 있다. 

 

TPU는 matrix operation을 더 효율적으로 하기 위해 on-chip high-bandwidth memory(HBM)을 탑재한다. HBM은 기존 메모리에 비해 대역폭이 훨씬 큰 새로운 구조의 메모리이다. TPU는 서로 연결된 수천 개의 multiply-accumulator가 physical matrix를 형성한다. 이를 systolic array architecture이라고 하는데, 각 processing element가 계산을 할 때마다 register에서 값을 참조해 오는 것이 아니라, 중간 ALU들은 이전 layer의 ALU로부터 계산된 값을 받아와 순차적으로 계산함으로써 register의 참조 수를 줄인다. 이는 data의 재사용이라고 볼 수 있다. 

Figure 2: 왼쪽이 CPU, GPU이고 오른쪽이 TPU이다.[2]

 

Cloud TPU v3에는 각 processor마다 2개의 systolic array로 구성된 $128\times 128$개의 ALU를 가지고 있다. 먼저 HBM에서 weight를 가져와 MXU(Matrix Multiplication Unit)에 저장한다. 

Figure 3: matrix weight의 저장[3]

 

그 후 data를 HBM에서 load한다. 각 multiplication이 수행된 이후 값은 next multiply-accumulator로 전달되어 바로 계산이 수행된다. 이 값이 accumulate되는 곳은 설정에 따라 다를 수 있고, 결과값은 mutliplication 값의 summation이 된다. 이때 이전 accumulator에서 바로 값을 받아 연산을 수행하므로 memory access는 최초 1번만 발생한다. 

 

Figure 4: matrix operation[3]

 

TPU의 활용

TPU는 8bit 연산*을 이용하여 GPU보다 연산 속도를 최적화한다. 이 때문에 low precision compuation에 적합하며, CNN 등에서 좋은 성능을 얻을 수 있다고 알려져 있다. 반면 GPU는 FC layer에 적합하고 CPU는 RNN에 이점이 있다고 알려져 있다.[4]

 

 

*8bit 연산의 디테일에 대해서는 추가해야 함

---

References

[1] Wikipedia contributors. (2023, November 6). Tensor processing unit. Wikipedia. https://en.wikipedia.org/wiki/Tensor_Processing_Unit

 

[2] Ogunfunmi, T., Ramachandran, R. P., Togneri, R., Zhao, Y., & Xia, X. (2019). A primer on deep learning architectures and applications in speech processing. Circuits, Systems, and Signal Processing, 38(8), 3406–3432. https://doi.org/10.1007/s00034-019-01157-3

 

[3] Introduction to Cloud TPU. (n.d.). Google Cloud. https://cloud.google.com/tpu/docs/intro-to-tpu

 

[4] Wang, Y. E., Wei, G. Y., & Brooks, D. (2019). Benchmarking TPU, GPU, and CPU platforms for deep learning. arXiv preprint arXiv:1907.10701.

 

 

 

 

Footnotes