Reproduce “LSH-MoE: Communication-efficient MoE Training via Locality-Sensitive Hashing” on OpenResearch

About this paper

Xiaonan Nie, Qibin Liu, Fangcheng Fu, Shenhan Zhu, Xupeng Miao, Xiaoyang Li, Yang Zhang, Shouda Liu and 1 more

Researchers from Peking University, Purdue University, and ByteDance developed LSH-MoE, a framework that uses Locality-Sensitive Hashing and residual-based error compensation to reduce the "all-to-all" communication bottleneck in distributed Mixture-of-Experts model training. The approach accelerates training convergence by up to 2.2x for T5-MoE and 1.6x for RoBERTa-MoE while maintaining model accuracy on various downstream tasks.

Key insights

• Input tokens in MoE models exhibit high similarity and clustering during all-to-all communication, presenting an opportunity for data compression.
• This observed redundancy can be exploited by transmitting compressed cluster centroids instead of individual tokens, directly addressing the communication bottleneck.
• A residual-based error compensation mechanism is crucial to maintain model accuracy and convergence quality when compressing intermediate activation values.

Summary by alphaXiv.