Hugging Face, TorchTPU 통합으로 PyTorch TPU 기본 지원

TPU 기본 지원의 시작

Hugging Face는 TorchTPU를 통합하여 PyTorch 사용자가 Google의 텐서 처리 장치(TPU)를 기본 장치로 활용할 수 있도록 지원한다. 이로써 TPU는 PyTorch의 또 다른 장치 백엔드로 추가되며, Hugging Face 생태계 내에서 TPU를 직접 지원하는 길이 열렸다. Transformers, Diffusers, Kernels와 같은 주요 라이브러리들이 이제 TPU에서 기본적으로 실행된다. 사용자는 허브에서 모델을 선택하고 장치 맵을 "tpu"로 설정하는 것만으로 TPU를 활용할 수 있다.

TPU의 역할과 기존 접근 방식

TPU는 Google이 자체 개발한 AI 가속기로, 신경망 연산의 대부분을 차지하는 행렬 계산에 특화된 칩이다. Google은 약 10년간 TPU를 설계해왔으며, Gemini를 포함한 자체 모델 학습 및 서비스에 이를 활용한다. TPU는 Google Cloud에서 단일 칩부터 수천 개의 칩으로 구성된 포드(pod) 형태로 제공되며, 대규모 학습이나 서비스에 필요한 컴퓨팅 파워를 제공할 때 강점을 보인다.

그동안 PyTorch 사용자가 TPU를 활용하려면 PyTorch/XLA(torch_xla)를 통해 코드를 포팅해야 했다. 이 방식은 작동했지만, 텐서가 지연 실행되는 특성, 작은 변경에도 그래프 재컴파일이 발생하는 문제, 그리고 디버깅의 어려움과 같은 제약이 있었다. Hugging Face는 이러한 불편을 해소하기 위해 Optimum TPU를 개발했으나, TorchTPU의 기본 지원으로 이제 Optimum TPU는 더 이상 필요하지 않다.

TorchTPU의 주요 기능과 활용

TorchTPU는 PyTorch에 TPU를 새로운 백엔드로 추가한다. 이를 통해 사용자는 일반적인 PyTorch 코드를 작성하여 TPU에서 실행할 수 있다. TorchTPU는 TPU에 최적화된 Pallas 커널을 포함하며, 분산 학습을 지원하고, 더 높은 성능을 위해 torch.compile과 통합된다. TPU VM에 TorchTPU를 설치할 때는 특정 인덱스 URL을 지정하는 설치 명령을 사용한다. 이 명령은 불필요한 CUDA 빌드 대신 CPU 빌드를 가져오도록 한다.

Hugging Face 라이브러리 통합

Hugging Face의 주요 라이브러리들은 TorchTPU 통합을 통해 TPU 지원을 강화했다.

Kernels Hub는 즉시 사용할 수 있는 TPU 커널을 제공한다. Kernels 라이브러리를 사용하여 커널을 로드하면 모델의 속도를 향상시킬 수 있다. TPU 커널은 하드웨어에 대한 저수준 제어를 제공하는 Pallas로 작성된다. Tokamax 라이브러리의 커널도 활용할 수 있다. 예를 들어, Tokamax의 ragged-dot 커널을 Mixture-of-Experts 모델에 적용할 수 있다.

성능 최적화 및 분산 학습

Diffusers는 단일 칩에 너무 큰 모델의 경우 TensorParallelConfig를 from_pretrained에 전달하여 모델 로드 시 트랜스포머를 여러 칩에 분할할 수 있다. 이를 통해 전체 모델이 한 곳에 들어갈 필요 없이 분산 처리된다. 많은 인기 파이프라인이 이 기능을 기본적으로 지원하며, 그렇지 않은 경우에도 tp_plan을 통해 지원할 수 있다.

성능 측정 결과는 다음과 같다 (Cloud TPU v6e-4, TP=4, bf16, 배치 1 기준): MiniMax-H3 (비디오 + 오디오) 모델은 1344x768 해상도, 49단계에서 컴파일 모드 188초, 이저 모드 262초가 소요되었다. Qwen-Image-2.1 모델은 1024x1024 해상도, 40단계에서 컴파일 모드 3.5초, 이저 모드 8.0초가 소요되었다. Qwen-Image-2.1 모델은 2048x2048 해상도, 40단계에서도 실행 가능하다.

Kernels의 경우, TPU v6e-8에서 배치 크기 64일 때 디코딩 처리량에서 큰 개선을 보였다. Qwen3-30B-A3B 모델은 기본 구현 대비 21.5배, Gemma 4 26B-A4B 모델은 12.8배 더 빨라졌다.

안정성 확보를 위한 노력

Hugging Face는 Transformers 라이브러리 테스트 스위트를 지속적으로 실행하여 TPU 지원을 개선하고 있다. TPU에서 야간 CI(Continuous Integration)를 추가하고 그 결과를 Transformers 대시보드에서 제공하기 위해 노력 중이다. 현재 TPU에서 실행되는 테스트 중 96.5%가 통과한다. 발견된 모든 실패는 TorchTPU 저장소에 재현 스크립트와 이슈로 보고되며, 현재까지 34개의 이슈가 보고되었고 그 중 19개가 이미 해결되었다. TPU에서 테스트를 실행하는 과정에서 Transformers 자체의 StaticCache 및 Qwen3 텐서 병렬 계획과 같은 몇 가지 버그도 발견되어 상위 프로젝트에 수정 사항이 반영되었다.

개발자 활용 및 향후 전망

이제 TPU는 PyTorch 사용자에게 더 이상 부가적인 작업이 아니다. TorchTPU를 설치하고 Hugging Face Hub에서 모델을 가져와 "tpu" 장치로 보내는 것만으로 TPU의 강력한 컴퓨팅 성능을 활용할 수 있다. Hugging Face는 개발자들이 TorchTPU를 통해 어떤 혁신적인 애플리케이션을 구축할지 기대하고 있으며, 문제 발생 시 이슈 보고를 통해 지속적인 개선을 약속한다.