invoke-ai / invoke-ai/InvokeAI
[bug]: First load of a model is extremely slow on ROCm — copies to VRAM from mmap-backed safetensors run at ~2 MB/s
- Dominant language
- Python
- Stars
- 28.2k
- Forks
- 3k
- Avg merge
- 6d 5h
- Merged PRs (30d)
- 19
Description
### Is there an existing issue for this problem?
- [x] I have searched the existing issues
### Install method
Invoke's Launcher
### Operating system
Linux
### GPU vendor
AMD (ROCm)
### GPU model
7900 xtx
### GPU VRAM
24
### Version number
6.13.7
### Browser
_No response_
### System Information
Environment: InvokeAI 6.13.7, Linux Mint 22, RX 7900 XTX (gfx1100), launcher-installed ROCm torch wheels.
### What happened
The first-ever load of a model onto the GPU takes minutes to hours (e.g. a 234 MB CLIP text encoder: 123 s; a 1.3 GB text_encoder_2: 193 s). One CPU core sits at 100% in driver ioctl calls while GPU and disk are idle. Later loads are fast.
Cause: _load_state_dict_with_fast_device_conversion (cached_model_with_partial_load.py) calls .to(device) on tensors that are still views into the memory-mapped safetensors file. On ROCm, host→device copies from file-backed memory are pathologically slow (~2 MB/s); the same copy from ordinary RAM runs at ~7 GB/s. A py-spy native dump shows the stall inside c10::hip::memcpy_and_sync → HSA runtime → ioctl. Reproducible outside Invoke:
python
from safetensors.torch import load_file
import torch
sd = load_file('')
sd_gpu = {k: v.to('cuda') for k, v in sd.items()} # crawls at single-digit MB/s
# but:
sd2 = {k: v.clone() for k, v in sd.items()} # detach from mmap (~1 s)
sd2_gpu = {k: v.to('cuda') for k, v in sd2.items()} # ~7 GB/s
Verified fix: clone CPU tensors before the device copy in _load_state_dict_with_fast_device_conversion (and the jit variant). Replace:
python
else:
state_dict[key] = state_dict[key].to(target_device)
with:
python
else:
t = state_dict[key]
if t.device.type == "cpu":
t = t.clone() # ROCm copies from file-backed memory are pathologically slow
state_dict[key] = t.to(target_device)
First loads drop from minutes/hours to seconds. Could be gated on torch.version.hip if the extra clone is unwanted on CUDA. Possibly the untracked second AMD problem mentioned in #9410.
### What you expected to happen
Expected it to be fast
### How to reproduce the problem
I guess have my hardware with this version of Invoke, maybe also the latest version of ROCm affects
### Additional context
_No response_
### Discord username
_No response_
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.