告别冷启动——LoRA推理速度提升300%的秘诀

发布日期:2023年12月5日
在 GitHub 上更新

简而言之:我们根据用户请求交换Stable Diffusion LoRA适配器,同时保持基础模型“热”启动,从而实现在多个用户之间快速进行LoRA推理。您可以通过浏览我们的LoRA目录并使用推理小部件来体验这一点。

Inference Widget Example

在本博客中,我们将详细介绍我们是如何实现这一目标的。

我们成功地大幅提升了基于公共Diffusion模型的公共LoRA在Hub上的推理速度。这使我们能够节省计算资源,并提供更快、更好的用户体验。

在给定模型上执行推理分为两个步骤:

  1. 预热阶段——包括下载模型和设置服务(25秒)。
  2. 然后是推理任务本身(10秒)。

通过这些改进,我们将预热时间从25秒缩短到3秒。现在我们能够用不到5个A10G GPU为数百个不同的LoRA提供推理服务,同时用户请求的响应时间从35秒减少到13秒。

让我们更深入地探讨如何利用Diffusers库中开发的一些最新功能,通过一个单一的服务以动态方式为许多不同的LoRA提供服务。

LoRA

LoRA是一种微调技术,属于“参数高效”(PEFT)方法家族,旨在减少受微调过程影响的可训练参数数量。它提高了微调速度,同时减小了微调检查点的大小。

我们不通过对所有权重进行微小更改来微调模型,而是冻结大部分层,只训练注意力块中的少数特定层。此外,我们通过将两个较小矩阵的乘积添加到原始权重来避免触及这些层的参数。这些小矩阵的权重在微调过程中得到更新,然后保存到磁盘。这意味着模型的所有原始参数都得以保留,我们可以使用适配方法在顶部加载LoRA权重。

LoRA(低秩适应)这个名称来源于我们提到的小矩阵。有关该方法的更多信息,请参阅这篇博文原始论文

LoRA decomposition

上图显示了两个较小的橙色矩阵,它们作为LoRA适配器的一部分被保存。我们稍后可以加载LoRA适配器并将其与蓝色基础模型合并,以获得黄色的微调模型。关键在于,卸载适配器也是可能的,因此我们可以在任何时候恢复到原始的基础模型。

换句话说,LoRA适配器就像基础模型的附加组件,可以按需添加和移除。而且由于A和B的低秩特性,它与模型大小相比非常轻量。因此,加载速度比加载整个基础模型快得多。

例如,如果你查看Stable Diffusion XL Base 1.0模型库(它被广泛用作许多LoRA适配器的基础模型),你会发现它的尺寸约为**7 GB**。然而,像这个典型的LoRA适配器仅占用**24 MB**的空间!

Hub上“蓝色”基础模型的数量远少于“黄色”模型的数量。如果我们能快速地从蓝色模型切换到黄色模型,反之亦然,那么我们就可以用少数几个不同的蓝色部署来服务许多不同的黄色模型。

有关LoRA的更详尽介绍,请参阅以下博客文章:使用 LoRA 进行高效 Stable Diffusion 微调,或直接查阅原始论文

优势

我们Hub上有大约**2500**个不同的公共LoRA。其中绝大多数(**约92%**)是基于Stable Diffusion XL Base 1.0模型的LoRA。

在实现资源共享之前,这意味着要为所有这些模型(例如,上图中所有黄色的合并矩阵)部署一个专用服务;释放并预留至少一个新的GPU。启动服务并使其准备好为特定模型提供服务的时间约为**25秒**,在此之上还有推理时间(在A10G上,1024x1024的SDXL推理扩散,25个推理步长,约为**10秒**)。如果某个适配器只是偶尔被请求,其服务就会被停止以释放被其他服务抢占的资源。

如果你请求一个不那么受欢迎的LoRA,即使它像目前Hub上绝大多数适配器一样基于SDXL模型,也需要**35秒**来预热并在第一次请求时获得响应(后续请求只需推理时间,例如**10秒**)。

现在:请求时间已从35秒减少到13秒,因为适配器将只使用少数几个不同的“蓝色”基础模型(例如,Diffusion模型只有两个重要的)。即使你的适配器不那么受欢迎,它的“蓝色”服务也很可能已经预热。换句话说,你很有可能避免25秒的预热时间,即使你不经常请求你的模型。蓝色模型已经下载并准备就绪,我们所要做的就是卸载上一个适配器并加载新的适配器,这只需要**3秒**,正如我们下面所看到的。

总的来说,这需要更少的GPU来服务所有不同的模型,尽管我们已经有办法在部署之间共享GPU以最大限度地利用其计算能力。在**2分钟**的时间内,大约有**10个**不同的LoRA权重被请求。我们不是启动10个部署并保持它们热启动,而是简单地用1到2个GPU(如果请求爆发,则更多)来服务所有这些请求。

实现

我们在推理API中实现了LoRA的共享。当我们的平台上一个可用模型被请求时,我们首先判断它是否为LoRA。然后我们识别LoRA的基础模型,并将请求路由到一个公共后端集群,该集群能够处理对所述模型的请求。推理请求通过保持基础模型“热”启动并在运行时加载/卸载LoRA来得到处理。通过这种方式,我们最终能够复用相同的计算资源来同时服务许多不同的模型。

LoRA结构

在Hub中,LoRA可以通过两个属性识别:

Hub

LoRA将具有`base_model`属性。这仅仅是LoRA所构建的模型,在执行推理时应该应用于此模型。

由于LoRA并非唯一具有此类属性的模型(任何复制的模型都将具有该属性),因此LoRA还需要一个`lora`标签才能被正确识别。

为Diffusers加载/卸载LoRA 🧨

请注意,使用peft库可以更无缝地实现本节所述的功能。有关更多详细信息,请参阅文档。原理与下文相同(从上图示意图中的蓝色框到黄色框,反之亦然)。


Diffusers库中使用4个函数来加载和卸载不同的LoRA权重。

load_lora_weightsfuse_lora 用于加载权重并将其与主层合并。请注意,在执行推理之前将权重与主模型合并可以使推理时间减少30%。

unload_lora_weightsunfuse_lora 用于卸载。

下面我们提供了一个示例,展示如何利用Diffusers库快速将多个LoRA权重加载到基础模型之上。

import torch


from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
)

import time

base = "stabilityai/stable-diffusion-xl-base-1.0"

adapter1 = 'nerijs/pixel-art-xl'
weightname1 = 'pixel-art-xl.safetensors'

adapter2 = 'minimaxir/sdxl-wrong-lora'
weightname2 = None

inputs = "elephant"
kwargs = {}

if torch.cuda.is_available():
    kwargs["torch_dtype"] = torch.float16

start = time.time()

# Load VAE compatible with fp16 created by madebyollin
vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix",
    torch_dtype=torch.float16,
)
kwargs["vae"] = vae
kwargs["variant"] = "fp16"

model = DiffusionPipeline.from_pretrained(
    base, **kwargs
)

if torch.cuda.is_available():
    model.to("cuda")

elapsed = time.time() - start

print(f"Base model loaded, elapsed {elapsed:.2f} seconds")


def inference(adapter, weightname):
    start = time.time()
    model.load_lora_weights(adapter, weight_name=weightname)
    # Fusing lora weights with the main layers improves inference time by 30 % !
    model.fuse_lora()
    elapsed = time.time() - start

    print(f"LoRA adapter loaded and fused to main model, elapsed {elapsed:.2f} seconds")

    start = time.time()
    data = model(inputs, num_inference_steps=25).images[0]
    elapsed = time.time() - start
    print(f"Inference time, elapsed {elapsed:.2f} seconds")

    start = time.time()
    model.unfuse_lora()
    model.unload_lora_weights()
    elapsed = time.time() - start
    print(f"LoRA adapter unfused/unloaded from base model, elapsed {elapsed:.2f} seconds")


inference(adapter1, weightname1)
inference(adapter2, weightname2)

加载数据

以下所有数字均以秒为单位

GPU T4 A10G
基础模型加载——未缓存 20 20
基础模型加载——已缓存 5.95 4.09
适配器1加载 3.07 3.46
适配器1卸载 0.52 0.28
适配器2加载 1.44 2.71
适配器2卸载 0.19 0.13
推理时间 20.7 8.5

每次推理增加2到4秒,我们就可以服务许多不同的LoRA。然而,在A10G GPU上,推理时间大幅减少,而适配器加载时间变化不大,因此LoRA的加载/卸载相对更耗时。

处理请求

为了处理推理请求,我们使用了这个开源社区镜像

您可以在 TextToImagePipeline 类中找到之前描述的机制。

当请求LoRA时,我们将检查已加载的LoRA,并仅在需要时进行更改,然后像往常一样执行推理。通过这种方式,我们能够为基础模型和许多不同的适配器提供服务。

下面是如何测试和请求此镜像的示例。

$ git clone https://github.com/huggingface/api-inference-community.git

$ cd api-inference-community/docker_images/diffusers

$ docker build -t test:1.0 -f Dockerfile .

$ cat > /tmp/env_file <<'EOF'
MODEL_ID=stabilityai/stable-diffusion-xl-base-1.0
TASK=text-to-image
HF_HUB_ENABLE_HF_TRANSFER=1
EOF

$ docker run --gpus all --rm --name test1 --env-file /tmp/env_file_minimal -p 8888:80 -it test:1.0

然后在另一个终端向基础模型和/或在Hugging Face Hub上找到的各种LoRA适配器发出请求。

# Request the base model
$ curl 0:8888 -d '{"inputs": "elephant", "parameters": {"num_inference_steps": 20}}' > /tmp/base.jpg

# Request one adapter
$ curl -H 'lora: minimaxir/sdxl-wrong-lora' 0:8888 -d '{"inputs": "elephant", "parameters": {"num_inference_steps": 20}}' > /tmp/adapter1.jpg

# Request another one
$ curl -H 'lora: nerijs/pixel-art-xl' 0:8888 -d '{"inputs": "elephant", "parameters": {"num_inference_steps": 20}}' > /tmp/adapter2.jpg

批量处理呢?

最近发表了一篇非常有趣的论文,描述了如何通过对LoRA模型执行批量推理来提高吞吐量。简而言之,所有推理请求将收集在一个批次中,与通用基础模型相关的计算将一次性完成,然后计算剩余的适配器特定乘积。我们没有实现这种技术(接近于text-generation-inference为LLMs采用的方法)。相反,我们坚持使用单个顺序推理请求。原因是,我们观察到批量处理对扩散模型不感兴趣:吞吐量不会随着批次大小显著增加。在我们进行的简单图像生成基准测试中,对于批次大小为8,吞吐量仅增加了25%,但延迟增加了6倍!相比之下,批量处理对于LLMs更有趣,因为您可以获得8倍的顺序吞吐量,而延迟仅增加10%。这就是我们没有为扩散模型实现批量处理的原因。

结论:时间

通过使用动态LoRA加载,我们成功节省了计算资源并提升了Hub推理API的用户体验。尽管卸载先前加载的适配器并加载我们感兴趣的适配器会额外花费时间,但由于服务进程通常已经启动并运行,整个推理响应时间大大缩短了。

请注意,若要让某个LoRA在Hub上受益于此推理优化,它必须是公开的、非限制访问的,并且基于非限制访问的公共模型。如果您将相同的方法应用于您的部署,请务必告诉我们!

社区

注册登录 发表评论