How about forcing the first and last block on device when groupoffloading is used?

[seed93]

Is your feature request related to a problem? Please describe.
When group offloading is enabled, the offload and onload cannot be streamed between steps and this is really a big time comsuming problem.

Describe the solution you'd like.
Is it possible to add an option that could make the first and last block forced on device to avoid offload and onload?

@a-r-r-o-w Could you please give some help? Thanks so much.

[seed93]

I have tried to make it a circle here: https://github.com/huggingface/diffusers/blob/v0.33.1/src/diffusers/hooks/group_offloading.py#L319
However it seems not work.

[a-r-r-o-w]

Hi @seed93, forcing particular block(s), or other internal layers, to remain on the device should be possible. I'll try to prototype something later this week. I'd imagine denoting what layers to keep on device will have to be determined via a regex, so that ModuleGroup's of remaining modules can be formed easily. Trying to prefetch the first layer from the last layer might also be possible 🤔

[seed93]

Thank you so much. I tried a circular prefetch as I said on the last comment but it seems not work. I checked this using nsys and the memcpy between steps still exists. As the image shows below, the offload and onload were captured by pre_forward and post_forward of GroupOffloadingHook. <img width="1790" height="989" alt="Image" src="https://github.com/user-attachments/assets/8eb24d6b-b459-4625-816f-c4d3305f53e4" />

[seed93]

diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py
index 36abf3c41..d1397e237 100644
--- a/src/diffusers/hooks/group_offloading.py
+++ b/src/diffusers/hooks/group_offloading.py
@@ -435,11 +435,12 @@ class LazyPrefetchGroupOffloadingHook(ModelHook):
             base_module_group_offloading_hook.next_group = group_offloading_hooks[0].group
             base_module_group_offloading_hook.next_group.onload_self = False
 
-        for i in range(num_executed - 1):
+        for i in range(num_executed):
             name1, _ = self.execution_order[i]
-            name2, _ = self.execution_order[i + 1]
+            next_i = (i + 1) % num_executed
+            name2, _ = self.execution_order[next_i]
             logger.debug(f"Applying lazy prefetch group offloading from {name1} to {name2}")
-            group_offloading_hooks[i].next_group = group_offloading_hooks[i + 1].group
+            group_offloading_hooks[i].next_group = group_offloading_hooks[next_i].group
             group_offloading_hooks[i].next_group.onload_self = False
 
         return output