TorchTPU: Google 규모에서 PyTorch를 TPU에서 네이티브로 실행하기

Google이 TorchTPU를 공개하며, PyTorch 워크로드를 TPU에서 네이티브하게 돌리는 통합 스택을 소개했다.

핵심 목표는 usability, portability, performance다. 기존 PyTorch 스크립트에서 초기화만 tpu로 바꾸는 수준의 최소 변경을 지향하며, TPU의 대규모 분산 인프라에 맞는 성능도 함께 노린다.

Eager First 접근을 채택해 개발 경험을 우선했다.

• Debug Eager: 연산마다 CPU와 동기화해 디버깅에 적합
• Strict Eager: 비동기 단일 연산 실행으로 기본 PyTorch 경험에 가깝게 동작
• Fused Eager: 연산 흐름을 자동으로 묶어 더 큰 덩어리로 실행하고, Strict Eager 대비 50%~100%+ 성능 향상을 주장

컴파일 경로는 torch.compile 기반의 전체 그래프 최적화도 지원한다. PyTorch Dynamo로 FX graph를 캡처한 뒤 Torch Inductor 대신 XLA를 백엔드로 사용하고, 연산은 StableHLO로 내려간다. 커스텀 연산은 Pallas와 JAX 커널을 지원하며, Helion 지원도 진행 중이다.

분산 학습 측면에서는 DDP, FSDPv2, DTensor를 기본 지원한다고 밝혔다. 기존 PyTorch/XLA가 순수 SPMD에 강했던 반면, TorchTPU는 rank별로 약간의 코드 차이가 있는 MPMD 실행도 다루도록 설계해 로깅 같은 비대칭 작업을 수용한다.

하드웨어 측면에서는 TPU의 효율을 살리기 위해 모델 구조도 조정이 필요하다고 짚었다. 예를 들어 attention head dimension은 64보다 128 또는 256이 TPU tensor core 효율에 더 유리할 수 있다고 설명했다.

2026년 로드맵도 제시했다.

• 공개 GitHub 저장소와 문서, 아키텍처 튜토리얼 공개
• Helion DSL 통합
• torch.compile의 동적 shape 지원 강화
• multi-queue 지원
• vLLM, TorchTitan 등 생태계 연동
• Pod-scale 인프라까지 선형 확장 검증