fMoE: Fine-Grained Expert Offloading for Large Mixture-of-Experts Serving

作者/机构: Hanfei Yu (Stevens Institute of Technology), Hao Wang (Rutgers University), Xingqi Cui (Rice University), Hong Zhang (University of Waterloo), Hao Wang (Stevens Institute of Technology)

A1 主要贡献

核心问题: 混合专家(MoE)架构虽然降低了大型语言模型(LLM)的训练成本,但在服务(serving)阶段存在严重的内存效率问题,因为即使是稀疏激活的专家(expert)也需要常驻GPU内存。现有的专家卸载(offloading)方法试图通过将不活跃的专家从GPU移至CPU内存来解决此问题,但这些方法通常是粗粒度的,导致在推理延迟和内存占用之间难以取得良好平衡,要么延迟过高,要么内存占用过大。

研究目标: 本文旨在解决MoE服务中的延迟-内存权衡困境,提出一个能够同时实现低推理延迟和高内存效率的专家卸载系统。

创新点:
* fMoE系统设计: 本文设计并实现了一个名为fMoE的细粒度专家卸载系统。该系统通过精确指导专家的预取、缓存和卸载决策,有效降低了推理延迟和模型内存占用。
* Expert Map数据结构: 提出了一种名为“专家地图”(expert map)的新数据结构。它能够以迭代(iteration)为单位,精细地记录门控网络(gate network)输出的专家选择概率分布,从而捕捉MoE模型细粒度的专家选择行为。
* 结合语义信息的专家地图搜索: fMoE利用输入提示(prompt)的语义嵌入(semantic embeddings)来增强专家地图的搜索过程。通过结合基于语义和基于专家激活轨迹(trajectory)的信息,fMoE能够更准确地找到最合适的历史专家地图来指导当前推理的专家卸载。

核心成果: fMoE系统在HuggingFace Transformers基础上进行了原型实现,并部署在一个六GPU的测试平台上。实验结果表明,与当前最先进的解决方案相比,fMoE能够将推理延迟降低47%,并将专家命中率提高36%。

A3 背景知识与动机

2.1 LLM服务

LLM服务的两个阶段。与传统的深度学习模型推理不同,大型语言模型(LLM)的服务过程包含两个连续的阶段:预填充(prefill)和解码(decode)。如图1a所示,在预填充阶段,LLM首先计算输入提示(prompt)中所有词元(token)的中间键值(KV)状态,填充KV缓存,然后生成第一个答案词元。在解码阶段,LLM以自回归的方式逐个词元地生成答案,其中先前生成的词元被用于生成下一个词元。

两个阶段的不同特性。预填充和解码阶段各有其独特的特点。预填充阶段只需要一次迭代,并行处理所有输入词元并生成第一个答案词元。解码阶段则跨越多次迭代,每次迭代生成一个词元,直到答案生成完毕。由于这些不同特性,近期研究【【38, Splitwise: Efficient Generative LLM Inference Using Phase Splitting, 2024, ISCA】, 【61, Distserve: Disaggregating prefill and decoding for goodputoptimized large language model serving, 2024, OSDI】】指出,预填充阶段是计算密集型(compute-bounded),而解码阶段是内存密集型(memory-bounded)。因此,人们通常使用不同的指标来量化LLM服务两个阶段的性能。对于预填充阶段,通常使用首词元时间(Time-To-First-Token, TTFT),衡量从收到用户请求到生成第一个答案词元之间的延迟。对于解码阶段,则使用每秒词元数(Tokens-Per-Second, TPS)或每输出词元时间(Time-Per-Output-Token, TPOT)来衡量LLM服务的生成速率。

2.2 基于MoE的LLM服务

MoE架构在LLM中的应用。通过在Transformer块【【51, Attention is all you need, 2017, NeurIPS】】中集成MoE层,MoE架构【【58, Twenty Years of Mixture of Experts, 2012, TNNLS】】已成为现代LLM(如Mixtral【【21, Mixtral of Experts, 2024, arXiv】】、Snowflake Arctic【【46, Snowflake Arctic: The Best LLM for Enterprise AI, 2024, https://www.snowflake.com/en/ data-cloud/arctic/】】和DeepSeek-MoE【【10, DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models, 2024, arXiv】】)的流行骨干。如图1a所示,在基于MoE的LLM中,前馈网络(FFN)模块被MoE层取代。每个MoE层由一个门控网络(gate network)和一组专家网络(expert networks)组成。在每个Transformer块内部,自注意力模块首先根据输入隐藏状态计算注意力,然后门控网络决定激活哪个(或哪些)专家来计算输出表示。与传统的密集型LLM相比,基于MoE的LLM在训练和推理期间只激活一部分参数,从而减少了计算开销,同时在参数数量相当的情况下提供了比密集型LLM更优的生成性能【【1, Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone, 2024, arXiv】, 【10, DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models, 2024, arXiv】, 【21, Mixtral of Experts, 2024, arXiv】, 【46, Snowflake Arctic: The Best LLM for Enterprise AI, 2024, https://www.snowflake.com/en/ data-cloud/arctic/】】, 【53, Announcing Grok, 2023, https://x.ai/blog/grok 】, 【57, Qwen2 Technical Report, 2024, arXiv】】。

MoE服务中的内存效率问题。尽管MoE在节省训练计算方面有优势,但基于MoE的LLM服务仍然存在GPU内存效率低下的问题,因为MoE推理需要将所有模型参数(包括那些不活跃的专家)加载到GPU内存中。例如,Mixtral-8×7B【【21, Mixtral of Experts, 2024, arXiv】】和DeepSeek-MoE【【10, DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models, 2024, arXiv】】由于MoE中专家激活的稀疏性,在推理过程中分别有72%和83%的参数处于不活跃状态,这导致了较低的内存效率和服务吞吐量。因此,为了高效地服务大型MoE模型,我们必须寻找解决MoE架构固有的内存效率低下问题的方案。

2.3 延迟-内存权衡

MoE服务效率的提升方法。最近,一些研究被提出来以提高基于MoE的LLM服务效率。图2描述了MoE服务的设计空间。现有的主要研究可分为两类:有损服务(Lossy serving)应用压缩【【39, Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing Policy, 2024, ICLR】】、剪枝【【26, STUN: Structured-ThenUnstructured Pruning for Scalable MoE Pruning, 2024, arXiv】】和量化【【23, Mixture of Quantized Experts (MoQE): Complementary Effect of Low-bit Quantization and Robustness, 2023, arXiv】】技术到原始MoE模型,以减少服务内存需求。然而,这类工作以牺牲生成质量为代价来换取服务效率。无损服务(Lossless serving)专注于将模型权重(参数【【4, DeepSpeed-Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale, 2022, SC22】, 【36, Get Up and Running with Large Language Models, https://ollama.com/】】或专家【 【16, Fast Inference of Mixture-of-Experts Language Models with Offloading, 2023, arXiv】, 【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】, 【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】)中在时间或空间上稀疏利用的部分从GPU内存卸载到CPU内存,旨在保持合理的推理延迟。具体来说,专家卸载旨在提前预测专家的激活,在推理期间仅在GPU内存中预取或缓存必要的专家。我们选择无损服务来设计fMoE,因为这类方法避免修改模型,从而保证了生成质量。

现有卸载方案的困境。然而,现有的卸载解决方案在服务基于MoE的LLM时,无法在延迟-内存权衡中达到最佳点。图1b比较了现有最先进(SOTA)卸载解决方案的性能(即推理延迟和内存占用),它们要么提供低推理延迟但内存占用大(例如,No-offload和MoE-Infinity【【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】),要么反之(例如,ProMoE【【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】】、Mixtral-Offloading【【16, Fast Inference of Mixture-of-Experts Language Models with Offloading, 2023, arXiv】】和DeepSpeedInference【【4, DeepSpeed-Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale, 2022, SC22】】)。这一困境的关键原因是,基于MoE的仅解码器(decoder-only)LLM具有均衡的专家路由(balanced expert routing)【【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】】,使得现有解决方案难以找到有效的模式来指导专家卸载。现有研究指出了造成这一困境的两个主要原因:首先,大多数基于MoE的LLM是仅解码器架构,与编码器-解码器MoE LLM相比,它们表现出均匀的专家激活模式和较低的专家访问偏斜度【【18, Lynx: Enabling Efficient MoE Inference through Dynamic Batch-Aware Expert Selection, 2024, arXiv】, 【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】】。其次,最近的基于MoE的LLM在训练时使用了一种独特的负载均衡损失(load balancing loss)【【1, Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone, 2024, arXiv】, 【10, DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models, 2024, arXiv】, 【21, Mixtral of Experts, 2024, arXiv】, 【46, Snowflake Arctic: The Best LLM for Enterprise AI, 2024, https://www.snowflake.com/en/ data-cloud/arctic/】】, 【53, Announcing Grok, 2023, https://x.ai/blog/grok】】,这强制门控网络在同一MoE层内平衡路由到每个专家的词元,确保在整个训练过程中没有专家是无关紧要的。这种均衡路由降低了专家模式的可预测性,从而使现有解决方案效果不佳 。


图1: 混合专家(MoE)大型语言模型(LLM)服务。


图2: 基于MoE的LLM服务的设计空间。

2.4 现有MoE卸载解决方案

粗粒度专家模式的低效性。现有的专家卸载方法【【16, Fast Inference of Mixture-of-Experts Language Models with Offloading, 2023, arXiv】, 【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】依赖于粗粒度的专家模式,这对于指导卸载是低效的。我们将粗粒度信息定义为在请求级别收集的专家模式,其中信息是在一个请求提示的多个迭代中聚合的。例如,MoE-Infinity【【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】跟踪请求级别的专家激活。细粒度信息则定义为在每个推理迭代中单独观察到的专家模式。图3a展示了Mixtral-8×7B【【21, Mixtral of Experts, 2024, arXiv】】的粗粒度和细粒度专家激活热图示例。热图记录了32个MoE层中的专家激活情况,每个层包含八个专家,并激活八个中的两个专家来计算表示。虽然细粒度(迭代级别)的热图显示出清晰的专家激活模式,但聚合的粗粒度(请求级别)热图却降低了可预测性。

熵分析证明细粒度模式更优。为了证明这一点,我们分析了三种流行MoE模型在每个MoE层专家激活的香农熵【【44, A mathematical theory of communication, 1948, The Bell System Technical Journal】】。熵是信息论中量化变量不确定性和不可预测性的重要指标。一个均衡的专家激活模式(例如,四个专家的概率分布为[0.25, 0.25, 0.25, 0.25])会导致高熵,这表明该模式更不可预测,难以选择专家。图3b展示了三种MoE模型(Mixtral-8×7B【【21, Mixtral of Experts, 2024, arXiv】】、Qwen1.5-MoE【【57, Qwen2 Technical Report, 2024, arXiv】】和Phi-3.5-MoE【【1, Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone, 2024, arXiv】】)在两个真实世界数据集(LMSYS-Chat1M【【60, LMSYS-Chat-1M: A Large-Scale Real-World LLM Conversation Dataset, 2023, arXiv】】和ShareGPT【【45, ShareGPT: Share Your Wildest ChatGPT Conversations, https://sharegpt.com/】】)上每层计算的平均熵。粗粒度专家模式的熵明显高于细粒度模式,这意味着粗粒度的专家模式在预测方面可能效果较差。图3c显示了在推理迭代过程中每层的平均熵。虽然在推理开始时熵较低,但随着迭代的进行,由于专家激活信息的聚合,熵逐渐增加,从而变得更不可预测 。

本文主张。与粗粒度的专家卸载解决方案相反,我们认为专家卸载应该由细粒度的设计来仔细指导:分析迭代级别的模式,理解模型的专家选择偏好,并利用请求提示的语义特征。

2.5 粗粒度卸载的问题

现有粗粒度专家卸载方案存在三个问题
1. 延迟-内存权衡不佳。现有解决方案以粗粒度方式预取和卸载专家,要么过分侧重于降低推理延迟而导致内存占用过大【【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】,要么以降低内存占用为代价而严重增加推理延迟【【4, DeepSpeed-Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale, 2022, SC22】, 【16, Fast Inference of Mixture-of-Experts Language Models with Offloading, 2023, arXiv】】。
2. 专家命中率低。现有解决方案采用粗粒度的专家模式跟踪方法(例如,MoE-Infinity中的专家激活矩阵【【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】),这产生的专家模式对于指导卸载决策效果不佳,导致专家命中率低和推理延迟高。
3. 忽略MoE模型和提示的异构性。现有解决方案在很大程度上忽略了不同MoE模型和输入提示的独特特征,并以一种“一刀切”的方式为它们提供服务【【4, DeepSpeed-Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale, 2022, SC22】, 【16, Fast Inference of Mixture-of-Experts Language Models with Offloading, 2023, arXiv】, 【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】, 【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】),这忽略了针对MoE服务中异构模型和提示进行细粒度优化的机会。

细粒度设计的优势。图4显示了在使用LMSYS-Chat-1M数据集【【60, LMSYS-Chat-1M: A Large-Scale Real-World LLM Conversation Dataset, 2023, arXiv】】为三种流行的基于MoE的LLM(Mixtral-8×7B【【21, Mixtral of Experts, 2024, arXiv】】、Qwen1.5-MoE【【57, Qwen2 Technical Report, 2024, arXiv】】和Phi-3.5-MoE【【1, Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone, 2024, arXiv】】)服务时,粗粒度和细粒度专家卸载设计在不同预取距离下的专家命中率。预取距离指的是在目标层激活其专家之前,预取指令提前发出的层数。通过利用细粒度的专家卸载,我们可以实现比粗粒度方法显著更高的专家命中率,并通过适应变化的预取距离来保持更好的性能。

(a) Mixtral-8×7B的粗粒度与细粒度专家热图。颜色越深表示专家激活次数越多。

(b) 三个MoE模型和两个数据集在粗粒度和细粒度专家模式下的每层平均熵。熵越高表示可预测性越低。 (c) 三个MoE模型和两个数据集在推理迭代过程中的每层平均熵。随迭代聚合专家模式会逐渐降低可预测性。

图3: 粗粒度(请求级别)和细粒度(迭代级别)的专家模式和可预测性分析。


图4: 在不同预取距离下,服务于Mixtral-8x7B、Qwen1.5-MoE和Phi-3.5-MoE时,粗粒度和细粒度专家卸载设计的专家命中率。

A2 方法细节

3. fMoE概览

3.1 目标与挑战

fMoE的设计目标。fMoE旨在实现以下三个目标:
1. 内存高效且推理延迟最小的MoE服务。我们已证明现有专家卸载方案【【16, Fast Inference of Mixture-of-Experts Language Models with Offloading, 2023, arXiv】, 【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】, 【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】未能在MoE服务中有效权衡延迟与内存(§2.3)。我们的目标是通过提出细粒度专家卸载,同时实现低内存占用和低推理延迟。
2. 最小化专家预取错误导致的专家未命中。专家预取涉及对未来专家激活的预测,是专家卸载方案中的关键步骤。然而,最近的研究【【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】】表明,由于预测错误导致的专家未命中会引起推理中高昂的按需专家加载延迟。我们应最小化专家未命中并减轻专家卸载中的预测错误。
3. 适应异构的MoE模型和提示。在真实场景中,MoE推理可能服务于异构模型【【10, DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models, 2024, arXiv】, 【21, Mixtral of Experts, 2024, arXiv】, 【46, Snowflake Arctic: The Best LLM for Enterprise AI, 2024, https://www.snowflake.com/en/ data-cloud/arctic/】】, 【53, Announcing Grok, 2023, https://x.ai/blog/grok 】, 【57, Qwen2 Technical Report, 2024, arXiv】】和变化的提示【【45, ShareGPT: Share Your Wildest ChatGPT Conversations, https://sharegpt.com/】 】, 【60, LMSYS-Chat-1M: A Large-Scale Real-World LLM Conversation Dataset, 2023, arXiv】】。现有解决方案采用“一刀切”的设计来处理不同的模型和提示,而我们应设计自适应的专家卸载以适应MoE服务中的异构性。

为实现目标需应对的挑战。为实现上述目标,我们必须解决三个关键挑战:
1. 如何在预取和卸载专家时最大化专家命中率? 专家命中率直接关系到推理延迟。命中越多的专家,需要按需加载的专家就越少。我们提出一种细粒度的专家卸载解决方案以实现高专家命中率。
2. 如何适应不同的MoE模型和提示? 异构的MoE模型和输入提示展现出独特的系统和语义特征。我们应通过细粒度的优化来打造我们的解决方案,以实现自适应性。
3. 如何在管理专家时避免额外的系统开销? 我们的设计在服务现有MoE LLM时不能引入额外的系统开销。我们在fMoE中应用了一系列系统优化,以确保服务效率并最小化额外开销。

3.2 架构与工作流

fMoE的架构。图5描述了fMoE的架构和工作流,它由三个主要组件构成:
* 专家地图存储(Expert Map Store):我们记录专家地图,这是fMoE中定义的一种新数据结构,用于追踪来自历史请求提示的细粒度专家激活模式。与现有粗粒度专家追踪方法(例如,MoE-Infinity中的专家激活矩阵【【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】)相比,专家地图提供了更细致的专家选择偏好。专家地图存储动态地保留最有用和最独特的专家地图,以供推理期间的实时查询。
* 专家地图匹配器(Expert Map Matcher):当一个请求提示到达时,fMoE在专家地图存储中搜索合适的专家地图,以在推理前指导专家预取。专家地图搜索通过计算两个方面的相似度得分来指导:语义相似度和轨迹相似度。
* 专家缓存(Expert Cache):接收到匹配的专家地图后,fMoE将专家从CPU内存预取到GPU内存以在推理中执行计算。如果超出专家缓存的容量,fMoE会驱逐并卸载低优先级的专家权重到CPU内存。


图5: fMoE的架构与工作流。

fMoE的工作流程。fMoE遵循以下五个步骤,以实现内存高效且推理延迟最小的MoE服务:
1. 步骤1:推理上下文收集。在每次推理迭代之前,fMoE收集必要的上下文,例如语义嵌入和先前的专家激活轨迹(§4.1),并将其提供给专家地图匹配器进行混合相似度匹配。
2. 步骤2:专家地图相似度匹配。接收到迭代级别的上下文后,专家地图匹配器通过将输入上下文数据与专家地图存储中的历史上下文数据进行比较,找到并提取最相似的专家地图(§4.2)。匹配的专家地图被转发到专家缓存,以指导专家预取和卸载决策。
3. 步骤3:引导式专家预取和卸载。我们动态计算专家选择阈值,以在搜索到的专家地图的指导下,确定在MoE模型中预取和卸载哪个(些)专家(§4.3)。然后,fMoE将专家权重从CPU预取到GPU内存,并在达到缓存限制时将缓存的专家从GPU卸载到CPU(§4.5)。
4. 步骤4:专家服务。整个推理过程包括预填充(Prefill)阶段的一次迭代和解码(Decode)阶段的多次迭代。对于每次迭代中的每个MoE层,如果相应的权重在GPU内存中可用(定义为专家命中),fMoE直接服务于门控网络所需的专家。否则,fMoE按需将专家权重从CPU加载到GPU以执行无损服务(定义为专家未命中)。
5. 步骤5:专家地图更新。fMoE观察每次迭代后产生的新专家地图,并在专家地图存储中更新它们(§4.4)。当达到存储容量(例如,1K个专家地图)时,fMoE通过识别并丢弃冗余的专家地图来对专家地图存储进行去重,以保持多样性,最大化为任何请求提示提供有效专家地图的可能性。

3.3 问题形式化

问题定义。我们考虑在一个GPU集群上为一个拥有$L$个MoE层的MoE-based LLM提供服务,每个MoE层有一个门控网络和$J$个专家。每层的门控网络选择前$K \in [1, J]$个专家进行计算。该MoE模型处理并为一个由$W$个独立请求提示组成的工作负载生成答案。每个请求提示$w \in [W]$包含在预填充和解码阶段处理的多次迭代,其中$[W]$是请求提示的集合。令$E_{l, j}^{(i)}$表示第$i$次迭代中第$l$层的第$j$个专家,其中$l \in [L]$,$j \in [J]$,并且$i \in [w]$。在每次迭代$i$中,我们最多可以做出$L \cdot J$个预取决策。令$E_{cache}$和$E_{activate}^{(i)}$分别表示缓存的专家集合和第$i$次迭代激活的专家集合。因此,我们表示一个专家$E_{l, j} \in E_{activate}$是命中(由$E_{cache}^{(i)}$服务)还是未命中(从CPU内存按需加载)的结果:

$$\begin{aligned} R_{l,j}^{(i)} = \begin{cases} 1, & \text{if } (E_{l,j}^{(i)} \in E_{activate}^{(i)}) \wedge (E_{l,j}^{(i)} \notin E_{cache}^{(i)}), \\ 0, & \text{otherwise}, \end{cases} \end{aligned}$$

其中$R_{l, j}^{(i)} = 1$意味着$E_{l, j}^{(i)}$是一次未命中。由于MoE模型中的所有专家通常被设计为具有相同的权重大小,我们假设专家的加载时间$T_e$和内存占用$M_e$是同质的。因此,总的按需加载延迟$T$是在整个推理过程中对每个专家的所有迭代求和,即$T := T_e \cdot \sum_{w \in [W]} \sum_{i \in [w]} \sum_{l \in [L]} \sum_{j \in [J]} R_{l, j}^{(i)}$。

优化问题。最后,利用上述定义,我们将MoE专家卸载问题形式化为一个整数线性规划(ILP)优化问题:

$$\min _{\{E_{l, j}^{(i)}\}}\left(T_{e} \cdot \sum_{w \in[W]} \sum_{i \in[w]} \sum_{l \in[L]} \sum_{j \in[J]} R_{l, j}^{i}\right)$$ $$\text{s.t. } |E_{cache}^{(i)}| \leq L \cdot J, \quad \forall i \in [w], \forall w \in [W],$$ $$|E_{activate}^{(i)}| = L \cdot K, \quad \forall i \in [w], \forall w \in [W],$$ $$|E_{cache}^{(i)}| \cdot M_{e} \leq M, \quad \forall i \in [w], \forall w \in [W].$$

目标是最小化按需加载延迟(理想情况下,通过完美预测使$T=0$),同时限制缓存专家的总内存占用以满足可用的GPU内存$M$。约束1表示预取的专家总数不应超过MoE模型中所有专家的总数。约束2表示激活的专家总数必须等于所有$L$层中top $K$专家数量的总和。约束3描述了预取专家的总内存占用必须受限于可用的GPU内存大小。需要注意的是,解决这个ILP问题本身就是NP难的【【9, Introduction to Algorithms, 2022, MIT press】】,而在现实中,预取专家总会有预测失误,这进一步使问题复杂化。因此,我们选择为fMoE采用一种基于启发式的设计。


图6: 由一张专家地图追踪的专家选择。

4. fMoE的设计

4.1 专家地图

一种新的数据结构。我们提出了一种新的数据结构,称为“专家地图”(expert maps),用于以细粒度的方式追踪专家激活模式。图6描绘了一张专家地图的结构。在第$i$次迭代中,第$l$个自注意力层首先计算注意力状态。门控网络接收注意力并计算一个关于第$l$层所有专家的概率分布$P_l^{(i)} \in R^J$:

$$\mathbf{P}_l^{(i)}:=\left\{p_{l, 1}^{(i)}, \ldots, p_{l, j}^{(i)}, \ldots, p_{l, J}^{(i)}\right\}, \quad \sum_{j \in[J]} p_{l, j}^{(i)}=1, \forall p_{l, j}^{(i)} \geq 0.$$

然后,从$P_l^{(i)}$中选择前$K \in [1, J]$个专家来计算第$l$层的表示。我们收集所有$L$层的概率分布$P_l^{(i)}$,以形成第$i$次迭代的专家地图:

$$map_i := \{\mathbf{P}_1^{(i)}, \dots, \mathbf{P}_l^{(i)}, \dots, \mathbf{P}_L^{(i)}\}, \quad l \in [L].$$

专家地图的优势。通过追踪专家地图,我们指导fMoE发现细粒度的专家模式——即通过概率分布体现的迭代级别的专家选择偏好。直观地说,分析概率分布使fMoE不仅能识别哪些专家被二进制地选择或忽略,还能从门控网络的角度评估分配给每个专家的置信度或偏好。专家地图的设计相比于现有的粗粒度专家追踪方法(例如,MoE-Infinity【【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】追踪请求级别的专家命中次数)有两个关键优势。首先,现有工作只关注聚合的请求级别的专家激活,而专家地图则以详细的专家选择来追踪单个迭代。其次,现有工作只记录专家命中次数,而我们追踪详细的概率分布。值得注意的是,专家地图可以通过对概率分布应用一个top $K$选择算子并聚合迭代间的专家计数,轻松恢复粗粒度信息,因此可以推广到现有的追踪方法。我们在§6.5中评估了fMoE与其他追踪方法的对比,以展示专家地图的有效性。

4.2 专家地图搜索

混合搜索方法。在为MoE模型预测和预取专家时,通常会定义一个预取距离以避免影响推理延迟【【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】】。预取距离指的是在目标层激活其专家之前,预取指令提前发出的层数,这与内存预取中的同名术语类似【【25, When Prefetching Works, When It Doesn’t, and Why, 2012, TACO】】。设$d$为要服务的MoE模型的预取距离。图7显示,fMoE采用两种细粒度的搜索方法来共同匹配专家地图,以指导专家预取。语义搜索通过比较输入嵌入与历史嵌入来找到具有相似输入的专家地图,而轨迹搜索则观察先前的专家轨迹(即概率分布)并匹配相似的专家地图。我们结合语义和轨迹特征来提高fMoE的地图匹配和专家卸载准确性。两种搜索方法的有效性在§6.5中进行了评估。

基于语义的专家地图搜索。对于初始层$l \in [1, d]$,由于预取距离$d$的存在,现有解决方案【【16, Fast Inference of Mixture-of-Experts Language Models with Offloading, 2023, arXiv】, 【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】, 【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】在目标层准备好激活专家之前,无法观察到专家模式来进行预测和预取。因此,它们通常为预取初始层定义粗粒度的规则。例如,MoE-Infinity【【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】预取所有历史数据点中最受欢迎的专家。相比之下,fMoE利用输入提示中的语义线索来搜索最有用的专家地图,这需要零知识的专家激活模式。在服务请求提示并记录其专家地图时,我们记录每次推理迭代的语义嵌入。由于现有的基于MoE的LLM都包含一个用于分析输入语义的嵌入层,因此使用模型原始嵌入层的输出来提取语义嵌入是自然而然的。对于任何输入提示,我们在语义嵌入$sem_{new} \in R^{B \times h}$与专家地图存储中历史语义嵌入的集合$sem_{old} \in R^{C \times h}$之间计算成对的余弦相似度得分$score^{sem} \in R^{B \times C}$:

$$score_{x,y}^{sem} := \frac{sem_x^{new} \cdot sem_y^{old}}{\|sem_x^{new}\| \cdot \|sem_y^{old}\|}, \quad x \in [B], y \in [C],$$

其中$B$是输入提示的批量大小,$C$是专家地图存储的容量,$h$是隐藏维度的大小。然后,对于提示$x$,选择得分最高的历史迭代$y$。我们使用所选迭代的部分专家地图,$\{P_1^{(y)}, \dots, P_d^{(y)}\} \in map_{old_y}$,来指导层$l \in [1, d]$的专家预取。


图7: 语义和轨迹专家地图搜索。

基于轨迹的专家地图搜索。与层$l \in [1, d]$不同,我们可以观察到前面$(l-1)$层的专家概率轨迹来为层$l \in [d+1, L]$搜索专家地图。与基于语义的搜索类似,我们计算观察到的轨迹$map_{new} \in R^{B \times (l-1)J}$与专家地图存储中历史专家地图集合$map_{old} \in R^{C \times (l-1)J}$之间的成对余弦相似度得分$score^{map} \in R^{B \times C}$:

$$score_{x,y}^{map} := \frac{map_x^{new} \cdot map_y^{old}}{\|map_x^{new}\| \cdot \|map_y^{old}\|}, \quad x \in [B], y \in [C].$$

我们选择得分最高的历史迭代。然后,我们使用所选专家地图中的$P_{l+d}^y \in map_{old_y}$来指导目标层$l+d$的专家预取,其中$d$是预取距离。通过结合这两种专家地图搜索方法,我们为MoE服务中的每次推理迭代精心定制了指导专家预取的地图。通过这种设计,专家地图搜索对端到端推理延迟引入的开销可以忽略不计,我们将在§6.7中对此进行论证。

4.3 专家预取

基于搜索结果的细粒度预取。给定为层$l \in [L]$搜索和定制的专家地图,我们以细粒度的方式指导专家预取。

相似度感知的专家选择。在迭代过程中收集到不同上下文的情况下,fMoE搜索到的专家地图也具有不同的相似度得分,这反映了搜索的置信度。为了量化相似度得分与专家命中率之间的相关性,我们使用§4.2中描述的方法,运行了三个MoE模型(Mixtral-8×7B、Qwen1.5-MoE和Phi-3.5-MoE)和两个数据集(LMSYS-Chat-1M和ShareGPT)。对于每次推理迭代,我们计算相似度得分,收集由搜索到的专家地图指导的专家命中率,并计算皮尔逊相关系数【【8, Pearson Correlation Coefficient, 2009, Noise Reduction in Speech Processing】】。皮尔逊系数通常用于衡量变量之间的相关性,系数接近1表示强正相关,接近0表示弱相关。图8显示了三个MoE模型和两个数据集中,相似度得分与专家命中率之间的皮尔逊系数。结果表明,如果使用相应的专家地图,高相似度得分可能与高专家命中率相关。因此,我们设计fMoE的专家预取为相似度感知的。


图8: 三个MoE模型和两个数据集中,语义和轨迹余弦相似度与专家命中率之间的皮尔逊相关系数。

动态专家选择阈值。对于一个得分为$score \in [-1, 1]$的层$l \in [L]$进行预取,我们首先动态计算一个专家选择阈值$\delta_l \in [0, 1]$,由下式给出:

$$\delta_{l}:=\operatorname{Clip}(1-score, 0, 1)=\max (0, \min (1-score, 1)),$$

其中$score$是在公式4和5中计算的余弦相似度得分。给定搜索到的$P_l$,我们通过从$P_l = \{p_{l,1}, \dots, p_{l,j}, \dots, p_{l,J}\}$中迭代地选择概率最高的专家,直到$E_{prefetch}$的概率总和超过$\delta_l$,来找到要预取的专家集合$E_{prefetch}$:

$$\min_{\{E_{l,j}\}} |E_{prefetch}|$$ $$\text{s.t. } \sum_{E_{l,j} \in E_{prefetch}} p_{l,j} \geq \delta_{l}, \ j \in [J], \ \forall l \in [L],$$ $$|E_{prefetch}| \ge K, K \le [J],$$

其中$K$是每层需要激活的专家数量(例如,Mixtral-8×7B每层激活两个专家)。约束7要求每层选定预取的专家总概率大于$\delta_l$。约束8表示选定的专家最小数量必须大于MoE模型要求的激活专家数量。直观地,我们为低分专家地图分配较高的$\delta$,以便预取更多专家来减轻预测失误,并为高分专家地图分配较低的$\delta$以减少内存占用。概率较高的专家被优先预取。

异步专家地图匹配与预取。现有研究【【16, Fast Inference of Mixture-of-Experts Language Models with Offloading, 2023, arXiv】, 【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】在推理过程中同步地预测和预取专家,严重影响了推理性能。例如,MoE-Infinity【【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】在每个MoE层完成专家预测和预取之前无法计算前向函数【【55, MoE-Infinity Codebase, https://github.com/TorchMoE/MoE-Infinity】】。为了最小化系统开销和推理延迟,我们使用异步的发布者-订阅者架构(图7)将地图匹配和专家预取从推理过程中解耦。专家地图存储是一个消息代理,它保存来自推理过程和专家地图匹配器的消息。随着推理的进行,fMoE的推理过程持续地发布并将推理上下文(即语义嵌入和专家概率分布)写入专家地图存储。同时,专家地图匹配器订阅上下文数据,根据新的上下文数据匹配专家地图,并以异步方式将专家预取到专家缓存中 。

4.4 专家地图存储管理

存储容量与去重。在实践中,我们设计fMoE的专家地图存储以维持一个容量为$C$的存储空间,用于存放独特的专家地图。为了有效地指导跨多样化提示的推理,识别和去重冗余的专家地图是有意义的。

专家地图去重。由于fMoE使用两种方法(即基于语义和基于轨迹)来计算相似度,我们统一这两种相似度得分来计算新迭代数据和历史迭代数据之间的成对冗余得分:$RDY_{x,y} := \frac{d}{L} \cdot score_{x,y}^{sem} + \frac{L-d}{L} \cdot score_{x,y}^{map}$,其中$x \in [B], y \in [C]$,$score_{x,y}^{sem} \in R^{B \times C}$和$score_{x,y}^{map} \in R^{B \times C}$是从公式4和5计算的基于语义和基于轨迹的成对相似度得分,$d$是预取距离,$L$是总层数,$B$是新交互数据的批量大小,$C$是专家地图存储的容量。直观地,如图7所示,基于语义和基于轨迹的相似度得分对搜索专家地图的贡献分别与$\frac{d}{L}$和$\frac{L-d}{L}$成正比。因此,我们遵循相同的比例来统一和计算冗余得分。每当新迭代的上下文数据到达专家地图存储时,我们计算成对冗余得分$RDY_{x,y}$来决定丢弃哪些旧的迭代。因此,我们在专家地图存储中用新的迭代$x$($RDY_{x,y}$中的相应行)更新旧的迭代$y$($RDY_{x,y}$中的列)。

理论分析。专家地图去重可以被形式化为一个最小球面覆盖问题【【17, The Minimum Covering Sphere Problem, 1972, Management Science】】。目标是最小化存储中专家地图的总数,其中每个专家地图是球体的向量表示,同时最大化专家激活空间中的球体覆盖范围。研究【【15, Covering Spheres with Spheres, 2007, Discrete & Computational Geometry】, 【41, On the Closest Packing of Spheres in N Dimensions, 1947, Annals of Mathematics】】已经证明,维持至少$2LJ$个专家地图可以保证75%的专家地图相似度下限(即我们可以找到一个与任何新迭代至少75%相似的专家地图),而保持$\frac{1}{2}LJ\ln(LJ)$个专家地图可以提供98%的相似度下限,其中$L$和$J$分别是MoE模型中的层数和每层专家数。鉴于现代基于MoE的LLM通常有$L \in [8,128]$和$J \in [24, 96]$,我们可以估算出专家地图存储的最大需求小于50K个专家地图,占用不到200 MB的CPU内存【【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】。

4.5 专家缓存与驱逐

缓存管理策略。与现有的专家卸载解决方案【【16, Fast Inference of Mixture-of-Experts Language Models with Offloading, 2023, arXiv】, 【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】, 【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】类似,我们设计fMoE在GPU上维护一个专家缓存,以便在服务不同请求提示时重用专家权重。根据§4.2中匹配的专家地图,我们指导fMoE的专家缓存为单个专家计算两个优先级分数:1) 预取优先级,用于决定搜索到的地图中专家的预取顺序;2) 驱逐优先级,用于决定专家缓存中专家的驱逐顺序。

专家预取优先级。回想一下,要预取的专家集合$E_{prefetch}$是在公式6中确定的。对于每个专家$E_{l,j} \in E_{prefetch}$,我们定义其预取优先级为:

$$PRI_{l,j}^{prefetch} := \frac{p_{l,j}}{l - l_{now}}, \quad l \in [L], j \in [J],$$

其中$p_{l,j}$是来自搜索到的专家地图的专家概率,$l_{now}$是推理过程当前所在的层。直观地,具有更高激活概率$p_{l,j}$的专家应该更早被预取,而距离当前层更近的专家(即较小的$l - l_{now}$)也应该被优先处理。

专家驱逐优先级。与MoE-Infinity【【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】类似,fMoE的专家缓存基于最不常用(LFU)算法。我们结合搜索到的地图来共同决定驱逐优先级。对于每个专家$E_{l,j} \in E_{cache}$,我们定义其驱逐优先级为:

$$PRI_{l,j}^{evict} := \frac{1}{p_{l,j} \cdot freq_{l,j}}, \quad l \in [L], \ j \in [J],$$

其中$freq_{l,j}$是缓存访问频率,$p_{l,j}$是专家$E_{l,j} \in E_{cache}$在搜索到的地图中的概率。直观地,当达到专家缓存限制时,我们希望首先驱逐那些命中频率较低且激活概率较低的专家。需要注意的是,与现有工作【【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】, 【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】类似,我们不考虑专家的最近使用情况,这与经典的最近最少使用(LRU)算法【【16, Fast Inference of Mixture-of-Experts Language Models with Offloading, 2023, arXiv】】相反。由于专家使用是逐层顺序的,即一层接一层,优先考虑最近使用的专家与MoE服务中顺序前向计算的性质相悖。

按需专家加载。专家预取的错误预测会导致专家缓存在中未命中,因为MoE模型找不到门控网络指定的可用专家。每当发生专家未命中时,fMoE会暂停所有专家预取任务,并立即将未命中的专家从CPU加载到GPU内存中,以实现快速服务。

5. fMoE的实现

原型系统。我们基于Huggingface Transformers框架【【52, HuggingFace’s Transformers: State-ofthe-Art Natural Language Processing, 2020, EMNLP】】和MoE-Infinity代码库【【55, MoE-Infinity Codebase, https://github.com/TorchMoE/MoE-Infinity】】构建了fMoE的原型。fMoE的实现细节如下 。

专家地图存储。专家地图存储(Expert Map Store)使用Python实现,并利用了PyTorch【【37, PyTorch: An Imperative Style, High-Performance Deep Learning Library, 2019, NIPS】】和NumPy【【19, Array Programming with NumPy, 2020, Nature】】库。我们使用ndarrays数据结构来存储所有的语义嵌入和专家地图,以实现高效的数组操作。这些数组被转换为张量(tensors)以计算相似度,用于专家地图匹配。

专家地图匹配器。专家地图匹配器(Expert Map Matcher)使用Python实现,并利用了PyTorch【【37, PyTorch: An Imperative Style, High-Performance Deep Learning Library, 2019, NIPS】】和TorchMetrics【【12, TorchMetrics: Measuring Reproducibility in Pytorch, 2022, JOSS】】库。我们使用TorchMetrics中的余弦相似度接口来实现成对计算,包括相似度(§4.2)和冗余度(§4.4)得分。我们使用Python的多线程库来实现异步的专家地图匹配和专家预取,其中线程与专家地图存储共享相同的内存空间,以实现高效的读写操作。

专家缓存。专家缓存(Expert Cache)基于MoE-Infinity代码库【【55, MoE-Infinity Codebase, https://github.com/TorchMoE/MoE-Infinity】】用C++实现。GPU中的专家管理通过CUDA运行时API【 【35, CUDA Runtime API :: CUDA Toolkit Documentation, 2024, https://docs.nvidia.com/cuda/ cuda-runtime-api/index.html】】实现。我们实现了fMoE的缓存逻辑,并修复了MoE-Infinity代码库中的关键错误以支持专家卸载。与MoE-Infinity一样,fMoE支持多GPU推理和专家并行,其中专家被映射到不同的GPU设备进行加载和卸载。我们使用哈希映射来将专家ID分配给不同的GPU,并在推理过程中检索它们。专家分配遵循轮询(round-robin)方式以平衡整体GPU负载。此外,我们在GPU空间中使用一个多线程任务池来调度和执行专家预取和按需加载任务。

A4 实验环境

硬件配置:
* GPU: 6块NVIDIA GeForce RTX 3090,每块拥有24 GB显存。所有GPU通过NVLink对连,并通过PCIe 4.0(32GB/s带宽)连接到CPU内存。
* CPU: 32核AMD Ryzen Threadripper PRO 3955WX。
* 内存: 480 GB CPU内存。

模型架构:
* 使用了三种流行的基于MoE的LLM进行评估:Mixtral-8×7B【【21, Mixtral of Experts, 2024, arXiv】】、Qwen1.5-MoE【【57, Qwen2 Technical Report, 2024, arXiv】】和Phi-3.5-MoE【【1, Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone, 2024, arXiv】】。
* 模型参数、MoE层数和每层专家数见下表1。
* 根据模型性能分析,预取距离$d$被优化设置为3。

表1: 评估中三个MoE模型的特性。

数据集与踪迹:
* 数据集: 使用了两个真实世界的提示数据集:LMSYS-Chat-1M【【60, LMSYS-Chat-1M: A Large-Scale Real-World LLM Conversation Dataset, 2023, arXiv】】和ShareGPT【【45, ShareGPT: Share Your Wildest ChatGPT Conversations, https://sharegpt.com/】】 。
* 数据划分: 对于大多数实验,数据集按7:3的比例划分,70%的提示上下文数据(语义嵌入和专家地图)用于填充fMoE的专家地图存储,30%的提示用于测试。
* 踪迹: 对于在线服务实验,专家地图存储被清空,使用微软Azure发布的真实世界LLM推理踪迹【【38, Splitwise: Efficient Generative LLM Inference Using Phase Splitting, 2024, ISCA】, 【48, DynamoLLM: Designing LLM Inference Clusters for Performance and Energy Efficiency, 2025, HPCA】】来设置输入和生成长度,并驱动调用。

软件配置与基线:
* 软件栈: fMoE原型基于HuggingFace Transformers【【52, HuggingFace’s Transformers: State-ofthe-Art Natural Language Processing, 2020, EMNLP】】和MoE-Infinity代码库【【55, MoE-Infinity Codebase, https://github.com/TorchMoE/MoE-Infinity】】实现 。
* 基线模型:
1. MoE-Infinity【【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】: 使用粗粒度的请求级专家激活模式和同步预取。
2. ProMoE【【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】】: 采用基于步幅的推测性专家预取(由于代码未开源,作者尽力复现了其原型)。
3. Mixtral-Offloading【【16, Fast Inference of Mixture-of-Experts Language Models with Offloading, 2023, arXiv】】: 结合了逐层推测性专家预取和基于LRU的专家缓存。
4. DeepSpeed-Inference: 采用与专家无关的逐层参数卸载方法,仅支持按需加载(作者在MoE-Infinity代码库中复现了其卸载逻辑并添加了专家缓存以公平比较)。

A4 实验结果

总体性能评估 (Prefill 和 Decode 阶段)

  • 实验内容: 评估fMoE及四个基线在服务三种MoE模型时的首词元时间(TTFT)和每输出词元时间(TPOT)。
  • 实验结果: 如图9所示,fMoE在TTFT、TPOT和专家命中率方面均显著优于所有基线。与DeepSpeed-Inference、Mixtral-Offloading、ProMoE和MoE-Infinity相比,fMoE在三个MoE模型上平均将TTFT降低了44%、35%、33%、30%,将TPOT降低了70%、61%、55%、48%。在专家命中率方面,fMoE相比这四个基线平均提升了147%、11%、34%和63%。
  • 分析结论: fMoE的细粒度卸载设计能有效降低推理延迟并提高专家命中率。尽管Mixtral-Offloading因其同步推测性预取(预取距离为1)获得了较高的命中率,但其延迟表现不佳。


图9: fMoE及其他四个基线在prefill和decode阶段的总体性能。

在线服务性能
* 实验内容: 在在线服务场景下(即专家地图存储初始为空),使用Azure LLM推理踪迹驱动的请求,评估fMoE和基线的端到端请求延迟。
* 实验结果: 如图10的累积分布函数(CDF)图所示,fMoE在三种MoE模型的在线服务场景中,其端到端请求延迟显著低于其他基线。
* 分析结论: fMoE即使在冷启动条件下也能快速适应并提供卓越的性能,有效降低了在线服务的请求延迟。


图10: MoE在线服务的请求延迟CDF。

专家缓存限制的影响
* 实验内容: 通过限制专家缓存的内存预算(从6 GB到96 GB),评估fMoE和基线的TPOT,以考察它们在延迟-内存权衡中的表现。
* 实验结果: 如图11所示,fMoE在所有不同的专家缓存限制下,TPOT始终优于其他基线。特别是在GPU内存有限(如6GB)的情况下,fMoE相较于DeepSpeed-Inference、Mixtral-Offloading、ProMoE和MoE-Infinity,TPOT分别降低了32%、24%、18%和18%。
* 分析结论: fMoE的细粒度专家卸载能够在维持较低GPU内存占用的同时,显著减少按需加载延迟,从而在MoE服务的延迟-内存权衡中取得了更好的平衡点。


图11: 在不同专家缓存限制下,fMoE及其他四个基线的性能。

消融研究
* 专家地图搜索有效性:
* 实验内容: 比较了五种专家模式追踪方法:Speculate(推测性预测)、Hit count(请求级命中计数)、Map (T)(仅轨迹相似度搜索)、Map (T+S)(轨迹+语义搜索)和Map (T+S+δ)(完整fMoE功能)。
* 实验结果: 如图12a所示,专家命中率随着fMoE功能的逐步增加而提高,完整的fMoE(Map (T+S+δ))达到了最高的命中率。粗粒度的命中计数方法性能最差。
* 分析结论: 专家地图、混合相似度搜索和动态专家选择阈值的设计对提高专家命中率至关重要。

  • 专家预取与缓存有效性:
    • 实验内容: 将fMoE的缓存策略与LRU(Mixtral-Offloading使用)和LFU(MoE-Infinity使用)进行比较。
    • 实验结果: 如图12b所示,fMoE的缓存策略在专家命中率上优于LRU和LFU。
    • 分析结论: fMoE结合了专家概率和访问频率的缓存策略比传统的LRU和LFU更适合MoE服务的场景。


图12: fMoE的消融研究。

敏感性分析

  • 预取距离: 如图13所示,对于fMoE,预取距离$d=3$时性能最佳。距离过小无法完全隐藏系统延迟,距离过大则导致命中率下降,从而影响性能。
  • 专家地图存储容量: 如图14a所示,随着存储容量增加,语义和轨迹相似度得分均提高,但在容量超过1K后增益减缓。因此,1K被选为平衡性能和内存开销的最佳容量。
  • 推理批量大小: 如图14b所示,随着批量大小从1增加到4,fMoE在大多数情况下都保持了最低的TTFT和TPOT。


图13: fMoE在不同预取距离下服务MoE模型的性能。


图14: fMoE的敏感性分析。

系统开销
* 延迟开销: 如图15所示,fMoE中同步操作(上下文收集、专家按需加载)引入的总延迟在所有模型中都低于30ms(占迭代时间的5%),可以忽略不计。异步任务(地图匹配、专家预取、地图更新)不影响端到端迭代延迟。
* 内存开销: 如图16所示,fMoE的专家地图存储的CPU内存占用非常小。即使容量达到32K个地图,所需内存也不到200MB,这对于通常拥有大量CPU内存的现代GPU服务器来说是微不足道的。


图15: fMoE单次推理迭代在三个MoE模型上的延迟分解。


图16: 不同容量下fMoE专家地图存储的CPU内存占用。

A7 补充细节

无损MoE服务。近年来,关于无损MoE服务的研究被广泛提出。DeepSpeed inference【【4, DeepSpeed-Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale, 2022, SC22】】在不感知专家的情况下逐层卸载参数。Mixtral-Offloading【【16, Fast Inference of Mixture-of-Experts Language Models with Offloading, 2023, arXiv】】采用LRU专家缓存,并引入推测性预测以实现专家预取。MoE-Infinity【【54, MoE-Infinity: Offloading-Efficient MoE Model Serving, 2024, arXiv】】提出了请求级的专家激活矩阵,以粗粒度方式指导卸载。SwapMoE【【42, SwapMoE: Serving Off-the-shelf MoE-based Large Language Models with Tunable Memory Budget, 2023, arXiv】】在GPU内存中维护一组关键专家,并根据工作负载变化进行调整,以最小化卸载开销。ProMoE【【47, ProMoE: Fast MoE-based LLM Serving using Proactive Caching, 2024, arXiv】】为每个MoE层训练预测器,以实现高推测性预测准确率和低推理延迟。Lina【【28, Accelerating Distributed MoE training and inference with Lina, 2023, USENIX ATC 23】】将不常用的专家移至主机内存,并更侧重于分布式MoE训练。Liu等人【【31, Optimizing distributed deployment of mixture-of-experts model inference in serverless computing, 2025, INFOCOM】】通过黑盒贝叶斯技术预测专家模式,在无服务器计算上服务MoE模型。与现有粗粒度卸载方案不同,fMoE从轨迹和语义两个方面追踪细粒度的专家模式,并超越了SOTA基线。

有损MoE服务。除了无损卸载,其他工作也提出了有损MoE服务。专家剪枝【【13, GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, 2021, ICLR】, 【50, Taskspecific expert pruning for sparse mixture-of-experts, 2022, arXiv】】通过移除未充分利用的专家来减少内存使用。知识蒸馏【【33, Efficient large scale language modeling with mixtures of experts, 2021, arXiv】, 【43, Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale, 2022, arXiv】】产生紧凑的稀疏MoE模型。ComPEFT【【56, Compeft: Compression for communicating parameter efficient updates via sparsification and quantization, 2023, arXiv】】展示了无精度损失的专家压缩,而MC-SMoE【【39, Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing Policy, 2024, ICLR】】进一步将合并的专家分解为低秩和结构化稀疏的替代品。Hobbit【【49, Hobbit: A mixed precision expert offloading system for fast moe inference, 2024, arXiv】】使用低精度来服务不太关键的专家。然而,有损服务可能会影响生成质量,并且与fMoE是正交的。

MoE重构。一些工作提出重新设计和重构当前的MoE架构。Pre-gated MoE【【20, Pregated MoE: An Algorithm-System Co-Design for Fast and Scalable Mixture-of-Expert Inference, 2024, ISCA】】利用抢占式专家选择来消除专家选择和执行之间的顺序依赖。SiDA-MoE【【14, SiDA: Sparsity-Inspired Data-Aware Serving for Efficient and Scalable Large Mixture-of-Experts Models, 2024, MLSys】】提出了一种受稀疏性启发、数据感知的推理系统,该系统将专家路由与推理解耦。READ-ME【【7, Read-ME: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System CoDesign, 2024, NeurIPS】】将预训练的密集LLM重构为专门化的MoE模型。与上述工作不同,fMoE服务于开源MoE模型时无需任何训练。

A5 结论

本文提出了fMoE,一个用于MoE服务的细粒度专家卸载系统,它在不引入显著模型内存占用的情况下实现了低推理延迟。fMoE使用“专家地图”(expert map)从MoE模型中追踪迭代级别的专家概率分布,并分析来自单个请求提示的输入语义嵌入。基于输入的语义和专家轨迹信息,fMoE搜索最准确的专家地图,以针对每个推理迭代精心指导专家预取、缓存和卸载决策。fMoE在HuggingFace Transformers之上进行了原型设计,并部署到一个六GPU的测试平台上。通过对开源MoE模型和真实世界工作负载进行的大量实验表明,与最先进的解决方案相比,fMoE将推理延迟降低了47%,并将专家命中率提高了36%。