Hacker News

Google open-sources TorchTPU, a native PyTorch back end for TPUs

Google, TPU용 네이티브 PyTorch 백엔드 TorchTPU 공개

Google이 TPU에서 PyTorch 모델과 연산을 실행하는 오픈소스 백엔드 TorchTPU를 공개했습니다. ATen 커널, XLA 컴파일 캐시, 분산 학습과 torch.compile()을 지원하며, 사용하려면 TPU 환경이 필요합니다.

에디터 노트

TorchTPU의 공개로 PyTorch 코드가 TPU에서 실행되는 경로는 구체적으로 마련됐지만, 실제 사용 범위는 TPU 접근성과 지원 범위에 달려 있습니다. 저장소는 ATen 커널과 분산 학습, 모델 예제를 제시하고 설치와 빌드 조건도 자세히 적습니다. 다만 성능 비교나 지원 연산의 범위는 제공된 설명에 없습니다. 댓글에서 DeepSeek의 추론 엔진과 TorchTPU를 같은 경쟁으로 보는 질문이 나왔지만, 다른 댓글은 TorchTPU를 PyTorch/XLA의 대체 경로이자 TPU 학습 생태계의 도구로 구분합니다. AWS Trainium용 PyTorch 지원과 Tenstorrent용 구현도 언급되며, 하드웨어별 백엔드 개발이 이어지는 사례로 볼 수 있습니다. TPU를 실제로 확보할 수 있는지, 기존 PyTorch/XLA와 어떤 차이를 보이는지는 이 자료만으로 판단하기 어렵습니다.

AI 요약

TorchTPU는 Google Tensor Processing Unit(TPU)에서 PyTorch 모델과 연산을 실행하도록 만든 백엔드입니다. PyTorch ATen 연산을 위한 TPU 커널과 XLA 연산의 실행을 최적화하는 컴파일 캐시를 포함합니다. 그래프 실행을 위한 torch.compile() 통합, 데이터 병렬화와 텐서 병렬화 기반 분산 학습도 지원합니다.

구성과 사용 환경

저장소에는 Llama, Qwen, ResNet, DLRM, minGPT, UNet 예제가 있으며 디버깅과 벤치마크 도구도 제공합니다. 문서는 TPU VM에서 실행하는 Marimo 노트북과 사용자 가이드로 구성됩니다. 프로젝트는 Python 3.11 이상을 지원하며, 공개된 PyPI 버전은 0.2.0입니다. 안내문은 TPU VM 중 v6e GCE 인스턴스를 가장 간편한 시작 환경으로 제시합니다.

빌드와 설치

소스 빌드에는 Bazel을 사용합니다. 휠 빌드 명령에는 --config=wheel_common이 필요합니다. 이 설정은 XLA/MLIR 백엔드를 하나의 libxla_base.so로 묶는 빌드 경로를 선택합니다. 설정을 빼면 버전별 백엔드가 중복 포함돼 torch를 불러올 때 정적 등록 중복으로 종료될 수 있다고 설명합니다. 기본 설정은 원격 빌드 실행과 캐시를 사용하며, RBE 인증 정보가 없으면 --config=no_rbe를 뒤에 붙여 끌 수 있습니다. 설치할 때는 CPU용 PyTorch를 받도록 PyTorch CPU 휠 저장소를 추가하라고 안내합니다.

Hacker News 반응

  • @yipinwong — AI 엔지니어가 아닌 사람의 어리석은 질문인데, 오늘 DeepSeek이 공개한 추론 엔진에 대한 Google의 대응인가요?
    • @Tanjreeve — 둘은 목적이 다르고 스택에서도 위치가 다릅니다. TorchTPU는 추론보다는 학습 쪽에 가깝고, Google TPU를 더 많은 사람이 익숙하게 쓰도록 하려는 움직임으로 보입니다.
    • @mathisfun123 — 이건 pytorch/xla를 대체하는 겁니다.
  • @npalli — 흥미롭네요. AWS Trainium용으로 Amazon도 비슷한 작업을 하고 있습니다.
  • @tuna74 — Linux 드라이버도 공개됐나요?
  • @pcmoritz — Tenstorrent 가속기용 오픈소스 libtpu를 부수 작업으로 만들고 있습니다. 이미 TorchTPU의 대부분을 실행할 수 있고, PyTorch를 타일 기반 가속기로 낮출 때 필요한 작업을 보여줍니다.
  • @ismailmaj — 예전에는 TPU에서 PyTorch를 돌리려면 XLA와 씨름해야 했는데, 이제 마침내 첫 번째 수준의 지원을 받게 된 것 같네요.

원문: GitHub / 번역·요약: Trawling

원문 보기 ↗토론 보기 ↗