GPU만으로는 부족할 때: 구글 TPU가 AI 학습을 바꾸는 이유

https://henryhmko.github.io/posts/tpu/tpu.html

AI 모델은 점점 더 크고 복잡해지고 있습니다. 이제는 단순히 빠른 연산만으로는 감당하기 어려운 수준에 도달했죠. 이런 변화 속에서, 구글은 GPU나 FPGA 대신 자체 개발한 **TPU(Tensor Processing Unit)**라는 전용 칩을 통해 대규모 AI 학습과 추론의 새로운 해법을 제시하고 있습니다.

이 글에서는 TPU가 어떤 칩인지, GPU와는 어떻게 다른지, 그리고 실제로 어떤 구조와 기술적 철학을 기반으로 만들어졌는지 자세히 살펴봅니다. 마지막에는 실제로 어떻게 활용되고 있는지도 확인할 수 있습니다. 복잡한 이론이 아닌, 구조와 목적, 실용성 중심으로 설명해드릴게요.

반응형

---

TPU란 무엇인가: 구글이 직접 설계한 AI 전용 칩

TPU는 구글이 2013년부터 시작해 개발한 ASIC(Application Specific Integrated Circuit) 기반 AI 특화 칩입니다. 기존의 GPU, FPGA는 범용성을 강조했지만, 구글은 AI 서비스의 폭발적인 확장에 대응하기 위해 특정 용도에 최적화된 전용 하드웨어가 필요하다고 판단했습니다.

특히 음성 검색, 번역, 추천 시스템 등에서 요구되는 빠른 응답성과 효율적인 연산 처리를 위해 TPU가 등장한 것이죠. 현재는 Gemini, Veo, PaLM 등 주요 AI 모델 학습에 핵심적인 역할을 하고 있습니다.

---

구조부터 다른 설계: TPU는 왜 GPU와 다르게 만들었나?

TPU의 가장 큰 차별점은 **시스톨릭 어레이(systolic array)**를 기반으로 한다는 점입니다. 이 구조는 행렬곱 연산처럼 반복적이고 병렬적인 연산을 아주 빠르고 효율적으로 수행합니다.

• 시스톨릭 어레이는 인접한 연산 유닛끼리 데이터를 전달하며 파이프라인 형태로 연산을 처리합니다.
• 이렇게 하면 외부 메모리 접근 없이 데이터가 내부에서 흘러가며 연산되므로, 병목을 줄이고 에너지 효율이 극대화됩니다.

또한 TPU는 선행 컴파일(AoT; Ahead-of-Time) 방식을 사용합니다. XLA 컴파일러를 통해 연산 그래프를 분석하고 최적화한 후, 실행 시점에는 미리 컴파일된 코드만 사용하게 됩니다. 이 덕분에 TPU는 불필요한 메모리 접근과 제어 흐름을 줄이며 높은 처리량을 유지할 수 있습니다.

---

TPU의 내부 구조: TensorCore와 고속 메모리의 조합

TPU는 연산에 필요한 요소들을 아주 정교하게 구성했습니다. 특히 TPUv4 기준 단일 칩에는 2개의 TensorCore가 포함돼 있고, 각 코어는 다음과 같은 구조로 이루어져 있습니다.

• MXU (Matrix Multiply Unit): 128x128 크기의 시스톨릭 어레이로 행렬곱 연산에 특화
• VPU (Vector Unit): 일반적인 원소 단위 연산 담당
• VMEM (Vector Memory): 32MiB로, 연산용 데이터를 준비하는 메모리
• SMEM (Scalar Memory): 10MiB로, 제어 흐름과 스칼라 연산을 처리
• CMEM (Central Memory): 각 코어당 128MiB의 온칩 메모리 제공

GPU보다 연산 코어 수는 적지만, 메모리 구조가 훨씬 크고 정교하게 구성돼 있어 특정 작업에선 월등한 성능을 보여줍니다.

---

확장성과 토폴로지: TPU가 대규모 학습에 강한 이유

TPU의 진짜 강점은 확장성입니다. 단일 칩을 넘어서 수천 개의 칩이 유기적으로 연결돼 초대형 AI 모델도 훈련할 수 있도록 설계됐습니다.

구조 단계별 정리

1. 트레이(Tray): 4개의 TPU 칩과 CPU 호스트 구성
2. 랙(Rack): 64개 TPU 칩이 4x4x4 형태의 3D 토러스 구조로 연결
3. 팟(Pod): 여러 랙을 하나로 묶은 최대 구성 단위. TPUv4는 최대 4096개 칩까지 지원
4. 멀티팟(MultiPod): 여러 팟을 DCN(데이터센터 네트워크)로 연결

이 구조에서 중요한 기술이 **OCI(Optical Circuit Interconnect)**와 **OCS(Optical Circuit Switching)**입니다. 이 덕분에 물리적으로 떨어진 노드들도 하나의 슬라이스처럼 구성할 수 있어 유연성과 유지보수성이 극대화됩니다.

또한 비틀린 토폴로지 같은 구성 방식은 연산 패턴에 따라 최적화된 데이터 흐름을 구성할 수 있어, 특정 병렬 처리에 큰 장점을 제공합니다.

---

실제 활용 사례: PaLM 훈련과 TPU

구글은 PaLM이라는 초대형 언어 모델을 TPU를 이용해 학습했습니다. 여기서 총 6144개의 TPUv4 칩, 즉 2개의 팟이 동시에 사용됐습니다. XLA 컴파일러가 슬라이스 간 통신 연산을 자동으로 삽입해, 코드 수정 없이 대규모 분산 학습이 가능했죠.

이런 구조는 단지 고성능만을 의미하는 것이 아니라, 효율적인 자원 운용과 자동화된 분산 처리를 가능하게 한다는 점에서 큰 의미가 있습니다.

---

728x90

TPU는 단순히 빠른 칩이 아니다

TPU는 단순히 GPU보다 빠른 칩이 아닙니다. TPU는 처음부터 AI 연산만을 위해 설계된 구조와 철학을 바탕으로 만들어졌습니다. 다음과 같은 점에서 GPU와 다른 방향을 걷고 있습니다.

• 연산 구조: 시스톨릭 어레이 기반으로, 행렬곱 연산에 최적화
• 컴파일 방식: XLA와 AoT 방식으로 메모리 효율 극대화
• 확장 구조: Pod, MultiPod까지 확장 가능한 유연한 시스템
• 통신 최적화: OCS 기술로 토폴로지 구성의 자유도와 성능 확보

AI 모델이 계속해서 커지고 복잡해지는 지금, GPU만으로는 한계가 있는 상황에서 TPU는 새로운 길을 제시하고 있습니다. 앞으로 TPU는 더 많은 연구자와 개발자들이 접할 수 있는 자원으로 확대되고 있으며, 구글의 TRC(TPU Research Cloud) 같은 프로그램을 통해 실험적 사용도 가능해지고 있습니다.

TPU가 AI 인프라에서 어떤 의미를 가지는지, 그 구조적 차이가 실제 성능과 효율로 어떻게 연결되는지를 이해하면, 더 전략적인 인프라 선택이 가능할 것입니다.

https://henryhmko.github.io/posts/tpu/tpu.html

TPU Deep Dive