Serving models at scale with LoRA

We apply this method to layers throughout the network. In a transformer architecture we can adapt the ‘Q’, ‘K’ and ‘V’ projection layers in each attention head, the linear layer merging attention heads together as well as the post-attention MLP layer. Going beyond transformers, with small modifications, LoRA can also be applied to convolutional layers.

Even though the number of trainable parameters can be as low as ~1% of the trainable parameters of the full model, in many problems LoRA fine-tuning reaches almost the same accuracy as fine-tuning the full model. Intuitively, this is because we can modify all parameters of the model throughout all its layers. The training molds the whole model into the right shape for the problem. We can contrast LoRA with a traditional “low parameter fine-tuning method”, where we fine-tune only the last few layers of the model, while keeping all other layers frozen; the number of trainable parameters might be comparable to LoRA, but their effect is much more limited, since most of the model remains rigidly frozen and only a small part is being adapted to the problem.

The collection of all low rank updates to a model is called an adapter and storing an adapter takes only ~1% of the memory required by the full model. While pre-LoRA we could fit only 4 full models on one GPU, we can now fit one base model together with several hundred adapters. The same GPU that previously could support inference for four document types, can now support inference for several hundred document types. In short, we have eliminated the GPU bottleneck from our system.

LoRA in practice

The practical implementation of LoRA is made easy using the PeFT library, short for Parameter-efficient Fine-Tuning) together with the Transformers library. . Adapting a model for LoRA training can be as simple as

import peft, timm

model = timm.create_model("convnext_base", nu)
config = peft.LoraConfig(
r=4,
target_modules=r".*\.conv_dw|.*\.mlp\.fc\d",
modules_to_save = ["head.fc"],
)
peft_model = peft.get_peft_model(model, config)

We use a regex to select, which layers should be targeted by LoRA training. In this case we target all depthwise convolutional and all fully connected layers. We also perform a full fine-tuning on the classification head.

>>> peft_model.print_trainable_parameters() trainable params: 4,341,770 || all params: 91,918,484 || trainable%: 4.7235

These choices result in a model that has only 4.7% of the trainable parameters of the full model.

After implementing LoRA training and training a first model, we may find that the accuracy of the LoRA model does not quite reach the level of a fully fine-tuned model. In other words, it might be time for some hyper-parameter tuning.

• Two key factors that influence both accuracy and the size of the adapter are the rank r, and the choice of which layers to adapt. A good rule of thumb is to start by adapting all linear layers and then tune r to get the desired accuracy/adapter size trade-off.
• The original LoRA paper includes a scaling parameter alpha, which acts on the update matrices viah = (W + \alpha / r AB)x. Empirically, alpha=2r, appears to be a good choice, but Sebastian Raschka has found that sometime tuning alpha can improve performance.
• We have observed that training a LoRA adapter requires about twice the number of epochs compared to regular fine-tuning to achieve the same accuracy. More recently, authors XYZ (arXiv:2402.12354) have applied theoretical insights on large width networks to show that by using different learning rates for the ‘‘A’ and ‘B’ matrices, we can reduce the training time to values comparable to regular fine-tuning.

If the need for even smaller adapters arises, then there are methods such as VeRA or LoRA-XS that reduce the trainable parameter count even further.