diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py index 10d3f0c245a1..c5a8f4fd8c76 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/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 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.