TPU 심층 분석

11 hours ago 2

TPU는 구글이 개발한 대규모 AI 학습 및 추론용 맞춤형 칩으로, GPU와는 다른 설계 철학을 가지고 있음
확장성과 에너지 효율성을 강조하며, 하드웨어(예: 시스템 온칩 구성, 대형 온칩 메모리)와 소프트웨어(** XLA 컴파일러**)를 함께 설계함
핵심 구조는 시스톨릭 어레이와 파이프라이닝, 선행 컴파일 방식으로, 대부분의 딥러닝 연산(특히 행렬곱)에 최적화됨
OCI 및 OCS 기술 덕분에 유연한 노드 구성, 고성능 병렬 처리, 다양한 토폴로지 선택이 가능함
TPU 시스템은 단일 칩에서 초대형 멀티팟까지 계층적으로 확장되어, 대규모 AI 모델 훈련과 효율적 자원 활용이 가능함

개요 및 배경

TPU는 구글이 AI 서비스의 대규모 확장에 대응하기 위해 개발한 ASIC 기반 칩임
초기에는 GPU와 FPGA, ASIC 도입을 저울질했으나, 2013년 음성 검색 등 딥러닝 활용 확대에 따라 맞춤형 하드웨어의 필요성이 증가함
현재 Gemini, Veo 등 주요 AI 서비스에 필수적 역할을 하며, 추천 모델(** DLRM**) 등에 폭넓게 활용되고 있음

TPU 단일 칩 구조

기본 구성

TPUv4 기준 단일 칩에는 2개의 TensorCore가 존재(추론 특화 TPU는 1개)
각각의 TensorCore는 CMEM(128MiB), HBM(32GiB) 메모리 장치와 연결됨

TensorCore 내부 구성

Matrix Multiply Unit (MXU) : 128x128 시스톨릭 어레이 기반, 행렬곱 전담
Vector Unit (VPU) : 일반적 원소 단위 연산 수행
Vector Memory (VMEM; 32MiB) : HBM에서 데이터를 복사해 작업 준비 메모리로 활용
Scalar Unit + Scalar Memory (SMEM; 10MiB) : 제어 흐름, 스칼라 연산, 메모리 주소 관리 담당

GPU와의 구조적 차이

TPU는 온칩 메모리(CMEM, VMEM, SMEM)이 GPU 대비 훨씬 큼
HBM 용량은 GPU가 더 크고, 연산 코어 수는 GPU가 훨씬 많음
TPUv5p 기준, 칩당 500 TFLOPs/sec, 전체 팟(8960칩) 기준 4.45 ExaFLOPs/sec 수준의 고성능 구현

TPU 설계 철학

1. 시스톨릭 어레이 및 파이프라이닝

시스톨릭 어레이는 처리 소자(PE)들이 배열되어 인접 소자에 연산 결과 전달
데이터 입력 이후 별도 제어 없이 연속 연산 가능, 메모리 읽고 쓰기가 오직 입·출력에만 발생
주로 행렬 곱셈, 컨볼루션에 최적화되어 있음
파이프라인 처리를 통해 연산·데이터 이동을 동시에 진행해 처리량 최적화

시스톨릭 어레이의 단점 - 희소성(sparsity)

모든 연산 소자가 항상 활성화되어 희소 행렬에는 적합하지 않음
앞으로 DL 모델이 불규칙 희소성을 추구할 경우 한계로 작용 가능

2. 선행 컴파일(AoT) 및 캐시 최소화

TPU-XLA 코디자인으로 불규칙한 메모리 액세스를 필요로 하는 캐시 의존도를 최소화, 에너지 사용 절감
XLA 컴파일러가 연산 그래프 분석을 통해 메모리 액세스 패턴을 미리 산출, 캐시 대신 스크래치패드 메모리 중심 운용
JAX의 @jit는 JIT와 AoT의 중간 형태로, 최초 실행 시 정적 그래프 생성 후 XLA로 AoT 컴파일 진행
입력 형태가 달라지면 재컴파일 필요, 동적 패딩/반복에는 비효율 존재
높은 에너지 효율성 달성하나, 유연성 부족이 단점

TPUv4 에너지 효율성

현대 칩은 HBM3 메모리를 사용해 에너지 소모 감소
연산 대비 메모리 작업의 에너지 소모가 수십~수백 배라, 메모리 액세스 최소화로 대폭 효율 향상 가능

TPU 다중 칩 구조

트레이 수준(Tray/Board; 4칩)

1트레이는 4개의 TPU칩(8개 TensorCore)과 CPU Host로 구성
Host↔Chip 연결은 PCIe, Chip↔Chip 연결은 Inter-Core Interconnect(ICI) 로, 더 높은 대역폭 제공

랙 수준(Rack; 4x4x4=64칩)

1랙은 64 TPU칩, 3D 토러스(4x4x4) 로 ICI, OCS(Optical Circuit Switching)로 연결
구글에서는 랙(Rack), 팟(Pod), 슬라이스(Slice)를 구분함
- 랙: 64칩 물리 단위(=큐브)
- 팟: ICI 및 OCS로 연결 가능한 최대 단위(예: TPUv4=4096칩=64랙)
- 슬라이스: 4칩~Superpod 사이 임의 구성 추상 단위

OCS의 이점

래퍼라운드(Wraparound) : 각 축을 링(1D 토러스)화하여 노드 간 최악 이동 홉 수 절감
유연한 논컨티구어스 슬라이스: OCS 스위칭 구조로, 물리적으로 떨어진 노드도 하나의 슬라이스로 구성 가능, 자원 활용·유지보수 용이
비틀린 토폴로지: 동일 칩 수(x,y,z고정)에서 연결 구조을 변경해 특정 연산 패턴에 속도 최적화 가능(예: twisted torus)

토폴로지 활용 사례

큐브형: 데이터/텐서 병렬에 적합(최대 대역폭)
직선형(시가형) : 파이프라인 병렬에 유리
비틀린 토러스: all-to-all 통신 필요시(예: tensor parallel) 속도 향상

슈퍼팟(Full Pod/Superpod; TPUv4: 4096칩/64랙)

여러 랙을 ICI 및 OCS로 연결해 초대형 시스템 구축
슬라이스 토폴로지 종류에 따라 통신 대역폭/병렬처리 성능 차별화
OCS 덕분에 논컨티구어스 슬라이스, 비틀린 토폴로지 등 유연성 확보

멀티팟 수준(Multi-pod/Multislice; TPUv4: 4096개 초과)

여러 팟을 데이터센터 네트워크(DCN) 로 연결, 대규모 학습 인프라 제공하나 ICI 대비 대역폭 낮음
PaLM 훈련은 2팟(6144 TPUv4) 활용, 총 6팟 자원관리로 진행
대규모 모델 개발 시, 연구자 설정값(병렬 차원 등)에 따라 XLA 컴파일러가 통신 패턴 최적화
XLA가 각 슬라이스·팟 간 통신 연산 삽입, 최소 코드 변경으로 대규모 분산학습 구현

실물 하드웨어와 도식 예시

TPU 랙: 4x4x4 3D 토러스가 한 유닛, 각 행마다 2트레이(8칩)
TPUv4 트레이: 실제로는 4개의 PCIe 포트(각 TPU 1개씩)
TPUv4 칩: 중앙의 ASIC, 주변 4개 HBM 스택(2 TensorCore 기준)
TPUv4i(추론형) 칩 플로어플랜: 1 TensorCore, 넓은 CMEM 면적 차지

마무리

구글 TPU Research Cloud(TRC) 에 연구 지원을 표함

참고자료

Read Entire Article