Unverified Commit 5bb38586 authored by Sayak Paul's avatar Sayak Paul Committed by GitHub
Browse files

[Core] fix offload behaviour when device_map is enabled. (#7919)

fix offload behaviour when device_map is enabled.
parent ec9e8813
......@@ -363,7 +363,7 @@ class LoraLoaderMixin:
is_model_cpu_offload = False
is_sequential_cpu_offload = False
if _pipeline is not None:
if _pipeline is not None and _pipeline.hf_device_map is None:
for _, component in _pipeline.components.items():
if isinstance(component, nn.Module) and hasattr(component, "_hf_hook"):
if not is_model_cpu_offload:
......
......@@ -419,6 +419,7 @@ class TextualInversionLoaderMixin:
# 7.1 Offload all hooks in case the pipeline was cpu offloaded before make sure, we offload and onload again
is_model_cpu_offload = False
is_sequential_cpu_offload = False
if self.hf_device_map is None:
for _, component in self.components.items():
if isinstance(component, nn.Module):
if hasattr(component, "_hf_hook"):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment