ME-Switch：适用于大语言模型的内存高效专家切换框架

24年6月来自澳大利亚Monash大学、商汤科技、北航和浙大的论文“ME-Switch: A Memory-Efficient Expert Switching Framework for Large Language Models”。

LLM 开发的典型过程包括在海量数据上预训练通用基础模型，然后在特定任务数据上进行微调以获得一系列专业专家。为这些专家提供服务可能带来重大挑战，因为将所有专家加载到设备上是不切实际的，并且根据用户请求在专家之间频繁切换会产生大量 I/O 成本，从而导致延迟和服务费用显著增加。以前的方法将专家权重分解为预训练模型的权重和残差增量权重，然后使用输出通道步长量化增量权重以减小模型大小。然而，这些方法忽略了一个事实，即增量权重的某些输入通道，在极低的位宽下会导致显著的量化误差。此外，现有方法假设预先知道适合用户请求的模型，而实际情况并非如此。

ME-Switch，一种专为 LLM 服务量身定制的内存高效专家切换框架。为了减少描述增量权重所需的位数，ME-Switch 采用了混合精度量化，该量化选择性地将增量权重的非显著输入通道量化为极低的位，同时保持显著通道不变，从而显著降低存储需求并保持性能。此外，开发一种路由方法，将模型选择问题转化为域分类问题，有效地将用户查询引导至最合适的专家。大量实验表明，ME-Switch 具有良好的内存效率和路由性能。例如，在为 Mistral-7B 系列的三个模型提供服务时，ME-Switch 将模型大小缩小 1.74 倍，并在指令、数学推理和代码生成任务上保持几乎无损的性能。此外，可以在 NVIDIA A100 GPU 上有效地为 Mistral-7B 系列的 16 个模型提供服务。

大语言模型 (LLM)，例如 GPT-4 [6] 和 Gemini [43]，在自然语言处理 (NLP) 方面取得了重大进展。这些模型在大规模文本数据集上进行训练，建立了广泛的语言理解基础，使其能够在需要常识性知识的任务上表现出色。为了获得特定于任务的知识，LLM 可以在专门的任务上进一步微调，从而使其适用于各种应用，例如交互智体 [45, 22]、代码生成 [34, 31] 和数学问题求解 [33, 31]，这显示了 LLM 的非凡多功能性。然而，鉴于没有一个模型可以同时掌握所有任务，因此为多个 LLM 提供服务（每个 LLM 都针对特定任务进行精确定制）对于获得良好的性能至关重要。

然而，为多个模型提供服务面临几个重大挑战。首先，即使模型数量相对较少，由于每个模型包含大量参数，存储需求也很大。例如，三个 LLaMA-2-70B 模型总共需要超过 384GB 的存储空间，按每个模型 128GB 乘以 3 计算。其次，这些模型的大量内存需求可能，使将它们全部同时加载到 GPU 内存中变得不切实际。虽然根据需要动态交换模型权重进出 GPU 内存是可行的，但模型的庞大规模使这个过程变得缓慢而低效，大大延迟了响应时间并对用户体验产生不利影响。

为了应对上述挑战，现有方法 [28, 50] 将微调模型的权重分解为预训练模型的权重和微调期间引入的增量权重。通过对这些增量权重应用每个张量量化 [35]，可以显著减少存储需求，从而实现在多个微调模型之间有效共享基础模型的存储。尽管如此，使用每个张量量化会忽略不同输入和输出通道之间的增量权重区别，从而导致大量的信息丢失。为了缓解这个问题，更细粒度的量化，例如输出通道量化 [48，47，27]，其中每个输出通道都分配有自己可学习的步长，仍然会忽略输入通道的变化。

为了进一步减少信息丢失，可以在量化之前重新缩放增量权重的输入通道 [26]。但是，对于极低位宽的量化，重新缩放只能提供有限的缓解。此外，现有方法假设预先知道响应用户请求时要使用的适当模型。实际上，动态确定使用哪个模型对于有效处理不同的用户查询至关重要。这一差距凸显了对智能模型选择机制的需求，该机制可以根据每个用户请求的具体要求自适应地选择最适合的模型。如图所示Speechless-Code-Mistral-7B 增量权重的输入通道最大值和最小值说明：输入通道之间的差异性突显出，某些突出通道（无论其大小如何）在使用超低位宽量化时都会导致显著的量化误差，这凸显了它们在保持性能方面的关键作用。

为了识别重要的增量权重输入通道，可以根据其大小进行选择 [13]，如图 (a) 所示，但这种方法无法捕捉到导致高量化误差的通道。相比之下，根据增量权重对输出激活中重建误差的影响来识别其显著的输入通道，如图 (b) 所示。

此外，开发一种路由方法，可根据用户请求选择最合适的 LLM，从而实现高效准确的模型切换。为了简化路由问题，假设每个 LLM 都专注于一个不同的领域，例如代码生成或数学推理。将模型选择视为多项选择问答任务，其中每个选项代表一个目标域，从而有效地将问题转化为域分类任务，如图所示。值得注意的是，对话 LLM 表现出强大的指令遵循能力。利用这一点，用一个小型预训练对话 LLM 作为模型级路由器来解决多项选择问答任务。然后，构建一个多项选择问答数据集并对路由器进行微调以提高其性能。

首先使用旨在引出域分类的模板，向路由器提示用户的问题。根据路由器的响应，动态加载所选域特定模型（例如数学模型）的相应压缩增量权重以生成输出。如下表是路由的模版提示：

由于路由器未针对查询域分类进行明确训练，因此其初始路由性能可能不是最优的。为了提高路由性能，构建一个针对路由问题量身定制的多项选择问答数据集。从各个域收集指令跟踪数据，并将查询插入提示模板。响应是通过考虑正确的域特定模型选择来构建的。然后，用监督微调使用构建的数据集对路由器进行微调，以进一步提高路由准确性。