Repository: huggingface/transformers
Pull request: [core] reuse unused reserved cuda memory when loading models

# What does this PR do?

TL;DR checks for unused reserved CUDA memory before preallocating more memory or deciding to do CPU offload.
Missing: benchmark whether this has a speed impact in `from_pretrained` on e.g. TP

### Context

(first commit containing the issue: https://github.com/huggingface/transformers/pull/36335)

There has been an issue with flaky model tests that is difficult to reproduce, and where resetting cuda memory was helping. E.g. if we remove the `tearDown` function with `torch.cuda.empty_cache` in `CacheHardIntegrationTest`, we might start getting failures (depending on the device).

Tracing down the issue, we can see that repeated `from_pretrained` calls may start offloading the model. More specifically, we can see that 
1. the reserved memory grows when we instantiate a second model, even when the first model is no longer actively allocating cuda memory
2. we're triggering CPU offload when there is plenty of memory for the model

On `main` + RTX 4090 (24GB), if we pick a 4B model in BF16 (~33% of device memory), we observe the following (see output below -- notice the CPU offload after the 3rd call):

```py
# How to reproduce the issue: pick a model/GPU/dtype combination such that the model takes >33% memory of the GPU.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "Qwen/Qwen3-4B"

def generate():
    tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.bfloat16)
    inputs = tokenizer(["Here's everything I know about cats. Cats"], return_tensors="pt").to(model.device)
    _ = model.generate(**inputs, do_sample=True, max_new_tokens=1, return_dict_in_generate=True, output_scores=True)

print("generate 1")
generate()
print("memory allocated (GB)", torch.cuda.memory_allocated(0) / 1024 ** 3)
print("memory reserved (GB)", torch.cuda.memory_reserved(0) / 1024 ** 3)

print("generate 2")
generate()
print("memory allocated (GB)", torch.cuda.memory_allocated(0) / 1024 ** 3)
print("memory reserved (GB)", torch.cuda.memory_reserved(0) / 1024 ** 3)

print("generate 3")
generate()
print("memory allocated (GB)", torch.cuda.memory_allocated(0) / 1024 ** 3)
print("memory reserved (GB)", torch.cuda.memory_reserved(0) / 1024 ** 3)

print("generate 4")
generate()
print("memory allocated (GB)", torch.cuda.memory_allocated(0) / 1024 ** 3)
print("memory reserved (GB)", torch.cuda.memory_reserved(0) / 1024 ** 3)
```

```
generate 1
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00,  2.94it/s]
memory allocated (GB) 0.0079345703125
memory reserved (GB) 8.224609375
generate 2
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00,  2.89it/s]
memory allocated (GB) 0.0079345703125
memory reserved (GB) 16.443359375
generate 3
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00,  4.00it/s]
Some parameters are on the meta device because they were offloaded to the cpu.
memory allocated (GB) 5.62136697769165
memory reserved (GB) 8.224609375
generate 4
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00,  3.17it/s]
memory allocated (GB) 5.62136697769165
memory reserved (GB) 16.443359375
```

### Solution

The solution is quite simple: when warming up memory or deciding whether to do CPU offload, let's check the memory available in the GPU *including* unused reserved memory.

After the fix in this PR, rerunning the script above we get

```
generate 1
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00,  2.77it/s]
memory allocated (GB) 0.0079345703125
memory reserved (GB) 8.22265625
generate 2
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00,  2.78it/s]
memory allocated (GB) 0.0079345703125
memory reserved (GB) 8.22265625
generate 3
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00,  2.92it/s]
memory allocated (GB) 0.0079345703125
memory reserved (GB) 8.22265625
generate 4
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00,  2.98it/s]
memory allocated (GB) 0.0079345703125
memory reserved (GB) 8.22265625
```


Changed files:
filename: src/transformers/modeling_utils.py
patch: @@ -1285,6 +1285,13 @@ def _get_device_map(
             max_memory = get_max_memory(max_memory)
         if hf_quantizer is not None:
             max_memory = hf_quantizer.adjust_max_memory(max_memory)
+
+        # CUDA: `max_memory` contains non-reserved memory. There may be *unused* reserved memory in the GPU, which we
+        # can use to allocate parameters.
+        for device_name in max_memory.keys():
+            if isinstance(device_name, int):  # it's a GPU device
+                unused_memory = torch.cuda.memory_reserved(device_name) - torch.cuda.memory_allocated(device_name)
+                max_memory[device_name] += unused_memory
         device_map_kwargs["max_memory"] = max_memory
 
         device_map = infer_auto_device_map(model, dtype=target_dtype, **device_map_kwargs)
@@ -5979,6 +5986,9 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: Dict,
             # Note that we use an absolute value instead of device proportion here, as a 8GiB device could still allocate too much
             # if using e.g. 90% of device size, while a 140GiB device would allocate too little
             byte_count = min(byte_count, max(0, int(device_memory - 1.2 * 1024**3)))
+            # If there is *unused* reserved cuda memory, we can skip/reduce the allocation.
+            unused_memory = torch.cuda.memory_reserved(index) - torch.cuda.memory_allocated(index)
+            byte_count = max(0, byte_count - unused_memory)
         # Allocate memory
         _ = torch.empty(byte_count // factor, dtype=torch.float16, device=device, requires_grad=False)