QLoRA Training does not give expected results
Author: karthsu
Pull request: no

### System Info

transformers version: 4.30.0
Platform: Linux-5.15.0-69-generic-x86_64-with-glibc2.29
Python version: 3.8.10
Huggingface_hub version: 0.15.1
Safetensors version: 0.3.1
PyTorch version (GPU?): 1.13.1 (True)
Tensorflow version (GPU?): not installed (NA)
Flax version (CPU?/GPU?/TPU?): not installed (NA)
Jax version: not installed
JaxLib version: not installed
Using GPU in script?: Yes
Using distributed or parallel set-up in script?: No

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Tried fine-tuning the [InstructCodeT5+](https://huggingface.co/Salesforce/instructcodet5p-16b) model using QLoRA and the loss is stuck at a particular value. I Followed the steps given in this example [notebook](https://colab.research.google.com/drive/1VoYNfYDKcKRQRor98Zbf2-9VQTtGJ24k?usp=sharing#scrollTo=FuXIFTFapAMI) for QLoRA. Code for the experiment:
```
import pandas as pd
import os
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, BitsAndBytesConfig
import torch
from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training, TaskType, prepare_model_for_kbit_training
from transformers import DataCollatorForSeq2Seq
import evaluate
import nltk
import numpy as np
from nltk.tokenize import sent_tokenize
from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments
from datasets import Dataset, DatasetDict
import argparse
import pickle
import json

parser = argparse.ArgumentParser(description='Options')
parser.add_argument('--dataset_dir', default='data', type=str, help="folder in which the dataset is stored")
parser.add_argument('--output_dir', default="lora-instructcodet5p", type=str, help="output directory for the model")
parser.add_argument('--results_dir', default="results", type=str, help="where the results should be stored")
args = parser.parse_args()

nltk.download("punkt")
tokenized_dataset = DatasetDict.load_from_disk(args.dataset_dir)
# Metric
metric = evaluate.load("rouge")
pad_tok = 50256
token_id="Salesforce/instructcodet5p-16b"

tokenizer = AutoTokenizer.from_pretrained(token_id)
# helper function to postprocess text
def postprocess_text(preds, labels):
    preds = [pred.strip() for pred in preds]
    labels = [label.strip() for label in labels]

    # rougeLSum expects newline after each sentence
    preds = ["\n".join(sent_tokenize(pred)) for pred in preds]
    labels = ["\n".join(sent_tokenize(label)) for label in labels]

    return preds, labels

def compute_metrics(eval_preds):
    preds, labels = eval_preds
    if isinstance(preds, tuple):
        preds = preds[0]
    for idx in range(len(preds)):
        for idx2 in range(len(preds[idx])):
            if preds[idx][idx2]==-100:
                preds[idx][idx2] = 50256
    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
    # Replace -100 in the labels as we can't decode them.
    labels = np.where(labels != pad_tok, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    # Some simple post-processing
    decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

    result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
    result = {k: round(v * 100, 4) for k, v in result.items()}
    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
    result["gen_len"] = np.mean(prediction_lens)
    return result
def get_dict(predicts):
    d = {}
    for num in range(len(tokenized_dataset['test'])):
        pred = tokenizer.decode([n for n in predicts[0][num] if n!=50256 and n!=-100])[1:]
        d[num+1] = {'Question':tokenizer.decode([n for n in tokenized_dataset['test'][num]['input_ids'] if n!=50256]),
                    'Ground truth solution':tokenizer.decode([n for n in tokenized_dataset['test'][num]['labels'] if n!=50256]),
                   'Prediction': pred if pred else None}
    return d

def find_all_linear_names(model):
    cls = torch.nn.Linear
    lora_module_names = set()
    for name, module in model.named_modules():
        if isinstance(module, cls):
            names = name.split('.')
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])


    if 'lm_head' in lora_module_names:
        lora_module_names.remove('lm_head')
    return list(lora_module_names)


def main():
    device = 'cuda'

    # huggingface hub model id
    model_id="instructcodet5p-16b"
    if not os.path.exists(model_id):
        model_id=token_id
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )
    # load model from the hub
    model = AutoModelForSeq2SeqLM.from_pretrained(model_id,
                                              # torch_dtype=torch.bfloat16,
                                              low_cpu_mem_usage=True,
                                              trust_remote_code=True, decoder_start_token_id=1, pad_token_id=pad_tok, device_map="auto", quantization_config=bnb_config)
    modules = find_all_linear_names(model)
    # Define LoRA Config
    lora_config = LoraConfig(
     r=16,
     lora_alpha=32,
     target_modules=modules,
     lora_dropout=0.05,
     bias="none",
     task_type=TaskType.SEQ_2_SEQ_LM
    )
    model = prepare_model_for_kbit_training(model, False)

    # add LoRA adaptor
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # we want to ignore tokenizer pad token in the loss
    label_pad_token_id = pad_tok
    # Data collator
    data_collator = DataCollatorForSeq2Seq(
        tokenizer,
        model=model,
        label_pad_token_id=label_pad_token_id,
        pad_to_multiple_of=8
    )
    output_dir=args.output_dir

    training_args = Seq2SeqTrainingArguments(
        output_dir=output_dir,
        per_device_train_batch_size=1,
        # per_device_eval_batch_size=1,
        predict_with_generate=True,
        weight_decay=0.05,
        # warmup_steps=200,
        fp16=False, # Overflows with fp16
        learning_rate=1e-4,
        num_train_epochs=5,
        logging_dir=f"{output_dir}/logs",
        logging_strategy="epoch",
        report_to="tensorboard",
        push_to_hub=False,
        # generation_max_length=200,
        optim="paged_adamw_8bit",
        lr_scheduler_type = 'constant'
    )

    # Create Trainer instance
    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=tokenized_dataset["train"],
        # eval_dataset=tokenized_dataset["validation"],
        # compute_metrics=compute_metrics,
    )

    # train model
    train_result = trainer.train()

if __name__ == '__main__':
    main()
```


### Expected behavior


Output using QLoRA and the generations are empty during the evaluation stage:

```
{'loss': 6.9007, 'learning_rate': 0.0001, 'epoch': 1.0}                 
                                                                                               
{'loss': 6.9007, 'learning_rate': 0.0001, 'epoch': 2.0}                
                                                                                                
{'loss': 6.9007, 'learning_rate': 0.0001, 'epoch': 3.0}       
                                                                                                         
{'loss': 6.9007, 'learning_rate': 0.0001, 'epoch': 4.0}     
                                                                                                           
{'loss': 6.9007, 'learning_rate': 0.0001, 'epoch': 5.0}    

```

The same is working in a setting where I use LoRA instead where the loss is reducing and the generations are much better:

```
{'loss': 0.8144, 'learning_rate': 0.0001, 'epoch': 1.0}                           
                                                                                                    
{'loss': 0.0745, 'learning_rate': 0.0001, 'epoch': 2.0}                 
                                                                                                                 
{'loss': 0.0391, 'learning_rate': 0.0001, 'epoch': 3.0}                     
                                                                                                        
{'loss': 0.0189, 'learning_rate': 0.0001, 'epoch': 4.0}                                    
                                                                                                   
{'loss': 0.007, 'learning_rate': 0.0001, 'epoch': 5.0}  
```

I understand that QLoRA can cause a little bit of performance drop but the results I get using this are close to nothing after fine-tuning. Any suggestions or help with this is greatly appreciated!

Comments:
cc @pacman100 

@karths8 did you manage to resolve the above issue?

No! Have not been able to solve it yet!

@karths8 @amdnsr 

I think the problem is rooted in the `find_all_linear_names` function, which is incompatible with `QLoRA` setting, the correct way to implement this function should be like following

```
def find_all_linear_names(args, model):
    cls = bnb.nn.Linear4bit if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
    lora_module_names = set()
    for name, module in model.named_modules():
        if isinstance(module, cls):
            names = name.split('.')
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])


    if 'lm_head' in lora_module_names: # needed for 16-bit
        lora_module_names.remove('lm_head')
    return list(lora_module_names)
```

Because under `QLoRA` setting, the `torch.nn.Linear` module is replaced by `bnb.nn.Linearkbit` to realize the quantilization, the reason why your training loss is always the same is that there is no target module in your model at all.

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.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.