[cosmicoptima]

This PR adds support for LoRA. It adds a model.lora config section with the fields enabled rank alpha dropout target_modules (modules adapted by the LoRA) trainable_modules (modules kept fully trainable).

The LoRA is merged into the base model on save for now because vLLM doesn't support loading a LoRA that keeps some modules fully trained. Support for it is being added though: vllm-project/vllm#11714.

---

GitHub Issue: #796

[cosmicoptima]

oops one sec im fixing the style check (edit: it's good now)

[cosmicoptima]

i'm able to fit 14b parameters into 2 trainer h100s with lora enabled, and 32b parameters into 3 trainer h100s with lora + liger kernel

[kalomaze]

some notes:

• if we want to standardize naming conventions with PEFT then trainable_modules should be modules_to_save
• this implementation is less invasive in terms of changing stuff at the expense of not being any faster at saving to disk (entire model is saved, not just the lora adapter), this is perfectly fine tho bc you cannot keep some layers like the output head or embeddings fully trainable otherwise because of how vllm's lora implementation works
• this implementation also implicitly adds support for freezing arbitrary layers or module types via regex, i.e someone could specify no target modules and leave only attention projs in trainable_modules, which would keep the MLPs frozen. this is slightly less direct than having an explicit configuration for freezing, but IMO is fine

[cosmicoptima]

if we want to standardize naming conventions with PEFT then trainable_modules should be modules_to_save

good idea

[mikasenghaas]

integration looks nice, i like it. will need some experiments and integration tests to confirm correctness

[cosmicoptima]

the gpu tests pass on my runpod and it looks like there has been the same OOM for unit tests on some other recent checks outside of this PR

[mikasenghaas]

ah yea, looks like a dangling process on our ci machine, cleaned it up and re-ran your tests

[samsja]

btw @cosmicoptima I forget to thank you for the great PR. We were not ready at the time to merge it. We are working on porting it here #1059 and will make sure you get the correct attribution for your work :)

[cosmicoptima]

cool to know it's being worked on @samsja ! also found a little trick whose implementation i don't think i put in this PR, and that i don't think i mentioned here--which is initializing a new lora every step and merging it into the base model at the end of the step

thereby approximating a higher rank over time while saving the same amount of compute

[mikasenghaas]

tysm for your contribution @cosmicoptima, super clean pr. we updated with latest main, ran some experiments to validate and merged this as an experimental feature now!

[mikasenghaas]

closing this now:)