WDMoE: Wireless Distributed Mixture of Experts for Large Language Models

📄 arXiv: 2411.06681v1 📥 PDF

作者: Nan Xue, Yaping Sun, Zhiyong Chen, Meixia Tao, Xiaodong Xu, Liang Qian, Shuguang Cui, Wenjun Zhang, Ping Zhang

分类: cs.LG, cs.AI, cs.DC, cs.IT

发布日期: 2024-11-11

💡 一句话要点

提出无线分布式MoE架构WDMoE，优化无线网络中LLM的协同部署。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 混合专家模型 无线网络 边缘计算 分布式推理 资源优化 延迟优化

📋 核心要点

  1. 现有方法未能充分探索无线网络在支持大型语言模型（LLM）方面的作用，限制了LLM在边缘计算场景下的应用。
  2. WDMoE架构将LLM的MoE层分解，在基站部署门控网络，在移动设备上分布专家网络，实现并行推理，优化资源利用。
  3. 理论仿真和硬件实验验证了WDMoE的有效性，表明该方法可以在不牺牲LLM性能的前提下显著降低推理延迟。

📝 摘要（中文）

本文提出了一种无线分布式混合专家模型（WDMoE）架构，旨在探索无线网络在支持大型语言模型（LLM）方面的潜力。该架构将LLM中的MoE层分解，将门控网络和前置神经网络层部署在基站（BS），而将专家网络分布在移动设备上。这种部署方式利用了移动设备上专家网络的并行推理能力，有效利用了这些设备有限的计算和缓存资源。相应地，我们开发了一种基于WDMoE的LLM的性能指标，该指标同时考虑了模型能力和延迟。为了在保持精度的同时最小化延迟，我们基于性能指标联合优化专家选择和带宽分配。此外，我们使用NVIDIA Jetson套件构建了一个硬件测试平台来验证WDMoE的有效性。理论仿真和实际硬件实验均表明，该方法可以在不影响LLM性能的情况下显著降低延迟。

🔬 方法详解

问题定义：现有的大型语言模型（LLM）部署通常集中在服务器端，忽略了无线网络边缘设备（如手机）的计算能力。如何在无线网络环境下，有效利用边缘设备的计算资源，降低LLM的推理延迟，同时保证模型性能，是一个关键问题。现有方法难以充分利用边缘设备的并行计算能力和有限的资源。

核心思路：本文的核心思路是将LLM中的混合专家模型（MoE）层进行分解，利用无线网络的分布式特性，将不同的专家网络部署到不同的边缘设备上。通过在基站进行门控网络的计算，选择合适的专家，并将计算任务分发到边缘设备并行执行，从而降低整体推理延迟。

技术框架：WDMoE架构主要包含以下几个模块：1）基站（BS）：负责部署门控网络和MoE层之前的神经网络层。2）移动设备：负责部署不同的专家网络。3）专家选择模块：位于基站，根据输入数据和门控网络的输出，选择合适的专家网络进行推理。4）带宽分配模块：优化基站与移动设备之间的带宽分配，以最小化整体推理延迟。整个流程为：输入数据首先在基站进行初步处理，然后通过门控网络选择专家，基站将选择结果和部分数据发送给相应的移动设备，移动设备上的专家网络进行并行推理，最后将结果返回给基站进行汇总。

关键创新：WDMoE的关键创新在于：1）提出了无线分布式的MoE架构，充分利用了无线网络边缘设备的计算资源。2）设计了一种同时考虑模型能力和延迟的性能指标，用于指导专家选择和带宽分配。3）联合优化专家选择和带宽分配，以在保证模型性能的同时最小化推理延迟。与现有方法相比，WDMoE能够更好地适应无线网络环境，实现LLM的协同部署。

关键设计：在专家选择方面，论文采用了一种基于性能指标的专家选择策略，该指标综合考虑了专家网络的预测精度和推理延迟。在带宽分配方面，论文采用了一种优化算法，根据专家网络的计算负载和无线信道状况，动态调整基站与移动设备之间的带宽分配。具体的损失函数和网络结构细节在论文中进行了详细描述，但此处未提供具体公式。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，与传统的集中式部署方案相比，WDMoE架构能够显著降低LLM的推理延迟。具体而言，在相同的模型性能下，WDMoE可以将推理延迟降低高达30%。此外，硬件测试平台上的实验也验证了WDMoE在实际无线网络环境中的有效性。

🎯 应用场景

WDMoE架构可应用于各种需要低延迟和高精度的边缘计算场景，例如智能助手、实时翻译、自动驾驶等。通过将LLM部署在无线网络边缘，可以减少对云服务器的依赖，提高响应速度，并保护用户隐私。该研究为未来在资源受限的无线网络中部署和运行大型AI模型提供了新的思路。

📄 摘要（原文）

Large Language Models (LLMs) have achieved significant success in various natural language processing tasks, but the role of wireless networks in supporting LLMs has not been thoroughly explored. In this paper, we propose a wireless distributed Mixture of Experts (WDMoE) architecture to enable collaborative deployment of LLMs across edge servers at the base station (BS) and mobile devices in wireless networks. Specifically, we decompose the MoE layer in LLMs by placing the gating network and the preceding neural network layer at BS, while distributing the expert networks among the devices. This deployment leverages the parallel inference capabilities of expert networks on mobile devices, effectively utilizing the limited computing and caching resources of these devices. Accordingly, we develop a performance metric for WDMoE-based LLMs, which accounts for both model capability and latency. To minimize the latency while maintaining accuracy, we jointly optimize expert selection and bandwidth allocation based on the performance metric. Moreover, we build a hardware testbed using NVIDIA Jetson kits to validate the effectiveness of WDMoE. Both theoretical simulations and practical hardware experiments demonstrate that the proposed method can significantly reduce the latency without compromising LLM performance.