서론

이 기사는 TPUv1, TPUv2, TPUv3를 「추론과 학습에 필요한 하드웨어 요구사항의 차이」를 통해 정리하는 기술 메모입니다.

주로 다음의 논문 및 1차 자료를 참조하고 있습니다.

• 논문: In-Datacenter Performance Analysis of a Tensor Processing Unit
• 논문: Benchmarking TPU, GPU, and CPU Platforms for Deep Learning
• 논문: A Study of BFLOAT16 for Deep Learning Training
• 공식 자료: TPU architecture
• 공식 자료: TPU v2
• 공식 자료: TPU v3
• 공식 자료: Improve your model's performance with bfloat16

상세한 배경 설명, 도해, 비교표는 개인 블로그에 정리해 두었습니다.

👉 완전판은 이쪽에서 확인하세요: https://kasblo.com/ai-paper/tpu-v1-v2-v3-training-inference/

3줄 요약

• TPUv1은 Jouppi 등의 ISCA 2017 논문에서 주로 데이터 센터 추론을 가속화하는 ASIC (특정 용도용 집적 회로)로서 분석되었습니다.
• TPUv2에서는 bfloat16 (float32에 가까운 값 범위를 가진 16-bit 부동 소수점 형식) 연산, float32 누적, HBM (High Bandwidth Memory, 고대역폭 메모리), 다중 칩 구성에 의해 학습을 주요 용도로 다루기 쉬워졌습니다.
• TPUv3는 TPUv2에서 확립된 학습 기반을 연산 성능, HBM 용량, 대역폭, Pod 규모 측면에서 강화한 세대입니다.

논문 정보

TPUv1은 무엇을 위해 설계되었는가

Jouppi 등의 논문에서 TPUv1은 Google의 데이터 센터에서 뉴럴 네트워크 (Neural Network) 추론을 가속화하기 위해 설계된 ASIC으로 설명됩니다.

TPUv1의 핵심은 65,536개의 8-bit MAC (Multiply-Accumulate, 곱셈-누산)으로 구성된 Matrix Multiply Unit (행렬 곱셈 유닛)입니다.

논문에서는 28 MiB의 software-managed on-chip memory (소프트웨어 관리형 온칩 메모리)를 가지며, 추론 워크로드에서 CPU/GPU보다 높은 성능과 전력 효율을 보였다고 보고하고 있습니다.

여기서 중요한 점은, 행렬 곱셈기가 빠르다는 것과 학습 가속기로서 실용적이라는 것이 동일한 의미는 아니라는 점입니다.

추론에서는 학습된 가중치 (Weight)를 사용한 순전파 (Forward Propagation)가 중심입니다.

반면, 학습에서는 순전파, 역전파 (Backpropagation), 경사 계산 (Gradient Calculation), 가중치 업데이트, optimizer state (최적화 알고리즘이 유지하는 보조 상태), 다중 칩 간의 동기화가 필요합니다.

학습에서 필요해지는 추가 요구사항

선형층 (Linear Layer)의 순전파는 다음과 같은 형태입니다.

가중치에 대한 경사 (Gradient)는 개념적으로 다음과 같이 쓸 수 있습니다.

SGD (Stochastic Gradient Descent, 확률적 경사 하강법)의 업데이트는 다음과 같은 형태입니다.

즉, 학습에서는 (XW)를 빠르게 하는 것만으로는 부족합니다.

activation (활성화 값), gradient (경사), weight gradient (가중치 경사), optimizer state (최적화 상태), 임시 버퍼를 유지하며 필요한 정밀도로 계속 업데이트해야 합니다.

이 때문에 학습에서는 메모리 용량과 메모리 대역폭이 모두 중요합니다.

용량은 값이 수용 가능한지에 영향을 미치고, 대역폭은 연산기로 충분한 속도로 공급할 수 있는지에 영향을 미칩니다.

TPUv1은 「학습 불가능」이 아니라 「학습에 최적화되지 않음」

TPUv1은 행렬 곱셈기 (Matrix Multiplier)를 가지고 있으므로, 이론적으로는 제한적인 저정밀도 학습이나 학습 프로세스의 일부를 구성할 여지는 있습니다.

따라서 "TPUv1에서는 절대로 학습할 수 없다"라고 단정 짓는 것은 지나친 표현입니다.

다만, TPUv1은 일반적인 DNN (Deep Neural Network, 심층 신경망) 학습을 주 목적으로 설계된 칩은 아닙니다.

int8 중심으로 학습을 구성할 경우, 양자화 스케일 (Quantization scale), 반올림 오차 (Rounding error), 포화 (Saturation), 그래디언트 (Gradient) 표현, 마스터 웨이트 (Master weight), 옵티마이저 상태 (Optimizer state) 등을 신중하게 다루어야 합니다.

흔한 오해를 정리합니다.

TPUv2에서 무엇이 바뀌었는가

TPUv2의 중요한 변화는 bfloat16 입력과 float32 누적을 사용하는 혼합 정밀도 (Mixed precision, 저정밀도 연산과 고정밀도 유지를 결합하는 계산 방식)입니다.

개념적으로는 다음과 같은 형태입니다.

bfloat16은 16-bit이지만, float32와 동일한 지수부 폭을 가지기 때문에 float32와 유사한 값의 범위 (Dynamic range)를 유지하기 쉬운 형식입니다.

학습에서는 그래디언트나 활성화 값 (Activation)이 층(Layer)이나 스텝(Step)에 따라 크게 변합니다.

따라서 단순히 비트 수를 줄이는 것뿐만 아니라, 값의 범위를 유지하기 쉬운 것이 중요해집니다.

Kalamkar 등의 논문은 다수의 모델에서 bfloat16 학습이 float32와 동등한 결과를 보여주었다고 보고하고 있지만, 무조건적인 보장은 아닙니다.

TPUv2의 HBM과 다중 칩 구성

Wang 등의 논문에서는 Cloud TPU v2의 1 board는 4개의 TPU package, 총 8개의 core로 구성되며, board 전체로 180 TFLOPS, 64 GB 메모리, 2400 GB/s의 총 메모리 대역폭을 가진다고 정리되어 있습니다.

1 core당으로는 8 GB HBM입니다.

HBM 용량은 모델, 활성화 값 (Activation), 그래디언트, 옵티마이저 상태 (Optimizer state)가 수용 가능한지에 영향을 미칩니다.

HBM 대역폭은 이것들을 MXU (Matrix Multiply Unit, 행렬 곱셈 유닛)로 충분한 속도로 공급할 수 있는지에 영향을 미칩니다.

또한, 다중 칩 학습에서는 각 칩이 서로 다른 미니 배치 (Mini-batch)를 처리한 후, 그래디언트 동기화가 필요합니다.

이 때문에 AllReduce (여러 디바이스의 값을 집계하여 모두에게 배포하는 통신)와 같은 집합 통신 (Collective communication)이 중요합니다.

TPUv3는 무엇을 강화했는가

TPUv3는 학습이 가능해진 첫 번째 세대라기보다, TPUv2의 학습 기반을 강화한 세대입니다.

Wang 등의 논문에서는 TPUv3가 TPUv2에 비해 MXU 수와 HBM 용량이 증가하였고, 피크 성능이 420 TFLOPS로 향상되었다고 정리되어 있습니다.

Google Cloud 공식 자료에서는 TPUv3 칩당 123 TFLOPS (bf16), 32 GiB HBM2, 900 GB/s, TPU Pod size 1024 chips, AllReduce bandwidth 340 TB/s 등이 제시되어 있습니다.

단순히 "TPUv3는 항상 TPUv2의 몇 배"라고 말할 수는 없습니다.

구현자를 위한 작은 추정 예시

학습에서는 가중치(Weight)뿐만 아니라 그래디언트와 옵티마이저 상태 (Optimizer state)가 메모리를 사용합니다.

다음은 Adam을 사용하는 학습에서, 파라미터 수로부터 대략적인 메모리를 추정하기 위한 간단한 예시입니다.

from __future__ import annotations
def estimate_adam_memory_bytes(
parameter_count: int,
...

실제 학습 메모리는 활성화 값 (Activation), 임시 버퍼 (Temporary buffer), 체크포인트 (Checkpoint), 병렬화 방식에 따라 더욱 달라집니다.

이 예시의 목적은 "가중치(Weight)만이 메모리를 사용하는 것이 아니다"라는 점을 확인하는 것입니다.

TPUv1·v2·v3 비교

학습에서 필요하게 되는 수치 표현, 누적 정밀도 (Accumulation Precision), 메모리 용량·대역폭 (Bandwidth), 다중 칩에서의 계산을 하드웨어와 시스템 양면에서 다룰 수 있도록 한 것입니다.

구현자 관점에서 눈에 띈 점

GPU의 혼합 정밀도 학습 (Mixed Precision Training), Tensor Core, TPU의 bfloat16 + float32 누적은 저정밀도 곱셈으로 연산량과 데이터 전송량을 줄이면서, 누적은 고정밀도로 유지한다는 방향성 측면에서 유사합니다.

반면, NPU나 카메라 내부 AI에서는 추론에 특화된 int8 설계가 유리해지기 쉽습니다.

DLNR (Deep Learning Noise Reduction, 딥러닝을 이용한 노이즈 저감), Demosaic (컬러 필터 배열로부터 RGB 이미지를 복원하는 처리), Super Resolution (초해상도)에서는 저지연 (Low Latency), 저전력, DRAM 대역폭, 지원 연산자 (Operator)가 강력한 제약 사항이 됩니다.

학습 환경에서는 메모리 용량, 대역폭, 분산 통신 (Distributed Communication)이 중요합니다.

카메라 내부 추론에서는 레이턴시 (Latency), 소비 전력, 지원 연산자, DRAM 대역폭이 중요합니다.

동일한 뉴럴 네트워크 (Neural Network)라도 학습과 추론에서는 하드웨어 요구 사항이 크게 다릅니다.

개인적인 소감

TPUv1, TPUv2, TPUv3의 흐름은 AI 가속기 (AI Accelerator)를 단순히 "행렬 곱셈이 빠른 칩"으로만 보지 않기 위한 좋은 소재라고 생각합니다.

행렬 곱셈은 중요하지만, 학습에서는 수치 형식, 누적, 메모리, 통신, 소프트웨어 스택 (Software Stack)이 하나로 묶여 작용합니다.

특히 TPUv1을 "학습 불가능"이라고 단정 짓지 않고, 설계 목적과 실용적인 학습 적합성을 구분하면 GPU나 NPU의 설계를 볼 때도 오해를 줄일 수 있습니다.

상세 버전

더 자세한 해설, 도해, 관련 연구와의 비교는 개인 블로그에 정리해 두었습니다.

👉 완전판은 여기에서 확인하세요: https://kasblo.com/ai-paper/tpu-v1-v2-v3-training-inference/

Discussion