Google, Meta TPU용 PyTorch 백엔드 TorchTPU 공개

Google과 Meta가 JAX 의존성을 제거하고 model.to("tpu") 명령만으로 실행 가능한 TPU용 네이티브 PyTorch 백엔드 TorchTPU를 공개함. 해당 프로젝트는 2026년 10월 PyTorch Conference에서 TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, TorchTitan 저장소를 오픈소스로 공개할 계획임.
Hugging Face 상위 5,000개 모델 중 4,917개를 검증하여 인기 모델 기준 98%, 수치 정확도 기준 97%의 통과율을 달성함.
1,024개 TPU v7x 칩에서 Llama 3 405B 모델을 500 tokens/sec/chip, MFU 50% 이상으로 실행하는 성능 벤치마크를 기록함.
SimpleFSDP, Scan Higher-Order Op, Pallas SplashAttention 최적화를 통해 FLUX.1-dev는 최대 3.80배, 어텐션 연산은 17.5% 속도 향상을 구현함.
분산 학습 시 2D+ DeviceMesh를 활용해 Tensor/Expert Parallelism은 고속 ICI, Data Parallelism은 저속 DCN에 할당하고 SparseCore 오프로딩으로 All-Gather/Reduce-Scatter 처리함.
TorchRec에 추천 모델용 네이티브 TPU 백엔드를 추가하여 DLRM v2에서 32칩 TPU v7x 클러스터 기준 JAX 베이스라인과 동일한 18.3ms 스텝 지연 시간을 달성함.
향후 로드맵에는 TPUv8용 FP8/FP4 양자화 지원, RL 루프용 Ray Core 오케스트레이션, Reasoning, Multimodal, Diffusion 모델 지원 확대가 포함됨.

핵심 내용