Add Riemannian-preconditioned LoRA optimizer

[smellslikeml]

Adds create_riemannian_optimizer for the r×r Riemannian preconditioning recipe from Riemannian Preconditioned LoRA for Fine-Tuning Foundation Models (arXiv:2402.02347v3). Before every base-optimizer step, the LoRA factor gradients are rescaled as:

g_A ← (BᵀB + reg·I_r)⁻¹ @ g_A
g_B ← g_B @ (AAᵀ + reg·I_r)⁻¹

Non-LoRA parameters are updated unchanged. The r×r preconditioner keeps storage and runtime overhead small in the LoRA rank.

Coordination

Filed as coordination issue #3380 before pushing. Original PR author @fangzhaozhang explicitly endorsed a revival (comment):

"Thanks for interests! I'm caught up by my school work after the initial implementation, so feel free to open a fresh PR if you'd like."

What this PR does

Revives the design accepted in principle in #1807, folding in round-1 review feedback preemptively:

• PEP 8 class naming, docstrings replace implementation-detail strings
• No assert in code — proper TypeError / ValueError / RuntimeError
• issubclass check on optimizer_cls — the wrapper works with any torch.optim.Optimizer subclass (AdamW, SGD, …)
• torch.linalg.pinv with damping for numerical stability on near-rank-deficient factors
• bf16-stable via float32 preconditioner computation, cast back to grad dtype
• DoRA-compatible — name-based lora_A / lora_B pair discovery skips the magnitude vector

Deviates from #1807 in one design choice: rather than copying and modifying AdamW.step(), this is a subclass-wrapper approach — a dynamic subclass of optimizer_cls that applies the preconditioner to .grad tensors before calling super().step(closure). Thin diff on top of the base optimizer's step, and works with any optimizer.

reg calibration (07-07 update)

The reg damping default is 1e-2, matching @fangzhaozhang's reference example in #1807 (create_riemannian_optimizer(..., reg=1e-2)) and the paper author's public repo. On the MetaMath benchmark (llama-3.2-3B, rank 32, 5000 steps, otherwise-default settings):

reg	GSM8K test acc.
1e-6 (initial ship)	0.45
1e-2 (this PR)	0.475
default LoRA r=32 baseline	0.49

The paper's reported gains are modest and rank-dependent — Table 1: +0.7 BLEU on GPT-2 medium at r=4; Table 2: +0.7pp GLUE avg on 4-bit Mistral 7B at r=16. MetaMath r=32 isn't a paper-benchmarked setting; the paper's stronger claim is LR-range robustness (Figures 1, 3) rather than fixed-LR gains at fixed rank.

Scope explicitly deferred from #1807's tree

• src/peft/tuners/lycoris_utils.py / poly/router.py changes — unclear necessity, dropped. Happy to add back if you confirm they're needed.
• examples/riemannian_lora/ — happy to add here or as a follow-up.
• The accidentally-committed events.out.tfevents.* file — dropped.

Files

• src/peft/optimizers/riemannian.py — factory + preconditioner
• src/peft/optimizers/__init__.py — one-line export
• tests/test_riemannian_lora.py — 9 tests
• method_comparison/MetaMathQA/utils.py — dispatch entry + TrainConfig validation for optimizer_type="riemannian"
• method_comparison/MetaMathQA/experiments/lora/llama-3.2-3B-rank32-riemannian/ — benchmark experiment (adapter + training params)

Test plan

• pytest tests/test_riemannian_lora.py — 9 passed (factory happy-path, subclass permissiveness on AdamW + SGD, DoRA compatibility, ValueError on no-LoRA model, TypeError on non-Optimizer class, exact paper-formula verification, bf16 gradient finiteness)
• make style — clean
• MetaMath benchmark on Colab A100 with the corrected reg=1e-2: 5000 steps completed, GSM8K test accuracy 0.475 (see reg calibration table above)

AI assistance

AI-assisted patch. Every changed line reviewed; local pytest and make style runs exercised the code paths. Following CLAUDE.md's AI-assisted contribution disclosure guidelines.

Credit for the original design + engagement with round-1 review goes to @fangzhaozhang (#1807), who is co-authored on the fix commit at their round-1 co-authorship suggestion.

Co-Authored-By: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com>

[BenjaminBossan]

Thanks @smellslikeml for reviving the Riemannian-preconditioned optimizer PR. I did a first review and I agree that this more generic approach is better. I have a couple of comments but generally, this PR is already quite mature.

One thing to add to the PR would be to extend the MetaMath benchmark to allow using this optimizer there. Check these lines:

peft/method_comparison/MetaMathQA/utils.py

Lines 280 to 286 in e4fe61b

	if optimizer_type == "lora+":
	optimizer = create_loraplus_optimizer(model, optimizer_cls=torch.optim.AdamW, **optimizer_kwargs)
	elif optimizer_type == "lora-fa":
	optimizer = create_lorafa_optimizer(model, **optimizer_kwargs)
	else:
	cls = getattr(torch.optim, optimizer_type)
	optimizer = cls(model.parameters(), **optimizer_kwargs)

After adding the Riemannian optimizer there, you would also have to add an experiment that uses it.

@fangzhaoz If you have the opportunity, it would be great if you could review the PR too. If @smellslikeml agrees, I would also suggest to add you as a co-author on the final commit.

[BenjaminBossan]

DoRA would only be one example. Let's reword this to make it clear it's not about DoRA specifically.

[BenjaminBossan]

Checking for ".lora_A."and ".lora_B." would be more precise, right?

[BenjaminBossan]

This type of comment can be removed. Instead, if the tests are not self-explanatory, add a comment at the start of the test.

[BenjaminBossan]

Let's put all tests inside a single test class that bundles them.

[BenjaminBossan]

Let's use a reasonable tolerance for allclose here and below.

[BenjaminBossan]

IMO, this test is not very useful, as it more or less mirrors the code used above to implement the optimizer. Maybe there are some edge cases that can be sanity checked without reimplementing the computation in the test, but as is, I think we can remove the test.

[BenjaminBossan]

Would it make sense to parametrize the test over bf16 + fp16?

[smellslikeml]

Thanks @smellslikeml for reviving the Riemannian-preconditioned optimizer PR. I did a first review and I agree that this more generic approach is better. I have a couple of comments but generally, this PR is already quite mature.

One thing to add to the PR would be to extend the MetaMath benchmark to allow using this optimizer there. Check these lines:

peft/method_comparison/MetaMathQA/utils.py

Lines 280 to 286 in e4fe61b

	if optimizer_type == "lora+":
	optimizer = create_loraplus_optimizer(model, optimizer_cls=torch.optim.AdamW, **optimizer_kwargs)
	elif optimizer_type == "lora-fa":
	optimizer = create_lorafa_optimizer(model, **optimizer_kwargs)
	else:
	cls = getattr(torch.optim, optimizer_type)
	optimizer = cls(model.parameters(), **optimizer_kwargs)

After adding the Riemannian optimizer there, you would also have to add an experiment that uses it.

@fangzhaoz If you have the opportunity, it would be great if you could review the PR too. If @smellslikeml agrees, I would also suggest to add you as a co-author on the final commit.

Agreed, thank you for your suggestions!
I've updated with changes and attribution

[BenjaminBossan]

Thanks, this looks pretty good from my side, just a few comments.

@fangzhaoz Please let me know if you have time to review the PR or not.

[smellslikeml]

@BenjaminBossan Just rebased onto main to clear the commit lag. Ready for re-review when you have a moment.

[BenjaminBossan]

Thanks for the updates, the PR generally looks good and is ready to be merged. Before doing that, I just have one question left.

[BenjaminBossan]

I'm just wondering: For something like LBFGS, which uses the closure, would the Riemanniannian preconditioner still make sense?

[smellslikeml]

For something like LBFGS, which uses the closure, would the Riemanniannian preconditioner still make sense?

No, LBFGS builds curvature from y = flat_grad - prev_flat_grad and gates on the secant condition; if P_k is recomputed each step, y stops being a gradient difference of any fixed function.

While looking at this I noticed the ordering had a separate issue: preconditioner.step() then super().step(closure) silently drops preconditioning for any closure-using optimizer, since torch optimizers evaluate the closure at the top of their own step (typical body is zero_grad(); loss.backward(), which overwrites .grad). Affects AdamW/SGD-with-closure too.

Pushed both fixes in 58648c6:

• Wrap the closure so preconditioning runs inside it, matching LoraFAOptimizer.step's convention.
• Guard closure-required optimizers at construction so we don't hardcode LBFGS:

  closure_param = inspect.signature(optimizer_cls.step).parameters.get("closure")
  if closure_param is not None and closure_param.default is inspect.Parameter.empty:
      raise ValueError(...)

  Matches the sibling TypeError on non-Optimizer classes and the AdaLoRA-style raise on unsupported combos.

Regression tests added to tests/test_riemannian_lora.py: closure-vs-no-closure identity check (verified to fail on the previous ordering) and an LBFGS construction-time rejection test.

[BenjaminBossan]

Thanks for investigating the closure case.

While looking at this I noticed the ordering had a separate issue: preconditioner.step() then super().step(closure) silently drops preconditioning for any closure-using optimizer, since torch optimizers evaluate the closure at the top of their own step (typical body is zero_grad(); loss.backward(), which overwrites .grad). Affects AdamW/SGD-with-closure too.

I did not understand this point. Could you please give a concrete example (not a full reproducer, just a snippet would be enough)? And are there real world cases of "AdamW/SGD-with-closure"? If this is a very hypothetical scenario, I'd also be fine with simple checking if a closure was passed and raising.

[smellslikeml]

def closure():
    optimizer.zero_grad()
    loss = model(x).sum()
    loss.backward()
    return loss

optimizer.step(closure)

Old step(closure) did:

1. preconditioner.step() writes P_k g_k into .grad
2. super().step(closure) calls e.g. AdamW.step(closure), which invokes closure() at the top, and zero_grad(); backward() overwrites .grad with raw g_k
3. AdamW reads raw g_k, preconditioning silently dropped

Fair, this is hypothetical for PEFT. LBFGS is the only shipping torch optimizer that requires a closure. Simplified per your suggestion in f90cc44: replaced the wrapper with a step-time raise on any closure, dropped the wrapper's regression test, kept the construction-time signature guard for LBFGS (fail-fast on construction) alongside the step-time check.

[BenjaminBossan]

Thanks for the last push. Since we don't know what happens within the closure, I agree it's better to fully reject it.