Trainer memory usage increasing a lot during evaluation steps when using Flan-T5 models with 8-bit/4-bit quantisation
Author: guilherme-pombo
Pull request: no

### System Info

Hello :) here are the details of my system:

Python: 3.8.10
Transformers version: 4.34.1
Bitsandbytes version: 0.41.2
Peft version: 0.6.2 
CUDA Version: 12.0  

When using T5-XXL with 8-bit or 4-bit quantisation and Hugging Face trainer the model memory consumption will be reasonable. However, as soon as the model does the evaluation steps, it runs out of memory since it's usage balloons up past the 48GB RAM that my NVIDIA RTX 6000 has.

This seems really odd since I expect the evaluation step to require less memory than the training step as it doesn't need to keep track of gradients.

This is not a problem when training in either FP32/FP16/BF16. Only happens when using 8 or 4-bit quantisation.

Any help would be greatly appreciated!

 @SunMarc  @pacman10

### Who can help?

 @SunMarc  @pacman10

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

This happens with any dataset, the code used for model loading and Trainer setup is as follows:

```
from transformers import AutoModelForSeq2SeqLM, BitsAndBytesConfig
from peft import prepare_model_for_kbit_training, TaskType, get_peft_config, get_peft_model
import transformers
from transformers import Seq2SeqTrainingArguments

'''
This happens with any dataset
'''

model = AutoModelForSeq2SeqLM.from_pretrained('google/flan-t5-xxl',
                                              device_map={"":0},
                                              torch_dtype=torch.bfloat16,
                                              load_in_8bit=True,
                                             )
                                             
model = prepare_model_for_kbit_training(model)

target_modules = ['.q', '.k', '.v', '.o', 'wi_0', 'wi_1', 'wo']
peft_config = LoraConfig(
    task_type=TaskType.SEQ_2_SEQ_LM,
    inference_mode=False,
    target_modules=target_modules,
    r=32,
    lora_alpha=32,
    lora_dropout=0.15
)

model = get_peft_model(model, peft_config)

trainer = transformers.Trainer(
        model=model,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        args=transformers.Seq2SeqTrainingArguments(
            per_device_train_batch_size=1,
            warmup_ratio=0.01,
            lr_scheduler_type='cosine',
            gradient_accumulation_steps=6,
            num_train_epochs=60,
            learning_rate=3e-4,
            eval_accumulation_steps=1,
            weight_decay=1e-2,
            bf16=True,
            logging_steps=10,
            optim="paged_adamw_8bit",
            evaluation_strategy="steps",
            save_strategy="steps",
            eval_steps=5,
            save_steps=5,
            output_dir='my_model',
            save_total_limit=1,
            load_best_model_at_end=True,
            ddp_find_unused_parameters=None,
            group_by_length=False,
        ),
    )

model.config.use_cache = False 
trainer.train()
```

### Expected behavior

When training T5-XXL in either FP32/FP16/BF16 there is no increase in memory usage during the evaluation step. However, when using 8 or 4-bit quantisation the memory usage balloons up to more than double. Since evaluation doesn't require keeping track of gradients, the expected behaviour should be that memory usage doesn't change.
 

Comments:
cc @BenjaminBossan 

Turns out this was happening because I was not defining the:

_per_device_eval_batch_size_

parameter when initialising the Trainer and it was being set to the default value on Trainer, which is 8. When I define it separately there is no issue now. Apologies for the confusion. Closing the issue