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

ai.euno.news 자동 뉴스 요약

PyTorch가 구글 TPU를 위한 네이티브 백엔드 TorchTPU를 공개했다. 구글과 메타가 공동 개발한 오픈소스 프로젝트로, 기존 PyTorch 코드의 model.to("cuda")를 model.to("tpu")로 바꾸는 정도의 수정만으로 TPU에서 모델을 실행하는 것이 목표다. 즉시 실행(eager execution)과 torch.compile()을 모두 지원하며, 구글은 이미 일부 대형 클라우드 고객의 추론·학습 환경에서 TorchTPU를 사용하고 있다고 밝혔다.¹

기존 선택지는 제약이 컸다. 2020년 등장한 PyTorch/XLA는 지연 텐서(Lazy Tensor)에 기반해 컴파일 영역을 코드에서 명시해야 했고, 컴파일 지연과 PyTorch의 즉시 실행 미지원, 최신 분산 처리 기능과의 격차가 문제로 지적됐다. TorchTPU는 런타임에서 JAX 의존성을 없애고, Torch Dynamo 기반 torch.compile()과 PyTorch 분산 API를 직접 지원하는 방향을 택했다. 현재 즉시 실행과 컴파일 모드, 연산을 모아 처리하는 DeferAndFuse 모드를 구현했으며 약 450개 ATen 연산자를 지원한다.¹

다만 코드를 한 줄 바꾼다고 TPU의 최고 성능까지 자동으로 얻는 것은 아니다. GPU에 맞춘 텐서 형태와 샤딩, 연산 구성은 TPU 구조에 맞게 다시 조정해야 한다고 PyTorch 측도 설명했다.¹ 이를 보완하는 것이 대규모 학습 프레임워크 TorchTitan 연동이다.

TorchTitan이 맡는 TPU 최적화#

구글은 메타의 TorchTitan을 TPU용으로 조정한 포크를 함께 공개했다. TorchTitan은 Llama 3, Qwen 3, FLUX 같은 밀집·전문가 혼합(MoE)·확산 트랜스포머의 사전학습 레시피를 제공한다. TorchTPU가 C10D 집합 통신 API와 DeviceMesh를 지원하므로 FSDP와 텐서 병렬화 같은 PyTorch 방식도 TPU 경로에서 활용할 수 있다.²

성능 최적화는 모델 본체나 공통 학습 루프를 고치는 대신 구성 트리를 통해 특정 블록을 교체하는 방식이다. TPU 레시피는 일반 어텐션을 Splash Attention으로 바꾸고, 최종 투영과 교차 엔트로피 손실을 결합해 전체 로짓을 고대역폭 메모리(HBM)에 기록하지 않는다. MoE 행렬곱에는 그룹형 Pallas 커널을 적용한다. 각 커널은 JAX 함수를 torch.tpu.pallas.jax_op 데코레이터로 PyTorch 사용자 정의 연산으로 감싼다.² TorchTPU가 코드 이식의 문턱을 낮춘다면, TorchTitan은 실제 분산 학습에서 필요한 샤딩과 TPU 전용 커널을 제공하는 역할을 맡는다.

읽기 목록은 이 브라우저에 저장됩니다.