Increased peak memory usage when upgrading to `transformers` v4.30 and inclusion of `safetensors`
Author: mariecwhite
Pull request: no

### System Info

- `transformers` version: 4.30.2
- Platform: Linux-5.19.0-1022-gcp-x86_64-with-glibc2.35
- Python version: 3.10.6
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): not installed (NA)
- Tensorflow version (GPU?): 2.12.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@gante @Rocketknight1 

From June 8, after upgrading to `transformers` version 4.30 which automatically installs `safetensors`, peak memory usage in BertLarge and T5Large (and possibly other models that we have not measured) increased to what appears to be a fixed value for smaller batch sizes.

![image](https://github.com/huggingface/transformers/assets/5143063/90197b6c-51e9-4479-9b74-74833883a484)



### Information

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

### Tasks

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

### Reproduction

Run the commands below on a CUDA enabled Linux machine:

```
# Clone benchmarking repo
git clone https://github.com/iree-org/iree-samples.git
cd iree-samples/iree-tf/benchmark

# Setup output file.
OUTPUT_PATH=/tmp/tf.json
echo "{\"trigger\": { \"timestamp\": \"$(date +'%s')\" }, \"benchmarks\": []}" > "${OUTPUT_PATH}"

# Setup virtual environment.
TENSORFLOW_VERSION=2.12.0 VENV_DIR=tf.venv ./setup_venv.sh
source tf.venv/bin/activate

# Run benchmark.
BENCHMARK_ID=47cb0d3a-5eb7-41c7-9d7c-97aae7023ecf-MODEL_BERT_LARGE-fp32-TF-384xi32-batch1
python benchmark_model.py --benchmark_id="${BENCHMARK_ID}" --device=gpu --output_path="${OUTPUT_PATH}" --iterations=5
```

Benchmark output will show `"device_memory_peak_mb": 4157.236992`

Now remove `safetensors` (which was installed with v4.30):

```
pip uninstall safetensors
python benchmark_model.py --benchmark_id="${BENCHMARK_ID}" --device=gpu --output_path="${OUTPUT_PATH}" --iterations=5
```

Benchmark output will show `"device_memory_peak_mb": 1591.090432`

### Expected behavior

Device peak memory usage should not have increased by 2.5x.

Comments:
Hey @mariecwhite 👋 

The script you shared with us is quite large -- any chance you could help us narrow it down? Ideally, we'd start from a short stand-alone script, unless we are unable to reproduce the issue :)

BTW, my immediate suspicion goes towards the `.from_pretrained()` function, as (de)serialization functions should be the only ones using `safetensors`.

Investigating and trying to make a minimal reproducer - in preliminary testing I do see some differences and I'd also guess that `safetensors` is the cause.

Confirmed that the issue is caused by loading `safetensors` weights with `from_pretrained` as @gante suspected. The difference in memory usage appears before training has even begun. It is transient and only occurs during weight loading - so unless your GPU goes OOM during weight loading itself, the rest of training will not be affected by this.

My guess is that the `safetensors` loading is creating two (or more?) copies of the weights on the GPU during the loading process before eventually cleaning them up. The most likely cause is that tensors in TF are created on-device by default, whereas in torch they are created on CPU and must be moved, so probably some of the code is accidentally creating some temporary variables on GPU? 

To test memory usage for loading in TF I used the following:
```python
import tensorflow as tf
from transformers import TFAutoModel

model = TFAutoModel.from_pretrained(repo_dir)
print(tf.config.experimental.get_memory_info("GPU:0"))
```

In my testing, peak GPU memory usage when loading `bert-large-cased` was 1.5GB when loading from TF `.h5` weights and 4.1GB when loading from `safetensors`, which matches @mariecwhite's benchmark.

cc @Narsil 

Further investigation: I think the cause is in the `safetensors` code [here](https://github.com/huggingface/safetensors/blob/main/bindings/python/py_src/safetensors/tensorflow.py#L130) - `tf.convert_to_tensor()` creates the tensor on the GPU by default if one is present, so the entire state dict is materialized on the GPU alongside the randomly initialized weights during loading.

Hi @mariecwhite, thanks again for the bug report! This is a significant issue and we really appreciate the warning. The PR to fix it is open at #24404 and will hopefully be merged soon. If you'd like to try using the PR branch before then, you can install it with
```python
pip install git+https://github.com/huggingface/transformers.git@tf_safetensors_reduced_mem_usage
```

Thank you for the quick follow-up!

No probs - it's our fault for missing this issue! The PR has now been merged, so you can just install from `main` to use it.
```
pip install git+https://github.com/huggingface/transformers.git
```
It'll be included in the next patch or full release of `transformers`, at which point you can go back to just `pip install transformers`. Thanks again for the clear bug report and the work you did tracing memory usage in different scenarios!