아마존에 이어 구글도 기존 AI 개발 도구를 자사 칩에서 쉽게 활용할 수 있도록 소프트웨어 생태계 확장에 나섰다. 개발자들이 익숙한 파이토치(PyTorch) 코드를 크게 손보지 않고도 자체 칩에서 실행할 수 있도록 지원해, 엔비디아 GPU와 CUDA 중심의 개발 환경에서 하드웨어 선택지를 넓히려는 움직임이다.
구글은 20일(현지시간) 열린 ‘파이토치 컨퍼런스 2026’에서 TPU용 파이토치 소프트웨어 스택인 ‘토치TPU(TorchTPU)’를 오픈소스로 공개했다.
아마존도 하루 앞서 자체 AI 칩 트레이니엄(Trainium)에서 기존 파이토치 코드를 실행할 수 있는 ‘토치뉴런(TorchNeuron)’을 발표했다. 아마존은 설정 하나만 바꾸면 코드를 다시 작성하지 않고 트레이니엄을 사용할 수 있다고 설명했으며, 20일부터 공개 베타로 제공한다.
파이토치는 AI 모델의 개발과 학습에 널리 사용되는 오픈소스 프레임워크다. 그동안 엔비디아 GPU와 CUDA에 맞춰 만들어진 코드와 도구가 많아 다른 칩으로 옮기려면 추가 작업이 필요했다. 두 회사는 전환에 따르는 부담을 낮춰 자사 칩의 활용 범위를 넓히려 한다.
메타의 파이토치 개발팀과 함께 만든 토치TPU는 기존 파이토치 API를 그대로 유지하면서 TPU에서 모델 학습과 추론, 배포를 지원한다. 개발자는 코드를 최소한으로 수정해 기존 애플리케이션을 TPU에서 실행할 수 있다.
일반 실행 모드에서 모델 오류를 확인하고, 이후 ‘torch.compile’을 적용해 성능을 최적화할 수 있다. 분산 학습 기능과 진단 도구도 기본으로 제공한다.
구글은 모델 실행에 그치지 않고 학습과 추론에 활용되는 주요 도구도 TPU에 연결했다. 허깅페이스와 협력해 트랜스포머스(Transformers) v5에서 ‘device_map="tpu"’ 설정을 이용하면 모델을 TPU에 배치할 수 있도록 했다.
이미지 생성 모델 개발 도구인 디퓨저스(Diffusers)는 TPU에서 일반 실행 모드와 컴파일 최적화를 지원한다. 여러 칩에 모델을 나눠 실행하는 텐서 병렬 처리도 가능하다. 허깅페이스 모델을 TPU 인프라에서 지속적으로 테스트하는 체계도 마련했다.
대규모 모델 추론 도구인 ‘vLLM’과 ‘SGLang’의 TPU 지원도 강화했다. 인퍼랙트(Inferact), 라딕스아크(RadixArk)와 각각 협력해 TPU용 백엔드를 구축했다. 기존 스케줄러와 API 서버 등을 유지하면서 TPU에 맞춰 연산을 최적화하도록 했다.
여러 서버에 추론 작업을 분산하는 쿠버네티스 기반 시스템 ‘llm-d’도 1.0 버전으로 공개했다. 요청 처리에 필요한 문맥이 이미 저장된 서버로 작업을 전달하고, 입력을 처리하는 프리필(prefill)과 답변을 생성하는 디코딩(decode)을 분리해 효율을 높인다. 구글 클라우드의 쿠버네티스 엔진(GKE)에서는 공통 인터페이스를 통해 GPU와 TPU를 사용할 수 있다.
학습 분야에서는 대규모 분산 학습 프레임워크 ‘토치타이탄(TorchTitan)’에 TPU 지원을 통합했다. 구글은 ‘큐원 3 8B’ 모델을 학습할 때 칩의 이론상 연산 성능을 실제 학습에 얼마나 활용하는지를 나타내는 모델 연산 활용률(MFU)이 최대 52%에 이르렀다고 밝혔다. 또 ‘아이언우드(Ironwood)’ TPU를 8192개까지 확장했을 때 칩 수 증가에 거의 비례해 성능이 향상됐다고 설명했다.
강화학습 프레임워크 ‘VeRL’에도 TPU 백엔드와 참고용 학습 구성을 추가했다. ‘레이(Ray)’를 이용하면 데이터 처리와 학습, 강화학습, 추론 등 각 단계에 CPU와 GPU, TPU를 배정해 하나의 작업 흐름 안에서 함께 사용할 수 있다.
개발자의 이전 작업을 지원하는 ‘TPU 에이전트 스위트(TPU Agent Suite)’는 깃허브에서 프리뷰 버전으로 제공된다. 플러그인 방식으로 개발 환경에 연결해 초기 설정과 오류 진단, 사용 관련 질문에 대한 안내, 추론 성능 조정을 돕는다.
아울러 온디바이스 AI 추론 도구 ‘라이트RT(LiteRT)’와 파이토치 모델을 연결하는 ‘라이트RT-토치(LiteRT-Torch)’도 공개했다. 파이토치 모델을 변환해 안드로이드와 iOS, 윈도, 맥OS, 리눅스, 웹, 사물인터넷(IoT) 기기에서 실행할 수 있도록 지원하며, 기기의 CPU와 GPU, 신경망처리장치(NPU)를 활용할 수 있다.