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.

[image]

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

[image]

Placement study (rank 8)

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.

[image]

Rank study (q + v)

[image]

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

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.