TPU: 구글이 NVIDIA GPU를 기다리지 않는 이유 (2/2)
Source: Dev.to
딥러닝을 위한 정밀도 포맷
전통적인 과학 컴퓨팅은 FP64(배정밀도) 또는 FP32(단정밀도)를 사용하며, 이는 매우 높은 정확도를 제공합니다.
하지만 딥러닝은 그 정도의 정밀도가 필요하지 않습니다. 그래서 구글은 bfloat16(Brain Floating Point)이라는 16‑비트 포맷을 만들었으며, FP32와 동일한 넓은 동적 범위(≈ 1e‑38 ~ 1e38)를 유지하면서 소수점 이하 정밀도만 일부 포기합니다.
- FP16은 범위가 제한적(~6e‑5 ~ 6e4)이라 학습이 불안정해질 수 있습니다.
- bfloat16은 FP32와 같은 범위를 유지하므로 AI 워크로드에 적합합니다.
NVIDIA는 이후 A100 및 H100 GPU에서 bfloat16을 채택했습니다.
TPU 팟과 칩 간 인터커넥트
단일 TPU 칩은 행렬 곱셈에 뛰어나지만 오늘날의 거대한 모델을 단독으로 처리하기엔 부족합니다. 구글은 칩을 계층적인 TPU 팟으로 묶습니다:
- TPU 칩 → TPU 보드
- 보드 → TPU 랙
- 랙 → TPU 팟
팟 하나에 최대 4,296개의 칩을 포함할 수 있으며, 소프트웨어는 이를 하나의 거대한 병렬 프로세서처럼 보여줍니다.
칩 간 인터커넥트 (ICI)
학습 중 지속적으로 요구되는 저지연 데이터 교환에 표준 이더넷은 너무 느립니다. TPU 팟은 CPU를 우회하는 전용 칩 간 인터커넥트 (ICI) 를 사용합니다. 칩들은 3‑D 토러스 토폴로지(도넛 형태)로 연결되어 있어, 어느 칩이든 몇 번의 홉만에 가장 먼 칩에 도달할 수 있습니다.
광회로 스위치 (OCS) – TPU v4
TPU v4는 광회로 스위치 (OCS) 를 도입해 전기‑광 변환 단계를 없앴습니다:
- MEMS 미러가 기울어져 빛을 실어 나르는 데이터 빔을 직접 지정하므로 거의 0에 가까운 지연을 제공합니다.
- 복원력: 일부 칩이 고장 나면 미러 각도를 다시 맞춰 트래픽을 즉시 재라우팅할 수 있습니다.
대규모 냉각
수천 개의 칩을 한데 모으면 엄청난 열이 발생합니다. 구글은 칩 직접 액체 냉각을 사용해 냉각 파이프를 칩 바로 위에 배치합니다—사실상 데이터 센터를 거대한 수족관처럼 만드는 방식입니다. 이 접근법은 NVIDIA가 최근 H100에 액체 냉각을 도입하기 전에 이미 사용되었습니다.
소프트웨어 스택: TensorFlow에서 JAX까지
| 구성 요소 | 역할 | 입력 | 출력 |
|---|---|---|---|
| JAX (프론트엔드) | 사용자 인터페이스 | Python 코드 | 중간 표현 (HLO) |
| XLA (백엔드) | 컴파일러 엔진 | HLO | TPU/GPU용 계산 그래프 / 바이너리 |
- 자동 미분 (grad) 및 벡터화 (vmap) 는 JAX가 담당합니다.
- XLA는 커널 융합을 수행해 메모리 접근을 줄이고 TPU의 systolic array를 최대한 활용합니다.
Ironwood TPU (TPU v7)
2025년 구글은 Ironwood라는 7세대 TPU를 발표했으며, 이는 LLM 추론 효율성과 대규모 학습 모두를 목표로 설계되었습니다.
- FP8 지원(네이티브 8‑비트 부동소수점 지원 최초 TPU).
- 연산량: 4,614 TFLOPS (FP8).
- 메모리: 칩당 192 GB HBM3E, 대역폭 7.37 TB/s(메모리‑대역폭 제한 워크로드).
- 팟 확장성: 팟당 최대 9,216개 칩.
- ICI 대역폭: 양방향 1.2 TB/s.
- 전력 효율: TPU v6 대비 약 2배 향상, 여전히 칩 직접 액체 냉각 사용.
Source: Google Cloud Blog – Ironwood TPU
GPU가 여전히 지배적인 이유
- 소프트웨어 생태계: NVIDIA의 CUDA는 2006년부터 지속적으로 발전해 왔으며, AI 연구자 중 90 % 이상이 CUDA용 코드를 작성합니다.
- 프레임워크 최적화: PyTorch는 CUDA에 최적화돼 있습니다. PyTorch가 TPU에서도 동작할 수는 있지만, 경험이 덜 성숙합니다.
- 접근성: GPU는 직접 구매해 온‑프레미스에 설치할 수 있지만, TPU는 Google Cloud Platform을 통해서만 이용할 수 있습니다. 이미 AWS나 Azure를 사용 중인 조직은 TPU 도입이 어렵습니다.
따라서 TPU가 하드웨어적으로 우수함에도 불구하고, 오랜 기간 구축된 소프트웨어 스택과 GPU의 접근성이 AI 개발에서 GPU를 여전히 선두에 서게 만들고 있습니다.