Goodbye cold boot – how we made LoRA Inference 300% faster

tl;dr: We swap the Stable Diffusion LoRA adapters per user request, while keeping the bottom model warm allowing fast LoRA inference across multiple users. You may experience this by browsing our LoRA catalogue and fidgeting with the inference widget.

On this blog we’ll go intimately over how we achieved that.

We have been capable of drastically speed up inference within the Hub for public LoRAs based on public Diffusion models. This has allowed us to avoid wasting compute resources and supply a faster and higher user experience.

To perform inference on a given model, there are two steps:

1. Warm up phase – that consists in downloading the model and establishing the service (25s).
2. Then the inference job itself (10s).

With the improvements, we were capable of reduce the nice and cozy up time from 25s to 3s. We are actually capable of serve inference for a whole lot of distinct LoRAs, with lower than 5 A10G GPUs, while the response time to user requests decreased from 35s to 13s.

Let’s talk more about how we will leverage some recent features developed within the Diffusers library to serve many distinct LoRAs in a dynamic fashion with one single service.

LoRA

LoRA is a fine-tuning technique that belongs to the family of “parameter-efficient” (PEFT) methods, which try to scale back the variety of trainable parameters affected by the fine-tuning process. It increases fine-tuning speed while reducing the dimensions of fine-tuned checkpoints.

As a substitute of fine-tuning the model by performing tiny changes to all its weights, we freeze a lot of the layers and only train a couple of specific ones in the eye blocks. Moreover, we avoid touching the parameters of those layers by adding the product of two smaller matrices to the unique weights. Those small matrices are those whose weights are updated through the fine-tuning process, after which saved to disk. Which means the entire model original parameters are preserved, and we will load the LoRA weights on top using an adaptation method.

The LoRA name (Low Rank Adaptation) comes from the small matrices we mentioned. For more information concerning the method, please discuss with this post or the original paper.

The diagram above shows two smaller orange matrices which might be saved as a part of the LoRA adapter. We are able to later load the LoRA adapter and merge it with the blue base model to acquire the yellow fine-tuned model. Crucially, unloading the adapter can also be possible so we will revert back to the unique base model at any point.

In other words, the LoRA adapter is like an add-on of a base model that will be added and removed on demand. And since of A and B smaller ranks, it is vitally light compared with the model size. Due to this fact, loading is way faster than loading the entire base model.

If you happen to look, for instance, contained in the Stable Diffusion XL Base 1.0 model repo, which is widely used as a base model for a lot of LoRA adapters, you may see that its size is around 7 GB. Nonetheless, typical LoRA adapters like this one take a mere 24 MB of space !

There are far less blue base models than there are yellow ones on the Hub. If we will go quickly from the blue to yellow one and vice versa, then we’ve got a way serve many distinct yellow models with only a couple of distinct blue deployments.

For a more exhaustive presentation on what LoRA is, please discuss with the next blog post:Using LoRA for Efficient Stable Diffusion Effective-Tuning, or refer on to the original paper.

Advantages

We’ve got roughly 2500 distinct public LoRAs on the Hub. The overwhelming majority (~92%) of them are LoRAs based on the Stable Diffusion XL Base 1.0 model.

Before this mutualization, this might have meant deploying a dedicated service for all of them (eg. for all of the yellow merged matrices within the diagram above); releasing + reserving not less than one latest GPU. The time to spawn the service and have it able to serve requests for a particular model is roughly 25s, then on top of this you’ve the inference time (~10s for a 1024×1024 SDXL inference diffusion with 25 inference steps on an A10G). If an adapter is just occasionally requested, its service gets stopped to free resources preempted by others.

If you happen to were requesting a LoRA that was not so popular, even when it was based on the SDXL model just like the overwhelming majority of adapters found on the Hub up to now, it might have required 35s to warm it up and get a solution on the primary request (the next ones would have taken the inference time, eg. 10s).

Now: request time has decreased from 35s to 13s since adapters will use only a couple of distinct “blue” base models (like 2 significant ones for Diffusion). Even in case your adapter isn’t so popular, there’s probability that its “blue” service is already warmed up. In other words, there’s probability that you just avoid the 25s warm up time, even if you happen to don’t request your model that usually. The blue model is already downloaded and prepared, all we’ve got to do is unload the previous adapter and cargo the brand new one, which takes 3s as we see below.

Overall, this requires less GPUs to serve all distinct models, although we already had a option to share GPUs between deployments to maximise their compute usage. In a 2min time-frame, there are roughly 10 distinct LoRA weights which might be requested. As a substitute of spawning 10 deployments, and keeping them warm, we simply serve all of them with 1 to 2 GPUs (or more if there’s a request burst).

Implementation

We implemented LoRA mutualization within the Inference API. When a request is performed on a model available in our platform, we first determine whether this can be a LoRA or not. We then discover the bottom model for the LoRA and route the request to a standard backend farm, with the power to serve requests for the said model. Inference requests get served by keeping the bottom model warm and loading/unloading LoRAs on the fly. This fashion we will ultimately reuse the identical compute resources to serve many distinct models directly.

LoRA structure

Within the Hub, LoRAs will be identified with two attributes:

A LoRA could have a base_model attribute. This is just the model which the LoRA was built for and needs to be applied to when performing inference.

Because LoRAs are usually not the one models with such an attribute (any duplicated model could have one), a LoRA may also need a lora tag to be properly identified.

Loading/Offloading LoRA for Diffusers 🧨

4 functions are utilized in the Diffusers library to load and unload distinct LoRA weights:

load_lora_weights and fuse_lora for loading and merging weights with the important layers. Note that merging weights with the important model before performing inference can decrease the inference time by 30%.

unload_lora_weights and unfuse_lora for unloading.

We offer an example below on how one can leverage the Diffusers library to quickly load several LoRA weights on top of a base model:

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()


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)
    
    model.fuse_lora()
    elapsed = time.time() - start

    print(f"LoRA adapter loaded and fused to important 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)

Loading figures

All numbers below are in seconds:

With 2 to 4 additional seconds per inference, we will serve many distinct LoRAs. Nonetheless, on an A10G GPU, the inference time decreases by so much while the adapters loading time doesn’t change much, so the LoRA’s loading/unloading is comparatively costlier.

Serving requests

To serve inference requests, we use this open source community image

You could find the previously described mechanism utilized in the TextToImagePipeline class.

When a LoRA is requested, we’ll take a look at the one which is loaded and alter it provided that required, then we perform inference as usual. This fashion, we’re capable of serve requests for the bottom model and plenty of distinct adapters.

Below is an example on how you may test and request this image:

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

$ cd api-inference-community/docker_images/diffusers

$ docker construct -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_XET_HIGH_PERFORMANCE=1
EOF

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

Then in one other terminal perform requests to the bottom model and/or miscellaneous LoRA adapters to be found on the HF Hub.

# Request the bottom 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 one other one
$ curl -H 'lora: nerijs/pixel-art-xl' 0:8888 -d '{"inputs": "elephant", "parameters": {"num_inference_steps": 20}}' > /tmp/adapter2.jpg

What about batching ?

Recently a extremely interesting paper got here out, that described the best way to increase the throughput by performing batched inference on LoRA models. In brief, all inference requests can be gathered in a batch, the computation related to the common base model can be done all of sudden, then the remaining adapter-specific products can be computed. We didn’t implement such a method (near the approach adopted in text-generation-inference for LLMs). As a substitute, we stuck to single sequential inference requests. The rationale is that we observed that batching was not interesting for diffusers: throughput doesn’t increase significantly with batch size. On the straightforward image generation benchmark we performed, it only increased 25% for a batch size of 8, in exchange for six times increased latency! Comparatively, batching is way more interesting for LLMs since you get 8 times the sequential throughput with only a ten% latency increase. That is the rationale why we didn’t implement batching for diffusers.

Conclusion: Time!

Using dynamic LoRA loading, we were able to avoid wasting compute resources and improve the user experience within the Hub Inference API. Despite the time beyond regulation added by the means of unloading the previously loaded adapter and loading the one we’re interested by, the undeniable fact that the serving process is most frequently already up and running makes the inference time response on the entire much shorter.

Note that for a LoRA to profit from this inference optimization on the Hub, it must each be public, non-gated and based on a non-gated public model. Please do tell us if you happen to apply the identical method to your deployment!