feat(lora): add LoRA+ support via lora_plus_lr_ratio

[ramkrishs]

Summary

Implements LoRA+ (Hayou et al., ICML 2024) by adding a --lora_plus_lr_ratio parameter to the LoRA finetune script.

Closes #963

---

What is LoRA+?

LoRA+ proposes applying a higher learning rate to lora_B parameters relative to lora_A. The theoretical motivation: lora_B lies in the output space and benefits from faster adaptation, while lora_A (input projection) should update more conservatively. The paper shows this asymmetry yields:

• 1–2% accuracy gains on downstream tasks
• Up to 2× faster convergence
• No additional compute cost

The recommended ratio is 16.0 for standard fine-tuning.

---

Usage

litgpt finetune lora \
    --checkpoint_dir checkpoints/meta-llama/Llama-3.1-8B \
    --lora_plus_lr_ratio 16.0

When lora_plus_lr_ratio is not set (default None), behaviour is identical to standard LoRA.

---

Implementation

litgpt/utils.py — new create_lora_plus_optimizer():

• Splits trainable parameters into two groups: lora_B at lr × ratio, everything else at base lr
• Works with any optimizer string or config dict (same interface as instantiate_torch_optimizer)

litgpt/finetune/lora.py:

• setup() and main() accept lora_plus_lr_ratio: float | None = None
• When set, calls create_lora_plus_optimizer() instead of instantiate_torch_optimizer()

tests/test_lora_plus.py — 11 unit tests:

• Two param groups created with correct LR ratio
• lora_B in high-LR group, lora_A in base-LR group
• All trainable params covered, no duplicates
• lora_B updates faster than lora_A per step (verifies the asymmetry works)
• Edge cases: no lora_B params, no trainable params, ratio=1.0

---

Reference

Hayou et al., LoRA+: Efficient Low Rank Adaptation of Large Models, ICML 2024
https://arxiv.org/abs/2402.12354

This PR was developed with AI assistance. All code has been reviewed, tested, and verified to follow LitGPT conventions.