Repository navigation
device_map="balanced" drops all but the last CPU-overflow component, which then loads without offload hooks and breaks inference #15204
Description
Describe the bug
When more than one pipeline component doesn't fit on the accelerator, _assign_components_to_devices overwrites the "cpu" entry for each one (device_id_component_mapping["cpu"] = [component]) instead of appending to it. Only the last overflowing component stays in the device map.
The dropped components are missing from pipe.hf_device_map, so they're loaded with no device_map at all: plain CPU modules without the accelerate hooks that the other overflow components get. Inference then fails with a device mismatch as soon as one of them runs.
Reproduction
import torch
from diffusers import StableDiffusionPipeline
device = 0 if torch.cuda.is_available() else "mps"
# In fp16: unet 1719 MB, safety_checker 608 MB, text_encoder 246 MB, vae 167 MB.
# With 700 MB on the accelerator, only safety_checker fits; unet, text_encoder and vae all overflow to CPU.
pipe = StableDiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
dtype=torch.float16,
variant="fp16",
device_map="balanced",
max_memory={device: "700MB"},
)
print("hf_device_map:", pipe.hf_device_map)
for name in ["unet", "text_encoder", "vae", "safety_checker"]:
component = getattr(pipe, name)
hook = getattr(component, "_hf_hook", None)
print(
f"{name:15s} in hf_device_map: {pipe.hf_device_map.get(name, 'MISSING'):<8} accelerate hook: {type(hook).__name__ if hook else None}"
)
image = pipe("a photo of an astronaut", num_inference_steps=2, height=256, width=256).images[0]
print("inference OK")Context
No response
Logs
On an NVIDIA B200:
hf_device_map: {'vae': 'cpu', 'safety_checker': 0}
unet in hf_device_map: MISSING accelerate hook: None
text_encoder in hf_device_map: MISSING accelerate hook: None
vae in hf_device_map: cpu accelerate hook: AlignDevicesHook
safety_checker in hf_device_map: 0 accelerate hook: AlignDevicesHook
File "src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 962, in __call__
File "src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 423, in encode_prompt
...
RuntimeError: Expected all tensors to be on the same device, but got index is on cuda:0, different from other tensors on cpu (when checking argument in method wrapper_CUDA__index_select)
Same on Apple Silicon (MPS), where it fails with `RuntimeError: Passed CPU tensor to MPS op`.
Expected: `unet`, `text_encoder` and `vae` are all listed under `"cpu"` and get the same offload hooks, so inference runs.
The tiny budget is only there to make the repro cheap. The same thing happens whenever two or more components overflow, for example a large pipeline loaded with `device_map="balanced"` on a GPU smaller than both its transformer and its text encoder.System Info
- 🤗 Diffusers version: 0.42.0.dev0 (main @ 3b56af0)
- Platform: Linux-6.11.0-1016-nvidia-x86_64-with-glibc2.39
- Python version: 3.12.3
- PyTorch version (GPU?): 2.14.1+cu130 (True)
- Huggingface_hub version: 2.3.0
- Transformers version: 5.20.2
- Accelerate version: 1.15.0
- Accelerator: NVIDIA B200, 183359 MiB
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No
Also reproduced on macOS (Apple M5 Pro, MPS), torch 2.14.1.
Who can help?
I have a three-line fix with a regression test, and I'm happy to open a PR if that sounds good.
Activity
sayakpaul commentedon Nov 7, 2026
Why would you do a
device_mapon B200 for Stable Diffusion?charchit7 commentedon Nov 7, 2026
Fair point, I was initially testing this locally on MAC, happened to have few credits left for B200 for some research use-case so tested on it parallely.
sayakpaul commentedon Nov 7, 2026
Feel free to open a PR then. Thanks!