[Hacker News 요약] 구글, PyTorch 워크로드를 TPU에서 네이티브로 실행하는 TorchTPU 공개로 AI 개발 가속화
30
설명
구글은 PyTorch 개발자들이 자사의 Tensor Processing Unit(TPU)을 더욱 쉽고 효율적으로 활용할 수 있도록 'TorchTPU'를 공개했습니다. 이 새로운 통합 스택은 PyTorch 모델을 구글 스케일의 TPU에서 네이티브로 실행할 수 있게 하여, 현대 AI 인프라의 복잡성을 줄이고 성능을 극대화하는 것을 목표로 합니다. TorchTPU는 유용성, 이식성, 그리고 뛰어난 성능을 핵심 가치로 삼아 개발되었습니다.
### 배경 설명
현대 인공지능(AI) 인프라 구축의 도전 과제는 근본적으로 변화했습니다. 최신 머신러닝 모델은 수천 개의 가속기를 아우르는 분산 시스템을 활용해야 하며, 모델이 수십만 개의 칩 클러스터에서 실행됨에 따라 소프트웨어는 성능, 하드웨어 이식성, 신뢰성에 대한 새로운 요구 사항을 충족해야 합니다.
구글의 TPU는 이러한 슈퍼컴퓨팅 인프라의 핵심입니다. 이 맞춤형 ASIC(주문형 반도체)은 구글 자체 AI 플랫폼(예: Gemini, Veo)과 클라우드 고객의 대규모 워크로드 학습 및 서비스를 지원합니다. AI 커뮤니티 전체가 TPU의 모든 기능을 쉽게 활용할 수 있어야 하며, 많은 잠재 사용자가 PyTorch로 모델을 구축하기 때문에, PyTorch가 TPU에서 네이티브하고 효율적으로 작동하도록 하는 통합은 매우 중요합니다. TorchTPU는 이러한 배경에서 탄생했으며, 기존 PyTorch 워크로드를 최소한의 코드 변경으로 마이그레이션하면서도 하드웨어의 모든 컴퓨팅 성능을 끌어낼 수 있는 API와 도구를 제공하고자 합니다.
### 설계 원칙: 유용성, 이식성, 성능
TorchTPU는 유용성, 이식성, 뛰어난 성능이라는 세 가지 핵심 원칙을 기반으로 설계되었습니다. TPU 시스템은 단순한 칩이 아니라, 호스트에 여러 칩이 연결되고 각 칩이 ICI(Inter-Chip Interconnect)를 통해 다른 칩과 연결되는 통합 네트워크입니다. 이 ICI는 칩들을 고효율 2D 또는 3D 토러스 토폴로지로 연결하여 전통적인 네트워킹 병목 현상 없이 대규모 확장을 가능하게 합니다. 각 칩 내에서는 TensorCores(밀집 행렬 연산)와 SparseCores(불규칙 메모리 접근 패턴)가 실행을 분담합니다. TorchTPU의 목표는 PyTorch 개발자가 기존 스크립트에서 초기화만 'tpu'로 변경하면 핵심 로직 수정 없이 학습 루프를 실행할 수 있도록, PyTorch와 같은 사용 경험을 제공하는 것입니다.
### TorchTPU 스택의 기술 아키텍처
TorchTPU는 'Eager First' 철학을 기반으로 유연성을 제공합니다. 개발자가 즉시 정적 그래프 컴파일을 요구하는 대신, PyTorch의 'PrivateUse1' 인터페이스를 사용하여 일반적인 PyTorch 텐서가 TPU에서 작동하도록 구현했습니다. 이를 통해 PyTorch의 Eager 실행 경험을 우선시합니다.
세 가지 Eager 모드를 지원하는데, Debug Eager는 디버깅에 유용하지만 느리고, Strict Eager는 비동기 단일 연산 디스패치를 통해 기본 PyTorch 경험을 모방합니다. 가장 혁신적인 Fused Eager 모드는 연산 스트림에 대한 자동 리플렉션을 사용하여 단계들을 더 크고 계산 집약적인 청크로 즉석에서 융합하여 TPU에 전달합니다. 이는 Strict Eager 대비 50%에서 100% 이상의 성능 향상을 제공합니다. 모든 모드는 단일 호스트 또는 다중 호스트 설정에서 작동하는 공유 컴파일 캐시의 지원을 받습니다.
최고 성능을 위해서는 `torch.compile` 인터페이스를 통한 전체 그래프 컴파일을 지원합니다. Torch Dynamo를 사용하여 FX 그래프를 캡처한 후, Torch Inductor 대신 XLA를 기본 백엔드 컴파일러로 활용합니다. XLA는 TPU 토폴로지에 최적화되어 있으며, ICI를 통한 밀집 계산과 집단 통신 간의 중요한 오버랩을 기본적으로 이해합니다. PyTorch 연산자는 XLA의 주요 중간 표현(IR)인 StableHLO로 직접 매핑되어 고도로 최적화된 TPU 바이너리를 생성합니다. 또한 Pallas 및 JAX로 작성된 사용자 정의 커널도 지원합니다.
### 분산 학습 및 MPMD(Multi-Program Multi-Data) 지원
TorchTPU는 대규모 환경에서 Eager 및 컴파일 모드의 유연성과 유용성을 유지하기 위해 PyTorch의 분산 API에 중점을 두었습니다. 현재 Distributed Data Parallel (DDP), Fully Sharded Data Parallel v2 (FSDPv2), PyTorch의 DTensor를 기본으로 지원하며, PyTorch의 분산 API를 기반으로 구축된 많은 서드파티 라이브러리가 TorchTPU에서 변경 없이 작동함을 검증했습니다.
이전 PyTorch/XLA의 주요 한계 중 하나는 순수 SPMD(Single Program Multiple Data) 코드만 지원했다는 점입니다. 그러나 PyTorch 입력의 현실은 종종 다른 랭크에서 실행되는 코드에 약간의 차이가 있다는 것입니다(예: '랭크 0' 프로세스가 로깅이나 분석을 위해 추가 작업을 수행하는 경우). TorchTPU는 이러한 이기종 실행(MPMD)을 신중하게 지원하도록 설계되었으며, 필요한 경우 통신 프리미티브를 격리하여 정확성을 최소한의 비용으로 유지합니다. 이 접근 방식은 기존 PyTorch 개발자에게 TPU 사용 경험이 가능한 한 자연스럽도록 보장합니다.
### TPU 하드웨어 인식 및 최적화
TPU는 매우 높은 성능과 효율성을 달성할 수 있지만, 최적의 모델 설계는 다른 하드웨어와 약간 다를 수 있습니다. 예를 들어, 현재 세대 TPU는 128 또는 256 차원에서 피크 행렬 곱셈 효율성을 달성하는 반면, 많은 모델이 어텐션 헤드 차원을 64로 하드코딩하는 경우가 많습니다. 모델을 128 또는 256 차원으로 수정하면 TPU 칩의 크고 밀집된 효율적인 텐서 코어를 더 잘 활용할 수 있습니다. TorchTPU는 먼저 올바른 실행을 확립한 다음, 곧 출시될 심층 가이드라인을 사용하여 최적화되지 않은 아키텍처를 식별하고 리팩토링하거나, 사용자 정의 커널을 삽입하여 최적의 하드웨어 활용을 돕는 계층적 워크플로우를 제공합니다.
### 가치와 인사이트
TorchTPU는 PyTorch 개발자들이 구글의 강력한 TPU 슈퍼컴퓨팅 인프라를 훨씬 더 쉽게 활용할 수 있도록 함으로써 AI 개발의 문턱을 낮춥니다. 기존 PyTorch 코드를 최소한의 수정으로 TPU에서 실행할 수 있게 하여, 개발자들이 새로운 하드웨어 아키텍처에 대한 깊은 이해 없이도 최첨단 AI 가속기의 이점을 누릴 수 있게 합니다. 이는 모델 학습 시간을 단축하고, 대규모 분산 환경에서의 효율성을 극대화하며, 궁극적으로 더 복잡하고 혁신적인 AI 모델의 개발을 가속화할 것입니다. 특히, MPMD 지원은 실제 PyTorch 워크로드의 유연성을 보장하여 개발자들이 하드웨어 제약보다는 모델 로직에 집중할 수 있도록 돕습니다. 이는 AI 연구 및 상용화에 있어 중요한 실무적 영향을 미칠 것입니다.
### 기술·메타
- PyTorch
- Google TPU (Tensor Processing Unit)
- XLA (Accelerated Linear Algebra)
- StableHLO
- Torch Dynamo
- Pallas
- JAX
- Helion DSL
- Distributed Data Parallel (DDP)
- Fully Sharded Data Parallel v2 (FSDPv2)
- PyTorch DTensor
- vLLM
- TorchTitan
### 향후 전망
TorchTPU 팀은 2026년 이후를 위한 야심찬 로드맵을 가지고 있습니다. 컴파일러 팀의 주요 초점은 동적 시퀀스 길이 및 배치 크기로 인해 발생하는 재컴파일을 줄이는 것입니다. XLA 내에서 고급 바운드 다이내미즘을 구현하여 컴파일 오버헤드 없이 형태 변화를 처리하는 것을 목표로 합니다. 또한 표준 연산을 위한 포괄적인 사전 컴파일된 TPU 커널 라이브러리를 구축하여 첫 실행 반복의 지연 시간을 대폭 줄일 예정입니다.
2026년에는 공개 GitHub 저장소 출시, PyTorch의 Helion DSL 통합, `torch.compile`을 통한 동적 형태 기본 지원, 분리된 메모리 및 컴퓨팅 스트림을 가진 비동기 코드베이스 마이그레이션을 용이하게 하는 네이티브 멀티큐 지원, 그리고 vLLM 및 TorchTitan과 같은 생태계 핵심 요소와의 깊은 통합 및 전체 Pod 크기 인프라까지의 선형 스케일링 검증이 계획되어 있습니다. 이러한 노력들은 TorchTPU를 PyTorch 생태계에서 마찰 없는 백엔드로 만들고, 다른 가속기 플랫폼과의 경쟁에서 우위를 점하며, 차세대 AI를 위한 TPU 슈퍼컴퓨팅 하드웨어의 접근성과 성능을 더욱 향상시킬 것으로 전망됩니다.
📝 원문 및 참고
- Source: Hacker News
- 토론(HN): [news.ycombinator.com](https://news.ycombinator.com/item?id=47881786)
- 원문: [링크 열기](https://developers.googleblog.com/torchtpu-running-pytorch-natively-on-tpus-at-google-scale/)
---
출처: Hacker News · [원문 링크](https://developers.googleblog.com/torchtpu-running-pytorch-natively-on-tpus-at-google-scale/)
신고 · 불법·유해·아동 안전(CSAE) 관련 콘텐츠
댓글 0
아직 댓글이 없습니다. 첫 댓글을 남겨 보세요.