TensorFlow TPU 학습 101

TensorFlow 2.0 모델을 Google Colab과 GCP의 TPU에서 학습시키는 과정을 기초 원리부터 코드까지 설명한다. TPU가 XLA로 연산 그래프를 컴파일하는 이유, TPUEstimator와 TPUStrategy의 차이, TPUClusterResolver 연결과 `strategy.scope()` 사용법을 다룬다. 무료 Colab의 불안정성과 GCP 과금 종료, 연구자용 TFRC 지원까지 운영상 주의점도 포함한다.

핵심 포인트

TPU는 선형대수 연산과 온칩 고대역폭 메모리에 최적화된 Google 전용 가속기이며, TPU v2 Pod가 V100 8개보다 ResNet-50 학습에서 약 27배 빨랐다는 비교를 소개한다.
XLA는 런타임에 TensorFlow 그래프를 JIT 컴파일해 중복 연산을 합치고 중간 메모리 할당을 줄인다.
TensorFlow 2.0에서는 `tf.distribute.Strategy` 기반 TPUStrategy로 TPU 연결과 복제 학습을 구성하는 방식을 권장한다.
Colab에서는 `COLAB_TPU_ADDR`로 gRPC 주소를 얻고 TPUClusterResolver 초기화 후 `strategy.scope()` 안에서 metric·compile·fit을 실행한다.
안정적인 학습은 GCP Cloud TPU가 낫지만 인스턴스와 TPU가 함께 과금되므로 작업 후 `ctpu delete`로 모두 종료해야 한다.

상세 정리

TPU 성격: Google이 머신러닝 학습과 추론의 대규모 선형대수를 가속하려고 만든 하드웨어다. Google 인프라에서만 Colab 무료 할당이나 GCP 유료 대여 형태로 접근한다.
데이터 이동: 매개변수를 온칩 고대역폭 메모리에 두고 in-feed queue의 batch를 읽어 GPU의 PCIe 경로보다 빠르게 공급하는 구조를 설명한다.
GPU와 차이: 단순히 빠른 GPU처럼 코드를 옮기면 안 된다. 모델 그래프를 TPU에 올리기 전에 지원 가능한 형태로 컴파일해야 하며 이 과정에서 TPU 특유의 오류가 발생할 수 있다.
XLA 최적화: softmax처럼 동일한 `exp(logits)`를 반복 계산하는 그래프를 분석해 연산을 결합한다. JIT로 네이티브 기계어를 만들며 중복 계산과 임시 tensor 메모리를 줄인다.
프레임워크 지원: XLA는 TensorFlow를 위해 시작했지만 PyTorch/XLA도 제공된다. 글은 당시 상대적으로 성능이 좋은 TensorFlow 2.0 경로에 집중한다.
두 API: TPUEstimator는 TensorFlow 1.x의 Estimator 확장이고, TPUStrategy는 TensorFlow 2.x의 분산 정책이다. 후자는 MirroredStrategy처럼 상태와 계산을 여러 TPU 코어에 복제한다.
Colab 연결: 런타임 가속기를 TPU로 바꾼 뒤 `COLAB_TPU_ADDR`를 gRPC URI로 만들고 TPUClusterResolver, cluster connect, TPU system initialize 순서로 실행한다.
학습 범위: resolver로 TPUStrategy를 만든 다음 metric 선언, `model.compile`, `model.fit`을 `strategy.scope()` 안에서 수행한다. 범위 밖에서 만들면 정상 실행되지 않을 수 있다.
Colab 제약: 무료 TPU는 선점형이고 연결이 불안정해 자세한 오류 없이 세션이 종료되기도 한다. 중요한 장기 학습이나 원인 분석 환경으로는 적합하지 않다고 경고한다.
GCP 생성: Cloud Shell에서 `ctpu up --zone=us-central1 --name test-tpu`를 실행하면 TPU와 연결된 Compute Engine이 함께 시작된다. TPU 이름은 `TPU_NAME` 환경 변수에서 얻는다.
GCP 비용: 글 작성 시점의 가장 저렴한 us-central1 v2-8이 시간당 약 4.50달러였다. 사용이 끝나면 같은 zone의 `ctpu delete`로 TPU와 Compute Engine을 함께 제거해야 추가 과금을 막는다.
연구 지원: TFRC는 승인된 연구자에게 30일 동안 on-demand v3 5개, on-demand v2 5개, preemptible v2 100개를 무료로 제공했다. 연구 결과 공유가 조건이다.