10월 20일, Google은 자사의 TPU를 위한 네이티브 PyTorch 백엔드인 TorchTPU를 GitHub에 공개했습니다. 이미 PyTorch로 TPU에서 실행하고 있는 사용자라면 TorchTPU FAQ의 다음 문구를 참고해야 합니다: "PyTorch/XLA는 사용 중단되었으며, 이제 공개된 TorchTPU로 대체되었습니다."
PyTorch/XLA README 역시 동일한 메시지로 시작합니다: "이 프로젝트는 사용 중단되었으며 TorchTPU로 대체되었습니다." 또한 이 저장소는 2026년 11월 15일에 아카이브될 것이라고 명시하고 있습니다. 이는 torch_xla 팀에게 프로젝트가 읽기 전용(read-only) 상태가 되기까지 4주 정도의 시간이 주어진다는 것을 의미합니다.
Google이 공개한 내용
Google은 PyTorch Conference 2026에서 TorchTPU를 발표했으며, 이는 Meta의 PyTorch 팀과 함께 구축되었다고 밝혔습니다 (Google Cloud blog). Google의 설명에 따르면, TPU는 이제 GPU 및 CPU와 함께 네이티브 PyTorch 대상(target)이 되며, 디버깅을 위한 eager 모드와 프로덕션 속도를 위한 torch.compile 기능을 지원합니다.
해당 저장소는 Apache 2.0 라이선스를 따르며, 최신 공개 버전은 PyPI에서 0.2.0입니다. Python 3.11 이상을 지원하며 (README), 기능 세트 페이지에서는 TPU v5e, v5p, v6e 및 v7이 검증되었다고 나열합니다. TPU v8i와 v8t는 현재 지원되지 않지만 가까운 미래 로드맵에 포함되어 있습니다.
서비스(serving) 측면도 함께 제공됩니다. Google은 Inferact로 구축된 vLLM 백엔드와 RadixArk로 구축된 SGLang 백엔드를 연결한 게시물을 올렸으며, Hugging Face Transformers v5가 device_map="tpu"를 사용하여 TPU에 배포된다고 밝혔습니다 (Google Cloud blog). 학습(training) 측면에서는 Google이 TorchTitan을 사용하여 Qwen 3 8B에서 최대 52%의 MFU를 달성했으며, 최대 8,192개의 Ironwood TPU에 걸쳐 거의 선형적으로 확장된다고 보고했습니다. 이는 독립적인 벤치마크가 아닌 Google 자체 수치입니다.
TechCrunch는 이번 출시를 Google이 Amazon과 함께 AI 개발자 시장에서 Nvidia의 지배력을 공략하는 행보로 보도했으며, 이는 AWS가 자체 PyTorch 백엔드를 Trainium에 발표한 지 하루 뒤의 일입니다. 저는 이전에 the TorchNeuron post에서 다룬 바 있습니다.
포트는 하나의 문자열로 시작합니다
FAQ는 기존 워크로드에 대한 변경 사항을 .to("cuda") → .to("tpu")로 설명합니다. README에서 기본 설치가 CUDA 빌드를 가져온다고 경고하므로, CPU 빌드의 torch를 사용하여 릴리스를 설치해야 합니다 (README):
pip install torch_tpu==0.2.0 \
--extra-index-url https://download.pytorch.org/whl/cpu
그런 다음 VM 시작 가이드의 코드를 간소화하여 사용합니다:
import torch
device = torch.device("tpu") # 이전에는 "cuda"였습니다.
...
이 가이드는 v6e 인스턴스를 예시로 TPU VM에서 실행됩니다. 또한 최근의 TPU 하드웨어는 VM 대신 GKE를 필요로 한다는 점을 언급합니다(guide).
기본값(Eager) 방식과 병목 현상이 발생하는 지점
실행 모델 페이지에서는 해당 한 줄의 변경으로 어떤 일이 발생하는지 설명합니다. 기본 엄격한 Eager 모드에서는 각 연산(op)이 단일 연산 그래프가 되어 -O1로 컴파일되고 캐싱되므로, 캐시가 빠르게 채워집니다. torch.compile은 전체 FX 그래프를 TorchTPU에 전달하여 -O2로 컴파일합니다. 이 사이에 두 가지 환경 변수가 존재합니다:
# 오류는 발생시킨 연산(op)에서 표시됩니다.
TORCH_TPU_LAUNCH_BLOCKING=1 python train.py
...
문서에 따르면, 모델 변경 없이 TORCH_TPU_DEFER_AND_FUSE를 사용할 경우 기본 Eager 방식 대비 최대 50%의 속도 향상을 주장하며, 동일한 페이지에서는 장치에서 병목 현상이 발생하는 지점들을 나열합니다. 예를 들어, if my_tensor.all()과 같은 데이터 의존적 Python 분기는 값을 CPU로 복사하고 큐를 고갈시키며, 컴파일되거나 융합된(fused) 모드에서도 그래프를 깨뜨립니다. torch.masked_select와 같이 데이터 의존적인 출력 형태를 가진 연산도 마찬가지입니다.
형태가 커지는 것(Growing shapes)이 매 단계마다 재컴파일을 강제하지는 않습니다. 동적 형태 페이지에서는 KV 캐시가 증가하는 디코딩 루프를 안내합니다. dynamic=False의 경우, 매 단계마다 재컴파일이 이루어집니다. 기본값인 dynamic=None의 경우, 증가하는 차원은 심볼릭(symbolic)으로 처리되며, 문서 예제에서는 처음 세 단계만 컴파일 시간이 소요됩니다. 내부적으로 TorchTPU는 128의 배수에 맞춰 버킷을 컴파일하고 그 안에 패딩합니다:
step = torch.compile(
decode_step, backend="tpu", dynamic=None
)
...
제한 사항은 동일합니다: 트레이스당 최대 8개의 버킷, dynamic=True 불가, torch.nonzero-스타일 출력 형태 불가, 해당 경로에서 커스텀 Pallas 커널 불가. 이러한 바운디드-다이나믹 실행 파일은 AOTAutograd 캐시에 저장되지 않으므로 프로세스가 시작될 때마다 재컴파일됩니다.
현재 PyTorch/XLA를 사용하고 있다면
FAQ는 두 가지 사이의 경계를 명확히 합니다: PyTorch/XLA는 JAX 런타임 컴포넌트와 자체 API에 의존했던 반면, TorchTPU는 JAX 런타임 종속성이 없는 네이티브 백엔드입니다. 또한 TorchTPU가 JAX를 대체하는 것이 아니라고 언급합니다.
먼저 코드가 이전 스택에 접촉하는 모든 지점을 찾아보세요:
git grep -n -e "torch_xla" -e "xla_device" -e "xm\."
그런 다음 각 히트(hit)를 기능 세트와 비교하세요. Eager 및 컴파일 실행, 순전파(forward) 및 역전파(backward) 패스, Pallas 커널, SimpleFSDP가 검증된 것으로 나열되어 있습니다. FSDPv2와 SymInts는 실험적이며, FSDPv2는 성능이 낮다고 표시되어 있습니다. 11월 15일이 지나면 PyTorch/XLA 레포지토리는 읽기 전용 상태가 됩니다: 코드를 계속 읽고 클론할 수는 있지만, 아무도 새로운 이슈나 풀 리퀘스트를 열 수 없습니다.
TPU 에이전트 스위트는 Claude Code 및 Antigravity와 연결됩니다
두 번째 발표는 TPU 에이전트 스위트이며, Google은 이를 미리 보기(preview)라고 부릅니다. 이는 코딩 에이전트를 위한 기술 및 서브에이전트 패키지이며, 온보딩, 프로파일링, 하드웨어 확인, 인용된 TPU 사실 및 커널을 위한 전문가에게 라우팅하는 하나의 진입점인 tpu-agent 스킬을 가지고 있습니다. Claude Code와 Antigravity가 지원되며, Codex와 Cursor는
실행하기 전에 설치 프로그램이 무엇을 쓰는지 읽어보세요. 설치 가이드에 따르면, Claude Code를 위한 프로젝트 범위(project-scope) 설치는 .claude/skills/ 아래에 스킬을 추가하고, 에이전트당 .claude/agents/ 아래에 프롬프트를 하나씩 생성하며, ~/.tpu-agent-suite 경로를 사전에 승인하는 .claude/settings.json 파일과 CLAUDE.md 내 관리 블록을 만듭니다. Antigravity의 경우 대신 AGENTS.md를 작성합니다. 이 블록을 건너뛰려면 --no-workspace-memory 플래그를 전달하세요. 브랜치에 설치하고, 무언가를 커밋하기 전에 git status와 git diff를 읽어보세요.
워크트리(worktrees)에서 에이전트를 실행할 때의 함정이 하나 있습니다: 새로 생성된 git worktree는 추적되는 파일만 체크아웃합니다. 커밋되지 않은 프로젝트 범위 파일은 설치한 폴더에 그대로 남아있습니다. 가이드에서 제시하는 --scope user 설치가 이를 방지합니다.
저는 Hermes IDE를 구축했으니, 이 점을 염두에 두세요. 현재 이 스위트가 지원하는 두 에이전트는 모두 Hermes 내장 에이전트이며, 각 작업은 자체 hermes/<name> 브랜치와 워크트리를 가질 수 있으므로, 동일한 TPU 포트에서 실행되는 Claude Code 실행과 Antigravity 실행의 파일들이 서로 영향을 주지 않습니다(README). 저는 직접 TPU 에이전트 스위트를 실행해 본 적은 없습니다. Hermes는 Google, Anthropic 또는 어떤 에이전트 벤더와도 제휴하지 않습니다.
The Hermes IDE New task sheet with its agent picker open: Claude Code installed and selected, Codex installed, and Antigravity CLI, Gemini CLI and other supported agents listed, plus Custom agent
Hermes 2.3.4의 에이전트 선택기: Claude Code 및 Codex 설치됨, Antigravity CLI와 기타 지원되는 에이전트 목록, 그리고 사용자 정의 에이전트.
이사할 건가요?
저는 TorchTPU를 실행해 본 적이 없으며, 위에 언급된 제한 사항들은 실제 포팅 과정에서 밝혀질 부분입니다. 사용 중단(deprecation)은 PyTorch/XLA에 의존하는 팀들이 기다려 볼 수 없다는 것을 의미합니다.
혹시 오늘 TPU에서 PyTorch를 사용하고 계신가요? 그리고 11월 15일 아카이브 날짜가 이미 로드맵에 포함되어 있나요? 댓글로 알려주세요.
유용했다면
만약 vLLM이나 SGLang을 사용하여 서빙(serving)을 TPU로 옮긴다면, 다음 질문은 그 이동이 효과가 있었는지 여부입니다: 지연 시간(latency), 요청당 토큰 수(tokens per request), 답변당 비용(cost per answer), 그리고 출력값이 변하지 않았음을 알려주는 평가(evals). LLM Observability Pocket Guide는 이러한 질문에 답하는 추적 및 평가 도구 선택 방법을 다룹니다.
Claude Code를 Antigravity나 Codex 옆에서 실행한다면, Hermes IDE가 소스 코드를 제공하며(BSL 1.1), macOS, Windows, Linux에서 실행됩니다. 소스는 GitHub에 있습니다.
LLM Observability Pocket Guide: Picking the Right Tracing & Evals Tools for Your Team
