EdgeShard：通过协作边缘计算实现高效的 LLM 推理

(2024-05-23) EdgeShard: Efficient LLM Inference via Collaborative Edge Computing (EdgeShard：通过协作边缘计算实现高效的 LLM 推理)

此文为机器辅助翻译，仅供个人学习使用，如有翻译不当之处欢迎指正

Abstract

大语言模型（LLMs）在自然语言处理和内容生成方面展现出了巨大潜力。然而，当前的大语言模型严重依赖云计算，这导致了延迟延长、带宽成本高昂以及隐私方面的问题。边缘计算有望通过在更靠近数据源的边缘设备上部署大语言模型来解决这些问题。一些研究尝试利用模型量化来减小模型大小，以适配资源受限的边缘设备，但这会导致精度损失。还有一些研究采用云边协作的方式，却受到网络连接不稳定的困扰。在本研究中，我们利用协同边缘计算来促进边缘设备和云服务器之间的协作，从而共同高效地执行大语言模型推理。我们提出了一个通用框架，将大语言模型划分为多个分片，并部署在分布式设备上。为了实现高效的大语言模型推理，我们构建了一个自适应的联合设备选择和模型分区问题，并设计了一种高效的动态规划算法，分别对推理延迟和吞吐量进行优化。在异构物理原型上对 Llama2 系列模型进行的实验表明，与基线方法相比，EdgeShard 的推理延迟最多可降低 50%，吞吐量最多可提高 2 倍。

引言

近年来，大语言模型（LLMs）的出现引起了公众、行业和学术界的广泛关注，它代表了人工智能（AI）领域的一项重大突破。许多参与者带着他们的先进模型进入这一领域，比如 OpenAI 的 GPT-4 [1]、Meta 的 Llama [2] 以及谷歌的 PALM [3]。大语言模型基于 Transformer 架构 [4] 构建，其特点是参数数量庞大，并且训练数据量巨大。大语言模型的参数规模通常高达数千亿，这使得模型能够捕捉语言和上下文的复杂模式，从而在生成连贯且符合上下文的回复方面表现出色。这种现象也被称为 “智能涌现” 。大语言模型的卓越能力使其在广泛的应用领域中具有很高的价值和出色的表现，从聊天机器人、内容生成（例如文本摘要和代码生成）到教育和研究辅助工具等。

然而，当前的大语言模型严重依赖云计算，面临响应时间长、带宽成本高以及隐私问题 [5]。首先，对云计算的依赖限制了模型在实时应用（如机器人控制、导航或勘探）中进行快速推理的能力，在这些应用场景中，即时响应至关重要。其次，大量数据（包括文本、视频、图像、音频和物联网传感数据）传输到云数据中心，这会导致大量的带宽消耗，给网络架构带来巨大压力。第三，基于云的大语言模型引发了严重的隐私问题，尤其是在处理医院和银行的敏感数据，以及手机上的文本输入和照片等个人数据时。

边缘计算有望通过在更靠近数据源的网络边缘设备（如边缘服务器、边缘网关和手机）上部署大语言模型（LLMs）来解决上述挑战[6]。然而，大语言模型计算量极大且资源需求高。例如，全精度的Llama2 - 7B模型推理至少需要28GB内存，这可能超出了大多数边缘设备的容量。一些研究利用模型量化技术[7]-[12]来减小模型大小，以适配资源受限的边缘设备。但这往往会导致精度损失。其他研究倾向于采用云边协作的方式[13]、[14]，即将大语言模型划分为两个子模型，并将部分计算工作负载卸载到配备高端GPU的强大云服务器上。然而，边缘设备与云服务器之间的延迟通常较高且不稳定。

另外，近年来我们见证了边缘计算能力的持续增长，大量边缘服务器和边缘云被部署在网络边缘，还有大量资源有待利用。因此，最近提出了协同边缘计算（CEC）[15,16]，旨在整合地理上分散的边缘设备和云服务器的计算资源，以实现高效的资源利用和性能优化。如图1所示，无处不在的分布式边缘设备和云服务器相互连接，形成一个共享资源池，共同提供即时数据处理和人工智能服务。协同边缘计算（CEC）与现有的边缘计算研究有所不同。现有的边缘计算研究侧重于云、边缘和终端设备之间的垂直协作，却忽视了边缘设备之间的水平协作，存在资源利用未优化、服务覆盖范围受限以及性能不均衡等问题。

受协同边缘计算（CEC）愿景的启发，我们提出了一种通用的大语言模型推理框架EdgeShard，以支持在分布式边缘设备和云服务器上进行高效的大语言模型协同推理。为简便起见，以下我们用计算设备指代边缘设备和云服务器。在一个由异构计算设备组成的网络中，EdgeShard将大语言模型划分为多个分片，并根据异构的计算和网络资源以及设备的内存预算，将它们分配给合适的设备。为了优化性能，我们构建了一个联合设备选择和模型分区问题，并设计了一种高效的动态规划算法，分别用于最小化推理延迟和最大化推理吞吐量。在实际测试平台上进行的大量实验表明，与设备本地推理和垂直云边协同推理方法相比，EdgeShard的推理延迟最多可降低50%，吞吐量最多可提高2倍。

我们的工作与那些在云数据中心将大语言模型进行划分并分配到多个GPU上的研究有所不同，比如Gpipe[17]和PipeDream[18]。在边缘计算环境中部署大语言模型与在云端部署有很大差异。其一，云服务器通常配备同质化的GPU，而边缘设备本质上具有异构的计算能力。其二，用于大语言模型的现代云GPU通常通过高带宽网络连接，如InfiniBand和Nvlinks，而边缘设备则通过异构的低带宽网络连接。例如，Nvlinks的带宽可达600GB/s，而边缘设备之间的带宽范围仅从几十Kbps到1000Mbps。为云数据中心设计的大语言模型部署方案忽略了异构且资源受限的边缘计算环境。

我们的贡献主要体现在三个方面。

• 首先，我们提出了一个通用的大语言模型推理框架，用于在边缘计算环境中部署大语言模型，该框架支持异构边缘设备和云服务器之间进行协同推理。
• 此外，我们定量研究了如何选择计算设备以及如何对大语言模型进行分区以优化性能。我们从数学角度构建了联合设备选择和模型分区问题，并分别提出了动态规划算法来优化推理延迟和吞吐量。
• 我们还在物理测试平台上，使用最先进的Llama2系列模型对EdgeShard的性能进行了评估。实验结果表明，EdgeShard的性能显著优于各种基线方法。

预备知识与研究动机

生成式大语言模型推理。 大语言模型通常是指基于解码器的拥有数十亿参数的Transformer模型。与基于编码器架构的BERT[19]不同，BERT的推理过程是单阶段的，而大语言模型的推理过程是迭代的，通常包括两个阶段：提示处理阶段和自回归生成阶段。提示处理阶段也被称为预填充阶段。

在提示处理阶段，模型将用户输入的初始词元$(x_1,...,x_n)$作为输入，并通过计算概率$P(x_{n+1}|x_1,...,x_n)$生成第一个新的词元$x_{n+1}$。

在自回归生成阶段，模型基于初始输入和目前已生成的词元，每次生成一个词元。这一阶段会通过多次迭代按顺序生成词元，直到满足停止条件，即生成序列结束（EOS）词元，或者达到用户指定的最大词元数量，又或者达到大语言模型所限制的最大词元数量。

如图 2 所示，假设大语言模型有 N 层，它会接收一系列输入词元，并运行所有层，逐个生成词元。在预填充阶段，模型一次性接收输入 “Today is a”，生成的第一个词元是 “good”。在自回归生成阶段，模型首先将 “Today is a good” 作为输入，生成下一个词元 “day”。然后，它将 “Today is a good day” 作为输入，生成下一个词元 “EOS”，这表明生成结束。由于生成的每个词元都由序列中它前面的所有词元决定，大语言模型利用键值缓存（KV 缓存）来避免重复计算，存储之前的计算结果以加快响应速度，从而减少计算工作量并缩短响应时间。预填充阶段生成一个词元的时间比自回归阶段要长得多（通常是10倍），因为预填充阶段需要计算所有输入词元的键值缓存（KV缓存）进行初始化。

大语言模型非常消耗内存。单个边缘设备可能没有足够的内存来容纳一个大语言模型。以最受欢迎的大语言模型之一——Llama2为例。如表1所示，Llama2有三种不同的版本，即70亿参数（7B）、130亿参数（13B）和700亿参数（70B）。从表中我们可以看到，Llama2-7B的全精度推理至少需要28GB的内存，而智能手机通常只有6到12GB的内存，英伟达Jetson Orin NX设备也只有8到16GB的内存。这些设备无法承担在本地进行大语言模型推理的任务。一些研究尝试使用低精度量化，例如8比特和4比特量化。然而，这仍然可能超过边缘设备的内存容量。例如，Llama2-70B的4比特推理至少需要35GB的内存，这是大多数边缘设备无法容纳的。此外，低精度推理会导致性能下降。

在这项工作中，我们利用了协同边缘计算，这是一种计算范式，在这种范式下，地理上分布的边缘设备和云服务器协同执行计算任务。基于这一理念，我们提出了EdgeShard，这是一个通用的大语言模型推理框架。该框架允许在分布式计算设备上进行自适应的设备选择和大语言模型分区，以解决高内存需求的问题，并利用异构资源来优化大语言模型的推理。

面向大语言模型的协同边缘计算

该框架包含三个阶段，分别是性能分析、任务调度优化以及协同推理。其工作流程如图3所示。

性能分析是一个离线步骤，它为优化步骤分析必要的运行时跟踪信息，并且只需要进行一次。这些跟踪信息包括：1）不同设备上每一层的执行时间；2）大语言模型每一层的激活数据大小和内存消耗；3）每个设备的可用内存以及设备之间的带宽。对于每一层的执行时间，我们分别分析预填充阶段和自回归阶段生成一个词元的时间，然后取平均值。对于那些可能没有足够内存来完整加载模型以进行性能分析的设备，我们采用了动态模型加载技术。在该技术中，模型的各层会被依次加载，以适配受限的内存。然后，这些性能分析信息将被用于支持智能的任务调度策略。

调度优化。在任务调度优化阶段，调度器通过确定让哪些设备参与、如何按层划分大语言模型以及应将模型分片分配给哪些设备，来生成一个部署策略。该策略全面考虑了异构资源、设备的内存预算以及隐私限制，随后会应用到选定的设备上，以实现高效的大语言模型推理。更多细节将在第四节中阐述。

协同推理。在获得大语言模型的分区和分配策略后，选定的设备将进行协同推理。我们会在每个参与的设备上预先为键值缓存分配内存空间。对于协同推理，我们考虑两种情况，即顺序推理和流水线并行推理。

在顺序推理中，设备轮流使用分配的模型分片来执行计算。如图 4（a）所示，假设大语言模型被划分为 3 个分片，并分别分配给设备 1、设备 2 和设备 3。设备 1 将首先处理输入数据，然后将激活数据 / 输出发送给设备 2，设备 2 处理数据后再传输给设备 3。顺序推理适用于为单个用户服务的场景，例如智能家居场景，在这种场景中，用户的个人设备（如平板电脑、手机和智能音箱）协同执行大语言模型推理。在这样的场景中，用户输入一个提示并得到回复，然后再输入另一个提示。我们的目标是将顺序推理的延迟降至最低。

然而，从系统的角度来看，顺序推理的资源利用效率并不高。当设备1在执行计算时，设备2和设备3处于闲置状态。因此，我们采用流水线并行的方式来提高资源利用率。对于像之前在云服务器上的Gpipe[17]和PipeDream[18]等工作中所采用的流水线并行推理，输入数据首先会被分割成微批次，然后输入到系统中。如图4(b)所示，设备1首先处理数据B1，然后将中间数据传输给设备2。在处理完数据B1后，设备1会立即开始处理数据B2。以这样的流水线方式运行，每个设备都处于忙碌状态，系统资源利用率很高。

优化大语言模型推理

我们考虑一个由异构设备和异构带宽连接组成的通用协同边缘网络。更具体地说，对于一组通过异构带宽连接的异构设备，EdgeShard旨在选择设备的一个子集，并将大语言模型划分为多个分片，然后将这些分片分配给选定的设备，以最小化推理延迟或最大化吞吐量。

系统模型。大语言模型通常具有分层架构，该架构由一个嵌入层、多个解码器层和一个输出层组成。各层的参数规模和激活值（即某一层的输出）大小各不相同。我们假设该模型有$N$层。$O_i$表示第$i$层（$0\le i\le N-1$）激活值的规模。第$i$层的内存消耗用$Req$表示。

我们考虑一个由$M$个边缘设备和云服务器组成的网络。这些设备具有异构的计算能力和内存能力，并且云服务器在计算能力方面比边缘设备强大得多。设备$3$的内存预算为$Me{m}_j$ 。这些计算设备是相互连接的。设备$k$与设备$j$之间的带宽为$B_{k,j}$ ，其中$0\le k\le M-1$ ，$0\le j\le M-1$ 。存在一个源节点，输入词元存储在该源节点中。在不失一般性的情况下，我们将源节点设为节点$0$。本文中使用的主要符号如表二所示。

A. 优化大语言模型（LLM）的推理延迟

问题表述。我们使用一个二元变量$X_{i,j}$来表示大语言模型（LLM）的分配策略。如果第$i$层被分配到节点$j$，则$X_{i,j}$等于$1$，否则$X_{i,j}$等于$0$。某一层将且仅会被分配到一个节点。因此，对于所有的$i$，有${\sum }_{j=0}^{M-1}{X}_{i,j}=1$ 。设$t_{i,j}^{\text{comp}}$表示第$i$层在节点$j$上的计算时间。假设第$i-1$层和第$i$层分别被分配到节点$k$和节点$j$。我们用$t_{i-1,k,j}^{\text{comm}}$表示将第$i-1$层的激活值从节点$k$传输到节点$j$的通信时间。数据传输时间由某一层的输出规模以及两个节点之间的带宽决定。如果第$i-1$层和第$i$层在同一个节点上，我们假设传输时间为$0$。因此，我们有：

t i − 1 , k