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

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


作者: Mingjin Zhang; Jiannong Cao; Xiaoming Shen; Zeyang Cui;
期刊: （发表日期: 2024-05-23）
期刊分区:
本地链接: Zhang 等 – 2024 – EdgeShard Efficient LLM Inference via Collaborative Edge Computing.pdf
DOI: 10.48550/arXiv.2405.14371
摘要: Large language models (LLMs) have shown great potential in natural language processing and content generation. However, current LLMs heavily rely on cloud computing, leading to prolonged latency, high bandwidth cost, and privacy concerns. Edge computing is promising to address such concerns by deploying LLMs on edge devices, closer to data sources. Some works try to leverage model quantization to reduce the model size to fit the resource-constraint edge devices, but they lead to accuracy loss. Other works use cloud-edge collaboration, suffering from unstable network connections. In this work, we leverage collaborative edge computing to facilitate the collaboration among edge devices and cloud servers for jointly performing efficient LLM inference. We propose a general framework to partition the LLM model into shards and deploy on distributed devices. To achieve efficient LLM inference, we formulate an adaptive joint device selection and model partition problem and design an efficient dynamic programming algorithm to optimize the inference latency and throughput, respectively. Experiments of Llama2 serial models on a heterogeneous physical prototype demonstrate that EdgeShard achieves up to 50% latency reduction and 2x throughput improvement over baseline methods.

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

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 ) (x_{1}, …, x_{n}) (x1,…,xn)作为输入，并通过计算概率 P ( x n + 1 ∣ x 1 , . . . , x n ) P(x_{n+1} | x_{1}, …, x_{n}) P(xn+1∣x1,…,xn)生成第一个新的词元 x n + 1 x_{n+1} xn+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 N N层。 O i O_{i} Oi表示第 i i i层（ 0 ≤ i ≤ N − 1 0 leq i leq N – 1 0≤i≤N−1）激活值的规模。第 i i i层的内存消耗用 R e q Req Req表示。

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

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

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

t i − 1 , k , j comm = { O i − 1 B k , j , if k ≠ j 0 , otherwise (1) t_{i – 1,k,j}^{ ext{comm}} = egin{cases} frac{O_{i – 1}}{B_{k,j}} , & ext{if } k
eq j \ 0, & ext{otherwise} end{cases} ag{1} ti−1,k,jcomm={
Bk,jOi−1,0,if k=jotherwise(1)
因此，总推理时间可以通过以下公式计算：

T t o l = ∑ i = 0 N − 1 ∑ j = 0 M − 1 X i , j × t c o m p i , j + ∑ i = 1 N − 1 ∑ j = 0 M − 1 ∑ k = 0 M − 1 X i − 1 , k × X i , j × t c o m m i − 1 , k , j (2) T_{tol }=sum_{i = 0}^{N – 1}sum_{j = 0}^{M – 1}X_{i, j} imes t_{comp }^{i, j}+sum_{i = 1}^{N – 1}sum_{j = 0}^{M – 1}sum_{k = 0}^{M – 1}X_{i – 1, k} imes X_{i, j} imes t_{comm }^{i – 1, k, j} ag{2} Ttol=i=0∑N−1j=0∑M−1Xi,j×tcompi,j+i=1∑N−1j=0∑M−1k=0∑M−1Xi−1,k×Xi,j×tcommi−1,k,j(2)

因此，最小化大语言模型推理延迟的问题可以表述如下，其中公式（4）是隐私约束条件。该条件表明，大语言模型的第一层应始终分配给节点0，节点0被设定为存储输入词元的源节点。在这种情况下，原始输入数据位于源节点，从而避免了在计算设备之间传输。公式（5）表明，分配给节点 j j j的所有层的内存需求不能超过其内存预算。 min T t o l (3) ext{min} T_{tol} ag3 min Ttol(3) X 0 , 0 = 1 (4) X_{0,0} = 1 ag4 X0,0=1(4) ∑ i = 0 N − 1 X i , j ∗ R e q i ≤ M e m j (5) sum_{i = 0}^{N – 1} X_{i, j} * Req_{i} leq Mem_{j} ag5 i=0∑N−1Xi,j∗Reqi≤Memj(5)

解决方案。为了最小化推理延迟，我们设计了一种动态规划算法。其思路是，前 i i i层的最小执行时间取决于前 i − 1 i – 1 i−1层，这意味着最优解可以由子问题的最优结果构建得出。该问题具有最优子问题性质，这促使我们采用动态规划方法。

令 D P ( i , j ) DP(i, j) DP(i,j)表示在第 i i i层被分配到节点 j j j之后，前 i i i层的最小总执行时间。状态转移方程表述为： D P ( i , j ) = { min ⁡ k ∈ M 1 ≤ i < N − 1 ( D P ( i − 1 , k ) + t c o m p i , j + t c o m m i − 1 , k , j ) min ⁡ k ∈ M i = N − 1 ( D P ( i − 1 , k ) + t c o m p i , j + t c o m m i − 1 , k , j + t c o m m i , j , 0 ) (6) DP(i, j)=left{egin{array}{l} min_{substack{k in M \ 1 leq i<N – 1}}left(DP(i – 1, k)+t_{comp}^{i, j}+t_{comm}^{i – 1, k, j}
ight) \ min_{substack{k in M \ i = N – 1}}left(DP(i – 1, k)+t_{comp}^{i, j}+t_{comm}^{i – 1, k, j}+t_{comm}^{i, j, 0}
ight) \ end{array}
ight. ag6 DP(i,j)=⎩
⎨
⎧mink∈M1≤i<N−1(DP(i−1,k)+tcompi,j+tcommi−1,k,j)mink∈Mi=N−1(DP(i−1,k)+tcompi,j+tcommi−1,k,j+tcommi,j,0)(6)
其中， D P ( i − 1 , k ) DP(i – 1, k) DP(i−1,k)表示如果将第 i − 1 i – 1 i−1层分配给设备 k k k，前 i − 1 i – 1 i−1层的最小执行时间。公式（6）表明， D P ( i , j ) DP(i, j) DP(i,j)的确定需要遍历上一层所有可能的节点，并选择能使前 i i i层执行时间最小的节点。此外，由于大语言模型具有自回归特性，生成的词元需要被送回源节点，以便进行下一轮的生成。因此，对于最后一层 N − 1 N – 1 N−1，通信时间不仅包括从第 N − 2 N – 2 N−2层传输数据的时间，还包括传输到源节点的时间 t N − 1 , j , 0 c o m m t_{N – 1,j,0}^{comm} tN−1,j,0comm。另外，考虑到隐私约束条件，我们按照公式（7）初始化 D P ( 0 , 0 ) DP(0, 0) DP(0,0)。 D P ( 0 , 0 ) = t c o m p 0 , 0 (7) DP(0, 0) = t_{comp}^{0,0} ag7 DP(0,0)=tcomp0,0(7) 基于公式（6）遍历每一层和每个节点，我们可以填充动态规划表 D P ( i , j ) DP(i, j) DP(i,j)，以跟踪到达每一层的最小总执行时间。最后，最后一层的最小总执行时间可以通过公式（8）计算得出。然后，通过回溯 D P ( i , j ) DP(i, j) DP(i,j)，我们可以得到每一层的最优节点分配方案。 min ⁡ j = 0 , … , M − 1 ( D P ( N − 1 , j ) ) (8) min_{j = 0, ldots, M – 1}(DP(N – 1, j)) ag8 j=0,…,M−1min(DP(N−1,j))(8) 这种方法简单而有效。利用动态规划，我们可以快速遍历解空间，找到最佳的大语言模型划分和分配策略。用于最小化推理延迟的最优大语言模型划分和分配策略的算法如算法1所示。

在算法1中，我们首先初始化动态规划表 D P ( i , j ) DP(i, j) DP(i,j)和选择表 c h o i c e ( i , j ) choice(i, j) choice(i,j)（第1 – 2行）。我们根据公式（7）初始化 D P ( 0 , 0 ) DP(0, 0) DP(0,0)。 c h o i c e ( i , j ) choice(i, j) choice(i,j)记录了承载第 i − 1 i – 1 i−1层的节点 k k k。它是公式（6）中的最优变量。然后我们从第1层开始遍历大语言模型的各层。对于每一层，我们遍历所有具有足够内存的计算设备，并计算推理时间（第3 – 19行）。在填充完动态规划表后，我们可以找到最小的 D P ( N − 1 , j ) DP(N – 1, j) DP(N−1,j)，它代表了执行大语言模型的最短时间，以及承载第 N − 1 N – 1 N−1层的最后一个节点。最后，通过回溯 c h o i c e ( i , j ) choice(i, j) choice(i,j)，我们得到模型划分和分配策略 R R R（第20 – 28行）。算法1的计算复杂度为 O ( N × M × M ) O(N ×M ×M) O(N×M×M)，其中 N N N是大语言模型的层数， M M M是网络中设备的数量。

B. 优化大语言模型（LLM）的推理吞吐量

问题表述。为了优化吞吐量，采用流水线并行来避免设备闲置。如前面所阐述的，第 i i i层在节点 j j j上的计算时间为 t i , j comp t_{i,j}^{ ext{comp}} ti,jcomp，并且如果从第 i i i层到第 m m m层都分配到节点 j j j，计算时间用 t i → m , j comp t_{i o m,j}^{ ext{comp}} ti→m,jcomp表示。第 i − 1 i – 1 i−1层的激活值从节点 k k k传输到节点 j j j的数据传输时间为 t i − 1 , k , j comm t_{i – 1,k,j}^{ ext{comm}} ti−1,k,jcomm。在流水线并行推理中，计算时间和通信时间可以重叠，以最大化吞吐量。因此，对于推理任务，节点 j j j的最大延迟可按如下方式计算： T j latency = max ⁡ { t i → m , j comp , t i − 1 , k , j comm } (9) T_{j}^{ ext{latency}} = maxleft{t_{i o m,j}^{ ext{comp}}, t_{i – 1,k,j}^{ ext{comm}}
ight} ag9 Tjlatency=max{
ti→m,jcomp,ti−1,k,jcomm}(9)
理想情况下，对于选定的设备而言，实现最大吞吐量等同于最小化最慢设备的延迟。我们用 S S S表示选定的设备，那么，最大化推理吞吐量的问题可以表述如下，其中 j ∈ S j in S j∈S。 min ⁡ { T l a t e n c y j ∣ j ∈ S } (10) min left{T_{latency }^{j} mid j in S
ight} ag{10} min{
Tlatencyj∣j∈S}(10)
解决方案。与最小化推理延迟类似，最大化吞吐量的问题也具有最优子问题性质。第一层到第 i i i层吞吐量的最大化可以通过解决第 i − 1 i – 1 i−1层的分配问题推导得出，这意味着整个问题的最优解可以由子问题构建而成。我们同样使用动态规划来解决这个问题。

令 g ( i , S , k ) g(i, S, k) g(i,S,k)表示在使用设备集合 S S S的情况下处理前 i i i层的最短时间，并且设备 k k k是所使用的最后一个节点， k ∈ S k in S k∈S。我们用 g ( m , S ′ , j ) g(m, S', j) g(m,S′,j)表示在使用设备集合 S ′ S' S′的情况下处理前 m m m层的下一状态，并且设备 j j j是所使用的最后一个节点，其中 0 ≤ i < m ≤ N − 1 0 leq i < m leq N – 1 0≤i<m≤N−1， j ∈ M S j in M S j∈MS（这里“ j ∈ M S j in M S j∈MS”可能表述有误，也许你想表达的是 j ∈ S ′ j in S' j∈S′之类的意思，你可以再检查一下原文）， S ′ = S ∪ { j } S' = S cup {j} S′=S∪{
j} 。

状态转移方程由公式（11）给出，其中 g ( m , S ′ , j ) g(m, S', j) g(m,S′,j)由前一状态 g ( i , S , k ) g(i, S, k) g(i,S,k)、设备 j j j的最大延迟决定，即由计算时间 t i − 1 , k , j comm t_{i – 1,k,j}^{ ext{comm}} ti−1,k,jcomm和通信时间 t i → m , j comp t_{i o m,j}^{ ext{comp}} ti→m,jcomp决定。最终的最优解 T t h r o u opt T_{throu}^{ ext{opt}} Tthrouopt是最小的 g ( N − 1 , S ′ , j ) g(N – 1, S', j) g(N−1,S′,j)，其中 S ′ ⊆ M S' subseteq M S′⊆M。 g ( m , S ′ , j ) ( S ′ = S ∪ { j } ) = min ⁡ 0 ≤ i < m ≤ N − 1 j ∈ M ∖ S max ⁡ { g ( i , S , k ) t i − 1 , k , j comm t i → m , j comp } (11) g(m, S', j) quad (S' = S cup {j}) = min_{substack{0 leq i < m leq N – 1 \ j in M setminus S}} max left{ egin{array}{l} g(i, S, k) \ t_{i – 1,k,j}^{ ext{comm}} \ t_{i o m,j}^{ ext{comp}} end{array}
ight} ag{11} g(m,S′,j)(S′=S∪{
j})=0≤i<m≤N−1j∈M∖Sminmax⎩
⎨
⎧g(i,S,k)ti−1,k,jcommti→m,jcomp⎭
⎬
⎫(11) 此外，在进行状态转移时我们有一些约束条件。它们是公式（12）所示的内存约束以及公式（13）中的隐私约束。 R e q i → m ≤ M e m j (12) Req_{i o m} leq Mem_{j} ag{12} Reqi→m≤Memj(12) g ( 1 , 1 , 0 ) = t 0 , 0 comp (13) g(1, 1, 0) = t_{0,0}^{ ext{comp}} ag{13} g(1,1,0)=t0,0comp(13) 算法2描述了找到最优解 T t h r o u opt T_{throu}^{ ext{opt}} Tthrouopt以及相应的模型划分和分配策略的伪代码。在算法2中，我们首先初始化动态规划表 g ( m , S ′ , j ) g(m, S', j) g(m,S′,j)和选择表 c h o i c e ( m , S , j ) choice(m, S, j) choice(m,S,j)，并将 t 0 , 0 comp t_{0,0}^{ ext{comp}} t0,0comp赋值给 g ( 1 , 1 , 0 ) g(1, 1, 0) g(1,1,0)（第1 – 2行）。然后我们从第1层开始遍历大语言模型的各层。对于每一层，我们遍历所有具有足够内存的计算设备，并计算最大延迟（第3 – 23行）。在填充完动态规划表后，我们可以找到最大延迟，基于此我们回溯选择表，最终得到模型划分和分配策略（第24 – 32行）。算法2的计算复杂度为 O ( N 2 × 2 M × M 2 ) O(N^2×2^M×M^2) O(N2×2M×M2)，其中 N N N是大语言模型的层数， M M M是网络中设备的数量。

流水线执行优化。需要注意的是，上述问题的表述和解决方案是基于理想情况，即在任何时候都不存在设备闲置的情况。一台设备处理一批数据后，无需等待即可继续处理另一批数据。然而，对于现实场景中的大语言模型推理来说，这种情况是不切实际的。

如图5（a）所示，与那些单阶段计算应用程序不同，基于解码器的大语言模型应用程序具有自回归特性，即会生成多个词元，并且当前词元的计算依赖于之前所有的词元。在获取到先前生成的词元之前，当前词元的计算无法开始。这会在流水线执行过程中产生气泡（即出现闲置或等待的情况）。

为了逼近理想情况并提高资源利用率以提升吞吐量，我们倾向于减少流水线执行中的气泡现象。我们提出了“EdgeShard-No-bubbles”方法，该方法允许在一次迭代中无需等待所有微批次结束即可立即生成词元。如图5(b)所示，在第一批的预填充阶段P1结束后，设备1会立即执行第一批的词元生成操作，如G1A所示。类似地，当G1A结束时，设备1会进入下一次词元生成迭代，如G1B所示。与“EdgeShard-Bubbles”方法相比，“EdgeShard-No-bubbles”方法通过减少设备闲置时间来减少bubbles现象，并且有望提高吞吐量。从图5的流水线执行图来看，我们可以看到，“EdgeShard-No-bubbles”方法能够在同一时间生成更多的词元。

实验评估

A. 实验设置

测试平台。我们使用各种边缘设备和云服务器来充当协作边缘计算中的异构计算设备。这些设备的规格列于表三。我们使用15台设备，包括12台英伟达 Jetson AGX Orin、2台英伟达 Jetson Orin NX以及一台云服务器来配置协作边缘网络。物理测试平台如图6所示。这些设备通过路由器和交换机相连。任意两台设备之间的带宽为1000Mbps。我们使用Linux系统的流量控制工具（TC工具）[20]来改变设备之间的网络带宽和通信延迟。

Benchmarks。我们使用一系列 Llama2 模型 [2]（包括 Llama2-7B、Llama2-13B 和 Llama2-70B）来测试 EdgeShard 的性能。Llama2 由 Meta 于 2023 年 7 月发布，是最受欢迎且功能强大的开源大语言模型之一，代表了人工智能和自然语言处理领域的一个开创性飞跃。对于模型推理，我们采用文本生成任务来测试其性能。我们使用来自 HuggingFace 的 WikiText-2 数据集 [21]。我们提取了输入词元长度为 32 的样本子集，并生成 96 个词元。在接下来的所有实验中，我们都使用全精度模型推理。

Baselines。我们根据延迟和吞吐量，将 EdgeShard 的性能与各种基线进行比较。（我们不将仅使用云服务的情况作为基线，因为这需要将输入词元传输到云服务器，可能会引发隐私问题）。

边缘单独部署（Edge-Solo）。在这种情况下，大语言模型（LLMs）在边缘设备上进行本地部署，不进行模型划分。
云边均匀划分（Cloud-Edge-Even）。在这种情况下，大语言模型被均匀划分为两部分。一部分分配到边缘设备，另一部分分配到云服务器。
云边优化划分（Cloud-Edge-Opt）。在这种情况下，大语言模型被划分为两个分片。一部分分配到边缘设备，另一部分分配到云服务器。对于大语言模型的划分策略，我们同样使用所提出的动态规划算法。不同之处在于，算法输入中只有两台设备。

B. 整体评估

我们将英伟达 Jetson AGX Orin 设为源节点，并将源节点与云服务器之间的带宽设置为1Mbps。其他计算设备之间的带宽设置为50Mbps，偏差为20%。为了测试吞吐量，我们将批量大小设置为参与设备所能支持的最大批量大小。大语言模型推理的延迟和吞吐量见表四。

我们有以下几点发现。首先，EdgeShard对于大语言模型的部署来说具有潜力且有益。对于Llama2-70B模型，其内存需求约为280GB，这远远超过了单独边缘部署以及云边协同部署的内存容量。我们有如下观察结果。首先，EdgeShard对于大语言模型的部署具有潜力且有益。对于Llama2-70B模型，其内存需求约为280GB，这远远超出了单独边缘部署和云边协同部署的内存容量，会出现内存不足问题（OOM）。然而，EdgeShard通过将大型模型拆分为多个分片并将它们分配到多个设备上，实现了协作式的模型推理，从而应对了这一挑战。其次，与基线方法相比，EdgeShard明显实现了更低的推理延迟和更高的推理吞吐量。对于Llama2-7B模型，EdgeShard实现了75.88毫秒的延迟，比“边缘单独部署（Edge-Solo）”和“云边优化划分（Cloud-Edge-Opt）”快约1.85倍，比“云边均匀划分（Cloud-Edge-Even）”快约3倍。在推理吞吐量方面，EdgeShard在最大批量大小为8的情况下达到了每秒52.45个词元，约为“边缘单独部署”和“云边优化划分”的2.2倍，约为“云边均匀划分”的7倍。在Llama2-13B模型上也观察到了类似的性能提升，EdgeShard的延迟分别比“云边均匀划分”和“云边优化划分”低45.7%和28.8%。此外，EdgeShard的吞吐量比“云边均匀划分”和“云边优化划分”分别高2.23倍和2.2倍。第三，我们还可以看到，对于Llama2-7B模型，“云边优化划分”在推理延迟和吞吐量方面的性能往往与“边缘单独部署”相同。这是因为在这个实验设置中，源节点与云服务器之间的带宽非常有限，即1Mbps。云边协同的最优部署策略是本地执行，这与“边缘单独部署”的情况相同。

C. 带宽的影响

我们将源节点设置为英伟达 Jetson AGX Orin，并将云服务器与源节点之间的带宽在1Mbps到50Mbps之间变化。大语言模型推理的延迟和吞吐量性能分别如图7和图8所示。

对于Llama2-13B模型，单个英伟达Jetson AGX Orin设备无法容纳整个模型。我们仅比较“云边均匀划分（Cloud-Edge-Even）”、“云边优化划分（Cloud-Edge-Opt）”和“EdgeShard”的性能。同样，由于内存限制，三种基线方法都无法部署Llama2-70B模型。相反，我们比较“EdgeShard”与其变体“EdgeShard-均匀划分（EdgeShard-Even）”的性能，在“EdgeShard-均匀划分”中，模型被均匀划分并部署到所有参与的计算设备上。我们选择11台英伟达Jetson AGX Orin设备和1台英伟达RTX 3090设备来部署Llama2-70B模型。

在延迟方面，除了“边缘单独部署（Edge-Solo）”之外，其他三种方法的延迟会随着带宽的增加而降低。这是因为这三种方法是基于协作的，延迟会受到数据传输时间的影响。带宽的增加会减少数据传输时间。我们还可以看到，对于这些协作方法，当云与源之间的带宽从1Mbps变化到10Mbps时，延迟会大幅降低，而从10Mbps变化到50Mbps时，延迟变化较小。这是因为此时带宽逐渐达到饱和状态，计算时间成为了瓶颈。

此外，我们可以看到，当带宽大于10Mbps时，云边协同方法的性能优于“边缘单独部署（Edge-Solo）”方法，因为云边协同方法引入了功能强大的云服务器来加速计算。然而，当带宽为1Mbps时，“云边均匀划分（Cloud-Edge-Even）”的性能比“边缘单独部署（Edge-Solo）”更差。这是因为在这种情况下数据传输成本很高。“云边优化划分（Cloud-Edge-Opt）”方法倾向于在本地部署大语言模型，这与“边缘单独部署（Edge-Solo）”方法相同。有趣的是，当带宽大于10Mbps时，“云边优化划分（Cloud-Edge-Opt）”和“EdgeShard”的延迟几乎相同。我们发现，“EdgeShard”生成的模型划分和分配策略与“云边优化划分（Cloud-Edge-Opt）”方法相同。差异来自模型执行过程中的微小波动。这表明“EdgeShard”的性能不会比“云边优化划分（Cloud-Edge-Opt）”差，并且“云边优化划分（Cloud-Edge-Opt）”方法是“EdgeShard”的一种特殊情况。对于Llama2-13B模型，也观察到了类似的模式。对于Llama2-70B模型，“EdgeShard”的性能优于其变体“EdgeShard-均匀划分（EdgeShard-Even）”，因为云服务器和边缘设备之间存在资源异构性，并且“EdgeShard”会在计算设备之间自适应地划分大语言模型。然而，性能提升并不十分明显，因为有11台计算能力相同的英伟达Jetson AGX Orin设备，且只有1台英伟达RTX 3090设备。

在吞吐量方面，对于Llama2-7B模型，也发现了与延迟评估类似的模式。不同且有趣的是，对于Llama2-13B模型，当带宽为10Mbps时，“EdgeShard”的性能与“云边优化划分（Cloud-Edge-Opt）”方法并不相近，而是有了显著提升，此时“EdgeShard”的吞吐量大约是“云边优化划分”方法的两倍。这是因为英伟达RTX 3090和作为源节点的英伟达Jetson AGX Orin的内存消耗较高。我们观察到，对于“云边优化划分”，这两台设备的内存消耗分别高达95%和98%，这使得最大批量大小仅为4。否则，计算设备上的键值（KV）缓存将没有足够的内存。然而，当带宽为10Mbps时，“EdgeShard”涉及多个边缘设备，单个设备的内存消耗大幅降低，允许更大的批量大小，在这种情况下为8。当带宽高于10Mbps时，“EdgeShard”的模型划分和分配策略往往与“云边优化划分”相同，这使得性能相近，就如同Llama2-7B模型的情况一样。对于Llama2-70B模型，“EdgeShard”的吞吐量有轻微提升，而“EdgeShard-均匀划分（EdgeShard-Even）”的吞吐量保持稳定，因为均匀划分策略不会随着云与源之间的带宽变化而改变。

D. 源节点的影响

我们还测试了源节点对推理延迟和吞吐量的影响，因为源节点可能具有不同的计算和内存容量，而EdgeShard会将大语言模型的第一层强制部署在源节点上，以避免原始数据传输。我们分别将英伟达Jetson AGX Orin和英伟达Jetson Orin NX设为源节点，并比较它们的性能。我们将源节点与云服务器之间的带宽设置为1Mbps。Llama2-7B模型推理的结果如图9所示。

我们发现，当源节点为英伟达Jetson Orin NX时，“边缘单独部署（EdgeSolo）”和“云边均匀划分（Cloud-Edge-Even）”方法会遇到内存不足（OOM）错误。这是因为英伟达Jetson Orin NX的内存相对较低，它无法容纳Llama2-7B模型，甚至连模型的一半都无法容纳。“云边优化划分（Cloud-Edge-Opt）”方法下的两种情况之间的差异比“EdgeShard”方法的差异要明显得多。对于“云边优化划分”方法，大约有60毫秒的差距，而对于“EdgeShard”方法，差距大约为5毫秒。这是因为在“云边优化划分”的情况下只有两台设备，并且它倾向于在源节点上放置更多的模型层。然而，英伟达Jetson AGX Orin在计算能力方面比英伟达Jetson Orin NX强大得多。“EdgeShard”倾向于涉及更多的设备，并且在源节点上放置更少的模型层，这可以弥补源节点之间在计算能力方面的差距。在吞吐量方面也观察到了类似的现象，对于“云边优化划分”方法，英伟达Jetson AGX Orin的吞吐量比英伟达Jetson Orin NX高6倍，而在“EdgeShard”方法下，英伟达Jetson AGX Orin的吞吐量仅比英伟达Jetson Orin NX高2倍。这表明“EdgeShard”可以充分利用网络中的计算资源来优化性能。

E. 流水线执行策略的影响

我们对两种流水线执行策略进行评估。我们将云服务器与源节点之间的带宽设置为1Mbps。结果如图10所示。

我们可以看到，对于所有方法而言，无气泡的EdgeShard（EdgeShard-No-bubble）的性能优于带有气泡的EdgeShard（EdgeShard-Bubble）。具体来说，对于Llama2-7b模型，对于“云边均匀划分（Cloud-Edge-Even）”和“EdgeShard”方法，无气泡的EdgeShard（EdgeShard-No-bubble）相较于有气泡的EdgeShard（Edgeshard-bubble），每秒分别能多处理约0.36个和6.96个词元。对于“云边优化划分（Cloud-Edge-Opt）”方法，在这种情况下它选择本地执行，不存在流水线执行，所以这两种方法的吞吐量是相同的。对于Llama2-13b模型，无气泡的EdgeShard（EdgeShard-No-bubble）相较于有气泡的EdgeShard（Edgeshard-Bubble），对于“云边均匀划分（Cloud-Edge-Even）”、“云边（Cloud-Edge）”和“EdgeShard”方法，每秒分别能多处理约1.69个、1.89个和5.21个词元。与有气泡的EdgeShard（EdgeShard-Bubble）相比，无气泡的EdgeShard（EdgeShard-No-bubble）无需等待一次迭代中所有微批次的完成，能够有效减少设备的空闲时间，从而实现更高的吞吐量。

讨论与未来工作

本节讨论一些可能会引起读者兴趣的开放性问题和未来的研究方向。

激励机制。在这项工作中，我们将大语言模型（LLM）划分为多个分片，并将它们分配到异构设备上。对于诸如智能家居、智能工厂等边缘计算场景，存在一组归单一利益相关者所有的可信设备。他们或许能够利用这些设备进行协同推理。然而，如果这些设备归不同的利益相关者所有，他们可能不愿意共享设备的计算资源。因此，需要进一步的激励机制来对资源共享给予奖励。

批量大小感知优化。较大的批量大小会增加内存使用量，并影响推理吞吐量。如实验所示，通过将大语言模型推理的工作负载划分到多个设备上，参与设备的内存使用量可以降低，从而能够采用更大的批量大小，进而提高吞吐量。然而，所设计的动态规划算法并未考虑批量大小的影响，这仍存在进一步优化的空间。

结论

在这项工作中，我们提出了EdgeShard，以实现大语言模型（LLMs）在协作式边缘设备和云服务器上的高效部署和分布式推理。我们构建了一个联合的设备选择和模型划分问题，分别对推理延迟和吞吐量进行优化，并使用动态规划算法来解决该问题。实验结果表明，EdgeShard能够在各种异构网络条件下自适应地确定大语言模型的划分和部署策略，以优化推理性能。EdgeShard并非旨在取代基于云的大语言模型推理，而是通过利用无处不在的计算设备提供一种灵活且自适应的大语言模型服务方法。实验还表明，当云带宽不足时，EdgeShard的性能优于云边协同推理方法，而当面对相对充足的云带宽时，它倾向于产生与云边协同推理方法相同的部署策略。

这是在协作式边缘计算环境中部署大语言模型的一项开创性工作。我们希望这项工作能够激发在这个充满前景的领域中产生更多的想法和进一步的研究。

文章版权归作者所有，未经允许请勿转载。如内容涉嫌侵权，请在本页底部进入<联系我们>进行举报投诉!

THE END