Google, TPU용 네이티브 PyTorch 백엔드 TorchTPU 오픈소스로 공개

GitHub README 앞부분을 기준으로 옮겼어요. 설치 명령과 사용자 가이드 본문은 다 읽지 않았어요.

무엇인가

TorchTPU는 Google의 TPU(Tensor Processing Unit)용 PyTorch 백엔드예요. 사용자가 PyTorch 모델과 연산을 TPU 장치에서 실행하게 해 줘요. 여러 PyTorch 연산에 대한 자체 ATen 커널, 실행을 최적화하는 컴파일 캐시, 분산 학습과 모델 추론 예제가 들어 있어요.

기능

• TPU용 PyTorch ATen 커널 구현
• XLA 계산 실행을 최적화하는 컴파일 캐시
• 분산 학습 지원(데이터 병렬, 텐서 병렬)
• Llama, Qwen, ResNet, DLRM, minGPT, UNet 등 모델 예제
• 그래프 모드 실행을 위한 torch.compile() 연동
• 디버깅·벤치마킹 도구

문서와 노트북

사용자 가이드는 torchtpu.readthedocs.io에 있어요. 대화형 노트북은 Jupyter와 비슷한 Marimo 형식이라 marimo edit file.py로 열어 브라우저에서 써요. TPU VM에서 돌려야 하고, v6e GCE 인스턴스가 가장 쉽다고 안내해요.

설치

Python 지원 정책은 PyTorch 릴리스 호환 표를 따르고, 지금은 Python 3.11 이상을 지원해요. 가상 환경(venv 또는 uv)을 만든 뒤 pip로 설치하며, PyPI 최신 공개 릴리스는 0.2.0이에요. 저장소 커밋 수는 4,464개로, 공개 전부터 오래 개발해 온 프로젝트로 보여요.