구글이 자사 AI 가속기인 TPU(Tensor Processing Unit) 위에서 PyTorch 모델과 연산을 실행할 수 있게 해주는 백엔드 'TorchTPU'를 오픈소스로 공개했다. 프로젝트 저장소는 google-pytorch 조직 아래에 있으며, 현재 PyPI를 통해 0.2.0 정식 릴리스가 배포돼 있다. 라이선스는 아파치 2.0이고, 기여를 위해서는 별도의 기여자 라이선스 계약(CLA) 서명과 코드 리뷰 과정을 거쳐야 한다.
TorchTPU의 핵심은 PyTorch 연산을 TPU 디바이스로 내려보내는 경로를 직접 구현했다는 점이다. 다양한 PyTorch 연산을 위한 커스텀 ATen 커널, 실행 최적화를 위한 컴파일 캐시, 그리고 분산 학습과 모델 추론 예제가 함께 들어 있다. ATen은 PyTorch의 텐서 연산이 실제로 호출되는 저수준 계층으로, 이 지점에 TPU용 커널을 끼워 넣었다는 것은 상위 모델 코드를 거의 바꾸지 않고도 TPU에서 돌리는 것을 목표로 한다는 의미다.
실행 환경과 사용 방법
이 백엔드를 쓰려면 TPU VM이 필요하며, 문서는 v6e GCE 인스턴스가 가장 손쉬운 선택지라고 안내한다. 즉 로컬 PC나 일반적인 GPU 서버에서 바로 체험할 수 있는 성격의 도구가 아니라, 구글 클라우드의 TPU 자원을 전제로 한다. 사용자 가이드는 torchtpu.readthedocs.io에 정적 문서 형태로 공개돼 있고, 원본은 Marimo 기반의 대화형 노트북으로 docs 폴더에 들어 있다. Marimo는 Jupyter와 유사한 노트북 환경으로, 'marimo edit 파일.py' 명령으로 실행한 뒤 브라우저에서 접근하는 방식이다.
파이썬 버전 정책은 PyTorch의 릴리스 호환성 매트릭스를 그대로 따르며, 현재 기준으로는 3.11 이상을 지원한다. 설치 시 주의할 점도 명시돼 있다. '--extra-index-url https://download.pytorch.org/whl/cpu' 옵션 없이 설치하면 CUDA 버전의 torch가 함께 깔리므로, TPU 환경에서는 CPU 빌드의 torch를 별도로 지정해 설치해야 한다는 것이다.
빌드 체계가 드러내는 설계 의도
빌드와 의존성 관리에는 Bazel을 사용한다. 설치는 Bazel을 직접 받는 대신 투명한 래퍼인 Bazelisk를 통해 이뤄지며, 저장소 루트에서 'bazel info'가 오류 없이 출력되면 설치가 정상임을 확인할 수 있다. OSS Bazel 빌드는 .bazelrc 설정에 따라 기본적으로 최적화 모드(-c opt)로 동작한다. fastbuild나 dbg 같은 비최적화 빌드는 성능을 떨어뜨릴 수 있어 휠(wheel) 빌드에는 쓰이지 않으며, C++ 디버깅 목적으로만 별도 플래그를 통해 허용된다.
휠 빌드에서 특히 눈에 띄는 대목은 '--config=wheel_common' 설정이다. 이 플래그가 없으면 빌드 자체가 실패하는데, 그 이유가 프로젝트의 내부 구조를 그대로 보여준다. 이 설정은 공유되는 XLA/MLIR 백엔드를 torch_version-reset 전환을 거쳐 하나의 libxla_base.so로 묶어낸다. 이 과정을 거치지 않으면 버전별 공통 모듈이 저마다 백엔드 사본을 담게 되고, 결과물은 'import torch' 시점에 중복 정적 등록 때문에 곧바로 중단된다. 여러 PyTorch 버전을 하나의 패키지 체계로 지탱하려다 보니 백엔드를 단일 공유 라이브러리로 강제하는 장치가 필요했던 셈이다.
빌드 설정에는 원격 빌드 실행(RBE)의 원격 캐시와 실행기도 엮여 있다. wheel_common이 이를 켜기 때문에 RBE 자격 증명이 없는 개발자는 '--config=no_rbe'를 덧붙여 해당 플래그를 걷어내야 한다. 이때 순서가 중요해, no_rbe는 반드시 wheel_common 뒤에 와야 한다. 두 설정이 같은 원격 캐시·실행기 플래그를 건드리므로 나중에 적용된 쪽이 최종값을 결정하기 때문이다. C++ 코드 탐색을 위해서는 clangd 확장을 쓰는 VSCode가 권장되며, Bazel 액션 그래프에서 compile_commands.json을 생성해 자동완성과 정의 이동 같은 기능을 쓸 수 있다. 다만 최초 생성에는 전체 빌드 때문에 30분 이상이 걸릴 수 있다.
실무자가 짚어야 할 지점
정리하면 TorchTPU는 그동안 사실상 JAX 중심으로 쓰이던 TPU 생태계에 PyTorch 사용자가 더 자연스럽게 진입할 수 있는 또 하나의 경로를 제시한다. 다만 아직 0.2.0이라는 초기 버전이고, TPU VM과 v6e 인스턴스라는 구글 클라우드 환경을 강하게 전제하며, CUDA torch가 기본으로 딸려 오는 설치 함정이나 민감한 빌드 플래그 순서처럼 운영상 주의할 지점이 적지 않다. 기존 PyTorch XLA 경로와의 관계나 구체적인 성능 수치는 공개된 저장소 설명만으로는 판단하기 어렵다. 따라서 당장 운영 파이프라인에 올리기보다는, TPU 도입을 검토 중인 팀이 호환성과 커널 지원 범위를 작은 규모로 가늠해 보는 단계에서 살펴볼 만하다.