headless-start/peft-lora-llm

LoRA Fine-Tuning of SmolLM2-360M on AG News

Trained checkpoints for github.com/headless-start/peft-lora-llm, a hand-written LoRA implementation on a frozen decoder-only model (HuggingFaceTB/SmolLM2-360M with a sequence-classification head). LoRA matrices wrap the attention q_proj and v_proj layers (alpha = 2r, B initialised to zero) and only they and the classification head are trained.

The repository holds every checkpoint behind the results in the GitHub README: the linear-probe / LoRA / full fine-tuning comparison, the placement study and the rank study. Code, training scripts and figures live on GitHub.

Results

Top-1 accuracy on the full AG News test split (7,600 articles, 4 topics: World, Sports, Business, Sci/Tech).

Headline run (LoRA rank 8 on q and v, 5 epochs on a 20,000-article training subset): 93.7% with 823K trainable parameters out of 362.6M (0.23%). File: checkpoints/best.pt.

Baselines

Method	Accuracy	Trainable parameters	Checkpoint file	Size
Linear probe	90.1%	3.8K (0.001%)	checkpoints/best_head.pt	0.02 MB
LoRA r=8, q+v	93.7%	823K (0.23%)	checkpoints/best_lora.pt	3.3 MB
Full fine-tuning	93.6%	361.8M (100%)	checkpoints/best_full.pt	1.4 GB

Placement study (rank 8)

Placement	Accuracy	Trainable parameters	Checkpoint file
q	92.7%	495K	checkpoints/best_r8_q.pt
k	92.5%	332K	checkpoints/best_r8_k.pt
v	92.9%	332K	checkpoints/best_r8_v.pt
q + k	92.8%	823K	checkpoints/best_r8_qk.pt
q + v	93.7%	823K	checkpoints/best_r8_qv.pt
q + k + v	93.4%	1.15M	checkpoints/best_r8_qkv.pt

The k and v projections are smaller than q because the model uses grouped-query attention, so LoRA on q costs more parameters at the same rank.

Rank study (q + v)

Rank	Accuracy	Trainable parameters	Checkpoint file
4	93.1%	413K	checkpoints/best_r4_qv.pt
8	93.7%	823K	checkpoints/best_r8_qv.pt
16	93.1%	1.64M	checkpoints/best_r16_qv.pt
32	93.2%	3.28M	checkpoints/best_r32_qv.pt

Notes on the numbers

• best.pt, best_lora.pt and best_r8_qv.pt are the same weights; the same run appears in the headline and in all three tables.
• Every number is a single run with seed 42. Each checkpoint is the epoch with the highest accuracy on the test split, which is also the split reported here, so the figures are best-epoch results rather than estimates from a held-out validation set.
• All checkpoints were re-evaluated on the test split before upload and reproduce the stored accuracies.

Files

checkpoints/
  best.pt              headline run, LoRA r=8 on q+v
  best_head.pt         linear probe (classification head only)
  best_lora.pt         LoRA r=8 on q+v, as used in the comparison tables
  best_full.pt         full fine-tuning (all weights)
  best_r8_<placement>.pt   placement study
  best_r<rank>_qv.pt       rank study
results/               the JSON results and figures from the GitHub repository

The LoRA and linear-probe checkpoints store only the trained tensors (LoRA matrices and head); the frozen backbone comes from the public SmolLM2-360M weights. best_full.pt stores the whole network. Every file is a PyTorch dictionary with the keys model, epoch and val_acc.

Usage

Clone the code, download a checkpoint and run the prediction script:

git clone https://github.com/headless-start/peft-lora-llm.git
cd peft-lora-llm
pip install -r requirements.txt

hf download headless-start/peft-lora-llm checkpoints/best.pt --local-dir .
python predict.py "Stocks rallied after the central bank held rates steady." --ckpt checkpoints/best.pt

For another LoRA checkpoint pass its rank and placement, for example --ckpt checkpoints/best_r16_qv.pt --lora-r 16 or --ckpt checkpoints/best_r8_k.pt --placement k.

In Python:

import torch
from huggingface_hub import hf_hub_download
from predict import load_model
from src.data import build_tokenizer

path = hf_hub_download("headless-start/peft-lora-llm", "checkpoints/best.pt")
model = load_model(path, "HuggingFaceTB/SmolLM2-360M", r=8, alpha_factor=2,
                   device=torch.device("cpu"), placement="qv")
tokenizer = build_tokenizer("HuggingFaceTB/SmolLM2-360M")

Text is tokenised with the SmolLM2 tokenizer, truncated to 128 tokens, with the end-of-sequence token used for padding.

Training setup

Setting	Value
Backbone	HuggingFaceTB/SmolLM2-360M, frozen for LoRA and the linear probe
Data	AG News, 20,000 articles sampled from the training split (seed 42); full test split for evaluation
Epochs	5
Optimiser	AdamW, learning rate 3e-4 (3e-5 for full fine-tuning), weight decay 0.05
Schedule	2 warmup epochs, then cosine decay to 1e-7
Batch size	32 (8 for full fine-tuning)
Other	mixed precision, maximum length 128 tokens

Licence

Released under the MIT licence, as is the code. SmolLM2-360M is Apache-2.0; check the original AG News terms before using the dataset or these weights beyond research.