X-Stage: An Overlooked Pipeline Stage for Communication-Computation Overlap in DiT Inference
작성자
카테고리
Empty
저자
Jianwen Xian, Zhiyuan Xu, Yuchen Li, Ziliang Lai, Kang He, Zhen Huang, Aichen Feng, Jinyan Chen, Yilin Zhang, Qinqin Chen, Chengru Song
💡 개요
본 논문은 분산 Diffusion Transformer (DiT) 추론 시 통신-연산 오버랩을 개선하기 위해 기존 시스템에서 간과되었던 'X-Stage'라는 새로운 파이프라인 단계를 식별합니다. X-Stage는 원격 저장(remote stores) 발행 후 완료되기 전까지의 진행 상황을 포착하여, 송신자 측 백프레셔를 예측하고 이를 완화하기 위한 새로운 모델과 최적화된 커널을 제안합니다. 제안된 기법들은 DeepGEMM MegaMoE와 Ulysses sequence-parallel attention에서 상당한 성능 향상을 입증했습니다.
🔑 시사점 및 한계
•
통신-연산 오버랩의 새로운 최적화 기회 발굴: X-Stage는 기존 시스템에서 간과되었던 통신 발행 후 연산 복귀 사이의 유휴 시간을 활용하여 성능을 향상시킬 수 있는 새로운 가능성을 제시합니다.
•
정교한 백프레셔 예측 및 관리: 제안된 Burst-Gap 모델은 통신 발행 후 발생하는 백프레셔를 예측하고 이를 완화하는 데 효과적이며, 이를 통해 시스템의 안정성과 성능을 높일 수 있습니다.
•
향후 연구 과제: 제시된 모델과 기법들이 더 다양한 하드웨어 아키텍처 및 DiT 모델 구성에서 일반화될 수 있는지에 대한 추가적인 검증과 최적화가 필요합니다.
