Hugging Face × TorchTPU: Transformers·Diffusers 모델을 `device_map="tpu"` 한 줄로 TPU에서
무엇인가요
Hugging Face가 PyTorch용 TPU 백엔드 TorchTPU 위에서 Transformers, Diffusers, Kernels 라이브러리를 돌리는 방법을 공개했어요. 예전 PyTorch/XLA 방식이나 Optimum TPU 없이 TPU를 그냥 하나의 장치로 쓰는 방식이에요. TPU VM에서 설치한 뒤 모델을 장치에 올리면 돼요.
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu모델은 device_map="tpu"로, 입력은 .to("tpu")로 옮겨요. 기본은 즉시 실행이고 torch.compile로 더 빠르게 돌릴 수 있어요.
왜 중요한가요
TPU를 쓰려면 XLA 전용 코드를 따로 짜야 해서 GPU 코드를 옮기기 부담스러웠는데, 그 장벽이 꽤 낮아졌어요. Hugging Face 측정으로 Qwen-Image-2.1 1024×1024 생성은 컴파일하면 3.5초, 즉시 실행은 8.0초였고업체 주장, MoE 모델 디코딩은 전용 커널로 최대 21.5배 빨라졌다고 해요업체 주장. GPU 수급이 빠듯한 팀이 대안 장비를 검토할 때 도움이 돼요.
완성도
Beta 베타베타 수준이에요. Transformers 테스트 중 TPU에서 도는 것의 96.5%가 통과하고업체 주장 야간 CI는 아직 다 초록이 아니라고 글에서 밝혔어요.
앞으로 이렇게 지켜보세요
- 이슈 추적: TorchTPU 저장소에서 보고된 34건 중 아직 열린 이슈와 수정 속도를 봐요.
- CI: Hugging Face가 야간 TPU CI를 완성했다고 알리는지 확인해요.
- 실측: 내 모델을 TPU v6e에서 GPU와 같은 조건으로 돌려 보고 가격 대비 성능을 비교해요.