구글이 PyTorch 모델을 TPU에서 실행하는 TorchTPU를 공개했다

구글이 PyTorch 모델과 연산을 텐서 프로세싱 유닛(TPU) 위에서 직접 돌리는 백엔드 TorchTPU를 오픈소스로 풀었다. 저장소는 google-pytorch/torch_tpu이며 아파치 2.0 라이선스를 따른다. 별도 프레임워크로 모델을 다시 작성하는 대신, PyTorch 코드를 TPU 디바이스에 그대로 올리는 것을 목표로 한다.

핵심은 두 층이다. 우선 여러 PyTorch 연산을 TPU에서 처리하기 위한 커스텀 ATen 커널을 직접 구현했다. 그 위로는 XLA 연산을 캐싱해 같은 그래프가 반복 실행될 때 컴파일 비용을 줄이는 컴파일 캐시가 붙는다. torch.compile()과도 연결돼 그래프 모드 실행 경로를 그대로 쓸 수 있다.

분산 학습 지원 범위도 함께 명시돼 있다. 데이터 병렬과 텐서 병렬 두 방식을 지원하며, Llama·Qwen·ResNet·DLRM·minGPT·UNet 등 규모와 성격이 다른 모델 예제를 저장소에 넣어 두었다. 디버깅과 벤치마킹을 위한 유틸리티도 같이 배포된다.

설치는 PyPI에 올라온 0.2.0 버전을 쓰면 된다. 다만 torch를 CPU 빌드로 받도록 별도 인덱스를 지정해야 하며, 지정하지 않으면 기본 CUDA 빌드가 끌려온다. 파이썬은 3.11 이상을 요구하고, 지원 정책은 PyTorch의 릴리스 호환 매트릭스를 그대로 따른다.

문서는 torchtpu.readthedocs.io에서 볼 수 있고, 인터랙티브 노트북은 Marimo로 작성돼 있다. 노트북을 열어 실행하려면 TPU VM이 필요하며, 안내에 따르면 v6e GCE 인스턴스가 가장 간편하다.

소스에서 휠을 직접 빌드할 계획이라면 Bazel(정확히는 Bazelisk) 기반 파이프라인을 이해해야 한다. wheel_common 설정이 필수이고, 이 설정이 공유 XLA/MLIR 백엔드를 하나의 libxla_base.so로 묶어 준다. 이 플래그 없이 빌드하면 버전별로 백엔드 사본이 따로 들어가 import torch 시점에 중복 등록으로 죽는다. 원격 빌드 자격 증명이 없다면 --config=no_rbe를 뒤에 붙여야 하는데, 순서가 바뀌면 적용되지 않는다.

TPU는 오랫동안 자체 스택 중심으로 쓰여 왔다. PyTorch가 사실상 표준 프레임워크로 자리 잡은 상황에서, TPU 쪽이 자체 백엔드를 내놓은 것은 하드웨어 선택지와 프레임워크 선택지를 분리하려는 시도로 읽힌다.

실무에서는 모델 코드를 유지한 채 학습 파이프라인의 디바이스 레이어만 바꾸는 그림이 가능해진다. torch.compile() 경로를 그대로 쓰고, 분산 학습은 저장소의 예제를 출발점으로 삼을 수 있다. 다만 아직 0.2.0이고, 모든 PyTorch 연산이 완비됐다는 언급은 원문에 없다. 커스텀 커널이 필요한 연산이 남아 있을 수 있으니 실제 모델에 붙이기 전에 벤치마킹 유틸리티로 먼저 확인하는 편이 안전하다.

전제 조건도 분명하다. 실행 환경은 TPU VM이어야 하고, 파이썬 3.11 이상이 필요하며, 휠을 직접 빌드할 경우 최적화 빌드가 기본값이라 디버깅용 비최적화 빌드는 별도 플래그를 줘야 한다. 기여하려면 CLA 서명과 코드 리뷰 절차를 거쳐야 한다.