구글이 PyTorch용 TPU 백엔드 오픈소스화
원제 Google open-sources TorchTPU, a native PyTorch back end for TPUs
195 포인트댓글 47
Key Point
PyTorch 사용자들이 NVIDIA GPU 대신 구글 TPU에서 직접 모델을 학습하고 추론할 수 있게 되어, 클라우드 인프라 선택지와 최적화 옵션이 확대됩니다.
핵심 요약
- 구글이 TorchTPU를 오픈소스로 공개했으며, 이는 PyTorch 모델과 연산을 구글 TPU 디바이스에서 직접 실행 가능하게 하는 네이티브 백엔드입니다.
- PyTorch ATen 커널 구현, XLA 컴파일 캐시 최적화, 데이터 병렬화와 텐서 병렬화를 포함한 분산 학습 지원, torch.compile() 통합 등의 기능을 제공합니다.
- Llama, Qwen, ResNet, DLRM, minGPT, UNet 등 다양한 모델의 예제가 포함되어 있으며, 디버깅과 벤치마킹을 위한 유틸리티도 제공됩니다.
- Python 3.11 이상에서 지원되며, PyPI를 통해 0.2.0 버전의 공개 릴리스를 pip로 설치하거나 소스에서 직접 빌드할 수 있습니다.
- 사용자 가이드는 https://torchtpu.readthedocs.io/에서 제공되며, Marimo 기반의 인터랙티브 노트북을 통해 TPU VM이나 v6e GCE 인스턴스에서 실행 가능합니다.
- 프로젝트는 Bazel을 사용해 빌드되며, Bazelisk 설치를 통해 의존성을 관리합니다.
- VSCode와 clangd 확장을 활용한 C++ 코드 네비게이션을 권장하며, 컴파일 데이터베이스를 자동으로 생성할 수 있습니다.
- Apache 2.0 라이선스로 배포되며, 기여자 라이선스 계약(CLA) 서명과 코드 리뷰 프로세스를 거쳐 기여를 받고 있습니다.