TPU 심층 분석 요약
카테고리
프로그래밍/소프트웨어 개발
서브카테고리
인프라/DevOps/보안
대상자
- 대상자: 소프트웨어 개발자, AI 연구자, 시스템 엔지니어
- 난이도: 중급~고급 (하드웨어 아키텍처 및 분산 컴퓨팅 지식 필요)
핵심 요약
- TPU는 구글이 ASIC 기반으로 개발한 AI 학습/추론 전용 칩으로, 시스톨릭 어레이와 XLA 컴파일러를 통해 행렬 곱셈, 컨볼루션 연산 최적화
- TPUv4 기준 500 TFLOPs/sec 성능, 8960칩 팟 기준 4.45 ExaFLOPs/sec의 초대규모 확장성 제공
- HBM3 메모리와 온칩 메모리(CMEM, VMEM, SMEM)를 통해 에너지 효율성과 메모리 액세스 최소화 달성
섹션별 세부 요약
1. TPU 개요 및 설계 철학
- TPU는 GPU와 달리 확장성, 에너지 효율성을 중점으로, 시스템 온칩(SoC)과 대형 온칩 메모리를 통합
- ASIC 기반으로, 2013년 이후 딥러닝 확대에 따라 맞춤형 하드웨어 필요성 증가
- Gemini, Veo, DLRM 등 주요 AI 서비스에 필수적 활용
2. 핵심 구성 요소 및 아키텍처
- TensorCore: CMEM(128MiB), HBM(32GiB) 연결, MXU(행렬 곱셈), VPU(일반 연산), VMEM(32MiB), SMEM(10MiB) 등으로 구성
- 시스톨릭 어레이: 128x128 구조, 행렬 곱셈 전용, 메모리 읽기/쓰기는 입출력만 발생
- XLA 컴파일러: 메모리 액세스 최적화, 캐시 의존도 감소, JAX의 @jit를 통해 정적 그래프 생성 및 AoT 컴파일
3. 확장성 및 네트워크 토폴로지
- 랙(Rack): 64칩, 3D 토러스(4x4x4) 구조, ICI(Inter-Core Interconnect) 및 OCS(Optical Circuit Switching) 연결
- 팟(Pod): 최대 4096칩(64랙), 슬라이스(Slice): 논컨티구어스 구성 가능
- 토폴로지 유형: 큐브형(데이터 병렬), 직선형(파이프라인 병렬), 비틀린 토폴로지(all-to-all 통신 최적화)
4. 성능 및 적용 사례
- TPUv5p: 500 TFLOPs/sec 성능, PaLM 훈련은 2팟(6144 TPUv4) 활용
- XLA 컴파일러가 병렬 차원에 따라 통신 패턴 최적화, 대규모 분산 학습 구현
- 구글 TPU Research Cloud(TRC) 및 GSPMD, TPUv4i 논문 지원
결론
- TPU는 에너지 효율성과 초대규모 확장성으로 대규모 AI 모델 훈련에 적합, XLA 컴파일러를 통해 메모리 액세스 최적화 필요
- 희소 행렬 처리 시 한계가 있으므로, 규칙적인 데이터 구조를 추천, OCS 및 ICI로 논컨티구어스 슬라이스 구성 가능
- TPUv4i(추론형)은 CMEM 확장으로 추론 성능 향상 가능, HBM3 메모리 도입으로 에너지 소모 감소 달성