Low-rank adaption (LoRA) is a technique to approximate the update to the linear layers in a LLM with a low-rank matrix factorization. This significantly reduces the number of trainable parameters and speeds up training with little impact on the final performance of the model. We demonstrate this method by instruction-finetuning Lit-GPT StableLM 3B on the Alpaca dataset on a single RTX 3090 (24GB) GPU.
The steps here only need to be done once:
- Follow the instructions in the README to install the dependencies.
- Download and convert the weights and save them in the
./checkpoints
folder. Weights can be downloaded following these instructions:
-
Download the data and generate the instruction tuning dataset:
python scripts/prepare_alpaca.py
(See this blog article for how to prepare and use custom datasets.)
python finetune/lora.py
The finetuning requires at least one GPU with ~24 GB memory (RTX 3090).
This script will save checkpoints periodically to the folder out/
.
Note: LoRA can be applied to not only
query
,key
orvalue
matrices, but also toprojection
,mlp
and classificationhead
. According to QLoRA paper (section 4): "LoRA on all linear transformer block layers are required to match full finetuning performance". By default LoRA is applied only to thequery
andvalue
matrices. In order to apply LoRA to other weight matrices - change the variables infinetune/lora.py
accordingly.
You can test the finetuned model with your own instructions by running:
python generate/lora.py --prompt "Recommend a movie to watch on the weekend."
Output:
I would recommend the movie The Martian (2015). It is a sci-fi movie starring Matt Damon that follows the story of...
If your GPU supports bfloat16
, you can additionally pass --precision bf16-true
to bring the memory consumption down to ~11 GB for StableLM-3B.
With only a few modifications, you can prepare and train on your own instruction dataset.
-
Create a json file in which each row holds one instruction-response pair. A row has an entry for 'instruction', 'input', and 'output', where 'input' is optional an can be the empty string if the instruction doesn't require a context. Below is an example json file:
[ { "instruction": "Arrange the given numbers in ascending order.", "input": "2, 4, 0, 8, 3", "output": "0, 2, 3, 4, 8" }, ... ]
-
Make a copy of
scripts/prepare_alpaca.py
and name it what you want:cp scripts/prepare_alpaca.py scripts/prepare_mydata.py
-
Modify
scripts/prepare_mydata.py
to read the json data file. -
Run the script to generate the preprocessed, tokenized train-val split:
python scripts/prepare_mydata.py --destination_path data/mydata/
-
Run
finetune/lora.py
by passing in the location of your data (and optionally other parameters):python finetune/lora.py --data_dir data/mydata/ --out_dir out/myexperiment
If you run into a CUDA error "Expected is_sm80 to be true, but got false", uncomment the line
torch.backends.cuda.enable_flash_sdp(False)
in the script below (see Lightning-AI/lit-llama#101).