多个大语言微调模型并行推断的潜力

随着开源预训练大型语言模型（Large Language Model, LLM ）变得更加强大和开放，越来越多的开发者将大语言模型纳入到他们的项目中。其中一个关键的适应步骤是将领域特定的文档集成到预训练模型中，这被称为微调。

通常情况下，来自领域特定文档的额外知识与预训练模型已经知道的相比微不足道。在这种情况下，低秩适应（Low-Rank Adaptation，LoRA ）技术证明是有价值的。

通过 LoRA，微调模型仅向预训练模型添加不到0.1%的参数。具体来说，这意味着 LoRA 微调模型仅增加了10~200 MB 的存储，具体取决于配置。从计算角度来看，考虑到与预训练模型相比参数的增加极少，额外的计算负载相对较小。

基于存储和计算的额外开销都很小这一点，我相信构建一个多租户的大语言微调模型的推断服务具有很大潜力。这个服务可以托管成千上万个 LoRA 模型，它们都共享相同的预训练大语言模型。在每个批次的执行中，每个用户请求都会调用一个独立的微调模型，从而分摊存储和计算成本到各种不同的模型中。

在我的上一篇博客文章中，我深入探讨了大语言模型推断中的批处理效应。在这篇文章中，我将详细介绍为什么多租户 LoRA 推断服务具有巨大的潜力。

背景知识：文本生成

文本生成服务，如ChatGPT，接受用户文本输入并提供文本响应。这个输入被称为“提示”（prompt）。在内部，当大语言模型处理文本时，它在一系列“词元”（token）上操作。我们可以大致将词元视为几个字符或一个单词。文本生成过程有两个主要阶段：

• 预填充阶段（或称“编码”，“初始化”）接受整个提示并生成随后的词元以及一个“键值缓存”(KV Cache)。
• 解码阶段处理新生成的词元和键值缓存，然后生成下一个词元，同时更新键值缓存。这个阶段不断重复，直到模型完成其输出。

有趣的是，尽管预填充阶段处理的词元数量比解码阶段多100倍，但它们的计算延迟是在同一数量级的。由于解码阶段会重复执行上百次，因此在这篇文章中，我将集中讨论如何优化解码阶段。

背景知识：大语言模型的架构和批处理

在其核心，大语言模型的架构非常简单。它主要包括多个 Transformer 层，所有层都共享相同的架构。每一层包括四个计算密集型组件： QKV 投影、自注意力、输出投影和前馈网络（Feed-Forward Network, FFN ）。

概括地说，其中包含了两种算子：

• 自注意力（Self-Attention，黄色标出）涉及矩阵-矩阵乘法。
• 密集投影（Dense Projection，绿色标出）涉及向量-矩阵乘法。

考虑到每个批次中每个序列只有一个词元，密集投影计算也非常微小，不足以充分利用 GPU 。因此，扩大批处理大小几乎不会影响密集投影的延迟，这使得大批处理大小对于构建高吞吐、低延迟的推断服务至关重要。

关于大语言模型推断中的批处理的更详细分析，请查看我的以前的博客文章。

背景知识： LoRA

给定形状为[H1, H2]的预训练参数矩阵W，LoRA微调训练一个形状为[H1, R]的小矩阵A和形状为[R, H2]的B。我们使用(W+AB)作为微调模型的权重。这里的R是 LoRA 微调指定的秩，通常远小于原始维度（>= 4096），通常在8~32之间。

这种方法背后的逻辑是，与原始权重相比，新增加的知识只占一小部分，因此 LoRA 将增量压缩成两个低秩矩阵。

与完全微调相比，LoRA 预训练显著降低了微调模型的存储和内存需求。

在大语言模型中，由于所有参数都位于密集投影中，LoRA 可以集成到 Transformer 层中的任何位置。虽然HuggingFace PEFT库仅将 LoRA 加到q_proj和v_proj，但一些研究，如QLoRA，主张将其包含在所有的密集投影中。

LoRA 延迟和批处理效应

尽管在存储方面，LoRA 矩阵明显小于原始权重矩阵，但计算的延迟并不成比例地减少。我们可以使用以下代码对骨干模型和 LoRA 增量的延迟进行基准测试：

h1 = 4096
h2 = 11008
r = 16
for bs in range(1, 33):
  w = torch.randn(h1, h2, dtype=torch.float16, device="cuda:0")
  a = torch.randn(bs, h1, r, dtype=torch.float16, device="cuda:0")
  b = torch.randn(bs, r, h2, dtype=torch.float16, device="cuda:0")
  x = torch.randn(bs, 1, h1, dtype=torch.float16, device="cuda:0")
  bench(lambda: x @ w)
  bench(lambda: x @ a @ b)

上图表明，LoRA 增量仅比骨干模型快2.5倍。

然而，显然 LoRA 批处理效应与骨干模型的效应相似。批处理大小的增加仅在较小程度上影响延迟。这个特性使得多租户 LoRA 非常可行。

在骨干模型中，所有一个批次中的所有请求都针对同一模型。而在多租户 LoRA 服务中，一个批次中的请求可能会调用不同的 LoRA 微调模型。数学表达如下：

挑战在于以批处理的方式将不同的 LoRA 增量应用于一个批处理中的各个输入，与此同时还要维持“白吃的午餐”式的批处理效应。

批处理 LoRA 算子

我们想要的批处理 LoRA 算子具有以下函数签名：

def add_lora(
    y: torch.Tensor,      # (batch_size, 1, out_features)
    x: torch.Tensor,      # (batch_size, 1, in_features)
    A: torch.Tensor,      # (num_loras, in_features, lora_rank)
    B: torch.Tensor,      # (num_loras, lora_rank, out_features)
    I: torch.LongTensor,  # (batch_size,)
):
  """Semantics: y[i] += x[i] @ A[I[i]] @ B[I[i]]"""
  raise NotImplementedError()

一个最简单的实现方法是在批处理维度上进行循环：

def lora_loop(
    y: torch.Tensor,  # (batch_size, 1, out_features)
    x: torch.Tensor,  # (batch_size, 1, in_features)
    A: torch.Tensor,  # (num_loras, in_features, lora_rank)
    B: torch.Tensor,  # (num_loras, lora_rank, out_features)
    I: torch.LongTensor,  # (batch_size,)
):
  for i, idx in enumerate(I.cpu().numpy()):
    y[i] += x[i] @ A[idx] @ B[idx]

让我们对循环版本进行基准测试。为了进行比较，我们可以包括一个“作弊”实现，其中我们假设批处理中每个请求的 LoRA 矩阵已经被合并在一起。这样，我们只测量批量矩阵乘法（Batched Matrix Multiplication, bmm）的延迟。

def lora_cheat_bmm(
    y: torch.Tensor,  # (batch_size, 1, out_features)
    x: torch.Tensor,  # (batch_size, 1, in_features)
    cheat_A: torch.Tensor,  # (batch_size, in_features, lora_rank)
    cheat_B: torch.Tensor,  # (batch_size, lora_rank, out_features)
):
  y += x @ cheat_A @ cheat_B


num_loras = 50
h1 = 4096
h2 = 11008
r = 16
A = torch.randn(num_loras, h1, r, dtype=torch.float16, device="cuda:0")
B = torch.randn(num_loras, r, h2, dtype=torch.float16, device="cuda:0")
for bs in range(1, 33):
  x = torch.randn(bs, 1, h1, dtype=torch.float16, device="cuda:0")
  y = torch.randn(bs, 1, h2, dtype=torch.float16, device="cuda:0")
  I = torch.randint(num_loras, (bs,), dtype=torch.long, device="cuda:0")
  cheat_A = A[I, :, :]
  cheat_B = B[I, :, :]

  bench(lambda: lora_loop(y, x, A, B, I))
  bench(lambda: lora_cheat_bmm(y, x, cheat_A, cheat_B))

可预见的是，循环版本明显较慢，并且失去了批处理效应。这是因为它逐个处理输入，而不是利用为批处理数据设计的高效 CUDA 核心。

然而，bmm方法提供了一个很好的启发。我们的目标变得清晰起来：首先将所有 LoRA 矩阵汇总到一个临时的张量中，然后使用bmm。经过一番挖掘，我发现了torch.index_select()函数，它可以高效地执行批量汇总(Batched Gather)。于是我们可以如下一个gbmm（ gather-bmm ）实现：

def lora_gbmm(
    y: torch.Tensor,  # (batch_size, 1, out_features)
    x: torch.Tensor,  # (batch_size, 1, in_features)
    A: torch.Tensor,  # (num_loras, in_features, lora_rank)
    B: torch.Tensor,  # (num_loras, lora_rank, out_features)
    I: torch.LongTensor,  # (batch_size,)
):
  a = torch.index_select(A, 0, I) # (batch_size, in_features, lora_rank)
  b = torch.index_select(B, 0, I) # (batch_size, lora_rank, out_features)
  y += x @ a @ b

BGMV 算子

虽然gbmm非常有效，但它并不是最终解决方案。我们没有必要仅仅是因为 bmm 需要连续的存储而将 LoRA 增量而汇总到一个连续的空间中。理想情况下，聚合可以在 CUDA 核心内部进行，于bmm操作同事进行。如果可能的话，这将消除与torch.index_select()相关的 GPU 内存读写操作。

我请叶子豪大牛帮忙，他是精通高性能 CUDA 核心的编写。经过几轮迭代，子豪开发了一个非常快的 CUDA 程序，把 LoRA 所需的计算分成两半。我们将这个算子命名为 BGMV (Batched Gather Matrix-Vector Multiplication)：

有了 BGMV 算子，add_lora()函数可以使用两次 BGMV 调用来实现：

def lora_bgmv(
    y: torch.Tensor,  # (batch_size, 1, out_features)
    x: torch.Tensor,  # (batch_size, 1, in_features)
    A: torch.Tensor,  # (num_loras, in_features, lora_rank)
    B: torch.Tensor,  # (num_loras, lora_rank, out_features)
    I: torch.LongTensor,  # (batch_size,)
):
  tmp = torch.zeros((x.size(0), A.size(-1)), dtype=x.dtype, device=x.device)
  bgmv(tmp, x, A, I)
  bgmv(y, tmp, B, I)

这是gbmm和bgmv的基准结果：

如图所示，gbmm非常有效。聚合过程相对于bmm增加了约20%的延迟，但保持了令人满意的批处理效应。而子豪编写的 bgmv 算子更加高效，甚至超过了bmm。

多租户 LoRA 文本生成性能

在bgmv算子的基础上，我开发了一个名为Punica的实验项目，支持多个 LoRA 模型。 Punica 的独特能力是能把不同 LoRA 模型的请求合并在一个批处理中。我对比了 Punica 和一系列知名系统的性能，包括HuggingFace Transformers、DeepSpeed、Faster Transformer和vLLM。

在测试中，每个请求都针对不同的 LoRA 模型。鉴于其他系统没有明确针对多租户 LoRA 服务进行优化，它们的批处理大小为1。我使用了HuggingFace PEFT库来把 LoRA 加到 HuggingFace Transformers 和 DeepSpeed 中。我尚未调整 vLLM 和 Faster Transformer，所以它们在没有 LoRA 的情况下运行。以下是结果：

DeepSpeed、vLLM 和 Faster Transformers 都有高度优化的 Transformer 实现，这些系统的吞吐量比标准的 HuggingFace Transformers 高出3倍。但由于批处理大小被限制在1，它们在多租户 LoRA 服务效率方面表现不佳。

相反，Punica 在批处理大小为16时比这些系统高出8倍，与普通的 HuggingFace Transformers 相比，甚至高达23倍。值得注意的是，Punica 的吞吐量几乎与批处理大小呈线性关系。

值得注意的是，Punica 仍然是一个早期的研究原型。它目前使用与普通的 HuggingFace Transformers 相同的 Transformer 实现，除了 LoRA 和自注意力运算符。一些已知的 Transformer 层的优化在 Punica 中尚未实现。这解释了在批处理大小为1时 Punica 与其他高度优化的系统之间的性能差距。

吞吐量不错，那么从延迟方面表现如何？

如图所示，Punica 中的批处理不会引入显著的延迟。

示例用途

在展示了多租户 LoRA 服务的高效性后，让我们设想一些潜在的应用场景：

• 用一本新的小说对 LoRA 模型进行细调，以帮助读者总结每个角色的旅程。
• 针对迅速发展的新闻对 LoRA 模型进行调整，以让读者了解最新动态。
• 基于网页内容对 LoRA 模型进行优化，提高读者的理解能力。

我将这种方法称为“即时细调”(Just-in-time Fine-tuning)，因为 LoRA 的训练速度非常快（在我的试验中，每个训练周期不到一秒）。

总结

本文展示了用批处理加速多个 LoRA 微调模型并行推断的可行性。我实现的 Punica 项目展现出了关于批处理大小几乎线性的吞吐量扩展，并且增加批处理大小并不显著增加延迟。

在 Github 上可以查看 Punica 的开源代码：https://github.com/punica-ai/punica

你可以在这个 Jupyter Notebook中访问本文中使用的基准代码。

这项研究仍然在进行中。我正在积极开展这项研究项目，预计很快会发布一个在线演示。我欢迎任何反馈或想法，请随时在评论部分或下面的按钮中分享。