From 14cd3f9519951319b0153e9be59874158b4e006f Mon Sep 17 00:00:00 2001 From: Daniel Gu Date: Fri, 4 Sep 2026 07:14:53 +0200 Subject: [PATCH 1/3] Add a group offloading introspection helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Group offloading leaves no way to ask what grouping it installed: the structure lives on the `GroupOffloadingHook` in each participating submodule's `_diffusers_hook` registry, several submodules share one `ModuleGroup`, and the groups hold module objects rather than names. So a device-mismatch failure names only the layer that tripped over it, never the group whose weights were never onloaded. Add `_get_group_offload_summary` next to the existing introspection cluster (`_get_top_level_group_offload_hook` and friends). It dedupes the groups and resolves the module objects back to qualified names, reporting each group against the module whose `forward` actually brings its weights over. That is not always the group's own leader: with `use_stream=True`, lazy prefetch sets `onload_self = False` and the group is onloaded by whichever group names it as `next_group`, so the prefetch chain is read off the per-module hooks rather than the shared group. That chain is wired only once a module's `forward` has completed once, which the docstring calls out — a module that had not run yet when the failure hit reports every group as onloading itself, and reading that as "this module is not prefetching" would be wrong. A group can reach modules that are not under the one being walked, so calling this on a submodule of an offloaded model — one component of a pipeline, say — finds the group its siblings belong to and cannot name them. Those are reported as `` rather than raising, since a debugging aid that dies when pointed at the thing being debugged is no use. It stays private. What a component reports would shift the day it gains `_group_offload_block_modules`, and reporting the grouping is all this can soundly do — whether a component is safe to offload at block level depends on which entry points the pipeline drives it through, which the module graph alone does not say. A component can group normally and still break: `PriorTransformer` splits its `transformer_blocks` into per-block groups yet reads `clip_mean` / `clip_std` from the root group in `post_process_latents`, and `AutoencoderKL` groups entirely into its `down_blocks` / `up_blocks` while `decode` bypasses `forward`. Use it in the pipeline group offload tests, where a `RuntimeError` from the offloaded run is now re-raised with the grouping of every offloaded component attached. Co-Authored-By: Claude Opus 5 (1M context) --- src/diffusers/hooks/group_offloading.py | 55 +++++++++++++++++++++++++ tests/pipelines/testing_utils/memory.py | 18 ++++++-- 2 files changed, 70 insertions(+), 3 deletions(-) diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py index 10d3f0c245a1..7ef77f1659ee 100644 --- a/src/diffusers/hooks/group_offloading.py +++ b/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 ``. + """ + name_of = {id(submodule): submodule_name or "" for submodule_name, submodule in module.named_modules()} + outside = "" + 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 "" + 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 diff --git a/tests/pipelines/testing_utils/memory.py b/tests/pipelines/testing_utils/memory.py index 7f811e2ee4cc..18858777dfb1 100644 --- a/tests/pipelines/testing_utils/memory.py +++ b/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. From e322f20e8369bf1ee967741c4ae097cbb5fb9c16 Mon Sep 17 00:00:00 2001 From: Daniel Gu Date: Fri, 4 Sep 2026 09:06:11 +0200 Subject: [PATCH 2/3] Test the group offloading summary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_get_group_offload_summary` is only read once something has already gone wrong, so a regression in it is invisible: a summary is still produced, it is just wrong, and it sends the reader to the wrong module. Nothing else in the suite would notice. Cover what it claims rather than how it formats it — the number of groups, and which module's `forward` brings each one over — asserting against the groups read straight off the hooks. Only "no group offloading applied" and "prefetched by" are matched as text, so the layout stays free to change. The two cases that carry their weight are prefetching, which is reported only once the first forward has wired the chain, and summarizing a submodule whose group reaches outside it. Both were live bugs in the helper: reverting either fix fails exactly its own test and no other. Co-Authored-By: Claude Opus 5 (1M context) --- tests/hooks/test_group_offloading.py | 104 ++++++++++++++++++++++++++- 1 file changed, 103 insertions(+), 1 deletion(-) diff --git a/tests/hooks/test_group_offloading.py b/tests/hooks/test_group_offloading.py index a903186aa6b4..ac97eef2e11d 100644 --- a/tests/hooks/test_group_offloading.py +++ b/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 From b967036ecbf942f2b5f214aa0c6401f3438f3436 Mon Sep 17 00:00:00 2001 From: Daniel Gu Date: Fri, 4 Sep 2026 09:53:23 +0200 Subject: [PATCH 3/3] make style and make quality --- src/diffusers/hooks/group_offloading.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py index 7ef77f1659ee..c5a8f4fd8c76 100644 --- a/src/diffusers/hooks/group_offloading.py +++ b/src/diffusers/hooks/group_offloading.py @@ -1032,9 +1032,9 @@ def _get_group_offload_summary(module: torch.nn.Module, name: str = "") -> str: 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. + 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 ``.