PaperlayerView exact revision on arXiv

Source-linked AI summary

MixLoRA: Enhancing Large Language Models Fine-Tuning with LoRA-based Mixture of Experts

Dengchun Li, Yingzi Ma, Naizheng Wang, Zhengmao Ye, Zhiyuan Cheng, Yinghao Tang, Yan Zhang, Lei Duan, Jie Zuo, Cal Yang, Mingjie Tang

arXiv:2404.15159v3cs.CLcs.AI

TL;DR

Existing PEFT methods reduce fine-tuning resources but can underperform in multi-task settings, while conventional MoE models remain difficult to run on consumer GPUs. MIXLORA builds a sparse MoE from LoRA experts, shared FFN weights, attention adapters, and load balancing, and reports higher accuracy with lower latency and memory use.

  • Problem

    LoRA reduces fine-tuning resource requirements but can underperform in multi-task learning, while MoE models remain resource-intensive for consumer-grade GPUs.

  • Method

    MIXLORA inserts multiple LoRA-based experts into a frozen dense model with a shared FFN, top-k routing, independent attention LoRAs, and auxiliary load balancing.

  • Results

    MIXLORA improves multi-task accuracy by 9.8% over LoRA and 9% over DoRA, while its framework reduces token computation latency by 30% and GPU memory usage by about 40%.

  • Takeaways & Limitations

    MIXLORA provides a parameter-efficient sparse MoE approach that improves downstream-task performance while reducing the computational and memory costs of training and inference.

Abstract

from arXiv · show

Fine-tuning Large Language Models (LLMs) is a common practice to adapt pre-trained models for specific applications. While methods like LoRA have effectively addressed GPU memory constraints during fine-tuning, their performance often falls short, especially in multi-task scenarios. In contrast, Mixture-of-Expert (MoE) models, such as Mixtral 8x7B, demonstrate remarkable performance in multi-task learning scenarios while maintaining a reduced parameter count. However, the resource requirements of these MoEs remain challenging, particularly for consumer-grade GPUs with less than 24GB memory. To tackle these challenges, we propose MixLoRA, an approach to construct a resource-efficient sparse MoE model based on LoRA. MixLoRA inserts multiple LoRA-based experts within the feed-forward network block of a frozen pre-trained dense model and employs a commonly used top-k router. Unlike other LoRA-based MoE methods, MixLoRA enhances model performance by utilizing independent attention-layer LoRA adapters. Additionally, an auxiliary load balance loss is employed to address the imbalance problem of the router. Our evaluations show that MixLoRA improves about 9% accuracy compared to state-of-the-art PEFT methods in multi-task learning scenarios. We also propose a new high-throughput framework to alleviate the computation and memory bottlenecks during the training and inference of MOE models. This framework reduces GPU memory consumption by 40% and token computation latency by 30% during both training and inference.

1 Introduction

MIXLORA combines LoRA-based experts with a shared FFN and routing to build a resource-efficient sparse MoE for downstream and multi-task learning. Its evaluations report improved accuracy and reduced computation and memory costs.

  • Architecture: MIXLORA adds independent self-attention LoRA adapters and an auxiliary load-balance loss to improve performance and address expert imbalance.The method combines attention adaptation with routing regularization.
  • Results: MIXLORA outperforms LoRA and DoRA in single-task and multi-task scenarios, with multi-task gains of 9.8% over LoRA and 9% over DoRA.The paper also reports single-task gains of 5.8% over LoRA and 2.3% over DoRA on LLaMA-2 7B.
  • Architecture: MIXLORA constructs multiple LoRA-based experts with a frozen shared FFN block from a pre-trained dense model.The design targets efficient sparse MoE construction under limited resources.
  • Efficiency: 30% lower token computation latency and 40% or more lower GPU memory usage are reported for training and inference on a consumer-grade 24GB GPU.The framework uses LLaMA2-7B in half precision for the stated memory result.
  • Results: 9.8% average accuracy improvement over LoRA is reported for multi-task learning, while the corresponding single-task improvement is 5.8% on LLaMA2-7B.The reported multi-task comparison is against LoRA across the evaluated benchmarks.

2 Related Works

Prior work uses MoE routing and LoRA adaptations to improve specialization and efficiency, while MIXLORA adopts a shared-FFN architecture with multiple LoRA-based experts. The related-work context emphasizes sparse experts as a way to increase model width without proportional computational demand.

  • Mixture-of-Experts: MoE replaces the transformer feed-forward sub-layer with sparsely activated experts, enabling greater model width without a corresponding surge in computational demand.This describes the modern MoE formulation discussed in the paper.
  • MIXLORA: MIXLORA’s architecture uses n experts formed from an original FFN sublayer combined with different LoRAs, sharing the FFN weights across experts.The shared FFN distinguishes the architecture’s expert construction.
  • MIXLORA: MIXLORA combines a sparse MoE block built from LoRA-augmented transformer components with a top-k router that assigns tokens to experts.The architecture is presented as a continuation of these routing-based LoRA-MoE approaches.
  • Prior LoRA-MoE Methods: Related LoRA-MoE methods route tokens to domain-specific or task-specific experts to address data conflicts, lifelong learning, or knowledge forgetting.The cited examples include LLaVA-MoLE, MoRAL, and LoRAMoE.

3 MIXLORA

MIXLORA builds a sparse MoE architecture by combining LoRA experts, shared frozen FFN weights, top-k routing, attention-layer adapters, and load balancing. Its optimized propagation shares computation across projections and supports efficient multi-model training and inference.

  • MIXLORA Architecture: MIXLORA replaces a dense model’s FFN with LoRA-based experts that share the pretrained FFN weights and are selected per token by a Top-K router.Each expert combines the shared original FFN with distinct LoRA parameters.
  • MIXLORA Architecture: A top-2 router selects the two most suitable experts for each input token, enabling different experts to process diverse tasks.The router dynamically activates experts according to token inputs.
  • Expert Load Balance: MIXLORA adds an auxiliary load-balancing loss using token dispatch fractions and router-probability fractions to reduce uneven expert utilization.The coefficient a = 10^-2 is intended to balance experts without overwhelming the primary cross-entropy objective.
  • Attention Adaptation: Independent LoRA adapters are added to the attention layer alongside the LoRA-based MoE FFN to improve MIXLORA’s fine-tuning capability.The adapters cover the q, k, v, and o projections.
  • Performance Optimization: Sharing W1 and W3 computation reduces MIXLORA token computation latency by approximately 30% while maintaining the same model performance.The optimized process computes W1 and W3 before routing and slices their outputs using router weights; W2 remains dependent on those outputs.
  • Performance Optimization: Multi-MIXLORA training and inference pack inputs from multiple models into a single batch, while frozen dense weights can be shared across models.This design targets higher throughput when operating two or more MIXLORA models.

4 Experiments

Experiments evaluate MIXLORA and MIXDORA across commonsense reasoning tasks, model sizes, single-task and multi-task settings, router configurations, ranks, expert loading, and computation efficiency. MIXLORA generally improves accuracy while maintaining balanced expert usage and manageable latency.

  • Experimental Setup: Experiments cover ARC, BoolQ, OpenBookQA, PIQA, SocialIQA, HellaSwag, and WinoGrande, with accuracy measured across datasets.The evaluated tasks include question answering, classification, science completion, and fill-in-the-blank settings.
  • Experimental Setup: MIXLORA and MIXDORA use rank 16, eight experts, a top-2 router, and LoRA applied to attention and feed-forward projections.The attention projections are q, k, v, and o; feed-forward weights are w1, w2, and w3.
  • Single-Task Results: MIXLORA and MIXDORA achieve strong single-task results, with LLaMA3 8B scores of 83.5% and 84.1% exceeding LLaMA2 13B scores of 83.0% and 83.1%.The corresponding LoRA and DoRA scores show a larger gap between LLaMA3 8B and LLaMA2 13B.
  • Multi-Task Results: In multi-task learning, LoRA and DoRA lose 1.9% and 6.7% average accuracy, whereas MIXLORA and MIXDORA maintain nearly the same average accuracy.The results mix ARC, BoolQ, OpenBookQA, and PIQA for training before separate evaluation.
  • Ablation Study: A router-loss coefficient of 1e-3 gives MIXLORA its highest average accuracy, while disabling router loss or increasing the coefficient lowers performance.At coefficient 0, MIXLORA trails MIXDORA by 1%, but at 1e-3 it leads by 0.9%.
  • Ablation Study: MIXLORA and MIXDORA balance expert workloads, with average standard deviations of 0.0223 and 0.0328, respectively.The lower MIXLORA deviation indicates more uniform loading, while both values are described as sufficiently small.
  • Ablation Study: Both methods perform well from rank 2 through rank 16, but average accuracy drops at rank 32 because of convergence difficulties.MIXDORA slightly trails MIXLORA across the tested ranks.
  • Computation Efficiency: MIXLORA achieves the best accuracy on Gemma 2B and LLaMA2 7B, while its token latency lies between LoRA and DoRA.The optimized implementation improves latency without changing model accuracy, especially for the larger 7B model.

5 Conclusion

The conclusion presents MIXLORA as a parameter-efficient MoE method that combines multiple LoRA experts with a frozen shared FFN, attention adapters, and load balancing. It reports higher accuracy than LoRA and DoRA alongside reduced computation and GPU memory use.

  • Method: MIXLORA combines multiple LoRA-based experts with a frozen shared FFN block to form a parameter-efficient MoE method.The design also applies LoRA adapters to self-attention and uses auxiliary load balancing.
  • Efficiency: The high-performance framework reduces computational complexity by 30% and GPU memory usage by about 40% during training and inference.These savings apply when processing multiple MIXLORA models.
  • Results: MIXLORA improves average accuracy by 5.8% over LoRA and 2.3% over DoRA in single-task learning on LLaMA2 7B.The comparison is reported for the single-task setting.
  • Results: In multi-task learning, MIXLORA surpasses LoRA by 9.8% and DoRA by 9% in accuracy.The conclusion reports these improvements as the multi-task outcome.

A.1 Hyperparameters and Implementation Details

The appendix records the hyperparameter configurations and hardware environments used for the experiments. Seven-billion-parameter experiments run on 24GB GPUs, while larger models use 48GB GPUs.

  • Hyperparameters: Table 3 specifies LoRA, DoRA, MixLoRA, and MixDoRA configurations for Gemma-2B, LLaMA2-7B/13B, and LLaMA3-8B.The configurations target commonsense reasoning tasks.
  • Implementation Details: Experiments use 24GB GPUs for 7B models and 48GB GPUs for 8B and 13B models.The 24GB hardware includes RTX 3090, RTX A5000, and RTX 4090; 48GB experiments use an RTX A6000.
  • Implementation Details: The implementation uses Python 3.10, Ubuntu 22.04, and x86-64 CPUs.

A.2 Datasets

The experiments use several commonsense reasoning datasets spanning question answering, classification, science completion, and fill-in-the-blank tasks. Dataset descriptions are provided in Table 4, and the data are downloaded through HuggingFace’s Datasets library.

  • Dataset Description: Table 4 describes the datasets used in the experiments, including task names, domains, train and test sizes, and task types.
  • Data Acquisition: All datasets are downloaded from HuggingFace using the Python Datasets library.

A.3 Multi-task Learning Evaluation Result using Gemma-2B

Table 5 compares PEFT methods for multi-task learning across various tasks using Gemma 2B, reporting accuracy under single-task and multi-task setups.

  • Table 5 evaluates different PEFT methods across various tasks with Gemma 2B as the base model.
  • Evaluation setups: The single-task setup trains and evaluates separate PEFT modules for each task.
  • Evaluation setups: The multi-task setup trains on mixed tasks before evaluating tasks separately.

A.4 Experimental Results of Performance Metrics.

The performance evaluation measures token computation latency and peak GPU memory for LLaMA-2 7B. MIXLORA is faster than DoRA, while its optimized variant further reduces latency and memory in multi-model settings.

  • Latency and memory: 535.2µs token computation latency lets MIXLORA outperform DoRA’s 659.4µs forward-propagation latency.Both methods use approximately 15GB peak GPU memory.
  • Latency and memory: 245.3µs token computation latency makes LoRA faster than MIXLORA, whose K = 2 routing sends each token to two experts.
  • Optimization: 30% lower token computation latency reduces MIXLORA from 535.2µs to 462.5µs with the optimized algorithm.
  • Optimization: Nearly 45% lower peak GPU memory reduces two-model training usage from 15.1GB to 8.8GB and inference usage from 13.7GB to 7.2GB.

A.5 Robustness of MIXLORA Towards Different Rank

The supplied materials identify Gemma 2B performance measurements and LLaMA2-7B accuracy comparisons for varying ranks and router losses, alongside MIXLORA’s optimized forward-propagation procedure.

  • Performance evaluation: Gemma 2B performance results are reported in Table 7 for comparing MIXLORA and related methods.
  • Rank comparison: LLaMA2-7B accuracy is compared for MIXLORA and MIXDORA across varying ranks on commonsense reasoning tasks.
  • Router-loss comparison: LLaMA2-7B accuracy is also compared for MIXLORA and MIXDORA under different router-loss configurations.
  • Forward propagation: The optimized forward pass allocates multi-task sequences to MIXLORA modules, routes tokens with normalized Top-2 scores, and concatenates task outputs.
  • Forward propagation: The algorithm reduces duplicated FFN calculations by precomputing projected token sequences before producing LoRA-expert outputs.
Loading 2404.15159v3…