Hacker News 2026.04.24 31

PyTorch가 드디어 TPU에서 '네이티브'로 돈다: TorchTPU가 바꿀 AI 학습 지형도

PyTorch 개발자에게 반가운 소식

GPU에 익숙한 분들이라면 구글의 TPU(Tensor Processing Unit)라는 이름은 들어봤지만 막상 써본 적은 없으실 거예요. TPU는 구글이 머신러닝 전용으로 만든 칩인데요, 엄청나게 빠르지만 한 가지 큰 단점이 있었습니다. PyTorch와 궁합이 별로 안 좋았다는 거예요. TPU를 제대로 쓰려면 JAX나 TensorFlow를 써야 했거나, PyTorch를 쓰더라도 torch_xla라는 징검다리 라이브러리를 거쳐야 했거든요. 이게 뭐냐면, PyTorch 코드를 XLA(가속기용 중간 표현 언어)로 번역해서 TPU에 돌리는 방식이었어요. 번역 과정에서 성능 손실도 있고, 디버깅도 까다로웠죠.

그런데 구글이 TorchTPU라는 걸 발표했습니다. PyTorch가 TPU 위에서 "네이티브로" 돈다는 게 핵심이에요. 네이티브라는 말은 중간 번역층 없이 PyTorch의 자연스러운 API와 동작 방식이 TPU 하드웨어에 바로 꽂힌다는 뜻이죠.

기술적으로 뭐가 달라진 걸까

기존 torch_xla 방식에서는 PyTorch의 동적 그래프(코드가 실행될 때마다 계산 경로가 달라질 수 있는 유연함)를 써도 내부적으로는 정적 그래프로 변환해야 했어요. 그래서 파이썬다운 코드, 예를 들어 if문이나 동적 shape를 쓰면 그래프가 계속 재컴파일되면서 속도가 떨어지는 문제가 있었습니다. 반복 학습 때마다 "어? 이번엔 모양이 다르네?" 하고 다시 컴파일하니까 낭비가 심했죠.

TorchTPU는 PyTorch 2.x에서 도입된 torch.compile과 Dynamo/Inductor 스택을 적극 활용해요. 쉽게 말하면 PyTorch 자체의 최신 컴파일러 인프라를 그대로 쓰면서, 타겟만 TPU로 바꾸는 거예요. 덕분에 GPU에서 쓰던 코드를 거의 그대로 TPU로 옮길 수 있고, 구글 내부에서 실제로 초대규모 학습을 돌려본 결과도 같이 공개됐습니다. Gemini 같은 프론티어 모델 학습 파이프라인에서 쓰이는 수준의 규모라고 하니, 장난감 수준이 아니라 프로덕션 검증을 거친 스택이라는 의미죠.

또 한 가지 중요한 점은 SPMD(Single Program, Multiple Data) 분산 학습 지원이에요. TPU는 수천 개 칩을 하나의 거대한 조직처럼 연결해서 쓰는 게 강점인데, 이걸 PyTorch 코드에서 자연스럽게 쓸 수 있게 해줍니다. GPU 클러스터에서 DDP나 FSDP를 쓰던 분들이 개념적으로 어렵지 않게 옮겨갈 수 있게 설계된 거죠.

경쟁 지형도

AI 하드웨어 전쟁은 지금 세 축으로 흘러가고 있어요. 첫째는 NVIDIA의 GPU + CUDA 생태계, 둘째는 구글의 TPU + JAX, 셋째는 AMD MI300이나 AWS Trainium 같은 신흥 가속기들이에요. PyTorch는 원래 첫 번째 축에서 탄생했고, CUDA와 가장 친한 프레임워크였습니다.

TorchTPU의 등장은 NVIDIA 독점에 금이 가는 또 하나의 사건이에요. Meta가 PyTorch를 만든 주인인데, 구글이 자사 하드웨어에 PyTorch를 1급 시민으로 대접한다는 건 "프레임워크는 중립, 하드웨어는 선택"이라는 메시지로 읽을 수 있거든요. 물론 JAX도 계속 밀겠지만, PyTorch 사용자층이 워낙 두껍다 보니 TPU를 쓸 이유 하나를 더 제공한 셈입니다.

한국 개발자에게 주는 시사점

한국에서 TPU를 직접 만져볼 일은 GCP를 쓸 때 말고는 드물어요. 하지만 이 소식이 의미 있는 이유는, 멀티 가속기 시대가 진짜로 오고 있다는 신호이기 때문입니다. 지금까지 "모델 짤 때는 PyTorch, 대규모 학습은 JAX" 같은 어정쩡한 분리가 있었는데, 앞으로는 PyTorch 한 벌로 GPU, TPU, 심지어 커스텀 NPU까지 다 돌릴 수 있는 흐름으로 갈 거예요.

당장 실무에서 써볼 수 있는지는 팀이 GCP를 쓰느냐에 달려 있어요. 국내 스타트업이 Gemini API 대신 자체 모델을 학습시키고 싶은데 H100을 구하기 어려워서 고민 중이라면, TPU v5 이상이 꽤 매력적인 대안이 될 수 있습니다. 특히 장기 학습 작업은 TPU의 시간당 비용이 H100보다 유리한 경우가 많거든요.