HybriMoE : Planification hybride CPU-GPU et gestion du cache pour une inférence efficace des MoE

HybriMoE: Hybrid CPU-GPU Scheduling and Cache Management for Efficient MoE Inference

April 8, 2025

Auteurs: Shuzhang Zhong, Yanfan Sun, Ling Liang, Runsheng Wang, Ru Huang, Meng Li

cs.AI

Résumé

L'architecture Mixture of Experts (MoE) a démontré des avantages significatifs en permettant d'augmenter la capacité du modèle sans augmentation proportionnelle des calculs. Cependant, la taille importante des modèles MoE introduit des demandes substantielles en mémoire, ce qui nécessite généralement un déchargement des experts sur des plateformes aux ressources limitées et entraîne un surcoût important. L'inférence hybride CPU-GPU a été proposée pour exploiter les calculs CPU afin de réduire le surcoût de chargement des experts, mais elle fait face à des défis majeurs : d'une part, les modèles d'activation des experts des modèles MoE sont très instables, rendant les stratégies de mappage fixes des travaux existants inefficaces ; d'autre part, le planning hybride CPU-GPU pour MoE est intrinsèquement complexe en raison de la diversité des tailles et structures des experts, ainsi que de la distribution inégale de la charge de travail, etc. Pour relever ces défis, dans cet article, nous proposons HybriMoE, un cadre d'inférence hybride CPU-GPU qui améliore l'utilisation des ressources grâce à un système novateur de planification CPU-GPU et de gestion du cache. HybriMoE introduit (i) une stratégie de planification intra-couche dynamique pour équilibrer la charge de travail entre le CPU et le GPU, (ii) un algorithme de préchargement inter-couche basé sur l'impact, et (iii) un algorithme de mise en cache basé sur des scores pour atténuer l'instabilité de l'activation des experts. Nous implémentons HybriMoE sur le cadre kTransformers et l'évaluons sur trois modèles de langage largement utilisés basés sur MoE. Les résultats expérimentaux montrent que HybriMoE atteint une accélération moyenne de 1,33 fois dans l'étape de préremplissage et de 1,70 fois dans l'étape de décodage par rapport au cadre d'inférence hybride MoE de pointe. Notre code est disponible à l'adresse : https://github.com/PKU-SEC-Lab/HybriMoE.

English

The Mixture of Experts (MoE) architecture has demonstrated significant advantages as it enables to increase the model capacity without a proportional increase in computation. However, the large MoE model size still introduces substantial memory demands, which usually requires expert offloading on resource-constrained platforms and incurs significant overhead. Hybrid CPU-GPU inference has been proposed to leverage CPU computation to reduce expert loading overhead but faces major challenges: on one hand, the expert activation patterns of MoE models are highly unstable, rendering the fixed mapping strategies in existing works inefficient; on the other hand, the hybrid CPU-GPU schedule for MoE is inherently complex due to the diverse expert sizes, structures, uneven workload distribution, etc. To address these challenges, in this paper, we propose HybriMoE, a hybrid CPU-GPU inference framework that improves resource utilization through a novel CPU-GPU scheduling and cache management system. HybriMoE introduces (i) a dynamic intra-layer scheduling strategy to balance workloads across CPU and GPU, (ii) an impact-driven inter-layer prefetching algorithm, and (iii) a score-based caching algorithm to mitigate expert activation instability. We implement HybriMoE on top of the kTransformers framework and evaluate it on three widely used MoE-based LLMs. Experimental results demonstrate that HybriMoE achieves an average speedup of 1.33times in the prefill stage and 1.70times in the decode stage compared to state-of-the-art hybrid MoE inference framework. Our code is available at: https://github.com/PKU-SEC-Lab/HybriMoE.