Google과 Meta, TPU에서 PyTorch를 네이티브로 실행하는 TorchTPU 공개
2025년 중반 Google에서 Meta의 PyTorch 팀과 함께 시작된 이 오픈소스 백엔드는 이미 Google Cloud의 대형 고객 일부가 프로덕션에서 사용하고 있다고 개발팀은 밝혔습니다.
AI 일러스트Google과 Meta는 PyTorch를 Google의 TPU 칩에서 네이티브로 실행할 수 있게 해 주는 오픈소스 백엔드 TorchTPU를 발표했습니다. 프로젝트 팀은 이 프로젝트가 2025년 중반 Google에서 Meta의 PyTorch 팀과 협력하여 시작되었으며 현재 활발히 개발 중이라고 밝혔습니다.
팀에 따르면 TorchTPU는 이미 Google Cloud의 대형 고객 일부가 추론 및 학습 워크로드에 프로덕션으로 사용하고 있습니다. TPU는 그동안 TensorFlow와 JAX 생태계에 묶여 있었다고 팀은 지적했습니다.
팀은 TorchTPU를 이전 시도와 결별한 것이라고 설명했습니다. Google은 2020년 Lazy Tensors를 기반으로 한 TPU용 첫 컴파일 방식 PyTorch인 PyTorch/XLA를 출시했습니다. 팀에 따르면 이 통합은 소스 코드에 침습적인 변경이 필요했고, 컴파일 시간 지연이 심했으며, PyTorch의 이거(eager) 실행 모드를 지원하지 않았고, 최신 분산 컴퓨팅 기능을 따라가지 못했습니다. 반면 TorchTPU는 이거 실행을 최우선 과제로 다루며 런타임에 JAX 의존성이 없습니다.
목표는 워크로드를 TPU로 옮길 때 코드에서 대상 디바이스를 GPU에서 TPU로 바꾸는 정도면 되도록 하는 것입니다. 팀은 모델의 셰이프, 샤딩, 연산이 GPU에 맞게 조정되어 있을 수 있으므로 이 변경만으로는 TPU의 최고 성능을 얻을 수 없다고 당부하면서도, 모델은 별도 작업 없이 실행되어야 한다고 말했습니다.
팀에 따르면 지금까지 이 프로젝트는 다음을 완료했습니다.
- DeferAndFuse라는 새로운 이거 모드를 포함해 이거 모드와 컴파일 모드 지원 추가
- 약 450개의 ATen 연산 지원 추가
- JAX/Pallas와 Helion DSL로 작성한 커스텀 커널 사용 가능
- 다운로드 수 상위 5,000개 Hugging Face 모델에서 별도 작업 없는 실행 검증
- vLLM, SGLang, TorchTitan, Ray, TorchRec과의 통합 검증
팀은 한계도 제시했습니다. TPU 디바이스에는 하나의 사용자 공간 프로세스만 접근할 수 있어 강화 학습 콜로케이션을 바로 지원할 수 없으며, 팀은 이를 우회할 방법을 조사 중이라고 밝혔습니다. 또한 기반 컴파일러와 런타임이 텐서 셰이프를 컴파일 시점에 알고 있다고 가정하기 때문에 제한된 동적 처리(bounded dynamism) 지원도 추가하고 있습니다.
정밀도와 관련해 TorchTPU는 bfloat16으로 곱셈하고 FP32로 누적하는 중간 모드를 기본값으로 사용합니다. 더 높은 모드는 약 3배 또는 약 6배의 연산 비용으로 FP32 정밀도를 더 많이 복원하며, 팀은 대부분의 딥러닝 워크로드에 기본값을 권장합니다.