研究背景
基于Transformer架构的深度学习大语言模型（LLM）在人工智能领域表现出了卓越的性能。尽管这些庞大的模型最初诞生于云数据中心，它们很快就扩展到了边缘设备，并在网络边缘落地了一系列智能应用，包括家庭或个人AI助手、智能机器人等。
图1：基于Galaxy系统的家庭AI语音助手示意图。
目前，大多数基于Transformer的智能应用严重依赖云服务，实际的大规模模型推理在云数据中心进行，边缘设备仅部署代理来转发用户请求。然而，这种基于云的方法面临三个重要问题：（1）用户数据隐私问题：传输敏感信息到云端引发隐私担忧。（2）服务质量问题：广域网连接不可靠导致响应延迟不稳定。（3）云数据中心的负荷问题：海量边缘应用请求对主干网络和云数据中心造成巨大压力。
为了解决上述问题，在边缘设备上直接进行数据处理和推理，不仅减少了对云端服务的依赖，而且确保了数据的本地存储，有效避免了网络传输的需求，这种方案展现出了巨大的应用潜力。然而，Transformer推理的高计算强度和高资源需求对资源有限且不可扩展的边缘设备构成了巨大挑战。与此同时，我们发现在智能家居等常见的边缘环境中，同一用户或组织通常拥有多台富裕或闲置的边缘设备（如手机、笔记本和智能家居设备）。这种思路启发我们，将周边的边缘设备视作资源的扩展，通过实施分布式协作计算，不仅可以实现在边缘设备上对Transformer模型的高效推理，还能确保推理过程的低延迟，这一方案具有显著的实用价值和应用前景。
然而，利用协同边缘设备加速Transformer推理面临几个关键挑战：（1）如何在多个边缘设备之间分配稀疏的Transformer推理任务，尤其是单样本请求。（2）如何根据异构边缘设备的资源预算来定制工作负载分配。（3）如何在带宽受限的边缘环境中减少协同推理的延迟。为了解决这些挑战，本工作设计并搭建了资源高效的面向边缘端Transformer大语言模型推理的分布式边缘智能系统Galaxy。该系统通过整合异构边缘设备的资源，实现低延迟的Transformer推理，从而支持实时的本地智能服务。
系统设计概述
Galaxy是一个资源高效的边缘端协同推理计算系统，旨在同时利用多个异构边缘设备，实现低延迟的本地Transformer推理。下图展示了Galaxy系统的工作流程，包括三个主要阶段：预处理阶段、并行规划阶段和执行阶段。
图2：Galaxy系统工作流概览。
预处理阶段：
Galaxy Profiler使用测试数据在真实的待部署边缘设备上进行模拟推理过程，记录设备的计算能力、可用内存、设备间的网络带宽、目标Transformer模型的参数量信息以及前向传播产生的中间值等。这一信息采集过程在部署前离线完成，所采集的信息可以被缓存并重复使用。
并行规划阶段：
Galaxy系统采用了一种新颖的混合模型并行架构，结合了张量并行和序列并行两种技术来并行利用分布式边缘设备。Galaxy Planner将Galaxy Profiler所收集的信息作为输入，生成并行规划配置。该配置综合考虑了设备的资源异构性、内存预算和设备间的网络带宽，以最大限度地发挥边缘环境中有限的计算资源。
图3：混合模型并行架构示意图。
推理执行阶段：
Galaxy系统将Galaxy Planner生成的并行配置应用于目标模型和边缘设备上，进行并行推理系统的部署。部署完成后，系统将能够进行资源高效的边缘协同推理。然而，分布式推理不可避免地涉及张量通信操作，这在低带宽网络环境下会成为性能瓶颈。为了解决这个问题，Galaxy系统设计了一种新颖的基于矩阵分块的细粒度通信计算重叠优化，有效减少了额外通信开销带来的推理性能下降。
图4：基于矩阵分块的细粒度通信计算重叠示意图。
研究结果
我们在六种包含不同边缘设备组合的边缘集群环境下进行了实验。我们评估了五种流行的基于Transformer的边缘大模型：DistilBERT、BERT、GPT-2 Large、OPT-Large和OPT-extra Large，输入序列来自GLUE数据集。我们将Galaxy系统与两种最先进的模型并行方法Megatron-LM和Sequence Parallelism进行了对比。评估结果显示，Galaxy在各种网络环境中均保持了高性能，与基准方法相比，延迟减少了最多达46%。我们的系统在异构环境中也展现了强大的鲁棒性。Galaxy在异构边缘环境中表现出色，相比其他并行方法，推理延迟最多减少了2.5倍。欢迎阅读原文[1]了解更多实验结果。
图5：Galaxy系统在不同实验设置下都表现出了突出的性能。
图6：Galaxy系统在异构边缘集群环境下表现出了很好的自适应性。
总结
我们设计并实现一个资源高效的面向边缘端Transformer大语言模型推理的分布式边缘智能系统Galaxy，创新性地利用了混合模型并行架构来协同多个边缘设备。Galaxy系统设计了一种并行规划算法来生成并行配置，该算法综合考虑了设备的资源异构性、内存预算和网络带宽环境。为了克服边缘环境中带宽受限的挑战，Galaxy系统采用了基于矩阵分块的细粒度通信计算重叠优化，有效减少了额外通信开销带来的推理性能下降。
拓展阅读
边缘协同（Edge Collaboration）是一种新颖的边缘智能系统设计范式，旨在通过利用本地化边缘设备和集群，实现经济实惠、资源高效和灵活可伸缩的个性化人工智能服务。除了本工作Galaxy外，我们还进行了另一项基于边缘协同架构的研究工作——Asteroid，该工作探索了在边缘环境中利用多个分布式边缘设备进行深度神经网络训练，可应用于多设备协同联邦模型训练、边缘大模型协同微调等新兴场景。该研究文章已被ACM MobiCom 2024录用发表，欢迎感兴趣的读者参考论文[2]了解详情。
论文信息
[1] Shengyuan Ye, Jiangsu Du, Liekang Zeng, Wenzhong Ou, Xiaowen Chu, Yutong Lu, Xu Chen*, “Galaxy: A Resource-Efficient Collaborative Edge AI System for In-situ Transformer Inference”.In IEEE International Conference on Computer Communications (INFOCOM), 2024. https://arxiv.org/abs/2405.17245.
[2] Shengyuan Ye, Liekang Zeng, Xiaowen Chu, Guoliang Xing, Xu Chen*, “Asteroid: Resource-Efficient Hybrid Pipeline Parallelism for Collaborative DNN Training on Heterogeneous Edge Devices”.In ACM Annual International Conference On Mobile Computing And Networking (MOBICOM), 2024. https://dl.acm.org/doi/abs/10.1145/3636534.3649363.