Papers
Topics
Authors
Recent
2000 character limit reached

Read-ME: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System Co-Design (2410.19123v1)

Published 24 Oct 2024 in cs.CL and cs.LG

Abstract: The proliferation of LLMs has led to the adoption of Mixture-of-Experts (MoE) architectures that dynamically leverage specialized subnetworks for improved efficiency and performance. Despite their benefits, MoE models face significant challenges during inference, including inefficient memory management and suboptimal batching, due to misaligned design choices between the model architecture and the system policies. Furthermore, the conventional approach of training MoEs from scratch is increasingly prohibitive in terms of cost. In this paper, we propose a novel framework Read-ME that transforms pre-trained dense LLMs into smaller MoE models (in contrast to "upcycling" generalist MoEs), avoiding the high costs of ground-up training. Our approach employs activation sparsity to extract experts. To compose experts, we examine the widely-adopted layer-wise router design and show its redundancy, and thus we introduce the pre-gating router decoupled from the MoE backbone that facilitates system-friendly pre-computing and lookahead scheduling, enhancing expert-aware batching and caching. Our codesign therefore addresses critical gaps on both the algorithmic and system fronts, establishing a scalable and efficient alternative for LLM inference in resource-constrained settings. Read-ME outperforms other popular open-source dense models of similar scales, achieving improvements of up to 10.1% on MMLU, and improving mean end-to-end latency up to 6.1%. Codes are available at: https://github.com/VITA-Group/READ-ME.

Definition Search Book Streamline Icon: https://streamlinehq.com
References (64)
  1. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 34:8583–8595, 2021.
  2. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1–39, 2022.
  3. Mixtral of experts. arXiv preprint arXiv:2401.04088, 2024.
  4. M3vit: Mixture-of-experts vision transformer for efficient multi-task learning with model-accelerator co-design. Advances in Neural Information Processing Systems, 35:28441–28457, 2022.
  5. Mod-squad: Designing mixtures of experts as modular multi-task learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11828–11837, 2023.
  6. Edge-moe: Memory-efficient multi-task vision transformer architecture with task-level sparsity via mixture-of-experts. In 2023 IEEE/ACM International Conference on Computer Aided Design (ICCAD), pages 01–09. IEEE, 2023.
  7. Serving moe models on resource-constrained edge devices via dynamic expert swapping. arXiv preprint arXiv:2308.15030, 2023.
  8. Fast inference of mixture-of-experts language models with offloading. arXiv preprint arXiv:2312.17238, 2023.
  9. Orca: A distributed serving system for {{\{{Transformer-Based}}\}} generative models. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 521–538, 2022.
  10. Sparse upcycling: Training mixture-of-experts from dense checkpoints. arXiv preprint arXiv:2212.05055, 2022.
  11. Moefication: Transformer feed-forward layers are mixtures of experts. arXiv preprint arXiv:2110.01786, 2021.
  12. Emergent modularity in pre-trained transformers. arXiv preprint arXiv:2305.18390, 2023.
  13. Exploiting transformer activation sparsity with dynamic inference. arXiv preprint arXiv:2310.04361, 2023.
  14. Prompt-prompted mixture of experts for efficient llm generation. arXiv preprint arXiv:2404.01365, 2024.
  15. Laszlo A. Belady. A study of replacement algorithms for a virtual-storage computer. IBM Systems journal, 5(2):78–101, 1966.
  16. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300, 2020.
  17. Robust mixture-of-expert training for convolutional neural networks. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 90–101, 2023.
  18. Openmoe: An early effort on open mixture-of-experts language models. arXiv preprint arXiv:2402.01739, 2024.
  19. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
  20. On the use of arxiv as a dataset, 2019.
  21. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020.
  22. Stack Excahnge. Stack exchange data dump, 2024.
  23. The lazy neuron phenomenon: On emergence of activation sparsity in transformers. arXiv preprint arXiv:2210.06313, 2022.
  24. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint arXiv:2401.06066, 2024.
  25. Accelerate — huggingface.co. https://huggingface.co/docs/accelerate/index. [Accessed 22-05-2024].
  26. Se-moe: A scalable and efficient mixture-of-experts distributed training and inference system, 2023.
  27. Moe-infinity: Activation-aware expert offloading for efficient moe serving, 2024.
  28. Pre-gated moe: An algorithm-system co-design for fast and scalable mixture-of-expert inference. arXiv preprint arXiv:2308.12066, 2023.
  29. Judging llm-as-a-judge with mt-bench and chatbot arena, 2023.
  30. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38–45, Online, October 2020. Association for Computational Linguistics.
  31. Attention is all you need. Advances in neural information processing systems, 30, 2017.
  32. Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020.
  33. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024.
  34. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
  35. Together Computer. Redpajama: An open source recipe to reproduce llama training dataset, 2023.
  36. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv e-prints, 2019.
  37. Wikimedia Foundation. Wikimedia downloads.
  38. Deepspeed- inference: Enabling efficient inference of transformer models at unprecedented scale. In SC22: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–15, 2022.
  39. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 611–626, 2023.
  40. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pages 2397–2430. PMLR, 2023.
  41. Openllama: An open reproduction of llama, May 2023.
  42. Sheared llama: Accelerating language model pre-training via structured pruning. arXiv preprint arXiv:2310.06694, 2023.
  43. Llm-pruner: On the structural pruning of large language models. Advances in neural information processing systems, 36:21702–21720, 2023.
  44. Slicegpt: Compress large language models by deleting rows and columns. arXiv preprint arXiv:2401.15024, 2024.
  45. Laco: Large language model pruning via layer collapse. arXiv preprint arXiv:2402.11187, 2024.
  46. Compresso: Structured pruning with collaborative prompting learns compact large language models. arXiv preprint arXiv:2310.05015, 2023.
  47. Huggingface tgi inference engine. https://github.com/huggingface/text-generation-inference. [Accessed 20-05-2024].
  48. An inter-reference gap model for temporal locality in program behavior. ACM SIGMETRICS Performance Evaluation Review, 23(1):291–300, 1995.
  49. Learn to be efficient: Build structured sparsity in large language models. arXiv preprint arXiv:2402.06126, 2024.
  50. Deja vu: Contextual sparsity for efficient llms at inference time. In International Conference on Machine Learning, pages 22137–22176. PMLR, 2023.
  51. Hire: High recall approximate top-k𝑘kitalic_k estimation for efficient llm inference. arXiv preprint arXiv:2402.09360, 2024.
  52. Optimizing dynamic neural networks with brainstorm. In 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23), pages 797–815, 2023.
  53. Zero-offload: Democratizing billion-scale model training. In USENIX Annual Technical Conference (USENIX ATC 21), pages 551–564, 2021.
  54. Atom: Low-bit quantization for efficient and accurate llm serving. arXiv preprint arXiv:2310.19102, 2023.
  55. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323, 2022.
  56. Lin Zhao. Awrq: Activation-aware weight reformulation quantizer for large language models.
  57. The emergence of essential sparsity in large pre-trained models: The weights that matter. Advances in Neural Information Processing Systems, 36, 2024.
  58. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024.
  59. Sida-moe: Sparsity-inspired data-aware serving for efficient and scalable large mixture-of-experts models, 2024.
  60. GitHub - NVIDIA/TensorRT-LLM: TensorRT-LLM provides users with an easy-to-use Python API to define Large Language Models (LLMs) and build TensorRT engines that contain state-of-the-art optimizations to perform inference efficiently on NVIDIA GPUs. TensorRT-LLM also contains components to create Python and C++ runtimes that execute those TensorRT engines. — github.com. https://github.com/NVIDIA/TensorRT-LLM. [Accessed 20-05-2024].
  61. Ekko: A {{\{{Large-Scale}}\}} deep learning recommender system with {{\{{Low-Latency}}\}} model update. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 821–839, 2022.
  62. Deepum: Tensor migration and prefetching in unified memory. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, pages 207–221, 2023.
  63. Sentinel: Efficient tensor migration and allocation on heterogeneous memory systems for deep learning. In 2021 IEEE International Symposium on High-Performance Computer Architecture (HPCA), pages 598–611. IEEE, 2021.
  64. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023.

Summary

We haven't generated a summary for this paper yet.

Ai Sparkles Streamline Icon: https://streamlinehq.com Summarize Now
Slide Deck Streamline Icon: https://streamlinehq.com

Whiteboard

Dice Question Streamline Icon: https://streamlinehq.com

Open Problems

We haven't generated a list of open problems mentioned in this paper yet.

Ai Sparkles Streamline Icon: https://streamlinehq.com Generate Now
Lightbulb Streamline Icon: https://streamlinehq.com

Continue Learning

We haven't generated follow-up questions for this paper yet.

Ai Sparkles Streamline Icon: https://streamlinehq.com Generate Now
List To Do Tasks Checklist Streamline Icon: https://streamlinehq.com

Collections

Sign up for free to add this paper to one or more collections.

User Circle Single Streamline Icon: https://streamlinehq.com Sign Up
X Twitter Logo Streamline Icon: https://streamlinehq.com

Tweets

Sign up for free to view the 1 tweet with 0 likes about this paper.

User Circle Single Streamline Icon: https://streamlinehq.com Sign Up for Free