TPU (Tensor Processing Unit)
텐서 처리 장치. Google이 설계한 AI 학습/추론 전용 ASIC으로, 대규모 딥러닝 워크로드에 최적화되어 있다.
역할
- 대규모 신경망 학습 (Training)
- 고속 추론 (Inference)
- 행렬/텐서 연산 집중 처리
- Google 내부 AI 서비스 및 GCP 제공
GPU와의 차이
| 항목 | GPU | TPU |
|---|---|---|
| 설계 목적 | 범용 병렬 연산 | 텐서 연산 특화 |
| 제조사 | NVIDIA / AMD 등 | Google (자체 설계) |
| 범용성 | 높음 | 낮음 (TensorFlow/JAX 위주) |
| 메모리 구조 | VRAM (GDDR/HBM) | HBM + 고대역폭 설계 |
| 가용 방법 | 자체 구매 / 클라우드 | GCP (Cloud TPU)만 |
| 전력 효율 | 중간 | 높음 (AI 특화) |
핵심 구조: Systolic Array
TPU의 핵심은 수축 배열(Systolic Array) 구조다. 행렬 곱셈을 파이프라인처럼 연속 흐름으로 처리하여 메모리 접근 횟수를 최소화한다.
데이터가 배열을 통해 파도처럼 흐름 → 캐시 미스 없이 연속 연산
→ 행렬 곱셈에서 GPU 대비 수십 배 효율
TPU 세대
| 세대 | 특징 |
|---|---|
| TPU v1 | 추론 전용, Google 내부 배포 |
| TPU v2/v3 | 학습 가능, 외부 공개 (GCP) |
| TPU v4 | 대규모 포드(Pod) 구성 가능, 광인터커넥트 |
| TPU v5e/v5p | 최신 세대, 비용 효율 및 성능 개선 |
활용 방법
GCP(Google Cloud Platform)의 Cloud TPU를 통해 사용:
import jax
import jax.numpy as jnp
# JAX는 TPU를 자동 감지하여 활용
devices = jax.devices("tpu")TensorFlow와 JAX가 기본 지원되며, PyTorch는 torch_xla로 TPU 활용 가능.
적합한 워크로드
- 대형 언어 모델(LLM) 학습
- 대규모 이미지 분류 모델 학습
- Google 내부: 검색 랭킹, 번역, AlphaGo/AlphaFold 등
개발자 관점에서의 고려사항
- GCP 비용 계획 시 TPU와 GPU 비용 비교 필요
- JAX 기반 코드가 TPU에서 가장 효율적
- TPU Pod: 수백~수천 개 TPU를 네트워크로 연결하여 초대형 학습 가능
- 온프레미스 사용 불가 (Google Cloud 전용)