PyTorch 모델을 device="tpu"로 실행하고 에이전트로 튜닝하기: Google TorchTPU와 TPU Agent Suite 해부

TPU 실리콘 가속기 위에서 PyTorch 연산 그래프와 네 종류의 전문 에이전트가 상태·수치 정합성을 점검하는 모습을 표현한 생성형 개념 이미지다. 실제 하드웨어 설계도나 콘솔 화면이 아니다.

Google이 공개한 TPU Agent Suite 소개 영상과 TorchTPU 아키텍처 투어 영상, 그리고 공식 오픈소스 발표문은 기존에 JAX·TensorFlow 중심으로 여겨지던 Google Cloud TPU 생태계에서 PyTorch 코드를 어떻게 수정 없이 네이티브로 실행하고, 코딩 에이전트를 통해 모델 온보딩과 커널 최적화를 자동화할 것인가를 다룬다. Google은 Meta의 PyTorch 팀과 공동 개발한 네이티브 백엔드 TorchTPU를 GitHub에 오픈소스로 공개하고, Claude Code 및 Antigravity에서 동작하는 TPU Agent Suite(프리뷰)를 함께 출시했다.

결론부터 말하면, 이번 발표의 핵심은 과거 지연 텐서(Lazy Tensor) 방식의 PyTorch/XLA가 요구하던 코드 수정과 긴 재컴파일 병목을 버리고 표준 PyTorch 디스패처·프로세스 그룹에 TorchTPU를 직접 연결해 device="tpu" 한 줄로 Eager 모드 디버깅과 torch.compile 최적화를 동시에 지원하며, 에이전트가 CPU 기준 모델과의 수치 정합성(Parity)까지 검증해 준다는 점이다. 다만 이 글은 Google Cloud 공식 영상과 기술 문서·논문을 교차 검증해 정리한 것이며, 실제 멀티호스트 TPU 슬라이스에서 커스텀 Pallas 커널 튜닝 효과나 MFU를 직접 측정한 벤치마크 결과는 아니다.

왜 기존 PyTorch/XLA 대신 Eager 우선(Eager-First) 아키텍처인가

2020년에 도입된 기존 PyTorch/XLA는 지연 텐서(Lazy Tensor)를 사용해 연산 그래프를 기록한 뒤 XLA 컴파일러로 넘기는 구조였다(TorchTPU 아키텍처 논문). 이 방식은 PyTorch를 억지로 JAX의 실행 방식에 끼워 맞춘 형태여서 개발자가 컴파일 영역을 표시하기 위해 모델 코드를 직접 수정해야 했고, 동적 크기(dynamic shape)가 바뀔 때마다 심각한 재컴파일 지연이 발생했으며, 파이썬 표준 디버거로 한 줄씩 텐서를 찍어보는 명령형(imperative) Eager 모드가 제대로 동작하지 않았다.

새로 공개된 TorchTPU 공식 문서는 PyTorch/XLA를 대체(deprecate)하며, C++ 백엔드를 PyTorch 표준 디스패처(dispatcher)와 분산 프로세스 그룹(ProcessGroup)에 직접 연결했다.

구분	기존 `PyTorch/XLA` (Lazy Tensor 기반)	차세대 `TorchTPU` (네이티브 디스패처 기반)
기본 실행 모델	지연 그래프 기록(Lazy Tensor) 후 일괄 컴파일	Eager-First: 표준 PyTorch와 동일한 1디바이스-1프로세스 명령형 실행 (`Strict` 비동기 기본 모드 및 텐서 경계 검사용 `Debug` 동기 모드 지원)
프로덕션 그래프 컴파일	컴파일 경계 수동 지정 및 잦은 그래프 재컴파일	`torch.compile(model, backend="tpu")` 한 줄로 전체 그래프를 `StableHLO` → `XLA`로 컴파일하고 프로세스·호스트·잡 재시작 간 컴파일 캐시 공유
분산 학습 및 서빙 생태계	최신 PyTorch 분산 API 지원 지연	표준 `ProcessGroup`에서 TPU 집합 통신(collectives)을 구현해 `DDP`, `FSDP2`, `DTensor`, `vLLM`, `SGLang`, `TorchTitan`, `Ray`가 코드 변경 없이 그대로 동작
커스텀 커널 확장	XLA C++ 수준 개입 필요	파이썬에서 작성한 `Pallas` 커널(칩 간 직접 통신 포함), `Tokamax` 라이브러리, 실험적 타일 기반 `Helion` 커널을 일반 PyTorch 함수처럼 호출

실제로 Hugging Face Transformers v5에서는 device_map="tpu"만 넘기면 바로 TPU에서 모델이 로드되며, TorchTitan을 활용한 대규모 학습에서는 최대 8,192개의 Ironwood TPU(v7x)에서 거의 선형에 가까운 확장성과 Qwen 3 8B 기준 최대 52% MFU(Model FLOPs Utilization)를 기록했다(Google Cloud 공식 블로그).

TPU Agent Suite의 4가지 전문 서브 에이전트와 수치 정합성(Parity) 검증

코드가 돌아가는 것(Make it run)과 하드웨어 성능을 끝까지 뽑아내는 것(Make it fast)은 별개의 문제다. Google이 함께 오픈소스로 공개한 TPU Agent Suite(프리뷰)는 Claude Code나 Google Antigravity 내부에 설치해 쓰는 멀티 에이전트 플러그인이다.

사용자가 단일 프론트엔드 에이전트에게 모델 온보딩이나 성능 튜닝을 지시하면, 오케스트레이터가 실행 계획을 먼저 보여준 뒤 다음 네 가지 전문 서브 에이전트로 작업을 라우팅한다.

TPU Model Onboarding: 대상 TPU 메모리에 모델이 들어가는지 확인하고, 먼저 vLLM-TorchTPU로 서빙을 시도한 뒤 지원되지 않는 구조면 네이티브 TorchTPU 서빙 코드를 작성해 배포 및 헬스 체크까지 마친다.

TPU Diagnostics: torch.profiler와 XProf 트레이스를 분석해 모델 실행 중 어디서 하드웨어 병목이나 컴파일 재시도가 발생하는지 찾아낸다.

TPU Kernel Development: 가장 비중이 큰 핵심 연산(hotspot)에 대해 커스텀 커널을 작성하고 튜닝한다.

TPU Knowledge Base: TPU 세대별(v5e, v5p, v6e, v7 / Ironwood) 하드웨어 사양과 소프트웨어 문서를 출처 링크와 함께 검색해 제공한다.

여기서 눈여겨볼 설계 원칙은 에이전트의 행동이 근거 기반(Grounded)이고, 횟수가 제한(Bounded)되며, 수치적으로 검증(Checked)된다는 점이다.

수치 정합성(Math Parity) 대조: TPU Model Onboarding 에이전트는 서빙 엔드포인트를 띄운 뒤 곧바로 완료 보고를 하지 않는다. 동일한 입력에 대해 원본 CPU PyTorch 모델의 출력값과 TPU에서 실행된 모델의 수학적 출력값을 대조(parity check)하여 정밀도 손실이나 오차가 없는지 확인한 뒤에야 엔드포인트를 반환한다.

제한된 커널 튜닝 루프(Bounded Tuning): 에이전트에게 커널 최적화를 맡기면 끝없이 루프를 돌며 크레딧과 시간을 소모할 위험이 있다. TPU Kernel Development는 한 번에 한 가지 시도씩 정해진 횟수만큼만 실행해 각각의 벤치마크 결과를 표로 기록한 뒤 멈추고, 최종 적용 여부는 개발자가 결정하게 한다.

GPU 워크로드를 Cloud TPU로 이전·검토할 때 확인할 점

기존 GPU 기반 PyTorch 학습·서빙 파이프라인을 TorchTPU로 옮기거나 에이전트로 최적화할 때는 다음 사항을 순서대로 점검하는 편이 안전하다.

기존 PyTorch/XLA 코드와 TorchTPU의 분리: 신규 프로젝트는 오는 11월 아카이브 예정인 레거시 PyTorch/XLA 대신 TorchTPU로 시작해야 하며, 기존 xm.mark_step() 같은 지연 텐서 전용 호출은 제거하고 표준 torch.compile(backend="tpu") 및 ProcessGroup 방식으로 전환해야 한다.

서빙 워크로드 특성에 맞는 엔진(vLLM vs SGLang)과 라우터(llm-d) 선택: 일반 고처리량 연속 배칭에는 vLLM-TorchTPU가 적합하고, 에이전트 워크플로처럼 반복되는 프롬프트 프리픽스의 KV 캐시를 RadixAttention으로 재사용해야 하는 구조에는 SGLang-TorchTPU가 유리하다. 또한 GKE 멀티노드 환경에서는 컨텍스트를 이미 보유한 레플리카로 요청을 보내는 CNCF 오픈소스 라우터 llm-d를 결합해 토큰 처리량을 높일 수 있다.

컴파일 후 CPU 대비 출력 정합성(Parity) 테스트 의무화: torch.compile이나 커스텀 Pallas 커널, 양자화를 적용한 뒤에는 에이전트의 온보딩 체크리스트처럼 반드시 소규모 검증 입력셋에서 CPU(또는 Eager 모드) 출력과 로짓(logit) 오차를 비교하는 자동화 테스트를 파이프라인에 넣어야 한다.

결국 하드웨어 다변화 시대에 중요한 것은 특정 칩 전용 프레임워크를 새로 배우는 것이 아니라, 익숙한 표준 PyTorch 코드를 유지한 채 컴파일러와 전문 에이전트로 하드웨어별 최적화 격차를 메우는 워크플로를 갖추는 일이다.