GMoE: Empowering LLMs Fine-Tuning via MoE Graph Collaboration

TL;DR

This paper introduces GMoE, a graph-based MoE framework that improves expert collaboration and stability in LLM fine-tuning, addressing load imbalance issues with novel routing and coordination strategies.

Contribution

GMoE presents a new graph router and coordination strategies for MoE, enhancing expert collaboration and stability during LLM fine-tuning with parameter-efficient methods.

Findings

GMoE outperforms baseline models on multiple benchmarks.

The graph routing improves expert collaboration.

Coordination strategies increase model stability.

Abstract

The sparse Mixture-of-Experts (MoE) architecture of large language models (LLMs) confronts an inherent issue of load imbalance arising from the simplistic linear router strategy, which ultimately causes the instability and inefficient learning of LLMs. To address this challenge, we introduce a novel MoE graph-based framework $\textbf{GMoE}$, aimed at enhancing the collaboration among multiple experts. In GMoE, a graph router function is designed to capture the collaboration signals among experts. This enables all experts to dynamically allocate information derived from input data by sharing information with their neighboring experts. Moreover, we put forward two coordination strategies in GMoE: the $\textit{Poisson distribution-based distinction strategy}$ and the $\textit{Normal distribution-based balance strategy}$, to further release the capacity of each expert and increase the model stability in the fine-tuning of LLMs. Specifically, we leverage a parameter-efficient fine-tuning technique, i.e., Low-Rank Adaptation (LoRA), to implement the graph MoE architecture. Extensive experiments on four real-world benchmark datasets demonstrate the effectiveness of GMoE, showing the benefits of facilitating collaborations of multiple experts in LLM fine-tuning. The code of experimental implementation is available at https://github.com/BAI-LAB/GMoE