TorchTPU PyTorch TPU 가이드: 2026년 설치와 마이그레이션 방법

PyTorch 코드를 cuda에서 tpu로 옮기는 작업 자체는 생각보다 간단합니다. 하지만 TPU에 제대로 접근할 수 있는지, 어떤 패키지를 설치했는지, 컴파일 방식과 텐서 shape가 적합한지가 실제 결과를 좌우합니다.

TorchTPU가 실제로 제공하는 것

TorchTPU는 Google TPU를 위한 PyTorch 네이티브 백엔드입니다. 호스팅 모델 API도, TPU VM을 대체하는 솔루션도 아닙니다. PyTorch 발표 자료에 따르면 eager 실행, torch.compile, 분산 프리미티브, 약 450개의 ATen 연산을 지원합니다. 공개 저장소에는 ATen 커널, 컴파일 캐시, 분산 실행 예제, torch.compile 통합 내용이 정리돼 있습니다.

마이그레이션 방식은 익숙한 PyTorch 코드와 크게 다르지 않습니다.

# Before
model.to("cuda")

# First TPU port
model.to("tpu")

다만 이 정도의 변경은 어디까지나 이식 가능성을 확인하는 수준입니다. 성능이 자동으로 보장되는 것은 아닙니다. PyTorch 발표 자료도 TPU에서 최고 성능을 내려면 모델 shape, 샤딩 방식, 하드웨어에 맞는 연산 구성이 중요하다고 설명합니다.

설치 전에 준비할 것

TorchTPU는 Google TPU가 할당된 환경에서만 제대로 활용할 수 있습니다. 저장소에서는 TPU VM에서 노트북을 실행하는 방식을 권장하며, 시작 환경으로는 v6e Google Compute Engine 인스턴스가 가장 간단하다고 안내합니다. CPU 전용 노트북에서도 패키지 설치는 가능하지만, TPU 실행 여부를 확인할 수는 없습니다.

공개 README는 Python 3.11 이상과 torch_tpu==0.2.0을 기준으로 합니다. 실험 환경이 CUDA 빌드나 다른 프로젝트의 의존성을 물려받지 않도록 별도의 가상 환경을 사용하는 편이 좋습니다.

CUDA 패키지를 끌어오지 않고 `torch_tpu` 설치하기

먼저 가상 환경을 만든 다음, 고정된 TorchTPU 패키지와 CPU용 PyTorch wheel을 함께 설치합니다.

python -m venv .venv
source .venv/bin/activate

pip install torch_tpu==0.2.0 \
  --extra-index-url https://download.pytorch.org/whl/cpu

README에 따르면 추가 인덱스를 생략하면 CUDA용 PyTorch wheel이 설치될 수 있습니다. TPU VM에서는 적합하지 않은 의존성이 들어오는 셈입니다.

첫 마이그레이션이라면 공식 릴리스 wheel을 사용하는 것이 좋습니다. 소스에서 빌드하려면 Bazel 또는 Bazelisk가 필요하며, 저장소에 명시된 --config=wheel_common과 필요한 경우 --config=no_rbe 플래그를 사용해야 합니다.

간단한 PyTorch 프로그램으로 백엔드 확인하기

다음 코드는 TPU VM에서 실행합니다.

import torch

x = torch.randn(2, 2, device="tpu")
y = torch.randn(2, 2, device="tpu")
z = x @ y

print(z)
print(z.device)

결과가 정상적으로 출력되고 디바이스가 tpu로 표시된다면 텐서 할당과 행렬 곱셈이 TPU 백엔드에서 실행된 것입니다. 디바이스를 사용할 수 없다면 모델 코드를 수정하기 전에 VM 런타임과 패키지 설치 상태부터 확인해야 합니다.

TorchTPU는 디버깅에 적합한 eager 실행과 컴파일 실행을 위한 torch.compile을 모두 지원합니다. PyTorch 발표 자료는 프로젝트 테스트에서 DeferAndFuse 모드가 strict eager 모드보다 50% 이상 빠를 수 있다고 소개합니다. 다만 이는 특정 프로젝트의 결과이므로 모든 모델에 적용되는 일반적인 벤치마크로 받아들여서는 안 됩니다.

실제 PyTorch 워크로드 옮기기

처음부터 TPU 전용 분산 설정까지 한꺼번에 바꾸기보다는, 가장 작은 변경으로 단일 디바이스 실행을 확인한 뒤 범위를 넓히는 방식이 안전합니다.

디바이스 선택을 torch.device("tpu")로 바꾸거나, 지원되는 환경에서는 torch.accelerator.current_accelerator()를 사용합니다.
모델과 입력 텐서를 같은 디바이스로 이동합니다.
컴파일이나 분산 실행을 활성화하기 전에 짧은 배치 하나를 실행합니다.
측정 중에는 배치 크기와 텐서 shape를 일정하게 유지합니다. shape가 바뀌면 컴파일 작업이 다시 발생할 수 있습니다.
대규모 마이그레이션을 결정하기 전에 대표 모델 경로와 실제 사용할 분산 패턴을 테스트합니다.

분산 CUDA 코드는 보통 NCCL을 사용합니다. 반면 TorchTPU에서 공개한 예제는 TPU용 프로세스 그룹 백엔드를 사용합니다.

import torch
import torch.distributed as dist

device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")
model = model.to(device)

정확한 분산 실행 명령은 TPU 환경과 워크로드에 따라 달라집니다. 디바이스 문자열만 바꾼다고 샤딩, 프로세스 배치, 입력 파이프라인까지 자동으로 구성되는 것은 아닙니다.

첫 실행과 이후 실행의 결과가 다를 수 있다는 점도 염두에 둬야 합니다. TorchTPU는 반복되는 그래프를 재사용할 수 있는 컴파일 캐시를 제공하므로, 워밍업 시간과 안정 상태 처리량을 분리해 기록하는 것이 좋습니다.

TorchTPU가 준비된 영역과 아직 다듬어지는 영역

Google은 vLLM, SGLang, TorchTitan, Hugging Face 도구와 TorchTPU의 통합을 문서화하고 있습니다. 추론 관련 발표에서는 KV 캐시 이동과 서빙 기능을 포함한 vLLM 및 SGLang의 네이티브 TPU 경로를 소개합니다. Google의 Cloud 발표는 TorchTitan에서 Qwen 3 8B가 52% MFU를 기록한 결과와 실제 운영 사례를 언급합니다.

다만 발표 자료의 인상적인 수치와 별개로, 실제 사용 시 확인해야 할 제약은 여전히 남아 있습니다.

영역	자료를 통해 확인되는 내용	직접 확인할 사항
설치	Python 3.11 이상, `torch_tpu==0.2.0`, CPU용 PyTorch 인덱스	사용 중인 VM 이미지와 wheel이 문제없이 의존성을 해결하는지
이식성	Eager 모드, `torch.compile`, `tpu` 디바이스, 분산 지원	지원되지 않는 연산과 모델별 폴백 동작
성능	약 450개의 ATen 연산과 TorchTitan의 Qwen 3 8B에서 프로젝트가 보고한 52% MFU 결과	사용할 모델의 워밍업 시간, shape, 배치 크기, 비용
개발 워크플로	공개 예제, 노트북, 디버깅 및 벤치마크 유틸리티	TPU 접근성, Bazel의 복잡성, 이용 가능한 관측성 도구
현재의 공백	기능 지원 범위가 계속 확장되는 단계	현재 워크로드에 필요한 기능이 아직 문서화되지 않았거나 지원되지 않는지

TorchTPU 도입을 결정하는 기준

이미 Google TPU를 사용할 수 있고 PyTorch 모델을 운영하며, 학습이나 추론에서 공통 프레임워크 경로를 원한다면 TorchTPU를 지금 테스트해볼 가치는 있습니다. 초기 도입에 가장 적합한 팀은 한 줄짜리 포팅만으로 GPU와 동일한 효율을 기대하기보다, 모델 shape를 프로파일링하고 조정할 준비가 된 팀입니다.

반대로 성숙한 호환성 매트릭스, 예측 가능한 콜드 스타트 동작, 또는 현재 스택에 꼭 필요한데 TorchTPU에서 아직 노출하지 않은 기능이 필요하다면 핵심 워크로드의 이전을 서두르지 않는 편이 좋습니다. 모델 하나, TPU 슬라이스 하나, 고정된 평가 세트 하나로 변경을 시작하고, 설치 시간·첫 컴파일·안정 상태 처리량·출력 정확도를 각각 측정하세요.

TorchTPU FAQ

TorchTPU는 PyTorch/XLA와 같은 것인가요?

아닙니다. TorchTPU는 eager 실행과 torch.compile을 중심으로 설계된 새로운 PyTorch 네이티브 백엔드입니다. PyTorch/XLA는 기존 TPU 통합 경로이므로, XLA 기반 코드를 옮길 때는 별도의 마이그레이션 검토가 필요할 수 있습니다.

TPU가 없는 노트북에도 TorchTPU를 설치할 수 있나요?

패키지 설치 자체는 가능하지만, 호환되는 Google TPU 인프라 없이는 TPU 실행을 검증할 수 없습니다. 공식 저장소의 실제 설정은 TPU VM을 전제로 합니다.

첫 TorchTPU 스텝이 느리게 실행되는 이유는 무엇인가요?

첫 실행에서는 그래프를 컴파일하고 캐시를 채우는 작업이 발생할 수 있습니다. 이후 반복 실행을 따로 측정하고, 재컴파일 여부를 확인하는 동안에는 shape를 일정하게 유지하세요.

`cuda`를 `tpu`로 바꾸면 같은 성능이 보장되나요?

아닙니다. 디바이스 교체는 우선 코드가 실행되는지 확인하는 이식 작업일 뿐입니다. 성능을 비교하려면 shape, 샤딩, 연산 지원 범위, 워밍업 시간을 모두 프로파일링해야 합니다.

TorchTPU는 프로덕션에 사용할 준비가 됐나요?

Google은 실제 운영 환경에서의 추론 및 학습 사례를 소개하고 있습니다. 하지만 공개 생태계의 구성 요소와 기능은 각각 성숙도가 다릅니다. 도입을 확정하기 전에 사용하려는 모델, 서빙 엔진, 분산 패턴을 그대로 테스트해야 합니다.

TPU VM에 고정 버전 wheel을 설치하고 2×2 검증 코드를 실행한 뒤, 현재 GPU 환경과 워밍업이 끝난 동일한 워크로드를 비교해보세요. 이렇게 하면 코드 이식성, 컴파일 워밍업, 연산 지원 범위, 하드웨어 비용을 서로 분리해 판단할 수 있습니다.