TPU 상세

개요

TPU(Tensor Processing Unit)는 Google이 자체 설계한 신경망 학습 전용 ASIC(Application-Specific Integrated Circuit)으로, 2015년부터 Google 데이터센터에서 운영되어 왔다. 범용 프로세서(CPU)와 그래픽 프로세서(GPU)가 일반적인 von Neumann 구조의 한계를 가지는 반면, TPU는 메트릭스 프로세서(Matrix Processor)로서 대규모 행렬 곱셈-누적(Multiply-Accumulate) 연산에 최적화된 systolic array 아키텍처를 사용한다.

Google은 2016년 Google I/O에서 TPU를 공식 발표했으며, TPU v1부터 최신 TPU v8까지 지속적으로 세대를 거듭해 왔다. TPU는 Google Cloud TPU 서비스를 통해 외부 사용자에게 제공되며, JAX, TensorFlow, PyTorch/XLA 등의 프레임워크를 지원한다. TPU의 핵심 목표는 기계 학습 워크로드에서 높은 처리량(throughput)과 에너지 효율을 달성하는 것이며, 이를 위해 메모리 대역폭 병목을 최소화하고 대규모 병렬 행렬 연산을 하드웨어 수준에서 직접 수행한다.

핵심 개념

Systolic Array 아키텍처

Systolic Array는 데이터가 곱셈-누산 유닛(Multiply-Accumulator) 격자를 순환적으로 흐르며 연산이 수행되는 아키텍처이다. CPU가 매 연산마다 메모리에서 operand를 가져오는 것과 달리, systolic array는 한 번 로드된 데이터가 인접 유닛으로 전달되면서 반복적으로 사용된다. 이를 통해 메모리 접근 횟수를 대폭 줄이고 에너지 효율을 높인다.

동작 원리:
- 가중치(Weight)는 MXU 내부에 프리로드(preload)된다.
- 입력 데이터(Activation)가 행 단위로 스트리밍 입력된다.
- 각 유닛에서 곱셈 후 결과가 옆 유닛으로 전달되며 누적된다.
- 최종 결과 벡터가 MXU에서 출력된다.
- 데이터가 격자를 통과하는 동안 별도의 메모리 접근이 불필요하다.

TPU MXU 크기 비교:

bfloat16 포맷

bfloat16(Brain Floating Point)는 Google Brain에서 개발한 16비트 부동소수점 포맷으로, IEEE 754 FP32의 지수 비트를 그대로 유지하고 mantissa를 7비트로 줄인 것이다. FP32와 동일한 동적 범위(10⁻³⁸ ~ 10³⁸)를 가지면서도 메모리 사용량과 대역폭을 절반으로 줄여 TPU의 연산 효율을 높인다.

TensorCore 구조

TensorCore는 TPU 칩 내부의 핵심 연산 단위로, 하나 이상의 MXU, 벡터 유닛(Vector Unit), 스칼라 유닛(Scalar Unit)으로 구성된다. 세대에 따라 칩당 TensorCore 수와 TensorCore 내부 MXU 수가 다르므로, TPU를 볼 때는 "칩당 TensorCore 수"와 "TensorCore당 MXU 수"를 함께 봐야 한다.

HBM(High Bandwidth Memory)

TPU는 HBM을 사용하여 MXU에 대량의 데이터를 빠르게 공급한다. TPU v2부터 HBM을 도입하였으며, 세대가 거듭될수록 용량과 대역폭이 확대되었다.

ICI(Inter-Chip Interconnect)

ICI는 TPU Pod 내 칩 간 고속 연결 네트워크로, 단일 칩의 메모리 대역폭을 넘어 대규모 모델 병렬 처리를 가능하게 한다. TPU v4부터는 큐브 간 연결에 광학적 회로 스위치(OCS: Optical Circuit Switch)를 사용해 대형 슬라이스의 장애 우회와 토폴로지 운영 유연성을 높였다.

확장 규모:

SparseCore

SparseCore는 희소 연산 전용 데이터프로세서로, 추천 모델(Recommendation Model)과 같은 대규모 임베딩 모델의 가속에 특화되어 있다. TPU v5p와 TPU7x에는 칩당 4개의 SparseCore가, TPU v6e에는 칩당 2개의 SparseCore가 내장된다.

비교/분석

TPU 세대별 비교

TPU vs GPU 아키텍처 비교

TPU VM 아키텍처

TPU VM은 물리적으로 TPU 하드웨어에 연결된 가상 머신으로, 사용자가 SSH를 통해 직접 접근할 수 있다. 이 방식은 초기 세대의 관리형 추상화보다 디버깅과 런타임 제어 범위를 넓혀 주며, JAX나 PyTorch/XLA 실행 환경을 VM 수준에서 직접 다룰 수 있게 한다.

동작 흐름:
1. TPU VM이 물리 호스트에 생성된다.
2. 사용자가 SSH로 TPU VM에 접속한다.
3. XLA 컴파일러가 모델을 TPU 칩에 최적화한다.
4. TPU 칩이 HBM에서 데이터를 로드하여 연산을 수행한다.
5. 결과가 TPU VM의 메모리에 기록된다.

동작 원리

TPU 칩 내부 구조

TPU 칩 하나는 하나 이상 TensorCore로 구성되며, 각 TensorCore는 MXU, 벡터 유닛, 스칼라 유닛을 포함한다. TPU v4와 v5p는 칩당 2개의 TensorCore를, TPU v6e는 칩당 1개의 TensorCore를, TPU7x는 칩당 2개의 TensorCore를 제공한다. SparseCore는 TPU v4가 아니라 v5p 이후 세대에서 등장했다.

데이터 흐름:
1. 호스트(Host) CPU가 TPU에 데이터를 전달한다.
2. TPU가 HBM에서 가중치를 MXU에 로드한다.
3. 입력 데이터가 행 단위로 MXU에 스트리밍된다.
4. systolic array 내에서 곱셈-누적 연산이 파이프라인으로 수행된다.
5. 결과가 HBM에 저장되거나 출력 큐로 전달된다.

Pod 스케일 아키텍처

대규모 모델 훈련을 위해 여러 TPU 칩이 ICI 네트워크로 연결되어 Pod를 형성한다. TPU v4, v5p, TPU7x 같은 대형 Pod 세대는 큐브 내부에서는 직접 링크를 사용하고, 큐브 간 연결에는 OCS를 사용해 대규모 3D 토폴로지를 구성한다.

토폴로지 예시:
- 단일 호스트: TPU v7 기준 4개 칩이 하나의 호스트에 연결
- 다중 호스트: 여러 호스트의 칩이 ICI로 연결
- Multislice: 여러 Pod가 DCN(Data Center Network)으로 연결되어 초대규모 병렬 처리 지원

3D 토폴로지:
- TPU v4/v5p는 3D 토폴로지(A×B×C)를 지원
- 64 칩 이상일 때 A, B, C는 4의 배수이어야 함
- 예: 4×4×4 = 64 칩 (1 cube), 8×8×8 = 512 칩

ICI Resiliency

ICI Resiliency는 광학 링크와 OCS의 오류를 견디는 기능으로, TPU v4, v5p, TPU7x의 큐브급 이상 슬라이스에서 기본 활성화된다. ICI 연결이 광학 OCS나 광학 ICI 오류를 우회하여 라우팅될 수 있어 슬라이스의 스케줄링 가능성이 향상되지만, 일시적으로 ICI 성능이 저하될 수 있다.

Multislice 병렬 처리

Multislice는 여러 슬라이스를 그룹화하여 ICI의 범위를 넘어 확장하는 기술이다. 각 슬라이스 내 데이터는 ICI로 전송되며, 슬라이스 간 데이터는 DCN으로 전송된다. 이를 통해 단일 슬라이스가 수용할 수 있는 것보다 더 많은 TPU 코어를 하나의 작업에 사용할 수 있다.

장단점

장점

단점

관련 기술

하드웨어 관련

• Tensor Core Architecture: NVIDIA 텐서 코어 아키텍처 상세 분석
• AI 가속기 개요: GPU/TPU/NPU 분류 체계 및 비교
• 혼합 정밀도 하드웨어: FP8, BF16, FP4 혼합 정밀도 기술
• HBM Generation Comparison: HBM 메모리 기술 비교

소프트웨어 관련

• Google JAX: TPU 최적화 수학 프레임워크, 함수형 프로그래밍 패러다임
• TensorFlow/XLA: TPU/GPU용 자동 컴파일러
• PyTorch/XLA: PyTorch의 TPU 지원 레이어
• OpenXLA: TPU/GPU 통합 컴파일러 인터페이스

참고 문헌

• Jouppi et al., "In-Datacenter Performance Analysis of a Tensor Processing Unit," ISCA 2017
• Jouppi et al., "In-Datacenter Performance Analysis of a Tensor Processing Unit," ISCA 2017
• Jouppi et al., "A Domain-Specific Supercomputer for Training Deep Neural Networks," Communications of the ACM, 2021
• Jouppi et al., "TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings," ISCA 2023
• Google Cloud TPU Documentation, "TPU architecture" (2026)
• Google Cloud TPU Documentation, "TPU v5p" (2026)
• Google Cloud TPU Documentation, "TPU v6e" (2026)
• Google Cloud TPU Documentation, "TPU7x (Ironwood)" (2026)

핵심 정리

1. TPU는 systolic array 아키텍처를 사용하여 메모리 접근을 최소화하고 대규모 행렬 연산의 에너지 효율을 극대화하는 Google의 AI 전용 ASIC이다.
2. bfloat16 포맷은 FP32와 동일한 동적 범위를 유지하면서 메모리 사용량을 절반으로 줄여 TPU의 연산 효율을 높인다.
3. ICI(Inter-Chip Interconnect)와 광학적 회로 스위치(OCS)를 통해 Pod 수준에서 수천 칩까지 확장하여 초대규모 모델 훈련을 지원한다.
4. TPU7x(Ironwood)는 192GB HBM3e와 7.38 TB/s 대역폭을 제공하며, 칩당 2,307 TFLOPS(BF16)와 9,216칩 Pod 확장을 지원한다.
5. TPU는 Google Cloud 전용이라는 접근성 제한이 있지만, 대규모 훈련에서의 비용 효율성과 확장성으로 경쟁력을 유지하고 있다.