Add LoRA+ support

[rockerBOO]

LoRA+ adds a hyperparameter to the optimizer for B or up in the LoRA to increase the ratio of that LR.

The hyperparameter should be adjusted depending on your dataset and how difficult it is to train. They suggest starting at 2^4 or 16 ratio to start.

The Text Encoder and UNet are currently separate ratios as with recommended starting values as follows.

Specify loraplus_lr_ratio with --network_args. Example: --network_args "loraplus_lr_ratio=16"
loraplus_unet_lr_ratio and loraplus_lr_ratio can be specified separately for U-Net and Text Encoder.

Example: --network_args "loraplus_unet_lr_ratio=16" "loraplus_text_encoder_lr_ratio=4" or --network_args "loraplus_lr_ratio=16" "loraplus_text_encoder_lr_ratio=4" etc.

I am willing to change how the code is implemented, name of arguments and so on as necessary.

https://arxiv.org/abs/2402.12354
https://github.com/nikhil-ghosh-berkeley/loraplus

[feffy380]

LoRA+ adds a hyperparameter to the optimizer for B or up in the LoRA to increase the ratio of that LR.

AFAIK B is lora_down
edit: nevermind I was mistaken. Damn notation

[feffy380]

LoRA+ doesn't make sense for LoRA-FA where one of the matrices is already frozen

[kohya-ss]

Thank you for this! This seems to be simple and nice. I'd like to merge this sooner, as soon as I have time.

[rockerBOO]

LoRA+ doesn't make sense for LoRA-FA where one of the matrices is already frozen

My inclination was to replicate the behavior for all the networks to make it consistent. Since it leaves B to be trainable I think it makes sense to still be able to utilize this same behavior for it. I think it will be less valuable for LoRA-FA but I don't think it's entirely wrong to support.

Do you think this assessment is ok? Ideally all the networks use the same prepare_optimizer_params to reduce code duplication in the future so I think this is ok?

[rockerBOO]

Bugs:

• Default LR not applied
• Logging LR shows incorrect labels (group 1 and group 2 will get labeled like lr/text_encoder and lr/unet)

[rockerBOO]

Made some fixes, including Default LR and invalid LR now properly working. Should now properly support single param group, only TE and only UNet and all LoRA+ and LR permutations.

The following test script I made to make sure it is working as expected. Willing to fix anything you find incorrect.

https://gist.github.com/rockerBOO/7e58b268d266c6f55a7f90085a1d1dfe

edit Additionally worked through all the logging for LoRA+. Added lr/all when it is set for a single param group (this happens with LyCORIS modules). Moved from possibly "lr" because of how grouping works so it is grouped together under "lr". Willing to adjust this if not desired to preserve compatibility historically.

all

LoRA+ with TE and Unet

LoRA+ with block weights

[kohya-ss]

Thank you for update! This looks promising. Please wait a little longer while I check the updated part.

[kohya-ss]

Thank you for this, and sorry for the dealy. I will add some backward compatibility to some scripts, and merge to dev!

[kohya-ss]

I notice while updating, I'd like to move loraplus_* args from the command line args to network_args for compatibility with LyCORIS. I appreciate your understanding.