Add Group Offloading Summary Utility

src/diffusers/hooks/group_offloading.py

@@ -1024,6 +1024,61 @@ def _get_group_onload_device(module: torch.nn.Module) -> torch.device:
     raise ValueError("Group offloading is not enabled for the provided module.")
 
 
+def _get_group_offload_summary(module: torch.nn.Module, name: str = "") -> str:
+    """Render the offload grouping of a module as a short block, for use in assertion messages.
+
+    The grouping lives on the `GroupOffloadingHook` in each participating submodule's `_diffusers_hook` registry, and
+    several submodules share one `ModuleGroup`, so this dedupes the groups and resolves the module objects back to
+    qualified names. Each group is reported against the module whose `forward` actually brings its weights over: a
+    group with `onload_self=False` is onloaded by whichever group names it as `next_group`, not by its own leader.
+
+    The grouping is reported as of the moment this is called. With `use_stream=True` the prefetch chain is wired only
+    once a module's `forward` has completed once, so a module that has not run yet reports every group as onloading
+    itself, exactly as it would without a stream.
+
+    A group can span modules that are not under the one being walked — calling this on a submodule of an offloaded
+    model reaches the group its siblings belong to — so those are reported as `<outside this module>`.
+    """
+    name_of = {id(submodule): submodule_name or "<root>" for submodule_name, submodule in module.named_modules()}
+    outside = "<outside this module>"
+    hooks, groups, seen = [], [], set()
+
+    for submodule in module.modules():
+        registry = getattr(submodule, "_diffusers_hook", None)
+        hook = registry.get_hook(_GROUP_OFFLOADING) if registry is not None else None
+        if hook is None:
+            continue
+        # The prefetch chain is wired onto individual hooks, while the modules of one group share a group object.
+        # Keep every hook for the chain, and let only the first hook of a group contribute a line.
+        hooks.append(hook)
+        if id(hook.group) not in seen:
+            seen.add(id(hook.group))
+            groups.append(hook.group)
+
+    if not groups:
+        return f"{name or type(module).__name__}: no group offloading applied."
+
+    prefetched_by = {
+        id(hook.next_group): name_of.get(id(hook.group.onload_leader), outside)
+        for hook in hooks
+        if hook.next_group is not None
+    }
+
+    lines = [f"{name or type(module).__name__}: {len(groups)} offload group(s)"]
+    for group in groups:
+        members = ", ".join(name_of.get(id(member), outside) for member in group.modules) or "<no modules>"
+        prefetcher = None if group.onload_self else prefetched_by.get(id(group))
+        onloaded_by = (
+            f"prefetched by {prefetcher!r}"
+            if prefetcher
+            else f"onloaded by {name_of.get(id(group.onload_leader), outside)!r}"
+        )
+        lines.append(
+            f"  {onloaded_by} forward: [{members}] (+{len(group.parameters)} params, +{len(group.buffers)} buffers)"
+        )
+    return "\n".join(lines)
+
+
 def _compute_group_hash(group_id):
     hashed_id = hashlib.sha256(group_id.encode("utf-8")).hexdigest()
     # first 16 characters for a reasonably short but unique name

tests/hooks/test_group_offloading.py

@@ -20,7 +20,8 @@
 import torch
 
 from diffusers import AutoencoderKL
-from diffusers.hooks import HookRegistry, ModelHook
+from diffusers.hooks import HookRegistry, ModelHook, apply_group_offloading
+from diffusers.hooks.group_offloading import _GROUP_OFFLOADING, _get_group_offload_summary
 from diffusers.models import ModelMixin
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.utils import logging as diffusers_logging
@@ -692,3 +693,104 @@ def test_conditional_modules_with_stream(self, offload_type: str):
             assert torch.allclose(out_ref_no_opt2, out_no_opt2, atol=1e-5), (
                 f"[{offload_type}] Outputs do not match on third pass (back to no optional_input)."
             )
+
+
+class TestGroupOffloadSummary:
+    """`_get_group_offload_summary` renders the installed grouping, for failure messages and manual inspection.
+
+    Its output is only read once something has already gone wrong, so a regression is invisible — a summary is
+    still produced, it is just wrong, and a reader is sent to the wrong module. These pin what it claims (how many
+    groups, and which module's `forward` brings each one over) and deliberately not how it formats them.
+    """
+
+    in_features = 64
+    hidden_features = 256
+    out_features = 64
+    num_layers = 4
+
+    def get_model(self):
+        torch.manual_seed(0)
+        return DummyModel(
+            in_features=self.in_features,
+            hidden_features=self.hidden_features,
+            out_features=self.out_features,
+            num_layers=self.num_layers,
+        )
+
+    @staticmethod
+    def installed_groups(module):
+        """Read the groups straight off the hooks, as an oracle for what the summary should describe."""
+        groups, seen = [], set()
+        for submodule in module.modules():
+            registry = getattr(submodule, "_diffusers_hook", None)
+            hook = registry.get_hook(_GROUP_OFFLOADING) if registry is not None else None
+            if hook is not None and id(hook.group) not in seen:
+                seen.add(id(hook.group))
+                groups.append(hook.group)
+        return groups
+
+    @staticmethod
+    def group_lines(summary):
+        return summary.splitlines()[1:]
+
+    @staticmethod
+    def reported_members(line):
+        """The members one summary line lists, so a test can check that none were dropped."""
+        return line[line.index("[") + 1 : line.rindex("]")].split(", ")
+
+    def test_reports_that_nothing_is_offloaded_when_offloading_is_not_applied(self):
+        assert "no group offloading applied" in _get_group_offload_summary(self.get_model())
+
+    def test_reports_one_line_per_installed_group(self):
+        model = self.get_model()
+        apply_group_offloading(
+            model,
+            onload_device=torch.device("cpu"),
+            offload_device=torch.device("cpu"),
+            offload_type="block_level",
+            num_blocks_per_group=1,
+        )
+        summary = _get_group_offload_summary(model)
+
+        assert len(self.group_lines(summary)) == len(self.installed_groups(model))
+        # Without a stream there is no prefetch chain, so every group onloads itself.
+        assert "prefetched by" not in summary
+
+    @pytest.mark.skipif(
+        torch.device(torch_device).type not in ["cuda", "xpu"],
+        reason="Test requires a CUDA or XPU device.",
+    )
+    def test_reports_prefetching_only_once_the_chain_is_wired(self):
+        model = self.get_model()
+        model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1, use_stream=True)
+
+        # The chain is wired by the lazy prefetch hook at the end of the first forward, so until then the groups are
+        # indistinguishable from the streamless case.
+        assert "prefetched by" not in _get_group_offload_summary(model)
+
+        model(torch.randn((4, self.in_features)).to(torch_device))
+
+        prefetched = sum(1 for group in self.installed_groups(model) if not group.onload_self)
+        assert prefetched > 0, "the first forward should have wired a prefetch chain"
+        summary = _get_group_offload_summary(model)
+        assert sum(1 for line in self.group_lines(summary) if "prefetched by" in line) == prefetched
+
+    def test_summarizes_a_submodule_whose_group_reaches_outside_it(self):
+        # With more than one block per group, a block's group holds its siblings too. Summarizing that block alone
+        # cannot name them, and a user inspecting one component should still get a summary rather than a KeyError.
+        model = self.get_model()
+        apply_group_offloading(
+            model,
+            onload_device=torch.device("cpu"),
+            offload_device=torch.device("cpu"),
+            offload_type="block_level",
+            num_blocks_per_group=3,
+        )
+        summary = _get_group_offload_summary(model.blocks[1])
+
+        groups = self.installed_groups(model.blocks[1])
+        lines = self.group_lines(summary)
+        assert len(lines) == len(groups)
+        # The group spans blocks 0-2, so from blocks[1] two of the three cannot be named — but all three are still
+        # its members, and a summary that drops the two it cannot name understates the group.
+        assert len(self.reported_members(lines[0])) == len(groups[0].modules) == 3

tests/pipelines/testing_utils/memory.py

@@ -18,6 +18,7 @@
 
 from diffusers import DiffusionPipeline
 from diffusers.hooks import apply_group_offloading
+from diffusers.hooks.group_offloading import _get_group_offload_summary
 
 from ...testing_utils import (
     assert_tensors_close,
@@ -336,6 +337,8 @@ def _enable_group_offload_on_components(self, pipe, **group_offloading_kwargs):
         if group_offloading_kwargs.get("use_stream"):
             self._assert_streams_took_effect(pipe, offload_names)
 
+        return offload_names
+
     def _assert_streams_took_effect(self, pipe, offload_names):
         """Guard against `use_stream=True` silently becoming a no-op.
 
@@ -367,9 +370,18 @@ def _run_group_offload_inference(self, base_pipe_output, expected_max_difference
         # difference under test. It stays on CPU here — the components are placed as they are hooked.
         pipe = self.create_pipe()
         self._skip_if_group_offloading_unsupported(pipe)
-        self._enable_group_offload_on_components(pipe, **group_offloading_kwargs)
-
-        assert_tensors_close(self.run_pipe(pipe), base_pipe_output, atol=expected_max_difference, rtol=1e-5, msg=msg)
+        offload_names = self._enable_group_offload_on_components(pipe, **group_offloading_kwargs)
+
+        try:
+            output = self.run_pipe(pipe)
+        except RuntimeError as error:
+            # A device mismatch here means some group's weights were never onloaded, and the bare error names only
+            # the layer that tripped over it. Say which module's `forward` brings each group over — for a
+            # prefetched group that is the group ahead of it in the chain, not its own leader.
+            summaries = "\n".join(_get_group_offload_summary(getattr(pipe, name), name) for name in offload_names)
+            raise RuntimeError(f"{error}\n\nOffload grouping under test:\n{summaries}") from error
+
+        assert_tensors_close(output, base_pipe_output, atol=expected_max_difference, rtol=1e-5, msg=msg)
 
     def _skip_if_streams_unsupported(self):
         # `apply_group_offloading` raises rather than degrading when `use_stream=True` has nowhere to put a stream.