Fix `save_pretrained` for partially offloaded models
Author: kylesayrs
Pull request: yes
Labels: name: bug

name: Accelerate

# What does this PR do? #
## Purpose ##
* Fix bug introduced by slight behavior change of the `get_state_dict_from_offload` introduced starting in accelerate 1.1.0
* This bug affects models which are executed on the GPU with partial cpu/disk offloading only

## Background ##
Starting in [accelerate 1.1.0](https://github.com/huggingface/accelerate/pull/3204), `get_state_dict_from_offload` moves state_dict tensors to the CPU by default.

In the case that a module has parameters on the GPU, `get_state_dict_from_offload` implements the following procedure:
* The module starts with a reference to the tensors on the GPU (1)
* When `get_state_dict_from_offload` is called, copies of the tensors are created on the CPU (2) which are returned later
* After the tensors are copied to the CPU, the reference to the GPU tensors (1) is dropped
* After the CPU tensors (2) are moved into the `state_dict`, the CPU tensors (2) are copied back to the GPU (3)

This procedure ensures memory efficiency due to how reference to GPU tensors (1) is dropped before GPU tensors (3) are allocated.

However, in the downstream case of transformers, `PretrainedModel.save_pretrained` keeps a copy of the original GPU tensors (1) for longer than they are needed. This long-lived reference to (1) results in references to both (1) and (3) being alive at the same time, leading to increased memory usage.

## Changes ##
* Remove references to tensors in the `state_dict` as they are loaded  by `get_state_dict_from_offload`
  * This ensures compatibility with accelerate versions which have been released prior to and after https://github.com/huggingface/accelerate/pull/3253

## Testing #
<details><summary>save_offloaded.py</summary>

```python3
import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("nvidia/Llama-3.1-Nemotron-70B-Instruct-HF", torch_dtype="auto", device_map="auto")
print(model.hf_device_map)
"""
{'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layer
s.4': 0, 'model.layers.5': 0, 'model.layers.6': 0, 'model.layers.7': 0, 'model.layers.8': 0, 'model.layers.9': 0, 'model.l
ayers.10': 0, 'model.layers.11': 0, 'model.layers.12': 0, 'model.layers.13': 0, 'model.layers.14': 0, 'model.layers.15': 0
, 'model.layers.16': 0, 'model.layers.17': 0, 'model.layers.18': 0, 'model.layers.19': 0, 'model.layers.20': 0, 'model.lay
ers.21': 0, 'model.layers.22': 0, 'model.layers.23': 0, 'model.layers.24': 0, 'model.layers.25': 0, 'model.layers.26': 0, 
'model.layers.27': 0, 'model.layers.28': 0, 'model.layers.29': 0, 'model.layers.30': 0, 'model.layers.31': 0, 'model.layer
s.32': 0, 'model.layers.33': 0, 'model.layers.34': 0, 'model.layers.35': 0, 'model.layers.36': 0, 'model.layers.37': 0, 'm
odel.layers.38': 0, 'model.layers.39': 0, 'model.layers.40': 0, 'model.layers.41': 0, 'model.layers.42': 'cpu', 'model.lay
ers.43': 'cpu', 'model.layers.44': 'cpu', 'model.layers.45': 'cpu', 'model.layers.46': 'cpu', 'model.layers.47': 'cpu', 'm
odel.layers.48': 'cpu', 'model.layers.49': 'cpu', 'model.layers.50': 'cpu', 'model.layers.51': 'cpu', 'model.layers.52': '
cpu', 'model.layers.53': 'cpu', 'model.layers.54': 'cpu', 'model.layers.55': 'cpu', 'model.layers.56': 'cpu', 'model.layer
s.57': 'cpu', 'model.layers.58': 'cpu', 'model.layers.59': 'cpu', 'model.layers.60': 'cpu', 'model.layers.61': 'cpu', 'mod
el.layers.62': 'cpu', 'model.layers.63': 'cpu', 'model.layers.64': 'cpu', 'model.layers.65': 'cpu', 'model.layers.66': 'cp
u', 'model.layers.67': 'cpu', 'model.layers.68': 'cpu', 'model.layers.69': 'cpu', 'model.layers.70': 'cpu', 'model.layers.71': 'cpu', 'model.layers.72': 'cpu', 'model.layers.73': 'cpu', 'model.layers.74': 'cpu', 'model.layers.75': 'cpu', 'model.layers.76': 'cpu', 'model.layers.77': 'cpu', 'model.layers.78': 'cpu', 'model.layers.79': 'cpu', 'model.norm': 'cpu', 'model.rotary_emb': 'cpu', 'lm_head': 'cpu'}
"""

torch.cuda.memory._record_memory_history()
model.save_pretrained("save_dir")
torch.cuda.memory._dump_snapshot(f"align.pickle")
torch.cuda.memory._record_memory_history(enabled=None)
```
</details>

### transformers <= 4.46.0 | accelerate < 1.1.0 ###
* Original references are not moved
<img width="1720" alt="Screenshot 2024-11-22 at 22 08 35" src="https://github.com/user-attachments/assets/7f21c040-ed8f-495b-a8b4-e7a195f4cc61">

### transformers <= 4.46.0 | accelerate == 1.1.0, 1.1.1 ###
* Results in bug
* Original references are not garbage collected
* Calling save_pretrained results in OOM (>80GiB)
<img width="1650" alt="without_change" src="https://github.com/user-attachments/assets/6b88675d-2577-4be3-b04e-16f3efaa7872">

### transformers > 4.46.0 (this branch) | accelerate == 1.1.0, 1.1.1 ###
* Original references are garbage collected and deallocated after each shard (note that peak represents tensors which are dereference but which have not been deallocated)
* Calling save_pretrained maintains existing memory usage (<68.9GiB)
<img width="1719" alt="Screenshot 2024-11-23 at 10 26 30" src="https://github.com/user-attachments/assets/616ccdff-2336-4fd4-82f9-a6ca77b47b99">

### transformers > 4.46.0 (this branch) | accelerate > 1.1.1 ###
* Original references are not moved
<img width="1717" alt="Screenshot 2024-11-24 at 23 28 27" src="https://github.com/user-attachments/assets/042f1b2b-82b0-4815-824d-56aeb232997e">

## Who can review?




Comments:
cc @muellerzr @SunMarc for accelerate

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_34890). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.