Deveoler Lee

PyTorch + TPU, 왜 지금 중요한가 대규모 AI 모델 학습의 시대가 본격화되면서, 수천~수만 개의 가속기를 효율적으로 활용하는 소프트웨어 스택의 중요성이 급격히 커지고 있어요. Google의 TPU(Tensor Processing Unit)는 Gemini, Veo 같은 자체 AI 플랫폼은 물론 Cloud 고객의 대규모 워크로드를 처리하는 핵심 인프라인데요. 문제는 ML 커뮤니티의 상당수가 PyTorch 생태계에 익숙하다는 점이에요. TorchTPU 는 바로 이 간극을 메우기 위해 설계된 스택으로, 기존 PyTorch 코드를 최소한의 변경만으로 TPU에서 네이티브로 실행할 수 있게 해줘요. 이 글에서는 TorchTPU의 핵심 아키텍처, Eager 모드 전략, 정적 컴파일 파이프라인, 그리고 2026년 로드맵까지 개발자 관점에서 깊이 살펴볼게요. TPU 하드웨어 아키텍처 이해하기 TorchTPU 를 제대로 이해하려면 먼저 TPU의 하드웨어 구조를 알아야 해요. TPU는 단순한 칩이 아니라 통합 네트워크 시스템이에요. ICI와 토러스 토폴로지 각 호스트에는 여러 개의 TPU 칩이 연결되어 있고, 칩들은 Inter-Chip Interconnect(ICI)를 통해 2D 또는 3D 토러스(Torus) 토폴로지로 연결돼요. 이 구조 덕분에 전통적인 네트워크 병목 없이 대규모 스케일업이 가능해요. GPU 클러스터에서 InfiniBand나 NVLink에 의존하는 것과는 근본적으로 다른 접근이에요. TensorCore와 SparseCore 칩 내부 실행은 TensorCore와 SparseCore로 나뉘어요. TensorCore : 단일 스레드 유닛으로 밀집 행렬 연산(dense matrix math)에 특화 SparseCore : 임베딩, gather/scatter 같은 불규칙 메모리 접근 패턴과 collective 오프로딩 처리 이 이중 구조가 TPU의 핵심 강점이고, TorchTPU는 이 두 코어를 모두 활용할 수 있도록 설계되었어요. Eager First 철학: ...

대규모 언어 모델(LLM) 학습에는 보통 수백, 수천 대의 GPU 클러스터가 필요해요. 100B 파라미터 모델 하나를 학습하려면 파라미터, 옵티마이저 상태, 그래디언트, 활성화 값까지 합쳐 테라바이트 단위의 메모리가 필요하기 때문이에요. 그런데 만약 단일 GPU로 1000억 파라미터 LLM 학습 이 가능하다면 어떨까요? Zhengqing Yuan 등이 발표한 MegaTrain 논문은 바로 이 문제를 정면으로 다루고 있어요. 풀 프리시전(Full Precision)을 유지하면서도 단 하나의 GPU에서 100B+ 규모 모델을 학습할 수 있는 시스템 아키텍처를 제안하고 있죠. 왜 단일 GPU 학습이 중요한가 현재 LLM 학습 생태계는 극심한 자원 불균형 상태에 있어요. NVIDIA H100 한 대가 약 6만 달러, 멀티 GPU 서버는 50만 달러 이상이에요. 대부분의 연구자와 중소 규모 조직은 이런 클러스터를 확보하기 어렵죠. 하지만 흥미로운 트렌드가 있어요. LLM 개발 워크플로우가 점차 사전학습 중심에서 사후학습(Post-training) 중심으로 이동하고 있다는 거예요. Instruction tuning, alignment, domain adaptation 같은 작업은 사전학습 대비 연산량이 훨씬 적어요. 원칙적으로는 단일 노드에서 충분히 수행할 수 있는 규모이지만, 기존 학습 시스템의 메모리 구조가 이를 가로막고 있었어요. 단일 GPU로 1000억 파라미터 LLM 학습 이 실현되면, GPU 클러스터 없이도 대형 모델을 커스터마이징할 수 있어요. 이건 LLM 민주화에 있어 결정적인 전환점이 될 수 있어요. GPU 메모리의 구조적 한계 100B 파라미터 모델을 BF16/FP32 혼합 정밀도로 학습한다고 가정해 볼게요. 파라미터 자체가 약 200GB(BF16), AdamW 옵티마이저 상태가 FP32로 파라미터의 3배(약 1.2TB), 그래디언트와 활성화 값까지 더하면 단일 GPU의 VRAM(80~192GB)으로는 어림도 없어요. 결국 핵심은 GPU 외부 메모리, ...