요약
구글이 TPU에서 PyTorch 모델을 코드 거의 그대로 돌리게 해 주는 TorchTPU를 공개했습니다.
구글과 메타의 PyTorch 팀이 함께 만든 오픈소스 백엔드이고, Apache 2.0 라이선스로 GitHub에 올라왔습니다. TPU(구글이 만든 AI 전용 칩)는 그동안 주로 JAX나 TensorFlow로 써 왔는데, 이제 PyTorch 이용자도 기존 코드를 조금만 고쳐 TPU에서 돌릴 수 있습니다.
왜 중요한가
- 장치 이름을 "cuda"에서 "tpu"로 바꾸는 정도로 옮길 수 있어, GPU용 PyTorch 코드를 TPU로 옮기는 부담이 줄어듭니다.
- JAX 런타임 없이 PyTorch에 직접 붙는 구조라서, TPU를 쓰려고 다른 프레임워크를 거칠 필요가 없습니다.
- 문서에 따르면 v5e, v5p, v6e, v7 TPU에서 검증했고 칩 약 8천 개 규모의 학습까지 지원해 실제 서비스 규모를 겨냥합니다.
핵심 내용
- 기본 즉시 실행(eager) 모드는 연산을 하나씩 비동기로 보내 GPU에서 PyTorch를 쓸 때와 비슷하게 동작합니다.
- torch.compile을 쓰면 그래프 전체를 StableHLO(XLA 컴파일러의 중간 표현)로 바꿔 XLA로 컴파일하고, TorchInductor는 쓰지 않습니다.
- 컴파일 결과는 프로세스, 같은 머신, 원격 저장소의 세 단계 캐시에 남겨 같은 컴파일을 되풀이하지 않습니다.
- Llama, Qwen, Gemma 같은 트랜스포머와 FLUX 같은 확산 모델, Hugging Face와 vLLM 연동이 검증된 기능으로 올라 있습니다.
- 현재 버전은 0.2.0이고, 입력 크기가 바뀌면 다시 컴파일해야 하는 점과 TPU 하나를 한 프로세스만 쓸 수 있는 점이 알려진 한계입니다.
HN 반응
- 예전에는 XLA와 씨름해야 했다며 공식 지원을 반기는 반응과 함께, TensorFlow가 설 자리를 잃었다는 평도 나왔습니다.
- 아마존도 Trainium용 PyTorch 백엔드를 내놓아 PyTorch 콘퍼런스에 맞춘 공개라는 지적과, 텐스토렌트 칩에서 TorchTPU 대부분을 돌렸다는 경험담이 나왔습니다.
관심 있는 분들을 위해 알려 드립니다. 저는 https://tenstorrent.com/ 가속기용 오픈소스 libtpu 구현을 사이드 프로젝트로 작업하고 있고, 이미 torchtpu의 대부분을 실행할 수 있습니다. https://github.com/pcmoritz/libtt (Tenstorrent 소프트웨어 스택 기반)를 보세요. torch를 타일 기반 가속기로 낮추는 과정(lowering)에 들어가는 요소가 전부 담겨 있습니다.
정말 멋지네요! 저는 Lean으로 PJRT를 구동하고 있는데, 이런 식입니다: https://github.com/brettkoonce/lean4-mlir 이론상으로는 당신의 코드로 플랫폼 하나를 더 열 수 있겠어요!
공유 감사합니다, 정말 멋지네요! 지금 제 정확성 검증 방식은 가능한 한 많은 테스트를 돌려서 고치는 건데(torchtpu 덕분에 쓸 수 있는 테스트가 엄청나게 늘었습니다), 형식 검증도 정말 멋진 방법이죠. 말씀하신 책 https://lean.brettkoonce.com/blueprint.pdf 를 살펴보는 중입니다.
참고로 Tenstorrent ISA는 문서화되어 있습니다. https://github.com/tenstorrent/tt-isa-documentation 하드웨어까지 포함해서 처음부터 끝까지 검증할 수 있을지 궁금하네요.
정말 멋져요!
응원 감사합니다. 하드웨어가 있어서 직접 써 보시게 되면 피드백을 꼭 듣고 싶어요 (저장소에 이슈로 남겨 주세요) :)
저도 그러고 싶은데 아쉽게도 지금은 하드웨어가 없어요 (아직은요?). 그런데 궁금한 게 있어요. 이 작업을 하려고 TT 장비를 직접 사셨나요, 아니면 Koyeb 같은 클라우드 업체를 쓰고 계신가요?
처음에는 써 보려고 직접 조립했고(정확한 하드웨어 구성은 https://relativity.cc/hardware/ 의 "one Blackhole card" 항목에 있습니다), 하드웨어가 마음에 들어서 quietbox 2를 샀습니다. 품질은 매우 좋지만 가격이 좀 나갑니다(그래도 다른 LLM 하드웨어에 비하면 저렴하고, 소비자·개발자 시장에서는 보기 드물게 인터커넥트가 훌륭합니다).
구글과 애플은 일반 소비자와 일반 기업 사용자를 엣지 추론으로 옮겨 놓을 겁니다. 이건 곧 있을 IPO에 대부분의 투자자가 아는 것보다 훨씬 큰 타격이 될 거예요.
구글, 애플, 그리고 우리의 친구 DeepSeek도 잊으면 안 되죠!
그쪽이 엣지 컴퓨팅으로 옮겨 가는 걸 보고 싶습니다. 일상적인 사용에서는 어느 정도의 자율성과 독립된 경험이 있어야 하니까요.
추론용으로 TPU를 살 수 있으면 좋겠네요.
살 수 있어요.. 이름이 Coral이에요. https://en.wikipedia.org/wiki/Tensor_Processing_Unit#Edge_TPU https://developers.google.com/coral https://physicalailab.io/en/google-coral-tpu-edge-tpu-status-coral-npu/
Google Coral TPU는 살 수 있는데, 그것도 쳐 주나요?
흥미롭네요, 아마존도 AWS Trainium용으로 같은 걸 하고 있습니다.
https://github.com/aws-neuron/torch-neuronx
오늘 PyTorch 컨퍼런스가 열리고 있어서 시기가 겹친 겁니다.
그래서 허깅페이스가 말한 Metal을 쓰는 맥에서의 5배 속도 향상이 설명되네요.
링크는요?
좋네요! XLA는 꽤 괜찮고, inductor보다 더 나은 컴파일러 인프라가 될 수도 있어서 모든 하드웨어의 성능 향상으로 이어질 겁니다.
그리고 이게 Tensorflow 관 뚜껑에 박는 마지막 못이 아닌가 싶네요.
놀라운 소식이네요. PyTorch를 TPU에서 돌리려고 XLA와 씨름하던 때가 기억나는데, 드디어 1st-party 지원이 생겼군요!
하드웨어 특성이 대부분 커널 계층(Pallas나 CUDA)에 있으니, 이제 Jax냐 PyTorch냐는 취향 문제에 더 가까워진 것 같은데, 맞나요?
무지해서 죄송한데, 구글 말고 TPU를 가진 곳이 또 있나요? 어디서 상업적으로 살 수 있는 건가요? 지금 구글 밖에서는 주로 어디에 쓰이나요?
클라우드에서 빌려 쓸 수 있고, 앤트로픽에는 직접 판매합니다. 앤트로픽이 (Trainium, GPU와 함께) TPU를 많이 쓰는 건 꽤 잘 알려져 있습니다.
구글은 클라우드 서비스를 통해, 그리고 일부 파트너에게는 실물로도 TPU라는 삽을 팔고 있습니다. 그러니 이해가 됩니다.
삽이 넘쳐나게 될 겁니다. Nebius는 이미 과열 최고점에서 거의 반토막이 났고요.
삽이 넘쳐날 조짐은 당장은 안 보이는데요.
다들 컴퓨팅이 엄청나게 부족하다고 하는데, 그게 대체 어디서 드러나나요? 누가 부족하다는 거죠?
저는 Claude에게 뭘 시킬 때마다 아무 문제 없이 처리되던데요. 다른 종류의 부족 사태에서 흔히 보이는 대기열 같은 걸 기다려 본 적이 없어요.
클라우드 업체에서 용량을 확보해야 하는데 못 구하는 기업(과 ISV)에서 드러납니다. GPU 인스턴스 종류가 나와 있을 때도 있지만, 필요하거나 돈을 내고 쓰고 싶은 것(예: TPU냐 V100이냐 B300이냐)이 아닌 경우가 많고, 스팟 인스턴스로 나오는 경우도 항상 있지는 않습니다(자주 없습니다).
하이퍼스케일러가 전부 수요 초과 상태라 기업에 약정(commit)을 요구합니다. 그래서 일시적인 수요만 있다면 [원하는 업체에서는] 구하지 못할 수도 있습니다. 게다가 GPU 서비스 업체라고 다 같은 건 아니에요. Semi Analysis에 그 이유를 잘 설명한 글이 여럿 있습니다.
삽도 금도 문제가 아닙니다, 문제는 인간이에요.
AI가 초고효율을 달성해서 실업으로 인한 체계적인 경제 수요 붕괴를 일으키거나, 기대에 못 미쳐서 대규모 기술 업종 자산 거품 붕괴를 일으키거나 둘 중 하나입니다.
AI가 기업 가치 평가만큼 대단하면 폭락하고, 그만큼 대단하지 않아도 폭락합니다. 계란을 전부 한 바구니에 담은 경제죠.
어떤 미래로 가든 엄청난 금융 폭풍이 옵니다.
AI는 모든 시장을 끝낼 기술로 가격에 반영돼 있어요. 평생 마지막으로 보게 될 강세장입니다.
성공해도 망하고 실패해도 망하는 이 피할 수 없는 외통수를 시장이 눈치채지 못한 게 기적일 정도입니다.
저는 벙커를 짓고 있습니다.
시장은 2008년에 배웠잖아요. Too Big To Fail(tm)이면 정부가 구제해 준다는 걸요.
2008년에 은행을 박살내지 않은 게 문제였어요. 미국 정부가 은행을 인수하고 투자자들을 싹 날려 버렸다면 지금쯤 반발이 훨씬 더 컸을 겁니다.
가능한 얘기입니다. 구글 AI 개요가 정확하다면, 과거의 투기적 광란과 어느 정도 비슷해 보입니다(예: 대공황 이전의 전기 관련 기업들). 규모는 더 큽니다(AI는 현재 35%, 전기 및 관련 기업은 25%). 신용이 더 많이 풀린 탓이라고 짐작합니다. 한편 금융위기 이전 시장에서 주택 부문 비중이 약 23%였는데 대공황만큼 경제에 큰 타격을 주지는 않았으니, AI 거품이 꺼져도 그렇게까지 나쁘지는 않기를 바랍니다.
AI는 지금 1880년대 철도 붐보다 큽니다.
"과잉 건설과 투기성 부채가 결국 1893년 공황으로 이어졌고, 미국은 극심한 경제 불황에 빠졌다."
구글 말로는 그렇지 않습니다. 1880년대 철도 거품은 시장의 60% 이상이었다고 합니다. 정확하지는 않을 수 있지만요.
"1880년대에 민간 철도 기업은 미국 주식시장 전체 시가총액의 60% 넘게 차지했으며, 1881년과 1885년에는 63%로 정점을 찍었다."