핵심 요약
구글이 HeyGen의 Avatar IV를 TPUs로 포팅하고, 8칩 메시 상에서 FSDP와 Ulysses 시퀀스 병렬화를 활용해 1.86배의 실시간 스트리밍 속도 향상을 달성했습니다.
구현 방법
- torchax와 XLA를 사용해 트릴리움 TPUs에 모델 배치
- 8칩 메시에서 FSDP 및 Ulysses 병렬성 구성
- 파이프라인화로 모든-대-모두(all-to-all) 컬렉트를 최적화
- 희소 어텐션 블록 크기 정렬로 마스크 패딩 제거
- softmax 직렬 의존성 우회: 미리 계산된 코시-슈왈츠 상한 사용
- Pallas 커널과 컴파일러 최적화 도입은 두 단계 품질 게이트 통과 후 배포
주요 결과
- 1.86배 속도 향상으로 실시간 스트리밍 가능
- 바이트 식별적 동등성 픽셀 출력을 보장하기 위한 2단계 품질 게이트를 거쳐 배포
