Trainer of AutoModelForSequenceClassification is saving the wrong score module (or trained parameters are in the wrong module)
Author: koen-dejonghe
Pull request: no
Labels: name: solved

### System Info

- `transformers` version: 4.33.1
- Platform: Linux-3.10.0-1160.71.1.el7.x86_64-x86_64-with-glibc2.35
- Python version: 3.10.13
- Huggingface_hub version: 0.17.1
- Safetensors version: 0.3.3
- Accelerate version: 0.22.0
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: MULTI_GPU
        - mixed_precision: bf16
        - use_cpu: False
        - debug: False
        - num_processes: 2
        - machine_rank: 0
        - num_machines: 1
        - gpu_ids: 1,2
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.0.1+cu117 (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?

@younesbelkada
@muellerzr
@pacman100

### 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

The trainer setup is like this.
The trainer saves checkpoints every 100 steps and does an evaluation of accuracy and F1.


```
q_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_quant_type="nf4",
   bnb_4bit_use_double_quant=True,
   bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForSequenceClassification.from_pretrained(
    "meta-llama/Llama-2-13b-hf",
    quantization_config=q_config,
    device_map="auto", 
    num_labels=n_labels,
)
model.config.pad_token_id = tokenizer.pad_token_id
model.config.use_cache = False 

peft_config = LoraConfig(
        r=16,  
        lora_alpha=64, 
        lora_dropout=0.1, 
        bias="none",
        task_type=TaskType.SEQ_CLS,
        target_modules=['v_proj', 'down_proj', 'up_proj', 'q_proj', 'gate_proj', 'k_proj', 'o_proj']
)

model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, peft_config)

training_args = TrainingArguments(...)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=ds_train,
    eval_dataset=ds_test,
    compute_metrics=compute_metrics,
    tokenizer=tokenizer,
)

trainer.train()
trainer.save_model("final-checkpoint")
```

I load the final checkpoint as follows. Note I've tried all the other possible ways to load the model as well. The problem is not in the loading.

```
model = AutoModelForSequenceClassification.from_pretrained(
    "final-checkpoint",
    device_map="auto", 
    num_labels=n_labels,
    quantization_config=q_config
)
```

When I do inference with this model on the same test dataset used during training, the loss, F1 and accuracy are really bad compared to the output of the training evaluator.

The modules in the trained model look like this:
```
(score): ModulesToSaveWrapper(
        (original_module): Linear(in_features=5120, out_features=647, bias=False)
        (modules_to_save): ModuleDict(
          (default): Linear(in_features=5120, out_features=647, bias=False)
        )
)
```

What is being saved to the checkpoints is `score.modules_to_save.default`. 

If I dump `score.original_module.weight` to a file and load it in a model instantiated from a checkpoint, I get the original loss and metrics. 
For example (skipped some steps):
```
# on trained model
orig_weights = trained_model.model.score.original_module.weight.cpu().detach()

# on checkpoint model:
checkpoint_model.score.load_state_dict({"weight": orig_weights})
```

See also https://discuss.huggingface.co/t/llama-2-sequence-classification-much-lower-accuracy-on-inference-from-checkpoint-compared-to-model/54910/2 


### Expected behavior

Metrics and loss during inference based off a checkpoint should be comparable to the evaluation during training with the same test dataset.

Comments:
Looks interesting, I am doing something similar with 0.17.2 , I noticed that the linear layers doesn't change, did you suffer from something like that

I have the same issues. `Peft` splits `score` into `original_module (no grad)` and `default (grad=true)`. Every checkpoint only `default` adapter layer is saved and without the `original_module` the evaluation metrics are off. I got better insights by going through #602 and #876. 

@natank1  your issue might be fixed with this [pull request](https://github.com/huggingface/peft/pull/755). 

Thanks

בתאריך יום ג׳, 17 באוק׳ 2023 ב-3:55 מאת Thein Oo ***@***.***>:

> I have the same issues. Peft splits score into original_module (no grad)
> and default (grad=true). Every checkpoint only default adapter layer is
> saved and without the original_module the evaluation metrics are off. I
> got better insights by going through #602
> <https://github.com/huggingface/transformers/issues/602> and #876
> <https://github.com/huggingface/transformers/issues/876>.
>
> @natank1 <https://github.com/natank1> your issue might be fixed with this pull
> request <https://github.com/huggingface/peft/pull/755>.
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/issues/26160#issuecomment-1765488409>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AB7W2Z7FZOV5ZVSC35ERNNTX7XJRZAVCNFSM6AAAAAA4X34AFKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTONRVGQ4DQNBQHE>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>


Sorry for re--discussing: in the version I work score(under the peftmodel..basemodel.model) is a linear and simply doesn't update : I work on 0.17.2 (HF) and 4.34.dev (trnaformers)

Note: the bug does not occur when you don’t specify `target_modules` in LoraConfig, so effectively only use default q & v of the attention blocks.

in my case I dont define target . Notsure whether it is due to different versions

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.

Hi,

I confirm this is severe a bug (probably only for 'multi-label-task') and the [workaround](https://discuss.huggingface.co/t/llama-2-sequence-classification-much-lower-accuracy-on-inference-from-checkpoint-compared-to-model/54910) does solve this issue.

I have fine-tuned LLAMA-2-7b using QLora for multi-label classification on 6k context using RoPE scaling. After training the model for 2 epochs I get F1 macro 0.51 and micro 0.65 on the evaluation dataset. However, When I tried to reproduce these scores in the inference pipeline on the same evaluation dataset I failed with F1 scores of 0.13 and 0.18, where all the predictions are random.

I tried training with and without flash_attention, and RoPE scaling but couldn't reproduce the
evaluation metrics. I also confirmed that the 'score' head is saved by setting 'module_to_save=[score]' but nothing worked. 


pinging @muellerzr and @pacman100 as it seems this issue still exists 

Hello @koen-dejonghe,

with the latest versions of Transformers, Accelerate and PEFT, when you call `trainer.save_model()`, only the adapter weights get saved and for inference you would need to do use `PeftModel.from_pretrained`. Please refer https://huggingface.co/docs/peft/quicktour#save-model for the details wrt saving and loading the adapter weights.

If you want to save the base model with the adapter weights merged in, please do the following:
```
model = trainer.model.merge_and_unload()
model.save_pretrained("final-checkpoint-transformers")
```

Then you can load the ckpt as usual using `AutoModelForSequenceClassification.from_pretrained `

Here is the minimal notebook example wherein I finetune tinyllama on mrpc seq cls task using qlora and targeting all linear layers. When I load the model for inference via `AutoPeftModelForSequenceClassification`, everything is working as expected. Please let us know if the recent releases have fixed this issue.

[tinyllama_qlora_seqcls.ipynb.zip](https://github.com/huggingface/transformers/files/14175992/tinyllama_qlora_seqcls.ipynb.zip)


@pacman100 this still seems to be an issue when doing multi-label classification. I just tried now with the latest updates for transformers, peft, and accelerate. I notice your notebook is binary classification not multi-label, which could account for the discrepancy?

Hello, could you please share the minimal example that we can run end-to-end to further deep dive?

I will try and write something appropriate asap. Your original notebook you posted could illustrate the problem if you use a multilabel problem instead of a binary classification problem. This huggingface dataset "victoriadreis/TuPY_dataset_multilabel" might be a good publicly available test.

When you reload the model it uses a random linear layer at the end instead of the one you have finetuned, giving random results. Other people have reported it doesn't properly save the final classification layer.


It seems this issue is solved now. I instruction fine-tuned the llama2 model with QLoRA and everything during training/validation and even inference was ok. 

```
# training args and trainer
training_args = TrainingArguments(output_dir=output_dir, ...)
trainer = Trainer(args=training_args, ...)
trainer.train()

trainer.save_model()
# load model and tokenizer
peft_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained(output_dir)
```


@shakibyzn have you tried it with AutoPeftModelForSequenceClassification head?

I've had no issues with CausalLM and Seq2Seq models.

I think the problems are relating to the linear classification head at the end not being saved properly and initialising a random classification head when reloading.

My personal problem is related to a multi-label classification problem using GatorTron medium, so I also don't know if that makes a difference. I want to write a script to demonstrate, but due to time pressures, and using confidential data in my project, I need to find the time to do it.

@stephenhbarlow My next task is to use the Llama2 model together with QLoRA for multi-class sequence classification. I'll let you know about the result by tomorrow.

@stephenhbarlow I did a single run and the result was what I expected. Here is a simplified version of my code:
I fine-tuned meta-llama/Llama-2-7b-hf as a sequence classification task.
```
# load your dataset
dataset = load_data(...)

# bnb config
bnb_config = BitsAndBytesConfig(...)

model = AutoModelForSequenceClassification.from_pretrained(model_name, quantization_config=bnb_config, num_labels=YOUR_NUM_LABELS)

tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
tokenizer.pad_token = tokenizer.eos_token

tokenized_datasets = dataset.map(YOUR_PREPROCESS_FUNC)

peft_config = LoraConfig(r=r, lora_alpha=lora_alpha, bias="none", task_type="SEQ_CLS")
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, peft_config)

training_args = TrainingArguments(output_dir=output_dir, ...)
trainer = Trainer(...)

trainer.train()
trainer.save_model()

# load model and tokenizer
peft_model = AutoPeftModelForSequenceClassification.from_pretrained(output_dir, id2label=id_to_label)
tokenizer = AutoTokenizer.from_pretrained(output_dir)

# evaluate 
evaluate_fn(peft_model, dataset['test'])
```

@shakibyzn interesting - must be related to multi-label classification for some reason. Still planning on making a notebook to demonstrate when I can.

@stephenhbarlow I wouldn't say it's related to that. Take a look at this [medium blog post](https://medium.com/@lukas.hauzenberger/multilabel-classification-using-mistral-7b-on-a-single-gpu-with-quantization-and-lora-8f848b5237f3
). It might be helpful. Does your model perform as expected on the validation set (output of trainer I mean)? Or is it only a problem with inference?

@shakibyzn  It's either that or related to the base model I'm using (Gatortron Medium). Basically it all works fine through training and as long as the trained model is still in memory it works great. The problem is with saving/reloading. The reloaded model initializes a new classification head so you get random results. I've tried the various things in the documentation like saving specific module and whatnot, but no success so far.

I've implemented it very similarly to the blogpost and results completely change when reloading (random basically cos the final layer is random).

Here's another person having a problem that linked me to this issue:

https://discuss.huggingface.co/t/llama-2-sequence-classification-much-lower-accuracy-on-inference-from-checkpoint-compared-to-model/54910


Perhaps this post can be beneficial  

(https://natan-katz.medium.com/codellama-classification-finetuning-28fa5546f64f




@stephenhbarlow I saw his post. Have you tried a newer version of transformers? e.g. 4.35?

I'm having the same issue, with esm (sequence classification + peft). Saving with merge+unload results in a garbage model upon reloading. 

@shakibyzn I'm using latest versions of Transformers, PEFT and Accelerate

I have been playing around again today and I think the issue I'm having might actually specific to the GatorTron medium model. For some reason it always uses a newly initialised classifier head when you save and reload and gives random results. I followed the exact protocol on my pc with the notebook @pacman100 provided and it worked fine and exactly like it did for him. Swap the base model for GatorTron and it doesn't.

(https://huggingface.co/UFNLP/gatortron-medium)

I am also struggling to get any sort of workaround operational (by saving the classifier layer separately and reloading etc...).

As this bug seems specific to a model I appreciate that it might not be something that can be fixed easily. If anyone has any insights or ideas for a workaround then that would be hugely appreciated. Thanks for all the tips and insights so far.

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.

I run into a similar issue. Resolved by adding modules_to_save = ["score"] in LoraConfig, then the score is saved. When I load the model again and evaluate the same dataset, the loaded model can give same evaluation result as the trained model

@FangyaoLiu , Hi, could you share your save and load code? I did the same thing, and it still doesn't work.