LSH-MoE: Communication-efficient MoE Training via Locality-Sensitive Hashing

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Nie, Xiaonan, Liu, Qibin, Fu, Fangcheng, Zhu, Shenhan, Miao, Xupeng, Li, Xiaoyang, Zhang, Yang, Liu, Shouda, Cui, Bin
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866910696387641344
author Nie, Xiaonan
Liu, Qibin
Fu, Fangcheng
Zhu, Shenhan
Miao, Xupeng
Li, Xiaoyang
Zhang, Yang
Liu, Shouda
Cui, Bin
author_facet Nie, Xiaonan
Liu, Qibin
Fu, Fangcheng
Zhu, Shenhan
Miao, Xupeng
Li, Xiaoyang
Zhang, Yang
Liu, Shouda
Cui, Bin
contents Larger transformer models always perform better on various tasks but require more costs to scale up the model size. To efficiently enlarge models, the mixture-of-experts (MoE) architecture is widely adopted, which consists of a gate network and a series of experts and keep the training cost constant by routing the input data to a fixed number of experts instead of all. In existing large-scale MoE training systems, experts would be distributed among different GPUs for parallelization, and thus input data requires additional all-to-all communications to access the target experts and conduct corresponding computations. However, upon evaluating the training process of three mainstream MoE models on commonly used GPU clusters, we found that the all-to-all communication ratio averaged around 45%, which significantly hinders the efficiency and scalability of training MoE models. In this paper, we propose LSH-MoE, a communication-efficient MoE training framework using locality-sensitive hashing (LSH). We first present the problems of scaling MoE training in existing systems and highlight the potential of exploiting token similarity to facilitate data compression. Then, we introduce an efficient LSH-based compression technique, which utilizes the cross-polytope hashing for rapid clustering and implements a residual-based error compensation scheme to alleviate the adverse impact of compression. To verify the effectiveness of our methods, we conduct experiments on both language models (e.g., RoBERTa, GPT, and T5) and vision models (e.g., Swin) for pre-training and fine-tuning tasks. The results demonstrate that our method substantially outperforms its counterparts across different tasks by 1.28x - 2.2x of speedup.
format Preprint
id arxiv_https___arxiv_org_abs_2411_08446
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle LSH-MoE: Communication-efficient MoE Training via Locality-Sensitive Hashing
Nie, Xiaonan
Liu, Qibin
Fu, Fangcheng
Zhu, Shenhan
Miao, Xupeng
Li, Xiaoyang
Zhang, Yang
Liu, Shouda
Cui, Bin
Distributed, Parallel, and Cluster Computing
Larger transformer models always perform better on various tasks but require more costs to scale up the model size. To efficiently enlarge models, the mixture-of-experts (MoE) architecture is widely adopted, which consists of a gate network and a series of experts and keep the training cost constant by routing the input data to a fixed number of experts instead of all. In existing large-scale MoE training systems, experts would be distributed among different GPUs for parallelization, and thus input data requires additional all-to-all communications to access the target experts and conduct corresponding computations. However, upon evaluating the training process of three mainstream MoE models on commonly used GPU clusters, we found that the all-to-all communication ratio averaged around 45%, which significantly hinders the efficiency and scalability of training MoE models. In this paper, we propose LSH-MoE, a communication-efficient MoE training framework using locality-sensitive hashing (LSH). We first present the problems of scaling MoE training in existing systems and highlight the potential of exploiting token similarity to facilitate data compression. Then, we introduce an efficient LSH-based compression technique, which utilizes the cross-polytope hashing for rapid clustering and implements a residual-based error compensation scheme to alleviate the adverse impact of compression. To verify the effectiveness of our methods, we conduct experiments on both language models (e.g., RoBERTa, GPT, and T5) and vision models (e.g., Swin) for pre-training and fine-tuning tasks. The results demonstrate that our method substantially outperforms its counterparts across different tasks by 1.28x - 2.2x of speedup.
title LSH-MoE: Communication-efficient MoE Training via Locality-Sensitive Hashing
topic Distributed, Parallel, and Cluster Computing
url https://arxiv.org/abs/2411.08446