Google TPU에 대한 단상

by Edward Yoon

Nov 28. 2025

16년 쯤 DataFlow, TensorFlow, 그리고 TPU의 출현!

나의 소프트웨어나 오픈소스 참여 열정을 빠르게 식힌 사건이라 잘 기억하고 있다.

뭔가 더 할 수 있는게 없는 느낌? 어쨌거나 그 시절 병렬처리나 분산시스템에 진심이었던 터라 TPU에 대한 이야기를 몇 자 적어본다.

PCIe 버스 통신의 병목

GPU 환경에서 추론은 다음과 같이 흘러간다:

1) CPU가 요청 받음

2) 입력 데이터를 CPU메모리에 준비

3) PCIe 버스를 통해 GPU에 업로드

4) GPU가 연산

5) 결과를 다시 PCIe로 CPU에 전송

여기서 문제는 PCIe 대역폭은 GPU 연산속도에 비하면 극도로 느리다는거다. 테라~페타 플롭스 연산 이후에 GB/s 수준으로 통신하거든. 따라서, 입력이 작아도 PCIe 왕복 latency가 연산 대비 상대적으로 너무 커서 batch=1 단일 인스턴스에 대한 추론을 예시로 아주 비효율적 기계가 되버린다.

ALU 가동률(?)

게다가, batch=1 같은 작은 배치에서는 스레드를 모두 채우지 못해 남는 ALU가 대기상태여서 느린 latency 뿐 아니라, 처리량도 낮아지는 현상이 발생한다. 이걸 warp underutilization이라고 한다.

그 반면, TPU는 행렬 곱 파이프라인에 맞춰 고정되고 입력이 작든 크든 파이프라인 단위로 연산이 흐르기 때문에 ALU 가동률이 최대치에 가까워진다. 이처럼 파이프라인 따라 계산만 흐른다 하여 TensorFlow 다. Computational Graph 최적화가 핵심이다.

Data locality, Weight-stationary 전략

주요 연산인 행렬 곱: Y=W⋅X 에서 W = weight (모델 파라미터), X = activation (입력) 중 어느 데이터를 만지작 거리면서 연산할지를 결정해야되는데, TPU는 데이터 처리 전략을 명확히 선택한 아키텍처를 가지고 있다. GFS, MapReduce 같은 구글의 분산 시스템에서 중요하게 생각하는 데이터 로컬리티 철학은 여기에서도 엿보인다.

GPU의 weight 가중치 처리 방식은 다음과 같다:

1) 모델 파라미터가 VRAM에 올라가있음

2) 연산할때 작은 chunk씩 register/Shared memory로 불러와 곱셈 후 accumulation

결과적으로, 작은 연산 단위마다 weight를 fetch하고 연산에 사용한다는 점에서, TPU와 달리 이동 비용이 발생한다. TPU는 내부 SRAM에 상주시켜놓고, 입력만 흐른다. 특히, 구글 번역 같은 즉 한 단어씩 들어갈때마다 추론하는 온라인 시퀀스 처리 즉 스트리밍 인풋에서 차이가 극명하게 벌어진다.

여기까지가 16~17년도 쯤 출현한 TPU v1의 추론 특화 프로세싱에 대한 이야기다. 온라인 서비스에 AI를 어떻게 효율적으로 처리할건지에 대한 고민을 하는구나라는 지점이 그들은 벌써 10년 전 얘기라는 것.

스케일 아웃 구조

GPU 세계에서 초기 대규모 학습을 위한 병렬 처리 구조는 역시나 통신 병목이 큰 허들이라 메모리를 늘리는 쪽으로, 그리고 여러 대에 동시 학습으로 gradient를 합치는 식으로 학습 속도 개선 하는데 더 주목하는 방향성을 잡았다 (GPU는 이때에도 throughput 중심). 서빙은 마이크로 배치로 묶어 처리하던지 뭐 그런식이었고. 이 방향성이 빠른 연구개발 환경, 프로토타입 개발에 도움 된다는 건 인정한다.

그러나 구글은 역시나 온라인 서비스와 스케일 문제에 진심이다.

DistBelief (대규모 CPU 클러스터 버전) 논문에서 부터 Jeff dean은 대형 모델을 처리하기 위한 데이터 병렬과 모델 병렬 두 마리 토끼를 모두 잡고 싶어하는 모습을 보였는데 그것이 TPU와 Tensorflow에서 고스란히 반영된 모습으로 진화했다.

즉, 모델을 쪼개서 여러 TPU에 배치시키고, TPU 보드에 interconnect 연결함으로써 결과적으로, 소프트웨어로 보면 MPI collective 방식을 하드웨어 자체 Collective 구조로 만들어서, weight-stationary 데이터 전략 + on-chip memory + 고성능 네트워크 통신으로 단일 샘플 latency, 그리고 대규모 학습 성능 둘 다 때려잡았다.