Google, TorchTPU 오픈소스로 공개해 TPU에서 네이티브 PyTorch 지원

Google은 2026년 10월 20일 PyTorch Conference 2026에서 자사의 Tensor Processing Unit(TPU)용 네이티브 PyTorch 스택인 TorchTPU를 GitHub에 오픈소스로 공개했다고 Google Cloud 블로그에서 밝혔다. 이 소프트웨어는 Meta의 PyTorch 팀과 공동 개발했으며, 기존 PyTorch 워크로드를 코드 변경을 최소화해 TPU에서 실행하도록 설계됐다.

Google에 따르면 이번 출시로 PyTorch에서 TPU도 GPU, CPU와 함께 네이티브 옵션으로 제공되며, vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face, LiteRT 등의 도구를 통해 학습, 서빙, 배포를 지원한다. 개발자는 이거 모드에서 모델을 디버깅한 뒤 torch.compile을 사용해 TPU 활용률을 높이고 메모리 오버헤드를 줄여 프로덕션 수준의 성능을 구현할 수 있다.

Meta의 PyTorch 엔지니어링 디렉터인 Jana van Greunen은 발표에서 TorchTPU를 통해 Meta가 핵심 머신러닝 시스템 팀과 함께 하드웨어를 유연하게 공동 개발하면서 주요 LLM, 광고 및 추천 시스템 워크로드를 지원할 수 있다고 밝혔다. 그는 “이를 통해 PyTorch 워크로드를 Google TPU에 배포하고, 엔지니어링 팀에 직접 피드백을 제공하며, 모두에게 필요한 생태계의 유연성을 구축할 수 있습니다”라고 말했다.

허깅 및 서빙 통합

Google은 Hugging Face와 협력해 Hugging Face Hub의 오픈 웨이트 모델을 TorchTPU에서 네이티브로 실행할 수 있도록 했다. Transformers v5에서는 인수를 전달하는 방식으로 TPU에 모델을 배포하며 device_map="tpu"Diffusers 라이브러리도 TPU에서 네이티브로 실행되며, 이거 모드, torch.compile, 여러 가속기에서의 텐서 병렬화를 지원한다. Google은 TPU 인프라에서 Hugging Face Transformers 제품군을 매일 밤 지속적으로 테스트하기 위한 전용 CI/CD 파이프라인도 구축했다고 밝혔다. Hugging Face의 제품 담당 부사장인 Jeff Boudier는 이번 협력이 하드웨어 선택권을 보장하겠다는 플랫폼의 약속을 실현한다고 말했다.

서빙 분야에서는 Inferact와 협력해 TorchTPU용 네이티브 PyTorch vLLM 백엔드를 개발했다. Google에 따르면 이 백엔드는 맞춤형 Pallas 커널을 사용해 TPU 하드웨어에서 연속 배칭, 양자화, 멀티칩 병렬 처리를 지원하면서도 업스트림 작업자 모델, 스케줄러, API 서버는 그대로 유지하므로 기존 추론 스택을 변경할 필요가 없다. Inferact의 CEO 겸 공동 창업자인 Simon Mo는 이번 통합으로 TPU에서 vLLM을 일급 지원하게 됐다고 말했다.

RadixArk와 함께 개발한 SGLang용 네이티브 PyTorch 백엔드는 구조화된 워크로드와 에이전트형 워크로드를 겨냥한다. 발표에 따르면 SGLang은 RadixAttention을 사용해 요청 간 KV 캐시를 재사용하므로 프롬프트 접두부가 반복되더라도 다시 계산할 필요가 없다. TorchTPU 버전은 HiCache의 I/O 백엔드 역할도 하는 비분리형 전송 백엔드를 사용한다. RadixArk의 CTO 겸 공동 창업자인 Banghua Zhu는 여러 종류의 실리콘에서 SGLang 지원을 확대하는 과정에서 TorchTPU가 회사의 작업을 가속했다고 말했다.

Google은 Red Hat, IBM, 500명이 넘는 기여자와 25개 조직이 참여하는 커뮤니티와 함께 llm-d 1.0도 Cloud Native Computing Foundation 샌드박스 프로젝트로 공개한다. 발표에서는 llm-d가 PyTorch 서빙 생태계와 Kubernetes 기반 프로덕션 배포를 연결한다고 설명한다. 라우터는 접두부 캐시를 고려한 라우팅을 사용해 각 요청의 컨텍스트를 이미 보유한 복제본으로 요청을 전달하고, 연산량이 많은 프리필과 메모리 사용량이 많은 디코드를 분리한다. 버전 1.0에는 멀티테넌트 흐름 제어, 라우터 고가용성, 안정화된 v1 API, 에이전트를 위한 세션 인식 라우팅이 추가된다. 또한 Google Kubernetes Engine에서는 공통 llm-d 엔드포인트를 GPU 또는 TPU에서 실행할 수 있다.

Zenteiq의 설립자인 Sashikumaar Ganesan 교수는 발표에서 llm-d 기반 GKE Inference Gateway를 사용해 회사의 합성 데이터 생성 추론 처리량이 총합 기준 4배를 넘게 증가했으며, TPU v6e 가속기 1,000개당 하루 생성 토큰 용량이 14억 7천만 개에서 59억 7천만 개로 늘었다고 밝혔다.

학습 및 오케스트레이션

분산 학습을 위한 PyTorch 네이티브 참조 구현인 TorchTitan에는 이제 Google이 프로덕션 등급이라고 설명하는 TorchTPU 통합 기능이 포함된다. 이 기능은 torch.compile, TorchTitan의 SimpleFSDP, 특화된 Pallas 커널, 그래프 수준 최적화를 기반으로 구축됐다. TorchTitan은 Llama 3 70B 및 405B, Qwen 3 8B 및 32B, MuseGlimmer 30B 같은 밀집 모델의 사전 학습과 Qwen 3 235B 전문가 혼합 모델, FLUX.1 12B 확산 트랜스포머를 위한 TorchTPU 레시피를 제공한다. Google은 이 통합 기능이 최대 8,000개의 Ironwood TPU(TPU v7x)에서 검증됐으며, TorchTitan을 통한 TorchTPU 학습에서 Qwen 3 8B의 모델 부동소수점 연산 활용률이 최대 52%에 달하고, 최대 8,192개의 Ironwood TPU에서 거의 선형적인 확장성을 보였다고 밝혔다.

Google은 강화학습 프레임워크인 VeRL에 TPU 백엔드와 엔드투엔드 참조 레시피도 업스트림으로 반영했다. 분리형 트레이너-샘플러 아키텍처를 기반으로 하고 수학적 추론 데이터셋에서 시연한 GRPO 레시피는 개발자가 기존의 하드웨어 비종속 모델 및 알고리즘 코드를 유지하면서 VeRL 워크로드를 TPU에서 실행할 수 있도록 설계됐다.

Ray는 단일 워크플로 내에서 CPU, GPU, TPU 전반에 걸쳐 PyTorch 데이터 처리, 학습, 강화 학습 및 서빙을 조율할 수 있도록 네이티브 TorchTPU 지원을 추가합니다. Anyscale의 수석 아키텍트인 Edward Oakes는 양사가 TPU를 Ray의 핵심 가속기로 만들기 위해 협력했다며 다음과 같이 말했습니다. “TorchTPU는 업계 표준 스택에 Google TPU를 직접 통합합니다.”

기존 아키텍처와 제공 현황

TorchTPU는 2026년 4월 7일, PyTorch의 PrivateUse1 인터페이스를 기반으로 하고 공유 컴파일 캐시를 지원하는 Debug, Strict, Fused 즉시 실행 모드를 갖춘 “Eager First” 설계를 설명한 Google Developers 블로그 게시물을 통해 공개적으로 처음 소개됐습니다. 당시 Google은 Fused Eager가 Strict Eager보다 50%에서 100% 이상의 성능 향상을 제공하며, TorchTPU에서 torch.compile을 사용하면 Torch Dynamo로 FX 그래프를 캡처하고 XLA와 StableHLO를 통해 하위 단계로 변환하며, 사용자 지정 커널은 Pallas와 JAX로 작성한다고 밝혔습니다. 게시물에 따르면 이 스택은 Distributed Data Parallel, FSDPv2, DTensor를 기본 지원하며, 랭크 간 실행이 달라지는 상황도 처리할 수 있도록 설계됐습니다. 또한 공개 GitHub 저장소의 출시도 TorchTPU의 2026년 로드맵 항목에 포함됐습니다.

torch_tpu 저장소는 이제 google-pytorch GitHub 조직에서 공개되어 있으며 Apache 2.0 라이선스를 따릅니다. README에는 사용자 지정 ATen 커널, XLA 연산용 컴파일 캐시, 데이터 병렬 처리 및 텐서 병렬 처리를 위한 분산 학습 지원, torch.compile 통합, 디버깅 및 벤치마킹 유틸리티와 Llama, Qwen, ResNet, DLRM, minGPT, UNet 예제가 나열되어 있습니다. 이 프로젝트는 Python 3.11 이상을 지원하며, torch-tpu 0.2.0 패키지는 Linux x86-64용 CPython 3.11~3.14 휠과 함께 2026년 10월 19일 PyPI에 출시됐습니다. torchtpu.readthedocs.io에서 사용자 가이드를 제공하며, README에는 프로젝트의 대화형 노트북을 TPU VM에서 실행해야 하고 v6e GCE 인스턴스가 가장 간편한 선택지라고 설명되어 있습니다.

같은 발표에서 Google은 이전에 TensorFlow Lite로 불렸던 LiteRT가 PyTorch 생태계에 합류한다고 밝혔습니다. LiteRT-Torch 라이브러리는 PyTorch 모델을 단 한 번의 litert-torch.convert() 호출로 변환하며, 변환된 모델은 Android, iOS, macOS, Windows, Linux, 웹 및 IoT 기기에서 실행됩니다. PyTorch 모델의 온보딩, 진단 및 추론 성능 튜닝을 자동화하기 위한 TPU 에이전트 도구 모음은 GitHub에서 초기 프리뷰로 제공되고 있습니다.

발표에서는 개발자가 TPU에서 PyTorch를 실행하려면 Google Cloud의 TPU 문서를 참고하도록 안내합니다.