@@ -405,12 +405,13 @@ def post_forward(self, module, output):
405
405
# if the missing layers end up being executed in the future.
406
406
if execution_order_module_names != self ._layer_execution_tracker_module_names :
407
407
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
+ )
414
415
415
416
# Remove the layer execution tracker hooks from the submodules
416
417
base_module_registry = module ._diffusers_hook
@@ -438,7 +439,8 @@ def post_forward(self, module, output):
438
439
for i in range (num_executed - 1 ):
439
440
name1 , _ = self .execution_order [i ]
440
441
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 } " )
442
444
group_offloading_hooks [i ].next_group = group_offloading_hooks [i + 1 ].group
443
445
group_offloading_hooks [i ].next_group .onload_self = False
444
446
0 commit comments