Moving to CUDA after loading is extremely slow

[francois-rozet]

Describe the bug

Moving a freshly loaded module to CUDA is extremely slow. Much slower than it should. This seems to be related to the way diffusers loads weights because cloning the weights before moving them is much faster.

Reproduction

import time
import torch

from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    variant="fp16",
    torch_dtype=torch.float16,
    low_cpu_memory_usage=False,
)

tic = time.time()

pipe.unet.to("cuda")

tac = time.time()

print(tac - tic, flush=True)  # 20s

If I manually clone the parameters and buffers first, the transfer is much faster.

import time
import torch

from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    variant="fp16",
    torch_dtype=torch.float16,
    low_cpu_memory_usage=False,
)

tic = time.time()

for p in (*pipe.unet.parameters(), *pipe.unet.buffers()):
    p.data = p.data.clone()

pipe.unet.to("cuda")

tac = time.time()

print(tac - tic, flush=True)  # 1s

Logs

System Info

- 🤗 Diffusers version: 0.35.2
- Platform: Linux-4.18.0-193.6.3.el8_2.x86_64-x86_64-with-glibc2.28
- Running on Google Colab?: No
- Python version: 3.11.11
- PyTorch version (GPU?): 2.6.0+cu118 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.36.0
- Transformers version: 4.57.1
- Accelerate version: 1.11.0
- PEFT version: not installed
- Bitsandbytes version: not installed
- Safetensors version: 0.6.2
- xFormers version: not installed
- Accelerator: NVIDIA GeForce RTX 2080 Ti, 11264 MiB
NVIDIA GeForce RTX 2080 Ti, 11264 MiB
NVIDIA GeForce RTX 2080 Ti, 11264 MiB
NVIDIA GeForce RTX 2080 Ti, 11264 MiB
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

Who can help?

@sayakpaul @DN6

[francois-rozet]

The issue seems to be linked to memory mapping. Patching the following code to never use memory mapping solves the issue.

diffusers/src/diffusers/models/model_loading_utils.py

Lines 174 to 177 in dcfb18a

	if disable_mmap:
	return safetensors.torch.load(open(checkpoint_file, "rb").read())
	else:
	return safetensors.torch.load_file(checkpoint_file, device=map_location)

[sayakpaul]

Thanks for the issue.

1. Why would you want to low_cpu_mem_usage=False?
2. The benchmark seems rather faulty to me, as the CUDA allocator would have been warmed up by then, and transfer time would give us a false sense of speedup. Could we ensure the timing calculation is run in freshly initialized environments?

[francois-rozet]

Hi @sayakpaul

1. I tried with both low_cpu_memory_usage=True and False. It doesn't matter.

2. The example I provide is with two different scripts, run independently. Nevertheless, re-running the entire script in a for loop (including the loading of the pipeline) leads to the same result.

The issue is with tensors mapped from safetensors memory. They are super slow to transfer to CUDA for some reason. With large models (e.g. Flux), the same test is even worse.

[sayakpaul]

If that's the case, I think this is better discussed in the safetensors repository as it's not diffusers specific

[francois-rozet]

I agree to some extent. But memory mapping is a feature of safetensors, it doesn't mean that diffusers should use it. At the very least there should be a way to disable it for the end user, which there isn't: the keyword disable_mmap cannot be provided to from_pretrained.

Also could you confirm whether this bug reproduces on your machine(s) ? This could be due to the filesystem of my cluster.

[francois-rozet]

@sayakpaul I have created an issue (huggingface/safetensors#672) in the safetensors repo. I still think the option to disable memory mapping would be very useful.