What if non_blocking=True?

Even with non_blocking=True, if a previous group onloaded this one on a side stream, we need a sync before the default stream uses the weights or we risk reading half-copied tensors. I’ve limited the sync to the record_stream=False case, when record_stream=True the tensors are tied to the consumer stream so we can safely skip the sync.

Why is this being deleted?

Rest of the diffs in this testing script are a bit difficult to follow honestly. Could we keep this cleaner?

Thanks for pointing this out. The class was not intentionally deleted. From the git history, this shows up as removed due to branch history / rebase artifacts while integrating changes (rather than a deliberate change to the test itself), which makes the diff noisier than it should be. I’m cleaning this up now: I’ll restore that block and reorganize the commits so the test diffs are more focused/atomic and easier to review.

Why do we have to set the default of default_stream?

I made it optional because the non-stream path calls _process_tensors_from_modules without a stream, there is nothing to record in that case, and record_stream is gated. None is a safety net for the record call, and it saves passing a placeholder from those call sites. If you prefer the stricter signature, I can keep it required and pass None explicitly where we don’t use streams. please do correct me thru my understanding if this is required to change

Let's stick to the existing implementation in this case i.e., a stricter signature.

Let's stick to a stricter signature.

This seems unnecessary. Explain?

originally added the empty root hook to tag the top module as offloaded when everything else was matched, but it did not change behaviour, the child hooks already mark the model as group-offloaded and the guardrails rely on those. It just added an empty group and potential extra files, so have removed it to simplify. Functionally nothing depends on it.

Let's remove this if it isn't relevant to pinning.

What's happening here?

It is the same thing as above, we prefix group_id with the parent name to avoid collisions (ids) when recursing into block_modules. Root stays empty to preserve existing ids, the prefix only appears when descending into children.

done

Let's stick to the existing implementation in this case i.e., a stricter signature.

Let's stick to a stricter signature.

Can you clarify this scenario in the comments as well? And provide a small example that justifies this change?

(nit): let's not get rid of the important comments.

done

Could I get a clarification on why this condition needs to be modified?

did not change the condition, those same four checks were already present in both branches. I consolidated them into one place to avoid duplication. We still only sync when the group did not onload itself, we are using a stream, there is no pending prefetch, and record_stream is not handling lifetime tracking.

This part of the diff reads very confusing to me and hence, a bit hard to confidently review. It seems to me, post_forward() was just brought up, _send_kwargs_to_device() was added (and I am not sure why) amongst other things. Possible to have a cleaner diff?

Thanks for the flag, I simplified the diff ( have removed _send_kwargs_to_device and the kwargs handling is back inline in pre_forward as before )

Are we just documenting pin_groups here? If so, we should remove that from here as apply_group_offloading() should already cover it:

https://github.com/bconstantine/diffusers/blob/335dca80fb2cdbdab7c2daa298c3ca934d7107b6/src/diffusers/models/modeling_utils.py#L542

removed redundant docstring

Let's also add a comment on how these modules were chosen to be included here.

sure, added comment

What is the test that's failing? If this change is to make a test pass and it is not related to pinning, let's revert it please. We can create a dedicated PR to address that issue.

Agreed. This eager disk offload was added to satisfy a test unrelated to pinning. have reverted it so offload remains lazy. If we still need eager materialization, I will open a separate PR scoped to that issue

Instead of using this loop to check all tensors, it would be better to introduce an _is_offloaded attribute to the ModuleGroup which changes to True/False during onload/offload. Use that to check if the group is one device or not.

? Change not needed I think? If block_modules is an empty set this won't evaluate to true.

reverted this check back to if name in block_modules since block_modules is already a set and the extra guard is redundant

Let's remove this if it isn't relevant to pinning.

What is the failing check that requires this change? It doesn't seem related to pinning?

My understanding here is that this logic is responsible for onloading/pinning the groups during the first inference pass. The first pass of offloading with streams will always be slow since we need to figure out execution order. The complexity introduced with this change is probably not worth the incremental performance benefit.

I would suggest we factor this into simpler

  def post_forward(self, module, output):
      # existing prefetch setup code unchanged

      if self.pin_groups is not None:
          self._mark_pinned_groups(group_offloading_hooks)

      return output

After marking, subsequent passes should just skip offloading. The group remains pinned after the second pass.

Also, what happens with pinning if we don't use streams? The prefetch hook doesn't run in that case.

refactored the pinning into a _mark_pinned_groups helper called from post_forward after the prefetch setup, matching the simplified structure you suggested. Also added a warning + doc note, pin_groups is only supported with use_stream=True, it is ignored otherwise since the lazy prefetch hook doesn’t run without streams