구글, TorchTPU를 TorchTitan에 통합한 TorchTitan TPU 공개… PyTorch 모델의 TPU 최적화 기법도 해설

TorchTitan TPU는 Meta의 TorchTitan을 포크한 것으로, https://github.com/google-pytorch/torchtitan/ 에서 공개되어 있다. 구글과 Meta가 공동으로 TPU 네이티브 실행 경로를 TorchTitan 안에 개발했다. 모델 정의는 CUDA 전용 코드를 포함하지 않는 단일 디바이스용 PyTorch(nn.Module)이며, 병렬화는 별도 레이어에서 끼워 넣는 구조다. 구성 트리를 바꾸면 어텐션은 Splash Attention으로, 최종 투영과 크로스 엔트로피는 융합 Pallas 커널로, MoE의 전문가 행렬곱은 그룹화된 Pallas 커널로 교체할 수 있으며, 모두 플래그로 전환한다.

성능의 핵심은 SimpleFSDP와 Graph Trainer다. SimpleFSDP는 파라미터 샤딩을 DTensor의 파라메트리제이션으로 표현하고, all-gather와 reduce-scatter를 트레이스된 그래프 안의 일반 노드로 만든다. FLUX에서는 이를 통해 JAX의 MaxDiffusion 대비 99%의 성능에 도달했다. Graph Trainer는 make_fx로 포워드, 손실, 백워드를 포함한 학습 스텝 전체를 하나의 XLA 프로그램으로 캡처하며, 레이어를 가로지르는 scan, 그래프 분할, 활성화 값의 호스트 메모리 오프로드를 가능하게 한다. functional AdamW를 그래프에 포함한 Qwen3-8B에서는 51~52% MFU에 도달했다.

최적화 절차는 플레이북으로 정리되어 있다. 베이스라인 확립(eager, compile, Graph Trainer), Pallas 커널과 타일 크기 조정, 배치 크기와 메모리 최적화, XLA와 libtpu 플래그 조정 및 프로파일링 순서이며, 대부분은 레시피 수준의 설정 변경으로 충분하다고 한다. FLUX.1 12B 모델을 TPU v6e-8에서 학습한 사례에서는 1스텝이 11.49초에서 2.78초로 4.1배 빨라졌다. Qwen3 235B는 64칩에서 8,192칩까지 weak scaling을 수행했으며, 칩당 처리량 저하는 8%에 그쳤다. 다만 이 규모의 데이터 병렬은 학습 레시피가 아니라 스케일링 실험이라고 명시되어 있다.

HuggingFace Transformers의 모델링 백엔드를 사용하면 Hub의 id로 모델을 구축하고, SimpleFSDP, 블록별 torch.compile, Splash Attention 등을 자동으로 적용할 수 있다. 지원하는 500개 이상의 모델을 같은 트레이너로 실행할 수 있지만, MoE 체크포인트는 전문가 병렬을 아직 지원하지 않는다. 새로운 모델 이식을 위해 코딩 에이전트용 onboarding 스킬도 저장소에 포함되어 있다. 로드맵에는 ML Goodput Measurement 및 Multi-Tier Checkpointing과의 통합, float8 레시피 확장, TPU 8t 및 8i 지원, 전문가 병렬과 장문맥 학습이 포함되어 있다.