Hugging Face, 네이티브 PyTorch 백엔드 TorchTPU로 Transformers와 Diffusers를 TPU 지원

기존에는 PyTorch를 TPU에서 실행하려면 PyTorch/XLA(torch_xla)를 사용해야 했다. 텐서가 지연 평가되기 때문에 코드가 예상한 시점에 실행되지 않거나, 작은 변경에도 그래프 재컴파일이 발생하고, 디버깅이 어려운 문제가 있었다고 한다. Hugging Face는 이를 보완하기 위해 Optimum TPU를 만들었지만, TorchTPU가 네이티브 지원을 제공하므로 Transformers와 Diffusers에서 Hugging Face 모델을 실행하는 데 Optimum TPU는 더 이상 필요하지 않다고 밝혔다.

TorchTPU는 PyTorch에 TPU를 새로운 백엔드로 추가한다. 일반적인 PyTorch 코드를 작성하기만 하면 되며, TPU에 맞게 조정된 Pallas 커널, 분산 학습 지원, 속도를 높이고 싶을 때 쓸 수 있는 torch.compile 연동이 함께 제공된다. Transformers에서는 `device_map="tpu"`를 전달하면 모델이 TPU에 직접 로드되며, 기본적으로 eager 모드로 동작한다. 예시로 Qwen/Qwen3.5-9B 코드가 제시되어 있다.

Diffusers에서는 TPU가 일급 장치가 되었다. 시험 제작과 디버깅에서는 eager 모드에서 `pipe.to("tpu")`만 호출하면 되고, 운영 환경에서는 컴파일 호출을 하나 추가하면 된다. 칩 하나에 들어가지 않는 대형 모델은 `TensorParallelConfig`를 from_pretrained에 전달하면 로드 시점에 칩 사이로 분할된다. Cloud TPU v6e-4(TP=4, bf16, 배치 1) 예시에서 Qwen-Image-2.1의 1024×1024, 40스텝 생성은 컴파일 시 3.5초, eager 시 8.0초가 걸렸다. MiniMax-H3의 동영상 및 음성 생성 예시는 컴파일 시 188초, eager 시 262초였다.

Kernels 허브에는 Pallas로 작성된 TPU용 커널이 있으며, Tokamax 라이브러리의 커널도 사용할 수 있다. Tokamax의 ragged-dot 커널을 Mixture-of-Experts 모델에 적용하면 배치가 클수록 효과가 커지며, 배치 크기 64에서 Qwen3-30B-A3B는 21.5배, Gemma 4 26B-A4B는 12.8배 빨라졌다.

품질 측면에서는 transformers 테스트 스위트를 TPU에서 실행하고 있으며, TPU에서 실행되는 테스트의 96.5%가 통과했다. 실패한 테스트는 재현 스크립트로 만들어 TorchTPU 저장소에 이슈로 보고했으며, 지금까지 34건을 보고해 19건이 수정되었다. TPU에서 실행하는 과정에서 StaticCache와 Qwen3의 텐서 병렬 플랜 등 Transformers 측 버그도 발견되어 상류(upstream)에서 수정했다. TPU용 nightly CI 추가도 진행 중이다.