[인공지능 하드웨어] 8 - Deep Learning HW Accelerator(TPU)

# 들어가며

TPU에 대해서 자세히 알아본다. 하드웨어 구현 수준은 아니지만 TPU가 어떻게 행렬 데이터를 다루는지, 그리고 TPU에 얽힌 스토리를 검토해본다. 컴퓨터 구조 관련한 배경지식이 있으면 이해하기 편하므로 글을 읽다가 이해가 되지 않는다면 관련 블로그 포스트나 유튜브를 찾아보는 것이 좋다. 이전 글은 다음 링크의 블로그 포스트를 참고하면 된다.

[인공지능 하드웨어] 7 - Winograd Convolution

---

# Google 그리고 TPU

Google은 자사 인공지능 학습에 NVIDIA의 GPU와 솔루션만 수동적으로 사용하는 것이 마음에 안들었었다. 그래서 2016년에 독자개발한 TPU를 출시한다. TPU는 Tensor Processing Unit으로, 이름에서 알 수 있듯이 텐서 연산에 최적화 되어 있다. 텐서 연산 중에서도 딥러닝 프로세스 그 중에서도 행렬곱에 최적화 되어있다.

Google TPU 실물 (왼쪽)과 AlphaGo의 실물 (오른쪽)

 

# TPU Architecture

TPU 구조도

 

노란색 Unit은 연산을 수행하고, 파란색 Unit은 data buffer로, 연산에 필요한 값들을 미리 가져와서 저장해놓고 연산할 때 제공해주거나 혹은 연산의 결과값을 저장한다. 연두색은 외부 메모리이다. 주요 유닛 하나씩 알아 보도록 하자.

 

Matrix Multiply Unit(MMU)

• 주요 연산이 일어나는 곳으로, 주요 연산은 DNN 연산에서 input과 weight값들을 행렬곱하고 그 값들의 합을 구하는 연산이다.
• 256x256 MACs로 구성되며, MACs는 Multiplication/Accumulation Units를 말한다. MACs는 PE(Processing Element)라고도 하며, Systolic Array에서 연산이 일어나는 Cell이라고 생각하면 된다. Systolic Array에 대한 자세한 내용은 이 링크를 참고하자.
• 8bit의 Weight과 Input Data의 곱셈으로 결과 값이 16Bit가 되는 Mixed-Precision 연산이 일어난다. 이때 8bit를 사용한 이유는 TPU v1의 구조가 Inference 용으로 설계 되었기 때문에 굳이 High-Precision 연산이 필요가 없었기 때문이다!
• Systolic Array로 구성되며, 그 중에서도 Weight Stationary 방식이다. 따라서 연산 전에 Weight를 Pre-load 하는 시간이 발생하며, Output은 아래로 흘러 Accumulator로 입력된다.

Systolic Array (Weight Stationary) 동작 방식

Unified Buffer(UB)

• CPU Memory에서 값을 읽어서 저장했다가 MMU에게 input 값을 넘겨주는 역할을 한다. 이때, MMU에 넘겨주는 방식으로는 Systolic Array에 맞추기 위해 Systolic Data Setup에 먼저 제공하게 된다.
• 당연하게도 중간 과정에서 Activation을 거친 결과값을 저장했다가 다시 제공하는 역할도 한다.

Systolic Data Setup

• UB로부터 데이터를 읽어와서 MMU에 Systolic Array에 맞춰서 입력 데이터를 전달하는 타이밍을 조절한다. Weight FIFO도 이와 비슷한 역할을 수행한다.

Weight FIFO

• Weight 즉 가중치들은 TPU 내에 저장되는 것이 아니라 Weight Memory라는 별도의 8GiB DRAM 칩에 저장되어 있는다.
• Weight FIFO는 이 값들을 읽어와서 MMU에 제공하는데, Queue로 읽어오기에 FIFO라는 이름이 붙었다.

Accumulators

• 16bit의 곱 결과들을 모아준다.
• 이때, Accumulator의 역할은 MMU에서 일어나는 Accumulation과 다른 것으로, 만약 1024x1024 크기의 행렬이 input으로 들어왔다면 TPU는 한번에 256 크기의 데이터밖에 처리하지 못하므로 Tiling과 같은 최적화를 하게 된다. 이때, 이 값들을 모았다가 나중에 한번에 더해야 하는 상황이 발생하는데 이 과정에서 Accumulator가 필요한 것이다.

Activation Unit(AU)

• Activation은 말그대로 활성화 함수층이다. 활성화 함수를 코드 수준에서 따로 연산하는 것이 아니라, 하드웨어에서 구현하여 쓸모없는 에너지 소모를 극한까지 줄인 모습을 볼 수 있다.

## TPU Instructions

TPU 내부적으로 정의된 명령어들은 다음과 같다.

• Read_Host_Memory: data를 CPU host memory로부터 Unified Buffer로 읽어온다.
• Read_Weights: weight를 Weight Memory에서 Weight FIFO(matrix multiply unit의 input)로 읽어온다.
• MatrixMultiply/Convolve: Unifed Buffer로부터 accumulator까지 matrix multiply unit이 matrix multiply 혹은 convolution을 수행하게끔 한다. matrix operation은 x*256개의 input을 matrix unit에 넣어 x pipeline동안 256x256의 constant weight input값들과 곱해져서 연산된다.(x값은 최대 256까지 가능)
• Activate: accumulator의 연산 결과에 Sigmoid나 ReLU와 같은 인공신경망의 비선형함수를 적용한 값을 unified buffer에 저장하고 convolution의 pooling 연산을 수행하기도 한다.
• Write_Host_Memory: unified buffer의 data를 CPU host memory로 쓴다

# CPU, GPU, TPU

CPU와 GPU, TPU는 매우 근본적인 차이들이 존재한다. 하나씩 대조하기엔 어색하지만 특징들을 보면 쉽게 알 수 있다.

CPU and GPU

 

• 범용적인(일반적인) 목적으로 만들어졌으므로, 가능한 모든 연산을 할 수 있다.
• 이 범용성을 위해서 연산을 수행하기 위한 값을 register에서 읽어와서 연산하고, 다시 register에 결과값을 쓰는 과정을 수행한다. 위의 그림을 참고하자. 따라서 여러 register에 접근하여 많은 cost가 발생한다.

TPU

• TPU는 systolic array 구조에 의해 한번 연산할 때 ALU가 값을 다른 ALU에 넘겨줌으로써 한 register의 값을 재사용한다. 따라서 register에 접근하는 횟수가 적고, 에너지를 덜 소모하게 된다.
• Matrix Multiplication의 에너지 효율성 측면에만 최적화된 구조로, 다른 일반적인 연산을 하기에는 부적합한 구조이다.

 

아래는 각 구조의 동작에 대한 매우 쉬운 이해이다.

 

CPU

 

GPU

 

TPU - Weight Preload

 

TPU-Weight Stationary Systolic Array Computation

 

# TPU v1의 문제점과 후속 모델

TPU v1에는 여러가지 문제가 있었다.

1. Inference에 최적화된 구조로 Training은 할 수 없다.
2. MMU를 풀가동 했을 때 DDR3에서 Weight FIFO로 공급되는 메모리 대역폭은 최대 179GB/s가 요구된다. 하지만 DDR3 메모리의 최대 대역폭은 30GB/s로 현저히 낮다. 이게 무슨말이냐면, 물 마실때 목이 너무 말라서 겁나 빨리 먹어야 하는데 물병 입구가 빨대라서 너무 느리게 갈증이 해소된다는 것이다.

그래서 이후 모델부터는 DDR이 아닌 HBM을 도입했다. 채널이 8개로, 물병 입구가 넓어졌으며 메모리 대역폭과 에너지 효율이 증가했다.

 

## TPU v3

 

TPU v3에서는 더 많은 작은 행렬 유닛 인스턴스를 사용한다. MMU 크기를 절반으로 줄였지만 연산에 16bit bfloat 포맷을 사용해서 정밀도를 높였다.

 

Brain Floating Point가 뭔가 하면 기존의 FP-16 즉 Half Precision 16비트가 가지고 있는 수를 표현하는 범위가 좁다는 단점을 해결하기 위한 포맷이다. 그러나 아래 그림을 보면 알겠지만 그냥 같은 16비트 상에 정밀도를 포기하고 범위를 늘린거라서 나름의 장단점이 존재한다. Brain이라는 이름이 붙은건 요새 딥러닝에 많이 쓰여서 그런것 같다.

---

국민대학교 권은지 교수님의 인공지능 하드웨어 강의를 수강하며 정리한 내용입니다.

Slides Provided

임베디드 시스템 프로그래밍 - Eunhyeok Park

딥러닝 최적화 - Eunhyeok Park

References

https://cloud.google.com/tpu/docs/intro-to-tpu