LM8-L3 — LoRA Fine-Tuning Experiment

Objective

Implement LoRALinear, add_lora, merge_lora and a training-memory estimate so that all five self-checks pass, and explain from the results what LoRA saves in trainable parameters and memory, and why the rank could not be seen to matter on this easy task.

Why a notebook? This exercise trains or runs a real PyTorch model. PyTorch does not run in the browser Python used by the site’s other labs, so you run this one on your own computer or in a hosted notebook service. The interactive browser labs of this module cover the surrounding ideas without PyTorch.

Setup and requirements

Hardware: CPU only, about 2 GB of RAM, a few minutes. CPU only; about 2 GB of RAM; the notebook takes about two to three minutes (building the base model, then three short adaptation runs). The model has about 113,000 parameters and works on characters. Nothing is downloaded.

Download the notebook (.ipynb) Reference solution (try the lab first)

On your own machine (Python 3.10 or newer):

python3 -m venv llm-env
source llm-env/bin/activate        # on Windows: llm-env\Scripts\activate
pip install torch numpy jupyterlab
jupyter lab lora-fine-tuning-experiment.ipynb

In a hosted notebook service such as Google Colab or Kaggle, upload the file and run it there; PyTorch is usually preinstalled. Matplotlib is used for optional plots when it is available; the notebook runs without it.

What you will do

1. LoRA layer. Write LoRALinear: frozen base layer plus (α/r)·x·Aᵀ·Bᵀ with B initialised to zero; check that it starts identical to the base layer.
2. Wrap a model. Freeze all parameters and wrap the attention and feed-forward layers; check the parameter count against r·(in+out).
3. Merge. Fold B·A into the base weights and verify that the function is unchanged.
4. Memory. Write the estimate 4P + 12·P_train and check it against GPT-2 small.
5. Experiment. Adapt to a new instruction with full fine-tuning, LoRA r=8 and LoRA r=2; read the table of parameters, memory, accuracy and forgetting.

Cells that you must complete are marked TODO in the notebook and have a self-check right after them (an assert with a message). Run the cells from the top; a failing check tells you what is wrong.

Expected behaviour

• Self-check 1: LoRALinear equals the base layer at initialisation, adds r·(in+out) parameters, and at the first step B receives a gradient while A receives exactly zero.
• Self-check 2: rank 8 trains 2·8·(256+128+320+320) = 16,384 parameters (14.5% of the base model) and the wrapped model reproduces the base model before training.
• Self-check 3: after merging, no wrapper remains, the parameter count is the original, and the outputs are unchanged.
• Self-check 4: the memory estimate gives about 1,991 MB for full fine-tuning of GPT-2 small.
• Self-check 5: the base model scores about 0 on the new instruction; full fine-tuning and LoRA with rank 8 learn it (at least 0.5 exact match on 30 held-out words) with LoRA training 14% of the parameters and needing less estimated memory. In the run used to prepare this lab all three variants, including rank 2 (3.6% of the parameters), reached about 1.0 on the new instruction and kept 0.9 to 1.0 on the old ones: this instruction is an easy change, so the rank did not matter. Estimated training memory was 1.82 MB for full fine-tuning and 0.72 MB for rank 8. Numbers vary with PyTorch version and seed.

How to know you are done

• All five self-checks print “passed” when the notebook is run from the top.
• You can derive the number of trainable parameters of a LoRA layer and of the whole model.
• You can say what the rank limits and why merging removes the inference cost.

The site cannot run or grade the notebook for you. The self-checks inside the notebook are your verification. The reference solution was executed with PyTorch on a CPU when this module was prepared. The model is tiny and the new instruction is a small change to a known one, a case chosen so that a low-rank update works (all ranks tried succeed, so the experiment cannot show where rank becomes limiting); results on large models and real tasks differ, and memory figures are estimates computed from a formula, not measurements.

If something goes wrong

• The wrapped model differs from the base at initialisation: B must be initialised to zeros, not random.
• Base weights change during training: call add_lora before creating the optimizer, and pass only parameters that require gradients.
• LoRA does not learn: use a larger learning rate than for full fine-tuning (3e-3 here) and enough steps (900).
• Merge changes the output: the update is scale * B @ A with scale = alpha / r; check shapes (out × r times r × in).

Theory for this lab

• LM8-T2: Low-Rank Adaptation (LoRA)
• LM8-T5: Choosing a Base Model and Managing Checkpoints

Reference material

• Build a Large Language Model (From Scratch), Appendix E: Parameter-efficient Finetuning with LoRA

Efficient Fine-Tuning and Model Adaptation labs

1. LM8-L1: Estimate Trainable Parameters and Memory
2. LM8-L2: LoRA Low-Rank Update Numerics
3. LM8-L3: LoRA Fine-Tuning Experiment: this lab.

Next lab: LM9-L1: Compare Two Checkpoints.