compute_loss takes a lot of extra memory after saving checkpoint and causes OOM
Author: HuiyingLi
Pull request: no

### System Info

- `transformers` version: 4.28.1 
- Platform: Linux-5.15.0-1033-aws-x86_64-with-glibc2.27
- Python version: 3.10.9
- Huggingface_hub version: 0.14.0
- Safetensors version: not installed
- PyTorch version (GPU?): 2.0.0 (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?: 
<img width="639" alt="image" src="https://github.com/huggingface/transformers/assets/1331543/1ade210a-8dd9-4d6a-a70b-b3d2982e1914">

- Using distributed or parallel set-up in script?: No

### Who can help?

@sgugger

### 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

Every time the trainer.py:_save() saves a checkpoint https://github.com/huggingface/transformers/blob/04ab5605fbb4ef207b10bf2772d88c53fc242e83/src/transformers/trainer.py#L2884
After saving the checkpoint and then resume training entering the training_step() function and executing compute_loss() https://github.com/huggingface/transformers/blob/04ab5605fbb4ef207b10bf2772d88c53fc242e83/src/transformers/trainer.py#L2731
There is a memory usage spike that causes OOM.

Without executing ` trainer.py:_save()` (a.k.a all previous normal training steps) however, the compute_loss() function does not allocate extra memory (except for the very first time of the forward backward pass happens which is expected). No memory increase is observed during `trainer.py:_save()` either. I have changed the save_steps to different numbers, the forward pass OOM is always triggered at the step right after saving checkpoint.

Minimal training script:
```
base_model_name="EleutherAI/pythia-6.9b"
model = transformers.AutoModelForCausalLM.from_pretrained(
            base_model_name,
            load_in_8bit=True,
            torch_dtype=torch.float16,
            device_map={'': 0}
        )
tokenizer = transformers.AutoTokenizer.from_pretrained(
            base_model_name,
            device_map={'':0}
        )
model = peft.prepare_model_for_int8_training(model)
model = peft.get_peft_model(model, peft.LoraConfig(
            r=lora_r,
            lora_alpha=lora_alpha,
            target_modules=["query_key_value", "xxx"],
            lora_dropout=lora_dropout,
            bias="none",
            task_type="CAUSAL_LM",
        ))
training_args = transformers.TrainingArguments(
        per_device_train_batch_size=8, 
        gradient_accumulation_steps=gradient_accumulation_steps,  
        num_train_epochs=epochs,  
        learning_rate=learning_rate, 
        fp16=True,  
        logging_steps=20, 
        output_dir=output_dir, 
        save_steps=5, #for debugging purpose
    )

trainer = transformers.Trainer(
    model=model, 
    train_dataset=data, 
    args=training_args, 
    data_collator=transformers.DataCollatorForLanguageModeling( 
        tokenizer,  
        mlm=False, 
    ),
)
model.config.use_cache = False
result = trainer.train(resume_from_checkpoint=False)
model.save_pretrained(output_dir)
```

### Expected behavior

I would expect saving action to not change the behavior of forward pass. I am wondering why there is the memory spike and whether it can be solved.

Comments:
There is a memory spike due to the model being in 8 bits probably, cc @younesbelkada 

Thanks for the prompt response. Do you have any insights why would it only happen after saving checkpoint?

Hmm, I just figured that since this is using lora, there is no need to save checkpoints anyways? 

> There is a memory spike due to the model being in 8 bits probably, cc @younesbelkada



Hi @HuiyingLi 
Maybe the default saving mechanism is the culprit, to be on the safe zone I suggest to save the adapters only, for that you should use a custom callback to properly save the adapter weights
Please have a look at the suggested solution here: https://discuss.huggingface.co/t/peft-lora-gpt-neox-loraconfig/35790 and let us know how it goes

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.