구글 TPU 들여보기

올해 4월초 구글에서 개발한 TPU(Tensor Processing Unit)와 관련된 ISCA 논문이 공개됐습니다. TPU는 딥러닝의 inferencing에 특화된 가속기로 2013년에 개발이 시작되어 2015년 부터 구글의 데이터센터에서 실제로 사용되기 시작했지만 성능과 관련된 데이터가 쌓이고 자세한 내용이 논문을 통해 공개된 시점은 2년 후인 2017년이 된 것이지요.

구글이 TPU 개발을 시작한 동기는 구글 포토, 음성 검색, 이미지 검색등의 서비스를 제공하려다보니 딥러닝 연산과 관련된 컴퓨팅이 증가하면서 CPU 만으로는 충분한 성능을 효율적으로 제공할 수 없다는 판단이 들었기 때문입니다. GPU도 생각할 수 있는 대안이긴 했지만 서비스 제공에 중요한 딥러닝의 inferencing 연산시에는 GPU도 그다지 효율적이지 않아 결국은 전용 가속기(TPU) 설계를 결정했다고 합니다.

딥러닝에는 대량의 데이터를 사용하여 각 NN(Neural Network) node 사이의 weight를 학습시키는 training 단계와 학습된 NN을 이용하여 입력 데이터에 대한 추론을 수행하는 inferencing 단계가 있습니다. Training 단계에서는 32-bit 부동수수점 연산이 필요하지만 inferencing 단계에서는 8-bit 부동소수점 또는 정수 연산만으로도 충분하기 때문에 TPU가 개발될 당시 32-bit/16-bit 부동소수점 연산에 촛점을 맞춰 설계되있던 GPU는 성능이 낭비되는 단점이 있었습니다. nVidia에서도 training과 inferencing시에 필요한 연산 특성이 다르다는 점에 주목해서 inferencing에 특화된 Tesla P40/P4 GPU를 발표하긴 했습니다만 이는 TPU 개발이 완료된 이후인 2016년의 이야기입니다. 참고로 P40/P4 GPU는 32-bit 부동소수점 대신 8-bit 정수 연산을 수행할 경우 한 clock에 4개의 정수연산을 동시에 수행하여 inferencing 성능을 높인 특징을 가지고 있습니다.

Training 유리한 P100 GPU와 Inferencing에 유리한 P40/P4 GPU (출처 – http://www.anandtech.com/show/10675/nvidia-announces-tesla-p40-tesla-p4)

TPU 전체 구조

구글이 발표한 논문에 의하면 TPU를 채용함으로서 딥러닝 inferencing 연산의 경우 2015년 당시의 GPU/CPU 대비 15~30배의 성능향상을 얻고 동일한 계산량에 대한 에너지 효율은 30~80배까지 향상시켰다고 하는데, 먼저 TPU의 전체 구조를 살펴보고 이를 inferencing 연산과 대응시켜보면서 어떻게 성능과 효율 두가지 모두를 향상시켰는지 살펴보도록 하겠습니다.

구글 TPU의 구조 (출처 – “In-Datacenter Performance Analysis of a Tensor Processing Unit,” ISCA, 2017)

위의 그림은 논문에 수록된 구글 TPU의 구조인데 색깔로 기능이 구분되어 있습니다. 노란색으로 표시된 부분은 실제 계산이 일어나는 로직이며 하늘색으로 표시된 부분은 계산에 필요한 입/출력값 및 Weight들을 DDR3 메모리와 TPU사이에서 옮기거나 중간결과를 임시로 저장하는 부분입니다. 빨간색은 명령어 수행을 위한 제어로직이고 초록색 부분은 TPU 외부와 연결되는 DDR3 메모리와 PCIe 인터페이스를 나타내고 있습니다.

TPU에서 가장 핵심적인 부분은 위의 그림에서 오른쪽 상당부분에 노란색 격자로 표시된 Matrix Multiply Unit인데, 65,536(=256×256)개의 곱셈기로 이루어져 있으며 각각의 곱셈기는 8-bit 정수 2개를 곱해 16-bit 결과를 만들어 냅니다. Matrix Multiply Unit 아래쪽에는 하나의 entry당 256개의 32-bit 레지스터를 가진 4096-entry Accumulator가 있는데 각 32-bit 레지스터는 세로방향을 따라 위치한 256개의 곱셈기가 계산한 16-bit 결과 값 256개를 누적하는데 사용됩니다. Accumulator의 4096개의 entry 중 어디에 계산 값이 누적될지는 TPU의 MatrixMultiply/Convolve 명령어에 포함된 accumulator address 필드에 의해서 결정됩니다. Accumulator의 결과는 Activation 유닛과 Normalize/Pool 유닛을 거쳐 Unified Buffer에 저장이 되는데, 계산된 값이 중간 결과일 경우 다시 Matrix Multiply Unit의 입력으로 보내지며 최종 결과일 경우 Host Interface와 PCIe 인터페이스를 거쳐 CPU의 메인메모리에 저장되게 됩니다.

Inferencing 연산

앞에서 살펴본 TPU의 구조가 inferencing에 유리한 이유를 이해하기 위해 간단히 inferencing에 필요한 연산을 살펴보겠습니다. 아래 그림은 layer당 3개의 노드를 가지는 NN의 inferencing 연산과정 일부를 보여주고 있는데 현재(N) layer의 각 노드값들을 계산하기 위해서는 이전(N-1) layer의 노드별 출력값과 두 layer 사이의 모든 노드쌍에 대한 weight 값이 필요합니다. 노드별 누적 값은 이전 layer의 노드별 출력 값에 연관된 weight를 곱해서 더하는 방식으로 계산되는데, 예를 들어 현재 layer의 맨 상단 노드의 값은 O_N-1,0xW_0,0 + O_N-1,1xW_1,0 + O_N-1,2xW_2,0 와 같은 식으로 계산이 됩니다. 각 노드별 계산식을 정리해보면 모든 노드에 대한 계산식은 아래그림의 하단과 같은 행렬 연산으로 표현할 수 있게 됩니다. 행렬곱의 결과로 계산된 노드별 누적 값은 ReLU(Rectified Linear Unit), Sigmoid, Tanh 함수 등의 입력으로 사용되어 최종적으로 현재 노드의 출력값이 계산되는데 이 과정을 Activation이라고 합니다. 아래 그림에서는 ReLU(y=max{0,x})함수가 activation을 위해 사용된 예제를 보여주고 있습니다.

Neural Network의 inferencing 연산

간단하게나마  inferencing 연산을 이해했다면 TPU는 inferencing 연산에 딱 필요한 회로만을 대량으로 가지고 있는 프로세서라는 특징을 발견할 수 있습니다. 행렬곱을 빠르게 수행하기 위한 65,536(=256×256)개의 곱셈기를 가지고 있고 이 곱셈결과들을 더해 출력 행렬의 원소를 계산하기위한 256개의 누산기(accumulator)도 갖추고있으며, 256개의 누적값에 대한 activation 또한 병렬로 계산할 수 있습니다. 이에 더해 CNN(Convolutional NN) 연산을 위한 전용 pooling 모듈까지도 갖추고 있습니다.

계산해야 할 행렬의 크기가 256×256 보다 클 경우에는 중간 결과를 Unified Buffer에 저장할 수 있는데 온 칩 메모리로는 상당한 용량인 24MB에 달하는 크기라 latency가 긴 DDR3 메모리를 접근하지 않고도 커다란 NN에 대한 inferencing 연산이 가능합니다. Unified Buffer의 경우 Matrix Multiply Unit에 데이터를 제공하기 용이하도록 256개의 8-bit 저장공간을 하나의 단위로 96K entry를 가지고 있습니다. (256 x 8-bit x 96K = 24MB)

Matrix Multiply Unit의 동작

TPU의 전체 구조가 inferencing에 유리한 이유를 살펴보았으니 이번엔 TPU 성능의 핵심인 Matrix Multiply Unit을 조금 자세히 살펴보도록 하겠습니다. TPU는 CISC 프로세서이기 때문에 Matrix Multiplication 또는 Convolution을 하나의 명령어로 실행할 수 있는데, Matrix Multiply Unit은 65,536(=256×256)개의 곱셈기를 가지고 있어 최대 256×256 크기의 weight matrix와 Bx256 크기의 입력 행렬을 곱하여 Bx256 출력 행렬을 한번에 계산할 수 있습니다(B의 최대값은 256). 여기서 한 번에 계산한다는 말은 한 싸이클에 행렬곱을 계산할 수 있다는 뜻이 아니라 atomic 하게 계산할 수 있다는 뜻인데 그 이유는 곧 설명할 Matrix Multiply Unit의 동작 방식과 관련이 있습니다.

65,536개의 곱셈기 전체에 대해서 동시에 계산할 operand를 공급하는 것은 메모리 대역폭에 부담이 되므로 Matrix Multiply Unit은 systolic array로 구성되어 동작되는데 아래 그림은 3×3 행렬 곱셈이 systolic array 상에서 어떤식으로 수행되는지를 보여주고 있습니다. 곱해질 두 행렬값이 systolic array의 윗쪽과 왼쪽에서 행과 열 별로 시간차를 두고 아래쪽과 오른쪽으로 shift 되면서 계산이 이루어지는데 보다 보다 자세한 내용은 링크를 참조하여 주시기 바랍니다.

Matrix Multiplication Unit 왼쪽의 Systolic Data Setup 모듈은 Unified Buffer로 부터 입력값 또는 중간 결과값을 불러와 위의 그림에서 설명된 순서로 입력 데이터를 전달하는 timing을 조절하는 역할을 하며 Weight FIFO 또한 비슷한 역할을 합니다. Systolic array의 장점은 곱셈기의 구조는 그대로 두더라도 입력데이터와 weight를 전달하는 순서를 바꾸면 matrix multiplication과 convolution 모두 수행할 수 있다는 점입니다. Systolic array를 사용할 경우 511(256×2 -1) 싸이클 동안 최대 65,536번의 곱셈과 65,536번의 덧셈을 할 수 있어 inferencing을 고속으로 수행할 수 있는 것이지요. 참고로 TPU의 동작속도는 700MHz라고 합니다.

TPU의 동작

TPU는 15개월이라는 짧은 개발기간 때문에 동작을 간단하게 하기위하여 PCIe 슬롯에 꽂히는 co-processor로 개발되었다고 합니다. TPU는 GPU와는 달리 스스로 수행할 명령을 호스트 CPU의 메인메모리로부터 fetch하지 못하고 호스트가 TPU의 메모리에 수행할 프로그램을 옮겨줘야 한다는 뜻이지요. TPU에서 동작할 프로그램이 PCIe 인터페이스를 통해 전달되는 관계로 프로그램 사이즈를 작게 하기위해 TPU는 CISC 프로세서로 설계되었고 주요 명령어로는 아래의 5개 명령어가 있다고 합니다.

1. Read_Host_Memory – 호스트의 메인메모리로부터 Unified Buffer로 입력 데이터를 이동
2. Read_Weights – TPU의 DDR3 메모리로 부터 Weight FIFO로 weight 데이터를 이동
3. MatrixMultiply/Convolve – Matrix Multiplication을 수행하여 Accumulator에 결과를 저장
4. Activate – Accumulator에 저장된 결과에 대해 Activation 함수를 적용한 후 Unified Buffer에 저장
5. Write_Host_Memory – Unified Buffer의 데이터를 호스트의 메인메모리로 저장

위의 5개 명령어를 순서대로 따라가면 TPU가 어떻게 동작하는지 그림이 그려지는데 첫번째와 두번째 명령어는 Matrix Multiply Unit이 계산할 입력과 Weight를 준비하는 과정입니다. Inferencing 연산시에는 weight 값들은 변하지않는 상수이므로 미리 DDR3 메모리에 로드되어 있다고 가정정했기 때문에 Host의 메인메모리로 부터 weight data를 옮기는 명령은 주요 명령어에 포함되지 않았습니다. 마찬가지 이유로 첫번째 그림에서 TPU 구조를 살펴보면 Unified Buffer와 DDR3 메모리 사이에는 직접 연결되는 data path도 보이지 않습니다.

다시 명령어의 흐름을 따라가면, 세번째 명령에서 행렬곱의 결과가 계산되어 4096-entry의 accumulator 중 1~256개의 entry에 저장이되고 네번째 명령은 다시 4096개의 accumulator entry 중에 N(?)개를 골라 activation 함수를 적용하여 Unified Buffer에 결과를 저장합니다. 계산해야할 행렬의 크기가 256×256보다 크거나 여러개의 NN layer에 대해 결과를 계산해야 한다면 세번째와 네번째 명령이 여러번 반복하여 실행될 것입니다. Inferencing 연산이 완료되면 마지막으로 다섯번째 명령을 수행하여 Unified Buffer에 저장된 마지막 노드의 출력 결과를 CPU의 메인메모리에 저장하게 됩니다.

구글 TPU의 보다 사진 (출처 – https://venturebeat.com/2017/04/05/google-opens-up-about-its-tpu-chips-for-ai/)

글을 마무리 하며

이번 글에서는 TPU의 구조가 어떻게 inferencing 연산에 도움이 되는지에 대해서만 집중해서 살펴보았는데, 공개된 논문에는 훨씬 많은 내용이 담겨져 있습니다. 특히 구글 데이터 센터에서 MLP(Multi Layer Perceptron), CNN, LSTM(Long Short Term Memory) 등 여러 종류의 딥러닝 NN이 어떤 비율로 사용되는지에 대한 데이터가 흥미로운데 자세한 내용이 궁금하다면 구글에서 발표한 논문을 읽어보시기 바랍니다.

개인적으로 흥미로웠던 내용은 처음에는 throughput 만을 고려해서 TPU를 설계했는데, 실제 데이터센터에 적용해보니 response time (<7ms) 또한 상당히 중요하다는 피드백을 받아 작업을 작은 단위로 쪼개서 처리하다보니 GPU 대비 성능향상이 두드러졌다는 부분이었습니다. TPU의 경우 범용성을 포기한 대가로 GPU 대비 단순한 구조를 가지는 덕분으로 본격적인 연산을 시작하기전에 준비할 작업의 양이 적기 때문으로 생각됩니다.

이 글을 마무리 하는 오늘 공교롭게도 후속버전인 TPU2가 공개되었는데 이 글에서 설명한 TPU가 지원하지않는 딥러닝의 training단계를 지원하고 multi-chip operation을 고려하여 설계되었다는 점이 중요한 향상점으로 보입니다. 얼마전에 출시된 nVidia의 새 GPU 아키텍처인 Volta에서도 4×4 행렬곱을 한 싸이클에 처리할 수 있는 Tensor Core를 추가하여 AI 성능의 엄청난 향상을 약속했는데, 구글의 TPU가 계속 살아남을지 다시 전문 GPU 회사의 솔루션을 도입하게 될지 지켜보는 것도 흥미로울 것 같습니다. TPU2에 대해서는 보다 자세한 자료가 공개되는대로 다시 한 번 다루도록 하겠습니다.