TorchTitan에서 1T MoE로 확장하기 [P]

요약

본 기사는 GH200 GPU를 활용하여 1T 파라미터 규모의 MoE 모델을 학습시키는 방법을 다룹니다. 저자는 병렬화 전략으로 FSDP와 EP 조합이 PP보다 효율적임을 제시하며, 특히 Full activation checkpointing을 통해 유휴 상태였던 GPU 자원을 재활용함으로써 성능 향상을 이끌어냈습니다.

핵심 포인트

• 1T MoE 모델 학습 시 GH200 GPU를 활용하여 높은 처리량(약 1,300 토큰/초) 달성 가능.
• 병렬화 전략은 FSDP + EP 조합이 PP보다 유리하며, 이는 네트워크 대기 시간을 줄이는 데 기여함.
• Full activation checkpointing을 통해 GPU의 유휴 컴퓨팅 자원을 활용하여 성능 향상을 얻음.
• InfiniBand와 같은 네트워크 통신 지연 시간이 전체 학습 시간에서 주요 병목 현상으로 작용함.