Tiny-TPU: 이유와 방법

아무도 TPU가 어떻게 작동하는지 진정으로 이해하지 못합니다… 저희도 마찬가지입니다! 그래서 저희는 완전한 초보자의 관점에서 이것이 어떻게 작동하는지 추측해 보고 시도해 보고 싶어서 이 프로젝트를 시작하게 되었습니다!

이 슬라이드쇼는 저희의 장난감 가속기(Accelerator)에서 진행되는 순전파(Forward pass)와 역전파(Backward pass)를 묘사합니다.

왜 이 프로젝트를 시작했나요?

저희는 우리가 마음먹은 것은 무엇이든 할 수 있다는 것을 스스로에게 증명하기 위해 매우 도전적인 일을 하고 싶었습니다. 특히 TPU를 구축하기로 선택한 이유는 상당히 간단합니다:

• ML(Machine Learning) 워크로드를 위한 칩을 만드는 것이 멋져 보였습니다.
• 추론(Inference)과 학습(Training)을 모두 수행하는 ML 가속기에 대해 잘 문서화된 오픈 소스 저장소(Repo)가 없었습니다.

저희 중 누구도 하드웨어 설계에 대한 실제 전문적인 경험이 없으며, 이는 어떤 면에서는 TPU가 얼마나 어려울지 정확히 예측할 수 없었기에 더욱 매력적으로 다가왔습니다. 프로젝트의 초기 단계에서 작업하면서 저희는 엄격한 설계 철학을 세웠습니다: 항상 해킹 방식(Hacky way)을 시도하라. 이는 외부 자료를 찾아보기 전에 먼저 머릿속에 떠오르는 "멍청한" 아이디어들을 먼저 시도해 보는 것을 의미했습니다. 이 철학은 저희가 TPU를 역공학(Reverse engineering)하는 것이 아니라, **재발명(Re-inventing)**하도록 도와주었으며, 이를 통해 TPU에 사용되는 많은 핵심 메커니즘을 스스로 도출해 낼 수 있었습니다.

또한 저희는 최근 약간의 어려움에 직면할 때마다 AI 도구에 의존하려는 본능이 생겼다고 느꼈기 때문에, AI가 대신 코드를 작성해 주는 것에 의존하지 않고 코딩하는 연습으로 이 프로젝트를 다루고 싶었습니다. 저희는 어려운 문제를 사고할 때 향후 어떤 노력에서도 가져가서 사용할 수 있는 특정한 **사고 방식(Style of thinking)**을 기르고 싶었습니다. 위대한 것을 만드는 데 허락은 필요하지 않다는 Sholto Douglas의 YouTube 쇼츠 메시지에서 영감을 받아, 저희는 계속해서 결과물을 공개(Shipping)하고 공개적으로 학습해 나갔습니다.[1]

이 프로젝트 전반에 걸쳐 저희는 딥러닝 (Deep Learning)의 기초, 하드웨어 설계 (Hardware Design), 그리고 알고리즘 생성에 대해 가능한 한 많은 것을 배우고자 노력했습니다. 저희는 이러한 분야를 배우는 가장 좋은 방법은 모든 것을 직접 그려보고 그것을 첫 번째 본능으로 만드는 것이라는 점을 깨달았습니다. 이 글을 읽으시면서, 저희의 설명이 이러한 철학에서 어떻게 영감을 받았는지 확인하실 수 있을 것입니다.

본론으로 들어가기에 앞서, 이 글이 무엇을 다루고 무엇을 다루지 않는지 명확히 하고 싶습니다. 이것은 TPU의 1대1 복제본이 아니라는 점에 유의하십시오. 이것은 저희가 직접 TPU를 재발명하려는 시도입니다.

TPU란 무엇인가?

TPU는 Google이 머신러닝 (ML) 모델의 추론 (Inference, 사용) 및 학습 (Training)을 더 빠르고 효율적으로 만들기 위해 설계한 주문형 집적 회로 (ASIC, Application Specific Integrated Circuit) — 기본적으로 맞춤형 칩 — 입니다. GPU가 프레임을 렌더링함과 동시에 ML 워크로드를 실행할 수 있는 반면, TPU는 수학 연산만을 수행할 수 있어 설계된 목적에 더 특화될 수 있습니다. 당연하게도, 여러 작업을 마스터하려고 노력하는 것보다 단일 작업을 마스터하는 것이 훨씬 쉽고 더 나은 결과를 낼 것이며, TPU는 이 철학을 강력하게 채택하고 있습니다.

하드웨어 설계에 대한 간단한 입문:

하드웨어에서 우리가 다루는 시간 단위는 클록 사이클 (Clock Cycle)이라고 불립니다. 이는 개발자로서 우리의 요구 사항을 충족하기 위해 설정할 수 있는 임의의 시간 단위입니다. 일반적으로 단일 클록 사이클은 수백 피코초 (ps)에서 1 나노초 (ns) 사이의 범위를 가질 수 있으며, 우리가 실행하는 모든 연산은 클록 사이클 사이 (BETWEEN)에 실행됩니다.

하드웨어를 설명하기 위해 사용하는 언어는 Verilog라고 불립니다. 이는 특정 하드웨어 모듈(소프트웨어의 함수와 유사함)의 동작을 설명할 수 있게 해주는 하드웨어 기술 언어 (HDL, Hardware Description Language)입니다. 하지만 프로그램처럼 실행되는 대신, 우리가 원하는 어떤 칩이라도 구축할 수 있는 디지털 로직을 만들기 위해 결합될 수 있는 불리언 논리 게이트 (Boolean Logic Gates: AND, OR, NOT 등)로 합성됩니다. 다음은 Verilog를 사용한 덧셈의 간단한 예시입니다:

module add (
input wire clk,
// 모듈을 리셋하기 위한 리셋 신호
...

위의 예시에서, 다음 클록 사이클(clock cycle)의 신호 b 값은 신호 a의 현재 값으로 설정됩니다. 소프트웨어 설계에서 볼 수 있는 즉각적인 업데이트와 달리, 대부분의 경우 신호(변수)는 순차적인 클록 사이클에 따라 업데이트된다는 것을 알 수 있습니다.

구체적으로, TPU는 행렬 곱셈(matrix multiplications)을 수행하는 데 매우 효율적입니다. 행렬 곱셈은 트랜스포머(transformers) 연산의 80-90%(매우 큰 모델의 경우 최대 95%)를 차지하며, CNN에서는 70-80%를 차지합니다. 각 행렬 곱셈은 MLP의 단일 레이어(layer)에 대한 계산을 나타내며, 딥러닝(deep learning)에는 이러한 레이어가 많이 존재하므로 모델이 커질수록 TPU의 효율성은 더욱 높아집니다.

우리의 "장난감" TPU를 어떻게 개발했나요?

이 프로젝트를 시작했을 때, 우리가 아는 것이라고는 $y = mx + b$라는 방정식이 신경망(neural networks)의 기초적인 빌딩 블록이라는 사실뿐이었습니다. 하지만 TPU의 다른 모듈들을 구축하기 위해서는 신경망 뒤에 숨겨진 수학을 완전히 이해(UNDERSTAND)해야 했습니다. 그래서 코드를 작성하기 전에, 우리 각자는 간단한 2 -> 2 -> 1 다층 퍼셉트론 (MLP, multi-layer perceptron)의 수학적 계산을 직접 수행했습니다.

왜 XOR인가요?

우리가 이 특정 네트워크를 선택한 이유는 XOR 문제(신경망의 "Hello World")에 대한 추론(inference)과 학습(training)을 목표로 했기 때문입니다. XOR 문제는 신경망이 해결할 수 있는 가장 간단한 문제 중 하나입니다. 다른 모든 게이트(AND, OR 등)는 입력값 중 어떤 것이 0에 해당하고 어떤 것이 1에 해당하는지를 분리하기 위해 단 하나의 선형 직선(하나의 뉴런)만 사용하여 입력을 통한 출력을 예측할 수 있습니다. 하지만 모든 XOR를 분류하기 위해서는 MLP가 필요한데, 이는 선형 방정식만으로는 달성할 수 없는 곡선 형태의 결정 경계(decision boundaries)가 필요하기 때문입니다. 기하학적이고 제1원리적인 접근 방식을 원하신다면, 무료 도서인 Understanding Deep Learning이 매우 훌륭합니다.

배치(Batching) 및 차원(dimensions)

이제 우리가 연속적인 추론 (Inference)을 수행하고자 한다고 가정해 봅시다 (예: 초당 여러 번의 예측을 수행하는 자율 주행 자동차). 이는 우리가 한 번에 여러 개의 데이터 조각을 보낸다는 것을 의미합니다. 데이터는 본질적으로 다차원적이며 많은 특징 (Features)을 가지고 있기 때문에, 우리는 매우 큰 차원을 가진 행렬 (Matrices)을 다루게 될 것입니다. 하지만 XOR 문제는 특징이 두 개(0 또는 1)뿐이고 가능한 입력 데이터 조각이 4개(0과 1의 네 가지 가능한 이진 조합)뿐이므로 차원을 단순화해 줍니다. 이는 4x2 행렬을 생성하며, 여기서 4는 행의 수 (배치 크기, Batch size)이고 2는 열의 수 (특징 크기, Feature size)입니다.

XOR 입력 행렬과 타겟 출력:

각 행은 가능한 네 가지 XOR 입력 중 하나를 나타내며, 출력 벡터는 기대되는 XOR 결과를 보여줍니다.

여기서 시스톨릭 어레이 (Systolic array) 예제를 위해 우리가 적용한 또 다른 단순화는 TPUv1에서 사용된 256x256 어레이 대신 2x2 어레이를 사용하는 것입니다. 하지만 수학적 원리는 여전히 충실하게 유지되므로, 실제로 내용을 낮춘 것이 아니라 규모를 축소 (Scaled down)한 것입니다.

방정식의 첫 번째 단계는 m과 x를 곱하는 것이며, 행렬 형태로는 $mx$가 됩니다.

더 공식적으로는:

여기서 $x$는 입력 행렬 (Input matrix), $w$는 가중치 행렬 (Weight matrix), $b$는 편향 벡터 (Bias vector)입니다.

하드웨어에서 행렬 곱셈을 어떻게 수행할 수 있을까요? 바로 시스톨릭 어레이 (Systolic array)라고 불리는 유닛을 사용할 수 있습니다!

시스톨릭 어레이 (Systolic array) 및 PE

TPU의 핵심은 시스톨릭 어레이 (Systolic array)라고 불리는 유닛입니다.[2] 이는 처리 요소 (Processing Elements, PE)라고 불리는 개별 빌딩 블록들로 구성되며, 격자 구조 (Grid-like structure)로 서로 연결되어 있습니다. 각 PE는 곱셈-누산 (Multiply-accumulate) 연산을 수행합니다. 즉, 들어오는 입력 $X$를 고정된 가중치 $W$와 곱하고[3], 이를 들어오는 누적 합 (Accumulated sum)에 더하는 과정을 모두 동일한 클록 사이클 (Clock cycle) 내에 수행합니다.

always_ff @(posedge clk or posedge rst) begin
if (rst) begin
input_out <= 0;
...

시스톨릭 행렬 곱셈 (Systolic matrix multiplication)

이 PE(Processing Element)들이 서로 연결되면, 시스톨릭(systolic) 방식으로 행렬 곱셈 (matrix multiplication)을 수행하는 데 사용될 수 있습니다. 즉, 매 클록 사이클(clock cycle)마다 출력 행렬의 여러 요소들을 계산할 수 있음을 의미합니다. 입력값들은 시스톨릭 어레이 (systolic array)의 왼쪽에서 들어와 매 클록 사이클마다 오른쪽의 인접한 PE로 이동합니다. 누적 합 (accumulated sums)은 첫 번째 PE 행의 곱셈 결과값에서 시작하여 아래로 이동하며, 각 후속 PE의 곱셈 결과에 더해지며, 마지막 PE 행에 도달하면 출력 행렬의 한 요소가 됩니다.

이 단일 유닛 덕분에 (그리고 행렬 곱셈이 모델에서 수행되는 연산의 대부분을 차지한다는 사실 덕분에), TPU는 어떤 모델이든 매우 쉽게 추론 (inference)하고 학습 (train)할 수 있습니다.

예제 실행 (Worked example)

이제 우리의 XOR 문제 예제를 살펴보겠습니다.

우리의 시스톨릭 어레이는 두 가지 입력, 즉 입력 행렬 (input matrix)과 가중치 행렬 (weight matrix)을 받습니다. 우리의 XOR 네트워크를 위해, 다음과 같은 가중치 (weights)와 편향 (biases)으로 초기화합니다.

레이어 1 파라미터 (Layer 1 parameters):

레이어 2 파라미터 (Layer 2 parameters):

입력 및 가중치 스케줄링 (Input and weight scheduling)

시스톨릭 어레이 내에 입력 배치 (input batch)를 입력하기 위해서는 다음 작업이 필요합니다:

• X 행렬을 90도 회전
• 입력을 지그재그로 배치 (STAGGER, 각 행을 1 클록 사이클만큼 지연)
  [4]

가중치 행렬을 입력하기 위해서는 다음 작업이 필요합니다:

• 가중치 행렬을 지그재그로 배치 (입력과 유사하게)
• 전치 (Transpose)!

회전과 지그재그 배치는 수학적인 의미를 갖지는 않으며, 단순히 시스톨릭 어레이가 작동하도록 하기 위해 필요하다는 점에 유의하십시오. 전치 역시 수학적인 관리를 위한 것일 뿐입니다. 신경망 도식 내에서 가중치 포인터 (weight pointers)를 설정한 방식 때문에 행렬 연산이 제대로 작동하도록 하기 위해 필요합니다.

지그재그 배치 및 FIFO (Staggering and FIFOs)

지그재그 배치 (staggering)를 수행하기 위해, 우리는 시스톨릭 어레이의 위쪽과 왼쪽에 각각 위치할 가중치 및 입력용으로 거의 동일하게 설계된 누산기 (accumulators)를 설계했습니다.

활성화 값 (activations)이 시스톨릭 어레이 (systolic array)로 하나씩 공급되기 때문에, 우리는 선입선출 큐 (FIFO)가 최적의 데이터 저장 옵션이 될 것이라고 생각했습니다. 하지만 전통적인 FIFO와 우리가 만든 누산기 (accumulators) 사이에는 약간의 차이가 있었습니다. 우리의 누산기에는 2개의 입력 포트가 있었습니다. 하나는 FIFO에 가중치 (weights)를 수동으로 쓰기 위한 용도였고, 다른 하나는 활성화 모듈 (activation modules)로부터 이전 레이어의 출력값을 입력 FIFO로 다시 (BACK) 쓰기 위한 용도였습니다 (이전 레이어의 출력은 현재 레이어의 입력이 됩니다).

우리는 또한 모든 레이어에 대해 유사한 방식으로 가중치를 로드해야 했으므로, 두 번째 포트가 없는 가중치용 FIFO 로직을 복제했습니다.

시스톨릭 어레이 행렬 곱셈 (Systolic array matrix multiplication)

편향 (Bias) 및 활성화 (Activation)

방정식의 다음 단계는 편향 (bias)을 더하는 것입니다. 이를 하드웨어에서 수행하려면 시스톨릭 어레이의 각 열 (column) 아래에 편향 모듈을 만들어야 합니다. 합계 값이 시스톨릭 어레이 내의 마지막 행을 통해 빠져나갈 때, 이를 즉시 편향 모듈로 스트리밍하여 사전 활성화 값 (pre-activations)을 계산할 수 있습니다. 우리는 이 값들을 변수 Z로 표기하겠습니다.

편향 벡터는 행렬의 모든 행에 걸쳐 브로드캐스트 (broadcast)됩니다. 즉, 각 행에 더해집니다.

이제 우리의 방정식은 고등학교 시절에 배웠던 것과 매우 흡사해 보이기 시작합니다. 다만 각 열이 스트리밍되어 나오는 다차원 형태이며, 시스톨릭 어레이에서 나오는 각 열은 자체적인 특징 (feature)을 나타냅니다!

다음으로 활성화 (activation)를 적용해야 하며, 이를 위해 우리는 Leaky ReLU를 선택했습니다.[5] 이 또한 편향과 마찬가지로 요소별 연산 (element-wise operation)이므로, 모든 편향 모듈 아래에 (결과적으로 시스톨릭 어레이의 모든 열 아래에) 활성화 모듈이 필요하며, 편향 모듈의 출력을 활성화 모듈로 즉시 스트리밍할 수 있습니다. 우리는 이 활성화 후의 값들을 H로 표기하겠습니다.

Leaky ReLU 함수는 다음과 같이 요소별로 적용됩니다:

여기서 $\alpha$는 우리의 누설 계수 (leak factor)입니다. 행렬의 경우, 이는 각 요소에 독립적으로 적용됩니다.

우리의 XOR 예시를 위해, Layer 1이 데이터를 어떻게 처리하는지 살펴보겠습니다. 먼저, 시스톨릭 어레이 (systolic array)가 다음을 계산합니다:

그 다음 편향 (bias)이 더해집니다:

마지막으로, LeakyReLU가 요소별 (element-wise)로 적용됩니다:

음수 값에는 0.5가 곱해지고, 양수 값은 변하지 않고 그대로 통과합니다.

편향 (bias) 및 LeakyReLU를 포함한 시스톨릭 어레이 (systolic array)

파이프라이닝 (Pipelining)

이제 여러분은 다음과 같이 질문할 수도 있습니다. 왜 편향 (bias) 항과 활성화 (activation) 항을 하나의 클록 사이클 (clock cycle)에 병합하지 않을까요? 그것은 바로 파이프라이닝 (pipelining)이라고 불리는 것 때문입니다! 파이프라이닝 (pipelining)을 사용하면 TPU의 서로 다른 단계 (stages)에 걸쳐 여러 연산을 동시에 실행할 수 있습니다. 즉, 하나의 연산이 완전히 끝날 때까지 기다렸다가 다음 연산을 시작하는 대신, 작업을 중첩될 수 있는 단계들로 나눕니다. 이를 조립 라인(assembly line)처럼 생각해보세요. 한 작업자 (활성화 모듈)가 부품을 처리하는 동안, 이전 작업자 (편향 모듈)는 이미 다음 부품을 작업하고 있습니다. 이를 통해 모든 모듈이 이전 단계가 완료되기를 기다리며 유휴 상태 (idle)로 있는 대신, 계속해서 바쁘게 작동할 수 있습니다. 이는 또한 우리가 TPU를 실행할 수 있는 속도에도 영향을 미칩니다. 만약 하나의 모듈이 단일 사이클 내에 너무 많은 연산을 욱여넣으려 한다면, 다른 모듈들은 그 단일 모듈의 속도만큼만 실행될 수 있기 때문에 우리의 클록 속도 (clock speed)는 해당 모듈에 의해 병목 현상 (bottleneck)이 발생하게 됩니다. 따라서 연산을 가능한 한 개별 클록 사이클 (clock cycle)로 나누는 것이 효율적이며 권장되는 방식 (best practice)입니다.