Use float32 for LoRA weights to avoid the risk of underflow and overflow.

[james77777778]

Description

As reported in keras-team/keras-hub#2629

We should use high precision (float32) for LoRA weights to stabilize the finetuning.

References:

• https://huggingface.co/docs/peft/developer_guides/troubleshooting (autocast_adapter_dtype)
• Adapters saved in float16 are loaded in float32 huggingface/peft#2421
• https://github.com/huggingface/peft/blob/c75485a2144a15a526e0290835a7439daefc3925/src/peft/tuners/lora/layer.py#L966-L968 (huggingface/peft impl)

Contributor Agreement

Please check all boxes below before submitting your PR for review:

• I am a human, and not a bot.
• I will be responsible for responding to review comments in a timely manner.
• I will work with the maintainers to push this PR forward until submission.

Note: Failing to adhere to this agreement may result in your future PRs no longer being reviewed.

[gemini-code-assist]

Code Review

This pull request updates the LoRA implementation across several layers—including Convolutional, Dense, EinsumDense, and Embedding—to ensure that LoRA weights are initialized as float32 to prevent numerical instability. It also introduces explicit casting to the appropriate variable or compute dtypes during kernel composition and forward passes. A critical issue was identified in the EinsumDense layer where a trailing comma incorrectly converts the LoRA update into a tuple, which will cause a TypeError during tensor operations.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.96%. Comparing base (e94cb07) to head (5822cef).

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #22559   +/-   ##
=======================================
  Coverage   82.95%   82.96%           
=======================================
  Files         596      596           
  Lines       69252    69259    +7     
  Branches    10814    10814           
=======================================
+ Hits        57451    57458    +7     
  Misses       8973     8973           
  Partials     2828     2828           

Flag	Coverage Δ
keras	82.78% <100.00%> (+<0.01%)	⬆️
keras-jax	58.72% <100.00%> (+<0.01%)	⬆️
keras-numpy	54.56% <69.23%> (-0.01%)	⬇️
keras-openvino	59.42% <69.23%> (-0.01%)	⬇️
keras-tensorflow	60.29% <100.00%> (+<0.01%)	⬆️
keras-torch	59.06% <100.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[james77777778]

PR rebased. The openvino test failure should be unrelated to this PR.

[amitsrivastava78]

Thanks for the PR, it's very well written, please check my minor comment about this.

[amitsrivastava78]

The PR itself is fine; just worth adding a note that users should merge LoRA weights before deploying for inference.

[james77777778]

Sure. The comments have been updated:

        # LoRA weights should be float32 to avoid the risk of underflow or
        # overflow during fine-tuning.
        # When deploying the model, these weights should be merged with the
        # original kernel while maintaining the original kernel's dtype.
        ...