TorchTPU: 구글 스케일에서 PyTorch를 TPU에서 네이티브로 실행하기

2026-04-07

요약

TorchTPU는 PyTorch를 구글 TPU에서 네이티브로 실행할 수 있는 스택입니다. 세 가지 Eager 모드(Debug, Strict, Fused)를 제공하며 Fused Eager는 자동 연산 융합으로 50~100% 이상 성능을 향상시킵니다. torch.compile을 통한 XLA 기반 정적 컴파일과 DDP, FSDPv2, DTensor 등 분산 API를 지원하여 대규모 분산 훈련에 최적화되어 있습니다. 또한 MPMD와 같은 분기 실행을 지원하고, 2026년에는 동적 형태 지원, Helion 커널, vLLM/TorchTitan 통합 등을 계획하고 있습니다.

인사이트

TorchTPU는 기존 PyTorch 코드를 최소한의 수정으로 TPU에서 실행할 수 있게 하여 개발자 편의성을 극대화했어요.
Debug Eager, Strict Eager, Fused Eager 모드를 제공하며, 특히 Fused Eager 모드는 자동 최적화로 성능을 50% 이상 향상시켜줘요.
torch.compile을 통한 XLA 기반의 정적 컴파일과 DDP, FSDPv2, DTensor 등 PyTorch 분산 API를 완벽 지원하여 대규모 분산 훈련에 최적화되어 있어요.

왜 중요한가

최신 AI 모델은 엄청난 규모의 분산 시스템과 수많은 가속기를 필요로 해요. 구글의 TPU는 이런 슈퍼컴퓨팅 인프라의 핵심인데, PyTorch 사용자들도 TPU의 모든 성능을 쉽게 활용할 수 있어야 하죠. TorchTPU는 PyTorch 개발자들이 TPU의 잠재력을 최대한 끌어내서 차세대 AI 모델을 더 빠르고 효율적으로 개발할 수 있도록 돕기 때문에 매우 중요합니다.

#PyTorch #TPU #GoogleCloud

TorchTPU: 구글 스케일에서 PyTorch를 TPU에서 네이티브로 실행하기 - 구글 개발자 블로그

검색

TorchTPU: 구글 스케일에서 PyTorch를 TPU에서 네이티브로 실행하기

2026년 4월 7일

클라우디오 바실레 Torch TPU 기술 리드

캣 코 TorchTPU 엔지니어링 리드

벤 윌슨 Torch TPU 소프트웨어 엔지니어

리 하우즈 Torch TPU 엔지니어링 리드

빌 지아 Core ML 엔지니어링 VP

조 패머 Core ML 엔지니어링 수석 디렉터

마이클 보즈네센스키 Torch TPU 엔지니어 리드

로버트 훈트 Core ML 우수 엔지니어

Image 3: Blog-A

현대 AI 인프라를 구축하는 일은 근본적으로 변화했어요. 최신 머신러닝 분야는 이제 수천 개의 가속기를 아우르는 분산 시스템을 활용해야 하죠. 모델의 규모가 수십만(O(100,000)) 개의 칩 클러스터에서 실행될 정도로 커지면서, 이런 모델들을 구동하는 소프트웨어는 성능, 하드웨어 이식성, 그리고 신뢰성에 대한 새로운 요구 사항을 충족해야만 해요.

구글에서는 Tensor Processing Units (TPUs)가 구글 슈퍼컴퓨팅 인프라의 핵심이에요. 이 커스텀 ASIC들은 구글 자체 AI 플랫폼인 제미나이(Gemini)와 베오(Veo)는 물론, 클라우드 고객들의 방대한 워크로드 훈련 및 서빙을 지원하고 있어요. 모든 AI 커뮤니티가 TPU의 모든 기능을 쉽게 이용할 수 있어야 한다고 생각하고, 또 많은 잠재적 사용자들이 PyTorch로 모델을 만들기 때문에, PyTorch가 TPU에서 네이티브로 효율적으로 작동할 수 있도록 통합하는 것이 매우 중요했죠.

TorchTPU의 등장이에요. 이 개발팀의 목표는 사용하기 쉽고, 이식성이 좋으며, 뛰어난 성능을 내는 스택을 만드는 거였어요. 개발자들이 기존 PyTorch 워크로드를 최소한의 코드 변경만으로 마이그레이션할 수 있게 하면서도, 구글의 하드웨어에서 최대한의 연산 능력을 끌어낼 수 있는 API와 도구를 제공하고 싶었답니다. 지금부터 TorchTPU를 이끄는 엔지니어링 원칙, 구축된 기술 아키텍처, 그리고 2026년 로드맵을 자세히 들여다볼게요.

사용성, 이식성, 성능을 위한 아키텍처 설계

TorchTPU를 이해하려면, 먼저 이 프로젝트 팀이 목표로 하는 하드웨어를 알아야 해요.

TPU 시스템은 단순히 칩 하나가 아니에요. 통합된 네트워크로 이루어져 있어요. 호스트 하나에 여러 칩이 연결되어 있고, 각 칩은 호스트와 다른 칩들에 구글의 ICI(Inter-Chip Interconnect)를 통해 연결되죠. 이 ICI는 칩들을 고효율의 2D 또는 3D 토러스(Torus) 토폴로지로 연결해서, 기존 네트워크 병목 현상 없이 엄청난 규모로 확장할 수 있게 해줘요. 각 칩 안에서는 실행이 TensorCores와 SparseCores로 나뉘어요. TensorCores는 조밀한 행렬 연산 전용의 단일 스레드 유닛이고, SparseCores는 임베딩, gather/scatter 연산, 그리고 콜렉티브 오프로딩(offloading collectives)처럼 불규칙한 메모리 접근 패턴을 처리해요.

이런 기능들 덕분에 TPU는 머신러닝을 위한 강력한 도구가 되는데요. 구글의 목표는 이런 독특한 기능을 완전히 활용하는 데 필요한 전문적인 지원을 제공하는 거예요. 여기서 PyTorch가 등장하는데요. PyTorch 툴체인은 이미 다른 장치 유형에 대해 일관성 있고 널리 사용되는 인터페이스를 제공하고 있어요.

사용성을 위한 구글 팀의 핵심 원칙은 간단해요. '파이토치처럼 느껴져야 한다'는 거죠. 개발자는 기존 PyTorch 스크립트를 가져와서 초기화를 "tpu"로 변경한 다음, 핵심 로직의 단 한 줄도 수정하지 않고 훈련 루프를 실행할 수 있어야 해요.

이것을 달성하려면 PyTorch가 TPU 컴파일러 및 런타임 스택과 상호 작용하는 방식에 완전히 새로운 접근 방식이 필요했어요.

TorchTPU 스택 엔지니어링: 기술적 현실

이거 퍼스트(Eager First): 유연성과 타협 없는 성능

개념에서 TPU의 네이티브 PyTorch 경험으로 나아가려면 실행 스택을 다시 생각해야 했어요. 이 프로젝트 팀은 "Eager First" 철학을 확립했어요. 개발자들에게 바로 정적 그래프 컴파일을 강요하는 대신, PyTorch의 "PrivateUse1" 인터페이스를 사용해서 TorchTPU를 구현했어요. 서브클래스나 래퍼 없이, 평범하고 익숙한 PyTorch 텐서가 TPU에서 작동하게 한 거죠. 이렇게 깊은 수준에서 통합함으로써, 개발자들이 PyTorch에서 기대하는 즉시 실행(eager execution) 경험을 완전히 우선시할 수 있었어요.

개발 수명 주기를 지원하기 위해 세 가지 명확한 즉시 실행 모드를 설계했어요.

첫 번째 즉시 실행 모드는 디버그 이저(Debug Eager)로, 한 번에 한 작업씩 디스패치하고 모든 실행 후에 CPU와 동기화해요. 본질적으로 느리지만, 형태 불일치, NaN 값, 메모리 부족 충돌을 추적하는 데는 매우 유용하답니다.

두 번째는 스트릭트 이저(Strict Eager)로, 단일 작업 디스패치를 유지하지만 비동기적으로 실행되며, 기본 PyTorch 경험을 반영하려고 해요. 이를 통해 CPU와 TPU가 사용자의 스크립트에서 동기화 지점에 도달할 때까지 동시에 실행될 수 있어요.

하지만 가장 큰 발전은 바로 퓨즈드 이저(Fused Eager) 모드예요. TorchTPU는 작업 스트림에 대한 자동 리플렉션(reflection)을 사용해서, 실행 중인 단계를 더 크고 계산적으로 밀도 높은 덩어리로 융합한 다음 TPU에 전달해요. TensorCore 활용도를 최대화하고 메모리 대역폭 오버헤드를 최소화함으로써, 퓨즈드 이저는 사용자 설정 없이도 스트릭트 이저에 비해 일관적으로 50%에서 100% 이상의 성능 향상을 제공하고 있어요.

세 가지 모드 모두 공유 컴파일 캐시(Compilation Cache)를 기반으로 작동하는데요. 이 캐시는 단일 호스트에서 작동하거나 여러 호스트 설정에서 영구적으로 구성할 수 있어요. 이는 TorchTPU가 워크로드를 학습할수록 컴파일에 드는 시간을 줄이고 실행에 더 많은 시간을 할애할 수 있다는 뜻이에요.

정적 컴파일: Dynamo, XLA, 그리고 StableHLO

TPU에서 최고 성능을 끌어내고 싶은 사용자들을 위해, TorchTPU는 전체 그래프 컴파일을 위한 torch.compile 인터페이스와 네이티브로 통합되어 있어요. Torch Dynamo를 사용해서 FX 그래프를 캡처하는 것부터 시작하는데요. 하지만 Torch Inductor를 통하는 대신, XLA를 기본 백엔드 컴파일러로 활용해요.

이것은 매우 의도적인 아키텍처 결정이었어요. XLA는 TPU 토폴로지에 대해 엄격하게 검증된 기술이거든요. 더 중요한 것은, XLA가 ICI 전반에 걸쳐 조밀한 계산과 집합적 통신(collective communications) 사이의 중요한 오버랩을 최적화하는 방법을 본질적으로 이해하고 있다는 점이에요. 이 프로젝트 팀의 변환 레이어는 PyTorch의 연산자들을 XLA의 주요 텐서 연산 IR(Intermediate Representation)인 StableHLO로 직접 매핑해요. 이는 PyTorch에서 XLA의 핵심 하위 경로(lowering path)로 직접 연결되어, 즉시 실행 모드에서 확립된 실행 경로를 재사용하면서도 고도로 최적화된 TPU 바이너리를 생성할 수 있게 해줘요.

커스텀 연산자를 작성하는 개발자들을 위해, 확장성이 성능을 저해하지 않도록 보장하고 있어요. TorchTPU는 Pallas와 JAX로 작성된 커스텀 커널을 네이티브로 지원해요. @torch_tpu.pallas.custom_jax_kernel로 JAX 함수를 데코레이션(decorate)하면, 엔지니어들은 구글 팀의 하위 경로와 직접적으로 상호 작용하는 저수준 하드웨어 명령을 작성할 수 있어요. Helion 커널도 지원하기 위한 작업이 진행 중이랍니다.

분산 훈련과 MPMD 과제

대규모에서 즉시 실행 모드와 컴파일 모드의 유연성 및 사용성을 유지하기 위해, PyTorch의 분산 API에 집중했어요. 현재 TorchTPU는 DDP(Distributed Data Parallel), FSDPv2(Fully Sharded Data Parallel v2), 그리고 PyTorch의 DTensor를 즉시 지원하고 있어요. PyTorch의 분산 API를 기반으로 하는 많은 서드파티 라이브러리들이 TorchTPU에서 변경 없이 작동한다는 것을 확인했답니다.

PyTorch/XLA(TorchTPU의 전신)의 주요 한계점 중 하나는 순수 SPMD 코드만 지원했다는 점이었어요. PyTorch 입력의 현실은 서로 다른 랭크(rank)에서 실행되는 코드에 약간의 차이가 자주 발생한다는 거예요. 예를 들어, '랭크 0' 프로세스가 로깅이나 분석을 위해 추가 작업을 하는 경우가 흔하죠. 이런 종류의 입력은 SPMD 최적화에 크게 맞춰진 TPU 스택에는 도전 과제가 된답니다. XLA는 시스템에서 실행되는 코드에 대한 전역적인 시야를 가질 때 가장 잘 작동하지만, 이것을 우회하려면 개발자가 불순한 동작을 신중하게 제거해야 하는 오버헤드가 발생해요.

TorchTPU는 분산 실행(MPMD)을 신중하게 지원하도록 설계되었으며, 정확성을 유지하기 위해 필요한 경우 통신 프리미티브를 최소한의 비용으로 분리할 예정이에요. 이 접근 방식은 PyTorch 개발자들이 TPU에서 PyTorch를 사용하는 경험이 가능한 한 자연스럽도록 보장하며, 동시에 XLA가 분산 TPU 배포의 전역적 시야를 통해 통신과 계산을 최대한 오버랩할 수 있도록 해줘요.

TPU 하드웨어 인식

TPU는 매우 높은 성능과 효율성을 달성할 수 있지만, 최적의 모델 설계는 다른 하드웨어와 약간 다를 수 있어요. 예를 들어, 어텐션 헤드(attention head) 차원을 64로 하드코딩하는 모델을 자주 보는데요. 하지만 현 세대 TPU는 128 또는 256 차원에서 최대 행렬 곱셈 효율성을 달성해요. 모델을 128 또는 256 차원으로 수정하면 TPU 칩의 크고 조밀하며 효율적인 텐서 코어를 더 잘 활용할 수 있죠.

이식성이 하드웨어 현실을 없애는 건 아니에요. 그래서 TorchTPU는 계층화된 워크플로를 제공해요. 먼저 정확한 실행을 확립한 다음, 곧 출시될 심층 가이드라인을 사용해서 최적화되지 않은 아키텍처를 식별하고 리팩터링하거나, 최적의 하드웨어 활용을 위해 커스텀 커널을 삽입할 수 있도록 할 거예요.

앞으로 나아갈 길: 2026년 이후

이 프로젝트 팀은 오늘날 훈련 및 서빙 지원 전반에 걸쳐 탄탄한 기반을 마련했고, TorchTPU를 PyTorch 생태계에서 마찰 없는 백엔드로 만들기 위해 몇 가지 열린 과제를 적극적으로 해결하고 있어요.

컴파일러 팀의 주요 초점은 동적 시퀀스 길이와 배치 크기로 인해 발생하는 재컴파일을 줄이는 거예요. XLA 내에서 고급 바운드 다이내미즘(bounded dynamism)을 구현함으로써, 컴파일 오버헤드 없이 형태 변경을 처리하는 것을 목표로 하고 있어요. 이는 반복적인 다음 토큰 예측과 같은 특정 워크로드에 중요한 기능이 될 수 있어요.

또한 표준 연산을 위한 포괄적인 사전 컴파일된 TPU 커널 라이브러리를 구축해서 첫 실행 반복의 레이턴시를 획기적으로 줄이는 작업도 하고 있답니다.

2026년 남은 기간 동안 구글 팀은 다음을 목표로 하고 있어요:

광범위한 문서와 재현 가능한 아키텍처 튜토리얼을 포함한 공개 GitHub 저장소 출시.
커스텀 커널 기능을 더욱 확장하기 위한 PyTorch의 Helion DSL과의 통합.
torch.compile을 통한 동적 형태에 대한 일등석 지원.
메모리와 컴퓨트 스트림이 분리된 비동기 코드를 쉽게 마이그레이션할 수 있도록 네이티브 다중 큐 지원.
vLLM 및 TorchTitan과 같은 생태계 핵심 요소와의 심층 통합, 그리고 전체 Pod 크기 인프라까지 검증된 선형 스케일링.

TorchTPU는 PyTorch 개발자들이 TPU 하드웨어에서 끊김 없고 고성능의 PyTorch 경험을 할 수 있도록 구글 팀이 쏟아부은 전념적인 엔지니어링 노력의 결과물이에요. PyTorch라는 여러분이 좋아하는 프레임워크와 차세대 AI에 필요한 TPU 슈퍼컴퓨팅 하드웨어 사이의 장애물과 마찰을 없애고 있답니다.

TorchTPU의 최신 업데이트에 대한 정보를 계속 받으려면 TPU 개발자 허브를 방문해 주세요.