Hugging Face, 네이티브 PyTorch TPU 백엔드 TorchTPU로 TPU 활용 간소화 및 성능 향상

TorchTPU는 PyTorch용 네이티브 TPU 백엔드로 TPU를 PyTorch의 새로운 디바이스로 추가해 기존에 복잡했던 PyTorch/XLA나 JAX를 거치지 않고도 TPU에서 모델을 직접 구동할 수 있도록 함.

Hugging Face 생태계 내 Transformers, Diffusers, Kernels 라이브러리들이 TPU를 네이티브로 지원하여 해당 라이브러리들을 TPU에서 바로 실행할 수 있게 됨으로써 TPU 사용 환경이 크게 확장되었어요.

추가로 TorchTPU는 TPUs를 위해 최적화된 Pallas 커널, 분산 학습 지원, 그리고 torch.compile과의 통합을 제공해 TPU 상에서의 성능 향상과 효율적인 분산 학습이 가능해요.

참고 자료

• https://huggingface.co/blog/hf-torchtpu

이미지: AI 생성 일러스트.