diff --git a/src/diffusers/loaders/lora_base.py b/src/diffusers/loaders/lora_base.py index 0009f04df9e2..db37290a233a 100644 --- a/src/diffusers/loaders/lora_base.py +++ b/src/diffusers/loaders/lora_base.py @@ -462,14 +462,16 @@ def _func_optionally_disable_offloading(_pipeline): if not isinstance(component, nn.Module): continue is_group_offload = is_group_offload or _is_group_offload_enabled(component) - if not hasattr(component, "_hf_hook"): - continue - is_model_cpu_offload = is_model_cpu_offload or isinstance(component._hf_hook, CpuOffload) - is_sequential_cpu_offload = is_sequential_cpu_offload or ( - isinstance(component._hf_hook, AlignDevicesHook) - or hasattr(component._hf_hook, "hooks") - and isinstance(component._hf_hook.hooks[0], AlignDevicesHook) - ) + for module in component.modules(): + if not hasattr(module, "_hf_hook"): + continue + hooks = getattr(module._hf_hook, "hooks", (module._hf_hook,)) + is_model_cpu_offload = is_model_cpu_offload or any( + isinstance(hook, CpuOffload) for hook in hooks + ) + is_sequential_cpu_offload = is_sequential_cpu_offload or any( + isinstance(hook, AlignDevicesHook) and hook.offload for hook in hooks + ) if is_sequential_cpu_offload or is_model_cpu_offload: logger.info( diff --git a/tests/lora/test_lora_loader_utils.py b/tests/lora/test_lora_loader_utils.py index d79fd77089b0..525b86239ca2 100644 --- a/tests/lora/test_lora_loader_utils.py +++ b/tests/lora/test_lora_loader_utils.py @@ -26,6 +26,7 @@ from diffusers.loaders.peft import PeftAdapterMixin from diffusers.models.modeling_utils import ModelMixin from diffusers.utils.import_utils import is_peft_available +from diffusers.utils.testing_utils import require_accelerate from ..testing_utils import CaptureLogger, require_peft_backend @@ -118,6 +119,35 @@ def test_local_directory_with_multiple_files_warns_and_uses_first(tmp_path, monk assert "contains more than one weights file" in cap_logger.out +@require_accelerate +def test_alignment_hook_does_not_enable_sequential_cpu_offload(): + from accelerate.hooks import AlignDevicesHook, add_hook_to_module + + component = torch.nn.Linear(1, 1) + add_hook_to_module(component, AlignDevicesHook(execution_device="cpu")) + pipeline = Mock(hf_device_map=None, components={"model": component}) + + offload_state = lora_base._func_optionally_disable_offloading(pipeline) + + assert offload_state == (False, False, False) + assert hasattr(component, "_hf_hook") + + +@require_accelerate +def test_sequential_cpu_offload_is_detected_and_disabled(): + from accelerate import cpu_offload + + component = torch.nn.Sequential(torch.nn.Linear(1, 1)) + cpu_offload(component, execution_device=torch.device("cuda")) + pipeline = Mock(hf_device_map=None, components={"model": component}) + + offload_state = lora_base._func_optionally_disable_offloading(pipeline) + + assert offload_state == (False, True, False) + assert not hasattr(component, "_hf_hook") + assert not hasattr(component[0], "_hf_hook") + + @require_peft_backend def test_unfuse_lora_partial_components_keeps_merged_adapters_in_sync(): """Regression test for gh-14214.