This repo contains the official implementation of the PoLoRA optimizer from the paper PoLoRA: A Preconditioned Orthogonalized LoRA Optimizer — a drop-in, PEFT-compatible optimizer for LoRA fine-tuning.
pip install -e . # the optimizer (torch only)
pip install -r requirements.txt # + transformers/peft/datasets for the training pipelinefrom peft import LoraConfig, get_peft_model
from polora import Polora
model = get_peft_model(model, LoraConfig(r=256, lora_alpha=256, target_modules="all-linear"))
model.to("cuda") # move to the target device BEFORE constructing the optimizer
opt = Polora(model, lr=1e-2)
# opt = Polora(model, lr=1e-2, compile=True) # torch.compile the optimizer kernels
# opt = Polora(pairs=[(A, B), ...], lr=1e-2) # non-PEFT use: pass factor pairs directly
for batch in loader:
model(**batch).loss.backward()
opt.step()
opt.zero_grad()The configs used in the experiments for the paper are in configs/:
python prepare_data.py --dataset openmath --model_name meta-llama/Llama-3.2-1B \
--max_train_samples 2_000_000 --max_eval_samples 1024 --tokenize_num_proc 8
python train.py --config configs/llama32_openmath_r256.json- Any config field can be overridden on the CLI, booleans included
(
--no-compile). --optimizer adamwruns the Adam baseline.- The
data/cache defaults to the repo folder; pass an absolute--outor--data_dirto keep it elsewhere.
@misc{ghosh2026polora,
title = {PoLoRA: A Preconditioned Orthogonalized LoRA Optimizer},
author = {Ghosh, Nikhil and Parshakova, Tetiana and Gower, Robert M.},
year = {2026},
eprint = {2607.17620},
archivePrefix = {arXiv},
primaryClass = {cs.LG},
url = {https://arxiv.org/abs/2607.17620}
}