[RFC] Remove automatic weight merging when training LoRA

[felipemello1]

Context:
Currently merging ckpt model + lora weights is the default in our recipes. We say that in our docs and assume it for generation. Our core users are used to it.

Problem:

1. IMO, this is a bad default. It is hidden from the user and can cause issues.

2. These issues are:

• 2.1: The user doesn't know that the weights are merged and deploy the merged trained model AND the adapter. This will silently have bad results;

• 2.2. The user trains a model that combines lora + full finetune (e.g. embeddings, llama vision), saves the merged weights, and now they can NOT resume from training, because, for example, there isnt a ckpt that has untrained transformers and trained embeddings;

• 2.3: The fact that we save both the adapter + model gives the false impression that they should be combined;

Solution 1:
Remove the weight merging default from recipes. If users want to have it, they can explicitly set a flag or call an utility.

Pros:

• No risk of deploying merged + lora
• No hidden behaviors
• Resuming from ckpt is straight forward

Cons:

• It changes our default behavior
• generation and docs have to be updated
• the user needs to run an extra step if they want it merged (unless we add a flag)

Solution 2:
Check if any layer, other than lora, requires grad. If they do, set merge = False.

Pros:

• Doesn't change default behavior for regular lora
• Can resume from ckpt if we add extra logic/safeguards

Cons:

• introduce another hidden behavior: sometimes its merged, sometimes it isnt.
• The user cant know or control if the model weights are merged or not merged
• bloated if/else code

[ebsmothers]

Thanks @felipemello1 for the RFC! This is a pretty important thing to get right so I'm glad we're having the discussion here. Imo the issue 2.2 is the most important reason why we should move to solution 1 in the long term. Because I don't see a clear way that solution 2 can solve the problem of a model that's part full finetune, part LoRA finetune cleanly.

That being said, I do not think we should land solution 1 without a full deprecation cycle. As you rightly point out, the fact that we merge the weights back into the original model is not the most obvious. But my biggest concern is that solution 1 is in some sense even worse than breaking BC of an API in a traditional sense. With a BC-breaking change of an API, I as a user will usually get an error message and then I know I need to go fix something. In this case we will output the same set of state dict keys, but the weights will change. So we will have a silent change in behavior, and without proper advance notice this can screw up a lot of people's finetuning -> inference workflows. Personally, I would recommend the following:

Release cycle 1: adopt solution 2 to unblock issue 2.2 (so that Llama vision resume from checkpoint works properly). Add hi-pri logging message that (a) Llama vision does not merge the weights and (b) we will be moving to unmerge by default (for all models) by release N
Release cycle 2: Add a flag to the config merge_lora_weights. Set it to True for all non-Llama-vision configs and False for Llama vision configs. Also do another round of comms with a script to unmerge weights for any merged checkpoints in case people need it
Release cycle 3: Switch the flag to False everywhere

[felipemello1]

thanks for answering! I was thinking we could start with cycle 2 (unless you think that this is too big of a lift for this release):

a) add flag with default = true
b) in recipe, check if any module requires grad except lora
c) if b is true, raise error

that way, it doesnt break backward compatibility, but it doesnt fail silently either.

[ebsmothers]

@felipemello1 I am fine on starting with cycle 2, but would maybe still keep it for two full releases. But yeah the main question is level of effort before this release. If the release cycle 1 version takes 1 hour (which is about what I expect), but the release cycle 2 version takes e.g. >= 2 days (E2E time), then I think it's not worth it to delay the release for that. So up to you based on how much effort you feel it is.