HybriMoE: Pianificazione Ibrida CPU-GPU e Gestione della Cache per Inferenza Efficiente di MoE

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

April 8, 2025

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

cs.AI

Abstract

L'architettura Mixture of Experts (MoE) ha dimostrato significativi vantaggi in quanto consente di aumentare la capacità del modello senza un incremento proporzionale del calcolo. Tuttavia, le grandi dimensioni dei modelli MoE introducono comunque notevoli richieste di memoria, che solitamente richiedono lo scaricamento degli esperti su piattaforme con risorse limitate e comportano un sovraccarico significativo. L'inferenza ibrida CPU-GPU è stata proposta per sfruttare il calcolo della CPU al fine di ridurre il sovraccarico di caricamento degli esperti, ma affronta importanti sfide: da un lato, i modelli di attivazione degli esperti nei modelli MoE sono altamente instabili, rendendo inefficienti le strategie di mappatura fisse presenti nei lavori esistenti; dall'altro, la pianificazione ibrida CPU-GPU per MoE è intrinsecamente complessa a causa delle diverse dimensioni degli esperti, delle loro strutture, della distribuzione disomogenea del carico di lavoro, ecc. Per affrontare queste sfide, in questo articolo proponiamo HybriMoE, un framework di inferenza ibrida CPU-GPU che migliora l'utilizzo delle risorse attraverso un innovativo sistema di pianificazione CPU-GPU e gestione della cache. HybriMoE introduce (i) una strategia di pianificazione intra-layer dinamica per bilanciare il carico di lavoro tra CPU e GPU, (ii) un algoritmo di prefetching inter-layer guidato dall'impatto, e (iii) un algoritmo di caching basato su punteggio per mitigare l'instabilità dell'attivazione degli esperti. Implementiamo HybriMoE sul framework kTransformers e lo valutiamo su tre LLM basati su MoE ampiamente utilizzati. I risultati sperimentali dimostrano che HybriMoE raggiunge un'accelerazione media di 1,33 volte nella fase di prefill e di 1,70 volte nella fase di decodifica rispetto al framework di inferenza ibrida MoE più avanzato. Il nostro codice è disponibile all'indirizzo: 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.