Google이 개발자가 널리 쓰이는 AI 모델 개발·실행 도구인 PyTorch를 Google의 텐서 처리 장치(TPU)에서 실행할 수 있도록 하는 소프트웨어 계층 TorchTPU를 오픈소스로 공개했다. Meta의 PyTorch 팀과 함께 개발한 이 도구는 기존 작업을 처음부터 다시 작성하는 대신, 비교적 적은 코드 수정만으로 TPU 하드웨어로 옮길 수 있도록 설계됐다.
Google에 따르면 TorchTPU는 디버깅에 유용한 즉시 실행 방식과 프로덕션 작업을 위한 컴파일 실행을 모두 지원한다. 개발자는 장치 설정을 “cuda”에서 “tpu”로 바꾸고 TPU 분산 백엔드를 이용할 수 있다. Google은 PyTorch의 가속기 감지 API를 대안으로 제시하기도 한다. TorchTPU는 GitHub에서 이용할 수 있지만, 동적 입력 형태, 컴파일 속도, 연산자 지원 범위, 강화학습, 양자화 등은 아직 개발 로드맵에 남아 있다.
개발자에게 달라지는 점은?
이번 출시로 PyTorch 개발자들이 이미 사용하는 소프트웨어에서 TPU를 활용할 수 있게 됐다. Google은 TorchTPU가 서빙 도구인 vLLM과 SGLang, 학습 및 오케스트레이션 프로젝트인 TorchTitan, Ray, TorchRec과 호환된다고 밝혔다. Hugging Face는 Transformers v5에서 device_map="tpu" 설정을 선택해 TPU를 지정할 수 있다고 설명했다. Hugging Face의 Diffusers 라이브러리도 컴파일 모드와 여러 가속기에 작업을 분산하는 기능을 포함해 TPU 실행을 지원한다.
이는 Google이 과거 PyTorch를 TPU에서 실행하기 위해 제공한 PyTorch/XLA와는 다른 접근 방식이다. Google은 기존 시스템에서 코드 변경이 더 많이 필요했고, 컴파일 시간이 길었으며, PyTorch의 즉시 실행 모드를 지원하지 않았다고 밝혔다. InfoQ에 따르면 Google은 500개가 넘는 아키텍처에 걸쳐 약 5,000개의 Hugging Face 모델을 평가하는 테스트 파이프라인을 구축했다. 이는 폭넓은 테스트를 거친다는 뜻이지만, 모든 모델이나 작업이 수정 없이 작동한다고 보장하는 것은 아니다.
다른 하드웨어 선택지와 비교하면?
Google은 즉시 실행과 컴파일 실행을 결합한 TorchTPU의 DeferAndFuse 모드가 Strict Eager 모드보다 최소 50% 빠르면서 가속기 메모리의 최대 사용량도 줄일 수 있다고 밝혔다. 별도의 테스트에서 Google은 표준 PyTorch로 매개변수 300억 개의 Muse Glimmer 모델을 실행했다. InfoQ에 따르면 시퀀스 길이가 5,120일 때 TPUv7 칩 1개에서 추론 처리량은 초당 25,700토큰, 10,240일 때는 초당 22,400토큰이었다. TPUv7 칩 4개를 사용한 학습 처리량은 각각 초당 5,600토큰과 4,600토큰이었다. 이는 Google이 공개한 결과로, 독립적으로 검증된 비교 결과는 아니다.
TechCrunch는 TorchTPU를 같은 주 발표된 Amazon의 TorchNeuron과 함께 소개했다. TorchNeuron은 Amazon의 Trainium 칩에서 PyTorch 작업을 실행하기 위한 도구다. 두 프로젝트 모두 Nvidia GPU의 대안을 더 쉽게 사용할 수 있도록 하는 것이 목표다. TorchTPU의 출시는 PyTorch 사용자에게 또 다른 하드웨어 선택지를 제공하지만, 작업을 옮기면 비용이 낮아지거나 현재 구성보다 성능이 나아진다는 점을 입증하는 것은 아니다.
왜 중요한가
이미 PyTorch를 사용하는 개발자는 TorchTPU를 통해 작업을 처음부터 다시 만들지 않고도 Google TPU에서 시험하거나 배포하기가 쉬워질 수 있다. 다만 특정 모델에서 전환이 비용이나 속도 면에서 이득인지는 아직 분명하지 않으며, 인용된 성능 수치는 Google 자체 테스트 결과다.
