facebook/opt-30b Cuda Allocation Error with version >= 4.50.0 code
Author: inf3rnus
Pull request: no
Labels: name: bug

### System Info

Ubuntu 24.04 LTS
torch==2.6.0
transformers==4.51.2

AWS EC2 g4dn.metal

### General Description

Hey all,

Figure these are some growing pains with the latest release.

On an aws ec2 g4dn.metal the model loads fine, no issues with version 4.49.0

Something has broke with loading sharded models for versions > 4.50.0 (I've tested up to the latest release, which as of this post is 4.51.2)

Note, I was observing memory utilization with `watch -n 0.1 nvidia-smi`

If it's helpful to know, it looks like there's an allocation step performed before weights are loaded, that fills each GPU with the same amount of data, I don't remember if it's 1GB or 10GB, but it's even for each GPU.

The old version loads weights one by one and appears to clear the GPUs of any preallocated memory before loading all the weights.

My theory is something is not being properly released before the weights are actually loaded.


This is with `torch==2.6.0`, I've not tested other torch versions, but I'm inclined to believe the issue is with model_utils.py and not with underlying code. 

Final note, I tried version 0.32 of accelerate and 1.6.0, but that had zero impact.

Best,
Aaron

### Who can help?

@Rocketknight1 

### Reproduction

To replicate, just run this code on a g4dn.metal instance. 

First with `transformers==4.49.0` which should work and then with `transformers==4.51.2` which should fail.


```py
from transformers import pipeline

# on a g6.48xlarge we got 33s in serial
if __name__ == "__main__":
    task = "text-generation"
    model_id = "facebook/opt-30b"

    pipe = pipeline(task=task, model=model_id, device_map="auto")
```

You should see that it loads with 4.49.0, and that you get a cuda OOM on 4.51.2

### Expected behavior

No OOM

Comments:
Hi @inf3rnus, thanks for the bug report! Firstly, can you check that it's still an issue using the latest version on `main`? If so, would you be willing to [git bisect](https://git-scm.com/docs/git-bisect) to find out exactly which PR caused it? That should make it a lot easier for us to track down the bug and assign someone to fix it.

Hey @Rocketknight1 updated my report, I had the wrong latest version specified, it's effecting the latest version on PyPi, so for clarity, as of this post, that's `4.51.2`.

With this said, the bug was introduced on version `4.50.0`.

When I have more time, can't guarantee that I will, but I'll see if I can bisect and narrow down to the specific commit between the two versions, although it may be faster to just run this through the debugger.

@Rocketknight1 

Found it, bad commit's here:

https://github.com/huggingface/transformers/pull/36742

Ran this test script to find it, with the commit ranges being 4.49.0 (last good) and 4.50.0 (end of range, bad commit)

```py
import subprocess
import sys
import os

# Install the current commit version of transformers
subprocess.run([sys.executable, "-m", "pip", "install", "."], check=True)

os.environ["HF_HOME"] = "/opt/dlami/nvme"


"""
# Optional but safe: reset to a clean working state
git reset --hard
git clean -fdx

# Start bisect
git bisect start
git bisect bad v4.50.0
git bisect good v4.49.0

# Run bisect using your script
git bisect run python3 bisect_test.py

"""
try:
    from transformers import pipeline

    task = "text-generation"
    model_id = "facebook/opt-30b"

    pipe = pipeline(task=task, model=model_id, device_map="auto")
    pipe("Hello, I'm a language model,", max_new_tokens=5)

    sys.exit(0)  # Good commit

except Exception as e:
    print(f"Bad commit. Error: {e}", file=sys.stderr)
    sys.exit(1)  # Bad commit

```

cc @cyrilvallez who wrote #36742, but let me know if you don't have bandwidth and I can take this one!

Hey! Could you provide the full traceback of the OOM you're encountering please? 🤗

I just checked locally, and the caching allocator is returning the exact byte size of the model, so I don't see how it could come from it. Are you sure the GPUs are not doing anything else in parallel? Because the pipeline will load in float32 by default, which leads to a model memory footprint of ~111.66 GiB. Your instance has 128 GiB available, so if the GPUs are used for something else in parallel, that could explain the issue!

@Cyrilvallez Hey, not running anything else in parallel. Something feels like it's being "double loaded"

If you run the code with version 4.49.0, everything loads just fine. If you load with 4.51.3, it explodes (anything > 4.50.0)...

I'm going to run the debugger on modeling_utils.py and see if I can hunt this down, but I'm time boxing it to 1 hour.

This is the stack trace:

```bash
root@ip-172-31-33-115:~# python3 load_model.py
Loading checkpoint shards:  14%|█████████████████████▎                                                                                                                               | 1/7 [00:08<00:51,  8.57s/it]
Traceback (most recent call last):
  File "/root/load_model.py", line 20, in <module>
    pipe = pipeline(task=task, model=model_id, device_map="auto")
  File "/usr/local/lib/python3.10/dist-packages/transformers/pipelines/__init__.py", line 942, in pipeline
    framework, model = infer_framework_load_model(
  File "/usr/local/lib/python3.10/dist-packages/transformers/pipelines/base.py", line 291, in infer_framework_load_model
    model = model_class.from_pretrained(model, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/models/auto/auto_factory.py", line 571, in from_pretrained
    return model_class.from_pretrained(
  File "/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py", line 279, in _wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py", line 4399, in from_pretrained
    ) = cls._load_pretrained_model(
  File "/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py", line 4833, in _load_pretrained_model
    disk_offload_index, cpu_offload_index = _load_state_dict_into_meta_model(
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py", line 824, in _load_state_dict_into_meta_model
    _load_parameter_into_model(model, param_name, param.to(param_device))
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 784.00 MiB. GPU 1 has a total capacity of 14.57 GiB of which 736.75 MiB is free. Including non-PyTorch memory, this process has 13.85 GiB memory in use. Of the allocated memory 13.02 GiB is allocated by PyTorch, and 747.95 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
```

Problem is line 4793, `caching_allocator_warmup()`

The pre allocation causes memory to be exhausted by the time weights are actually loaded... Currently seeing if there's a way to fix it easily.

@Cyrilvallez Alright, I have confirmed it's `caching_allocator_warmup()`

The problem is too much memory is being preallocated. I don't have time to understand why there's memory overhead in loading tensors onto the meta device resulting in OOM, but reducing the total memory to preallocate to 90% solves the problem.

```py
def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: Dict, factor=2):
    """This function warm-ups the caching allocator based on the size of the model tensors that will reside on each
    device. It allows to have one large call to Malloc, instead of recursively calling it later when loading
    the model, which is actually the loading speed botteneck.
    Calling this function allows to cut the model loading time by a very large margin.

    A few facts related to loading speed (taking into account the use of this function):
    - When loading a model the first time, it is usually slower than the subsequent times, because the OS is very likely
    to cache the different state dicts (if enough ressources/RAM are available)
    - Trying to force the OS to cache the files in advance (by e.g. accessing a small portion of them) is really hard,
    and not a good idea in general as this is low level OS optimizations that depend on ressource usage anyway
    - As of 18/03/2025, loading a Llama 70B model with TP takes ~1 min without file cache, and ~13s with full file cache.
    The baseline, i.e. only loading the tensor shards on device and adjusting dtype (i.e. copying them) is ~5s with full cache.
    These numbers are reported for TP on 4 H100 GPUs.
    - It is useless to pre-allocate more than the model size in this function (i.e. using an `allocation_factor` > 1) as
    cudaMalloc is not a bottleneck at all anymore
    - Loading speed bottleneck is now almost only tensor copy (i.e. changing the dtype) and moving the tensors to the devices.
    However, we cannot really improve on those aspects obviously, as the data needs to be moved/copied in the end.
    """
    # Remove disk and cpu devices, and cast to proper torch.device
    accelerator_device_map = {
        param: torch.device(device) for param, device in expanded_device_map.items() if device not in ["cpu", "disk"]
    }
    if not len(accelerator_device_map):
        return

    tp_plan_regex = (
        re.compile("|".join([re.escape(plan) for plan in model._tp_plan]))
        if _torch_distributed_available and torch.distributed.is_initialized()
        else None
    )
    total_byte_count = defaultdict(lambda: 0)
    for param_name, device in accelerator_device_map.items():
        param = model.get_parameter_or_buffer(param_name)
        # The dtype of different parameters may be different with composite models or `keep_in_fp32_modules`
        param_byte_count = math.prod(param.shape) * param.element_size()

        if tp_plan_regex is not None:
            generic_name = re.sub(r"\.\d+\.", ".*.", param_name)
            param_byte_count //= torch.distributed.get_world_size() if tp_plan_regex.search(generic_name) else 1

        total_byte_count[device] += param_byte_count

    # This will kick off the caching allocator to avoid having to Malloc afterwards
    for device, byte_count in total_byte_count.items():
        if device.type == "cuda":
            index = device.index if device.index is not None else torch.cuda.current_device()
            device_memory = torch.cuda.mem_get_info(index)[0]
            # Allow up to 90% of max device memory <--- I changed this from 95% to 90%, and it now works fine
            byte_count = min(byte_count, int(0.90 * device_memory))
        # Allocate memory
        _ = torch.empty(byte_count // factor, dtype=torch.float16, device=device, requires_grad=False)
```

However, with that said, I ran a few benchmarks with fresh uncached files (deleted them from the file system, and re downloaded them and then timed `_load_pretrained_model()`), and for at least `facebook/opt-30b`, there appears to be no tangible speed up with pre allocating.

Results for `facebook/opt-30b`:

```
With caching_allocator_warmup() on RAM disk:
41.2s
40.64s

With caching_allocator_warmup() on local SSD:
41.55s

Without caching_allocator_warmup() on RAM disk:
41.51s
40.96s

Without caching_allocator_warmup() on local SSD:
41.51s
```

Because of this, I'd recommend this enhancement be iced until there is a script or test that details the exact environment that shows the speed up, I've been doing a lot of benchmarking myself and depending on the hardware being utilized you can get very different results. e.g. using instances stores, ram disks, vs EBS drives. Caching speed ups seem to be wildly different. EBS drives tend to be much more sensitive to filesystem caching. 

Local instance stores and RAM disks I have found show almost no difference in load times when files are cached by the OS vs not.

Running these tests in Docker also effects load times, as it has its own interaction with the OS that can complicate the predictability of optimizations.

Final note, I know you all are slammed right now, especially with .xet and this latest release, but if someone could please help me push this thing along https://github.com/huggingface/transformers/pull/36835, I would be greatly appreciative! 

This speeds up model loading by ~50% when everything is cold. Depends on the model and your hardware, but it's a massive increase in performance for cold loading.

> torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 784.00 MiB. GPU 1 has a total capacity of 14.57 GiB of which 736.75 MiB is free. Including non-PyTorch memory, this process has 13.85 GiB memory in use. Of the allocated memory 13.02 GiB is allocated by PyTorch, and 747.95 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

Alright, this traceback is extremely useful: what happens is that for that very particular model and hardware config, indeed the 95% memory reservation is just very slightly too much, because we end up with 13.02 GiB already allocated, and 748 MiB still reserved (warmed up). However, next param is 784 MiB > 748 MiB still reserved, so Torch tries to plainly allocate that much on the GPU, without using the 748 that are still reserved, which ends up blowing the GPU as it's larger than the 736 MiB that are still free. We could argue that it is somewhat a default of Torch internal memory reservation system.

I will make sure to immediately slightly lower the 95% threshold, to avoid such edge cases. 🤗

As for the caching allocator itself, it is only warming up cuda, so it has nothing to do with the OS filesystem caching. It will make things much faster when cuda is cold (i.e. you just started a new python process, and nothing was yet initialized on cuda by torch), which explains why you can't measure speedups probably. If you want to measure it, you should make sure to measure timings while launching new python processes everytime. Also, note that some hardware are quite slow to copy data. Here, as you are loading the model in float32 but the weights are 16 bits precision (not sure if bf16 or fp16, but it's 16 bits), every parameter needs to be upcasted on-the-fly, which requires a copy. This can sometimes become the loading bottleneck depending on the hardware, in which case warming-up cuda does not provide much speedups, as all the time is taken not to allocate, but to copy the data. But on more recent hardware such as A100/H100 and such, it does provide very large speedups.

Finally, no worries I will try to review your PR asap! It is still very much welcome!!! 🤗😍

Thanks @Cyrilvallez... Interesting, good to know. And much appreciated for the extra eyes!

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.