Skip to content

set is_sequential_cpu_offload = True only when some component is on cpu and has AlignDevicesHook simultaneously #8750

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

set is_sequential_cpu_offload = True only when some component is on cpu and has AlignDevicesHook simultaneously #8750

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7ee9465
modify _optionally_disable_offloading()
zhangvia Jul 1, 2024
0b86e90
Merge branch 'main' of https://github.com/zhangvia/diffusers into mod…
zhangvia Jul 1, 2024
ecc3032
modify all is_sequential_cpu_offload
zhangvia Jul 4, 2024
eebed22
Merge branch 'main' of https://github.com/zhangvia/diffusers into mod…
zhangvia Jul 4, 2024
9c87a4a
fix conflict
Jul 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 6 additions & 5 deletions src/diffusers/loaders/textual_inversion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -419,15 +419,16 @@ def load_textual_inversion(
if isinstance(component, nn.Module):
if hasattr(component, "_hf_hook"):
is_model_cpu_offload = isinstance(getattr(component, "_hf_hook"), CpuOffload)
is_sequential_cpu_offload = (
is_sequential_cpu_offload = component.device.type=='cpu' and (
isinstance(getattr(component, "_hf_hook"), AlignDevicesHook)
or hasattr(component._hf_hook, "hooks")
and isinstance(component._hf_hook.hooks[0], AlignDevicesHook)
)
logger.info(
"Accelerate hooks detected. Since you have called `load_textual_inversion()`, the previous hooks will be first removed. Then the textual inversion parameters will be loaded and the hooks will be applied again."
)
remove_hook_from_module(component, recurse=is_sequential_cpu_offload)
if is_model_cpu_offload or is_sequential_cpu_offload:
logger.info(
"Pipeline offload enabled and Accelerate hooks detected. Since you have called `load_textual_inversion()`, the previous offloadhooks will be first removed. Then the textual inversion parameters will be loaded and the hooks will be applied again."
)
remove_hook_from_module(component, recurse=is_sequential_cpu_offload)

# 7.2 save expected device and dtype
device = text_encoder.device
Expand Down
12 changes: 6 additions & 6 deletions src/diffusers/loaders/unet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -373,16 +373,16 @@ def _optionally_disable_offloading(cls, _pipeline):
if not is_model_cpu_offload:
is_model_cpu_offload = isinstance(component._hf_hook, CpuOffload)
if not is_sequential_cpu_offload:
is_sequential_cpu_offload = (
is_sequential_cpu_offload = component.device.type == "cpu" and (
isinstance(component._hf_hook, AlignDevicesHook)
or hasattr(component._hf_hook, "hooks")
and isinstance(component._hf_hook.hooks[0], AlignDevicesHook)
)

logger.info(
"Accelerate hooks detected. Since you have called `load_lora_weights()`, the previous hooks will be first removed. Then the LoRA parameters will be loaded and the hooks will be applied again."
)
remove_hook_from_module(component, recurse=is_sequential_cpu_offload)
if is_model_cpu_offload or is_sequential_cpu_offload:
logger.info(
"Pipeline offload enabled and Accelerate hooks detected. Since you have called `load_lora_weights()`, the previous offload hooks will be first removed. Then the LoRA parameters will be loaded and the hooks will be applied again."
)
remove_hook_from_module(component, recurse=is_sequential_cpu_offload)
Comment on lines +381 to +385
Copy link
Member

@sayakpaul sayakpaul Nov 20, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks alright to me.


return (is_model_cpu_offload, is_sequential_cpu_offload)

Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/pipelines/pipeline_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -376,7 +376,7 @@ def module_is_sequentially_offloaded(module):
if not is_accelerate_available() or is_accelerate_version("<", "0.14.0"):
return False

return hasattr(module, "_hf_hook") and (
return hasattr(module, "_hf_hook") and hasattr(module,'device') and module.device.type == "cpu" and (
Copy link
Member

@sayakpaul sayakpaul Nov 20, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why the expansion?

Copy link
Author

@zhangvia zhangvia Mar 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

manually add a hook to the model and attempting to move it to another GPU does not mean that sequential_cpu_offload is enabled. but the original code will still return true

isinstance(module._hf_hook, accelerate.hooks.AlignDevicesHook)
or hasattr(module._hf_hook, "hooks")
and isinstance(module._hf_hook.hooks[0], accelerate.hooks.AlignDevicesHook)
Expand Down