Skip to content

Commit 66e17e6

Browse files
committed
update
1 parent bccda52 commit 66e17e6

File tree

1 file changed

+9
-7
lines changed

1 file changed

+9
-7
lines changed

src/diffusers/hooks/group_offloading.py

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -405,12 +405,13 @@ def post_forward(self, module, output):
405405
# if the missing layers end up being executed in the future.
406406
if execution_order_module_names != self._layer_execution_tracker_module_names:
407407
unexecuted_layers = list(self._layer_execution_tracker_module_names - execution_order_module_names)
408-
logger.warning(
409-
"It seems like some layers were not executed during the forward pass. This may lead to problems when "
410-
"applying lazy prefetching with automatic tracing and lead to device-mismatch related errors. Please "
411-
"make sure that all layers are executed during the forward pass. The following layers were not executed:\n"
412-
f"{unexecuted_layers=}"
413-
)
408+
if not torch.compiler.is_compiling():
409+
logger.warning(
410+
"It seems like some layers were not executed during the forward pass. This may lead to problems when "
411+
"applying lazy prefetching with automatic tracing and lead to device-mismatch related errors. Please "
412+
"make sure that all layers are executed during the forward pass. The following layers were not executed:\n"
413+
f"{unexecuted_layers=}"
414+
)
414415

415416
# Remove the layer execution tracker hooks from the submodules
416417
base_module_registry = module._diffusers_hook
@@ -438,7 +439,8 @@ def post_forward(self, module, output):
438439
for i in range(num_executed - 1):
439440
name1, _ = self.execution_order[i]
440441
name2, _ = self.execution_order[i + 1]
441-
logger.debug(f"Applying lazy prefetch group offloading from {name1} to {name2}")
442+
if not torch.compiler.is_compiling():
443+
logger.debug(f"Applying lazy prefetch group offloading from {name1} to {name2}")
442444
group_offloading_hooks[i].next_group = group_offloading_hooks[i + 1].group
443445
group_offloading_hooks[i].next_group.onload_self = False
444446

0 commit comments

Comments
 (0)