peft icon indicating copy to clipboard operation
peft copied to clipboard

Ineffective Fine-Tuning Bug: Using `get_peft_model()` Before Loading LoRA Produces Outputs Identical to the Base Model

Open Hoper-J opened this issue 1 year ago • 2 comments

System Info

transformers==4.41.2 peft==0.11.1

Who can help?

No response

Information

  • [ ] The official example scripts
  • [ ] My own modified scripts

Tasks

  • [ ] An officially supported task in the examples folder
  • [ ] My own task or dataset (give details below)

Reproduction

Solution: To prevent ineffective fine-tuning when using LoRA, do not use get_peft_model() before loading LoRA weights with PeftModel.from_pretrained(). Instead, simply load the LoRA weights using PeftModel.from_pretrained() on the original model.

# Correct way
model = PeftModel.from_pretrained(model, PATH)

# Incorrect way
model = get_peft_model(model, lora_config)
model = PeftModel.from_pretrained(model, PATH)

If you’re interested, here’s the background story on how I discovered the issue:

It took me three hours of eliminating all other potential issues before I found the root cause. The problem arose while I was writing instructional code for instructional purposes, converting an example from using load_state_dict to PEFT for learning purposes.

# Original project code (correct):
# Apply LoRA configuration to text_encoder and unet
text_encoder = get_peft_model(text_encoder, lora_config)
unet = get_peft_model(unet, lora_config)

# If set to resume training, load the weights from the last model checkpoint
# You can modify the model_path to specify another path
if resume:
    # Load only the weights from the last training,
    # not the whole model as in model = torch.load(...)
    text_encoder = torch.load(os.path.join(model_path, "text_encoder.pt"))
    unet = torch.load(os.path.join(model_path, "unet.pt"))

Converted to PEFT format:

# Apply LoRA configuration to text_encoder and unet
text_encoder = get_peft_model(text_encoder, lora_config)
unet = get_peft_model(unet, lora_config)

# If set to resume training, load the weights from the last model checkpoint
if resume:
    # Load the LoRA model using PEFT's from_pretrained method
    text_encoder = PeftModel.from_pretrained(text_encoder, os.path.join(model_path, "text_encoder"))
    unet = PeftModel.from_pretrained(unet, os.path.join(model_path, "unet"))

At first, everything seemed fine—the model didn’t throw any errors—but it behaved exactly the same as if LoRA hadn’t been applied at all. It was such a tricky bug, and I couldn’t help but feel frustrated by the time it took. But I'm glad I found it and can now share it with you all.

Let’s break it down with a simple example using a linear layer for clarity:

import torch
import torch.nn as nn
from torch.optim import Adam
from copy import deepcopy
from peft import get_peft_model, LoraConfig, PeftModel

# Set a fixed random seed for reproducibility
torch.manual_seed(42)

# Define a simple linear model
class LinearModel(nn.Module):
    def __init__(self, input_size, output_size):
        super(LinearModel, self).__init__()
        self.linear = nn.Linear(input_size, output_size)

    def forward(self, x):
        return self.linear(x)

# Instantiate the linear model
model = LinearModel(input_size=10, output_size=1)

# Deepcopy the original model before applying LoRA to ensure a fair comparison later
original_model = deepcopy(model)

# Configure LoRA parameters
config = LoraConfig(
    inference_mode=False,
    r=4,
    lora_alpha=16,
    target_modules=['linear'],
)

# Apply LoRA to the model
lora_model = get_peft_model(model, config)

# Define a simple loss function and optimizer
criterion = nn.MSELoss()
optimizer = Adam(lora_model.parameters(), lr=1e-3)

# Generate some simulated training data
input_data = torch.randn(100, 10)  # 100 samples, each with 10 features
target_data = torch.randn(100, 1)  # Corresponding target values

# Train for one epoch
lora_model.train()
for epoch in range(1):  # Train for 1 epoch
    optimizer.zero_grad()
    outputs = lora_model(input_data)
    loss = criterion(outputs, target_data)
    loss.backward()
    optimizer.step()

# Save LoRA weights after training
lora_model.save_pretrained('linear_lora_model')

# Method 1: Use get_peft_model before loading LoRA weights
model1 = PeftModel.from_pretrained(get_peft_model(deepcopy(original_model), config), 'linear_lora_model')

# Method 2: Directly load LoRA weights
model2 = PeftModel.from_pretrained(deepcopy(original_model), 'linear_lora_model')

# Generate the same input data to compare outputs
test_input = torch.randn(1, 10)

# Compare outputs from the four models (Original, LoRA, Method 1, Method 2)
def compare_model_outputs(input_data):
    # Original model
    original_output = original_model(input_data)
    print("Original model output:", original_output.detach().numpy())

    # LoRA model (trained before saving)
    lora_output = lora_model(input_data)
    print("Trained LoRA model output:", lora_output.detach().numpy())

    # Method 1: Use get_peft_model before loading LoRA
    output1 = model1(input_data)
    print("Method 1 (use get_peft_model before loading LoRA) output:", output1.detach().numpy())

    # Method 2: Directly load LoRA
    output2 = model2(input_data)
    print("Method 2 (directly load LoRA) output:", output2.detach().numpy())

    if torch.allclose(original_output, output1):
        print("The original model and Method 1 produce the same output.")
    if torch.allclose(lora_output, output2):
        print("The trained LoRA model and Method 2 produce the same output.")


# Compare the parameters between two models
def compare_params(m1, m2):
    for (n1, p1), (n2, p2) in zip(m1.named_parameters(), m2.named_parameters()):
        if n1 != n2 or not torch.allclose(p1, p2):
            print(f"Parameter mismatch: \n{n1}\n{n2}")
            return False
    return True

# Compare outputs from the four models
compare_model_outputs(test_input)

# Check if the parameters from Method 1 and Method 2 are the same
if compare_params(model1, model2):
    print("The parameters of Method 1 and Method 2's LoRA models are consistent!")
else:
    print("The parameters of Method 1 and Method 2's LoRA models are inconsistent!")

Output:

Original model output: [[-0.03600371]]
Trained LoRA model output: [[-0.03428639]]
Method 1 (use get_peft_model before loading LoRA) output: [[-0.03600371]]
Method 2 (directly load LoRA) output: [[-0.03428639]]
The original model and Method 1 produce the same output.
The trained LoRA model and Method 2 produce the same output.
Parameter mismatch: 
base_model.model.base_model.model.linear.base_layer.weight
base_model.model.linear.base_layer.weight
The parameters of Method 1 and Method 2's LoRA models are inconsistent!

From the output, you can see that Method 1 (which uses get_peft_model() before loading LoRA) produces the exact same output as the original model, meaning that LoRA was not applied effectively.

In contrast, Method 2 (which directly loads the LoRA weights using PeftModel.from_pretrained()) produces an output that matches the output of the trained LoRA model. This confirms that Method 2 successfully loads the LoRA fine-tuning and applies the pre-trained weights as intended.

Additionally, the parameter mismatch between Method 1 and Method 2 confirms that using get_peft_model() before loading LoRA interferes with the structure of the model, leading to ineffective fine-tuning. This is why the parameters in Method 1 and Method 2 are different, even though they should be consistent if LoRA were applied correctly.

I hope this solves your issue and saves you time! 😄

Expected behavior


Hoper-J avatar Sep 30 '24 11:09 Hoper-J

Thanks for this very detailed report. Indeed, one should not use get_peft_model and PeftModel.from_pretrained together. The first is intended for creating a fresh PEFT adapter with the intent of training it from scratch. The second is for loading an already trained PEFT adapter, with the intent of training it further or doing inference.

Since you spend a lot of time debugging this, I assume that you checked our docs and didn't find good information. Could you please suggest where you looked, so that we can improve the docs at that location?

Other than that, I decided to create a PR to raise a warning if there are missing keys, see #2118. LMK what you think.

BenjaminBossan avatar Sep 30 '24 15:09 BenjaminBossan

Thanks for this very detailed report. Indeed, one should not use get_peft_model and PeftModel.from_pretrained together. The first is intended for creating a fresh PEFT adapter with the intent of training it from scratch. The second is for loading an already trained PEFT adapter, with the intent of training it further or doing inference.

Since you spend a lot of time debugging this, I assume that you checked our docs and didn't find good information. Could you please suggest where you looked, so that we can improve the docs at that location?

Other than that, I decided to create a PR to raise a warning if there are missing keys, see #2118. LMK what you think.

In fact, your docs are quite good, and this issue arose because I performed an unanticipated action while adapting some pre-existing code: I directly replaced the previous load operation with the PEFT methods

I appreciate the PR to raise a warning for missing keys—that will certainly help others avoid a similar issue. I'll keep an eye on #2118 and provide feedback if needed. Thanks again for your quick response and for submitting the PR!

Hoper-J avatar Oct 01 '24 03:10 Hoper-J

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

github-actions[bot] avatar Oct 30 '24 15:10 github-actions[bot]

Should be resolved via #2118. If not, feel free to re-open.

BenjaminBossan avatar Oct 30 '24 15:10 BenjaminBossan