핵심 요약
구글이 TorchTPU를 PyTorch에 네이티브로 통합하고 TorchTitan TPU 포크를 통해 SimpleFSDP, Graph Trainer, Pallas 커널 등을 활용하여 TPU에서 대규모 모델을 재작성 없이 학습하고 최적화하는 방법을 공유합니다.
구현 방법
- torch.device("tpu") 백엔드의 네이티브 통합
- SimpleFSDP, Graph Trainer, Pallas 커널로 MFU를 높이는 최적화
- Dense 및 MoE 아키텍처에 대한 최적화 플레이북과 Hugging Face 백엔드 활용
주요 결과
- Dense/MoE에서 MFU 달성으로 TPU 활용도 향상
- 대규모 모델 학습의 확장성과 효율성 개선
