TorchTPU: PyTorch 코드를 Google TPU에서 직접 실행하기

배경과 문제

TPU는 머신러닝의 행렬 곱셈을 빠르게 처리하도록 설계된 전용 칩이다. 그러나 칩이 빠르다는 것만으로 기존 PyTorch 프로그램을 쉽게 옮길 수 있는 것은 아니다. 연산을 TPU가 실행할 형태로 바꾸고, 메모리·통신·컴파일을 연결하는 소프트웨어가 필요하다. 익숙한 개발 방식이 달라지면 작은 실험이나 오류 추적부터 어려워진다.

TorchTPU는 2025년 중반 Google에서 시작해 Meta의 PyTorch 팀과 협력한 프로젝트다. 표준 PyTorch 연산을 쓰는 모델이 장치 이름을 cuda에서 tpu로 바꾸는 정도의 적은 수정으로 실행되도록 설계됐다. 이는 이식의 출발점을 낮춘다는 뜻이다. GPU에 맞춰 작성한 연산이나 샤딩 방식까지 자동으로 TPU의 최고 성능을 낸다는 뜻은 아니다.

기존 접근의 한계

공식 글은 이전 PyTorch/XLA가 Lazy Tensor와 컴파일 중심의 방식이어서 코드 수정, 컴파일 지연, 새로운 분산 기능 지원에 부담이 있었다고 설명한다. TorchTPU가 우선하는 것은 PyTorch처럼 연산을 호출하며 확인하는 eager 실행이다. 이후 속도가 필요할 때 여러 연산을 묶어 컴파일하는 흐름으로 이어진다.

여기에는 단순 번역 이상의 문제가 있다. PyTorch에서는 여러 텐서가 같은 저장 공간을 바라보고 그 값을 제자리에서 바꿀 수 있다. 반면 TPU 컴파일에 사용하는 StableHLO는 새 값을 만드는 함수형 표현이다. TorchTPU는 공유된 데이터의 일부를 보는 뷰와 그 변경을 추적·변환해, 익숙한 PyTorch 동작을 유지하면서 컴파일 가능한 형태로 바꾼다.

구현 방식: 연산부터 컴파일·분산 실행까지

먼저 PyTorch가 연산을 디스패치하면 C++로 만든 TPU 백엔드가 이를 받는다. Strict는 기본 eager 모드로 연산을 하나씩 TPU에 보내되 CPU와 비동기로 진행한다. Debug는 실행을 동기화해 문제 지점을 찾기 쉽게 한다. DeferAndFuse는 연산을 잠시 모았다가 결과를 CPU로 가져오는 등 값이 실제 필요한 지점에 함께 실행해, 작은 연산을 여러 번 처리하는 부담과 중간 저장 공간을 줄인다.

모아 둔 연산은 StableHLO로 표현하고 XLA가 TPU 실행 코드로 컴파일한다. torch.compile을 쓰면 더 큰 계산 그래프를 묶을 수 있다. 반복 실행은 프로세스 안, 같은 머신, 원격 저장소의 캐시로 재컴파일을 줄인다. 원격 영속 캐시는 별도 활성화가 필요하며, 처음 보는 그래프나 입력 크기에서 생기는 컴파일 비용까지 사라지는 것은 아니다.

분산 실행은 PyTorch의 Process Group 인터페이스와 연결돼 FSDP 같은 도구를 이용한다. 최고 성능이 필요한 연산에는 Pallas 커널을 선택해 쓸 수 있다. 기본 백엔드가 JAX 런타임을 필수로 쓰지 않는다는 것과, 선택적으로 JAX 기반 Pallas 코드를 연결할 수 있다는 것은 양립한다. 모든 사용자가 다른 프레임워크로 모델 전체를 다시 작성해야 하는 구조는 아니다.

결과와 숫자의 의미

모델 호환성 검사에서는 4,917개 중 4,821개, 약 98%가 오류 없이 실행됐다. 이는 합성 가중치로 실행·그래프 변환·컴파일을 확인하는 검사이며 모델의 업무 정확도 98%가 아니다. 수치 검사는 별도 대표 모델 657개 중 638개, 약 97%가 통과했다. LLM은 GPU 기준 퍼플렉서티의 상대 차이 0.10 이내, 다른 모델은 정해진 평균제곱오차 기준을 사용했다. 통과는 허용 오차 안이라는 뜻이지 모든 출력이 비트 단위로 같다는 뜻이 아니다.

대규모 학습은 Qwen3 MoE 235B, TPU v7, 길이 4,096의 입력, TensorCore당 배치 16 조건이다. 1,024칩에서 8,192칩으로 늘릴 때 전체 배치도 32,768에서 262,144로 늘렸다. TensorCore당 초당 토큰 수는 2,236에서 2,104, MFU는 28.8%에서 27.1%, 한 단계 시간은 29.3초에서 31.2초였다. 장비 수만큼 일을 더 맡겨도 개별 장비의 처리 효율이 비슷한 약한 확장성(weak scaling) 결과다. 같은 고정 작업이 8배 빨라졌다는 실험은 아니다.

이 235B 실험에는 Pallas 커널, 활성값 오프로딩과 재계산, 여러 층을 묶는 scan, 통신과 계산의 겹치기 등 최적화가 들어갔다. 별도의 순수 PyTorch 코드 예시와 같은 조건으로 합칠 수 없다. 처리량 그래프는 규모 확대를 보여주지만, 손실 그래프는 첫 12단계만 비교한다. 따라서 최종 학습 수렴까지 같다고 입증한 결과로 읽어서는 안 된다.

구체적인 예시

정밀도는 실제 이식에서 확인할 문제다. 기본 medium 모드의 TPU 행렬 곱셈은 입력을 BF16으로 바꾸고 합을 FP32로 누적한다. 따라서 FP32 텐서를 쓴다는 사실만으로 다른 장치와 같은 계산 정밀도를 기대할 수 없다. high는 3회, highest는 6회로 계산을 나눠 더 높은 정밀도를 근사한다. 약 3배·6배 비용은 해당 행렬 연산의 추가 계산량에 대한 설명이며 전체 프로그램 지연이 반드시 그만큼 늘어난다는 뜻은 아니다.

원문에는 FLUX 모델을 단계적으로 최적화한 실제 예도 있다. eager 기준 학습 단계 시간이 11.49초였고, DeferAndFuse, 샤딩, 전용 attention 커널, torch.compile 등을 차례로 적용한 뒤 최종 2.78초가 됐다. 하나의 설정만 바꾼 효과가 아니라 여러 변경의 누적 결과다. 이 예는 코드를 옮겨 실행하는 단계와 하드웨어를 충분히 활용하는 단계를 구분해야 하는 이유를 보여준다.

한계와 오해하기 쉬운 점

호환성 및 성능 수치는 개발팀의 보고를 확인한 것이며, 이번 해설에서 직접 실행해 재현한 것은 아니다. 공식 글이 인용하는 외부 추론 비용 벤치마크도 특정 모델·소프트웨어·입출력 길이 조건의 결과다. 이를 상용 API의 실제 요금표나 모든 GPU보다 저렴하다는 결론으로 확대할 수 없다.

동적인 입력 크기 지원에는 일부 시험 결과가 있지만 완전히 끝난 기능으로 볼 수 없다. 글은 bounded dynamism의 운영 안정화, 연산 범위 확대, 향후 TPU의 FP8·FP4 지원, 강화학습 실행과 진단 도구 등을 앞으로 보강할 계획으로 설명한다. 현재 제공하는 기능과 다음 6~12개월의 개발 목표를 구분해야 한다.

실제로 어떻게 활용할까

실무에서는 기존 CPU나 GPU에서의 출력, 입력 크기, 자료형을 기준으로 정하고 TPU eager 실행 결과를 먼저 비교하는 방식으로 접근할 수 있다. 이 단계에서 오류와 허용 오차를 확인한 다음 컴파일과 캐시를 적용하면, 어떤 변경이 결과와 성능에 영향을 주는지 구분하기 쉽다. 이는 이 글의 구조에서 도출한 활용 방법이며 새로 수행한 실험 결과는 아니다.

성능을 판단할 때는 입력을 처음 처리하는 속도, 다음 토큰 생성 속도, 학습 단계 시간, 메모리 사용량을 따로 재야 한다. 모델에 따라 Pallas 커널이나 vLLM·SGLang·TorchTitan 같은 스택을 추가로 검토할 수 있다. TorchTPU의 가치는 기존 PyTorch 개발 흐름에서 TPU를 시험할 길을 넓힌 데 있으며, 모든 연산 지원이나 특정 서비스의 비용 절감을 보장한 데 있지는 않다.