구글, PyTorch를 TPU에서 실행하는 백엔드 공개

핵심 요약

• Google이 PyTorch 모델을 Google의 TPU(Tensor Processing Unit) 장치에서 실행할 수 있게 하는 torch_tpu 백엔드를 오픈소스로 공개했다.
• PyTorch 연산을 위한 ATen 커널 구현, XLA 계산 최적화를 위한 컴파일 캐시, 데이터 병렬화와 텐서 병렬화를 지원하는 분산 학습 기능을 포함한다.
• Llama, Qwen, ResNet, DLRM, minGPT, UNet 등 다양한 모델의 예제와 torch.compile()을 통한 그래프 모드 실행 통합, 디버깅 및 벤치마킹 유틸리티를 제공한다.
• Python 3.11 이상을 지원하며, PyPI에서 0.2.0 버전을 설치할 수 있다.
• 소스 코드에서 빌드할 때는 Bazel 빌드 시스템을 사용하며, VSCode와 clangd를 이용한 C++ 코드 내비게이션 설정을 지원한다.
• 공식 사용 가이드와 Marimo 기반 대화형 노트북이 torchtpu.readthedocs.io에서 제공되며, TPU VM 또는 v6e GCE 인스턴스에서 실행할 수 있다.
• Apache 2.0 라이선스로 배포되며, 기여자 라이선스 협약(CLA) 서명이 필요한 오픈소스 프로젝트다.