Punica：多 LoRA 模型统一服务框架

python examples/tui-multi-lora.py

低秩适配（Low-Rank Adaptation，LoRA）是一种参数高效的方法，可以为预训练 LLM 添加新知识。尽管预训练 LLM 需要数百 GB 的存储空间，但经过 LoRA 微调的模型只会增加 1% 的存储和内存开销。Punica 能以运行一个模型的成本，同时运行多个经过 LoRA 微调的模型。

假设形状为 [H1, H2] 的 W 是预训练模型的权重，LoRA 会添加两个小矩阵：形状为 [H1, r] 的 A，以及形状为 [r, H2] 的 B。在微调后的模型上处理输入 x，可以表示为 y := x @ (W + A@B)，这等价于 y := x@W + x@A@B。

当存在 n 个 LoRA 模型时，就会有 A1、B1、A2、B2、……、An、Bn。给定一个输入 batch X := (x1,x2,...,xn)，其中每个输入分别映射到对应的 LoRA 模型，输出为 Y := X@W + (x1@A1@B1, x2@A2@B2, ..., xn@An@Bn)。等式左侧这一项负责在预训练模型上计算输入 batch，效率非常高。得益于强大的 batching 效应，其延迟与只有一个输入时几乎相同。

我们找到了一种高效计算等式右侧部分（即 LoRA 附加部分）的方法。我们将这一操作封装到一个 CUDA kernel 中，称为分段聚集矩阵-向量乘法（Segmented Gather Matrix-Vector Multiplication，SGMV），如下图所示。

在下面的微基准测试图中，可以观察到预训练模型强大的 batching 效应。LoRA 的朴素实现速度较慢，如橙色曲线所示。通过 SGMV 实现的 LoRA 非常高效，并且保留了这种强大的 batching 效应。

下图展示了 Punica 与其他系统在文本生成吞吐量方面的对比，其中包括 HuggingFace Transformers、DeepSpeed、FasterTransformer 和 vLLM。该基准测试考虑了 LoRA 模型热度的不同分布情况。Distinct 表示每个请求都使用不同的 LoRA 模型；Identical 表示所有请求都使用同一个 LoRA 模型；Uniform 和 Skewed 则介于两者之间。与当前最先进的系统相比，Punica 实现了 12 倍的吞吐量。

阅读我们的论文以了解更多内容：Punica: Multi-Tenant LoRA Serving。

你可以通过二进制包安装 Punica，也可以从源码构建。

从二进制包安装

优点：无须编译，安装速度快。

缺点：可能与你的 CUDA 版本、CUDA 架构、PyTorch 版本或 Python 版本不匹配。

当前预编译版本：CUDA：11.8、12.1；Python：3.10、3.11；TORCH_CUDA_ARCH_LIST：8.0 8.6 8.9+PTX

TORCH_CUDA_ARCH_LIST：8.0 8.6 8.9+PTX

pip install ninja torch
pip install punica -i https://punica-ai.github.io/whl/cu121/ --extra-index-url https://pypi.org/simple
# Note: Change cu121 to your CUDA version.

# Please install torch before punica
pip install ninja numpy torch

# Clone punica
git clone https://github.com/punica-ai/punica.git
cd punica
git submodule sync
git submodule update --init

# If you encouter problem while compilation, set TORCH_CUDA_ARCH_LIST to your CUDA architecture.
# export TORCH_CUDA_ARCH_LIST="8.0"

# Build and install punica
pip install -v --no-build-isolation .

同时提供多个 LoRA 模型服务

微调、转换为 Punica 格式，并使用 Punica 提供服务。

参见 examples/finetune/。

文本生成基准测试

python -m benchmarks.bench_textgen_lora --system punica --batch-size 32

@misc{punica,
    title={Punica: Multi-Tenant LoRA Serving},
    author={Lequn Chen and Zihao Ye and Yongji Wu and Danyang Zhuo and Luis Ceze and Arvind Krishnamurthy},
    year={2023},
    eprint={2310.18547},
    archivePrefix={arXiv},
    primaryClass={cs.DC}
}