From c5872769dee8b78fb46602cb7dede191d8a30c0e Mon Sep 17 00:00:00 2001 From: Daniel Gu Date: Fri, 28 Aug 2026 07:18:16 +0200 Subject: [PATCH 1/3] tests: split block/leaf group offloading and derive the offloaded components `GroupOffloadTesterMixin.test_group_offloading_inference` ran both offload levels in one test body, so a pipeline that failed at one level had to skip both. Split it into `test_group_offloading_inference_block_level` and `test_group_offloading_inference_leaf_level`, sharing the helpers the test body used to define inline, and compare against the class-scoped `base_pipe_output` rather than rebuilding a baseline per level. The set of components to offload was a hardcoded list of eight names, so a pipeline with a component under any other name had it silently left on CPU for the forward pass to trip over. Derive the set instead: every `nn.Module` component is offloaded unless the config lists it in `group_offloading_leaf_level_exclude_modules`, the new `group_offloading_exclude_modules`, or `group_offloading_onload_component_names`. A name in either exclusion list that matches no component on the pipeline fails as the typo it is. Ideogram4's `unconditional_transformer` was one of the silently dropped components, which is why its group offload test was skipped; it now needs no declaration at all, and `text_encoder` picks up block-level coverage it never had. LTX2's skip goes the same way, with `audio_vae` declared alongside the other VAEs the tests keep on the accelerator. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/references/testing.md | 3 +- .../ideogram4/test_pipeline_ideogram4.py | 20 +- tests/pipelines/ltx2/test_ltx2.py | 16 +- tests/pipelines/testing_utils/common.py | 22 ++- tests/pipelines/testing_utils/memory.py | 171 +++++++++--------- 5 files changed, 120 insertions(+), 112 deletions(-) diff --git a/.ai/references/testing.md b/.ai/references/testing.md index 9dd0fd7280c7..e942d90b4738 100644 --- a/.ai/references/testing.md +++ b/.ai/references/testing.md @@ -30,7 +30,8 @@ Follow the style introduced in [#14113](https://github.com/huggingface/diffusers - `MemoryTesterMixin` — CPU offload, group offload, layerwise casting. - Cache mixins — `PyramidAttentionBroadcastTesterMixin`, `FasterCacheTesterMixin`, `FirstBlockCacheTesterMixin`, `TaylorSeerCacheTesterMixin`, `MagCacheTesterMixin`. Guidance-distilled models override the cache config (e.g. `FASTER_CACHE_CONFIG = {... "is_guidance_distilled": True}`). Don't introduce caching related tests in the first iteration. These tests are added on a case-by-case basis. - In the first pass, just add tests related to `PipelineTesterMixin` and `MemoryTesterMixin`. -- **Declare a component that can't be offloaded — don't hand-write a skip.** Leaf-level offloading hooks only the supported leaf types (`nn.Linear`, `nn.Conv*`, `nn.Embedding` — see `_GO_LC_SUPPORTED_PYTORCH_LAYERS` in `src/diffusers/hooks/_common.py`) and onloads each on its own `forward`, so any code that reads a leaf's `.weight` instead of calling the leaf bypasses that leaf's hook and computes against offloaded weights. Which fix applies depends on who owns the component. For a diffusers model, set `_supports_group_offloading = False` on the `ModelMixin` subclass (as `HunyuanDiT2DModel` does) — both offload mixins honor the flag and skip themselves, so the gap is declared on the model instead of buried in a test file. For a third-party component you can't annotate, such as a `transformers` encoder, list it in `group_offloading_leaf_level_exclude_modules` on the config class; `enable_group_offload` keeps excluded components on the accelerator, so every other component stays covered — including the VAE, which the component-scoped `test_group_offloading_inference` leaves out. Block-level offloading is usually unaffected, hence the level in the name — a component that fails at both levels does need a skip. +- **Declare a component that can't be offloaded — don't hand-write a skip.** Leaf-level offloading hooks only the supported leaf types (`nn.Linear`, `nn.Conv*`, `nn.Embedding` — see `_GO_LC_SUPPORTED_PYTORCH_LAYERS` in `src/diffusers/hooks/_common.py`) and onloads each on its own `forward`, so any code that reads a leaf's `.weight` instead of calling the leaf bypasses that leaf's hook and computes against offloaded weights. Which fix applies depends on who owns the component. For a diffusers model, set `_supports_group_offloading = False` on the `ModelMixin` subclass (as `HunyuanDiT2DModel` does) — both offload mixins honor the flag and skip themselves, so the gap is declared on the model instead of buried in a test file. For a third-party component you can't annotate, such as a `transformers` encoder, list it in `group_offloading_leaf_level_exclude_modules` on the config class; excluded components are kept on the accelerator, so every other component stays covered. Block-level offloading is usually unaffected, hence the level in the name, and `test_group_offloading_inference_block_level` still covers the component — a component that fails at both levels goes in `group_offloading_exclude_modules` instead, with a comment saying why. + - **Every `nn.Module` component is group offloaded unless a config list names it** — `group_offloading_leaf_level_exclude_modules`, `group_offloading_exclude_modules`, or `group_offloading_onload_component_names` (the VAE and friends, kept on the accelerator because tiling breaks stream tracing). A pipeline that adds a second denoiser or an extra encoder therefore gets it exercised without touching the shared mixin, and losing coverage takes naming the component. A name in an exclusion list that matches no component on the pipeline fails the test as a typo. - `torch.nn.MultiheadAttention` is the common instance: it passes `self.out_proj.weight` straight to `torch.nn.functional.multi_head_attention_forward` instead of calling `self.out_proj`, so the hook on `out_proj` never fires. `SiglipVisionModel`'s attention pooling head wraps one — see `tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py`, whose `image_encoder` is excluded for this reason. - `HunyuanDiTAttentionPool` (`src/diffusers/models/embeddings.py`) shows the same failure without an MHA module: a plain `nn.Module` that hands its `q_proj` / `k_proj` / `v_proj` / `c_proj` weights to `torch.nn.functional.multi_head_attention_forward`, so all four projections stay offloaded rather than just one. `HunyuanDiT2DModel` opts out of group offloading entirely with `_supports_group_offloading = False`. - Before adding a skip or an exclusion, confirm the failure still reproduces — several existing skips are stale, having outlived the upstream cause. diff --git a/tests/pipelines/ideogram4/test_pipeline_ideogram4.py b/tests/pipelines/ideogram4/test_pipeline_ideogram4.py index caea8ca36edb..c8a1b986217f 100644 --- a/tests/pipelines/ideogram4/test_pipeline_ideogram4.py +++ b/tests/pipelines/ideogram4/test_pipeline_ideogram4.py @@ -48,8 +48,10 @@ class Ideogram4PipelineTesterConfig(BasePipelineTesterConfig): required_input_params_in_call_signature = frozenset(["prompt", "height", "width", "guidance_scale"]) batch_input_params = frozenset(["prompt"]) output_shape = (3, 16, 16) - # `encode_prompt` drives the Qwen3-VL decoder layers directly instead of calling `text_encoder.forward`, so the - # offloading hooks would leave its inputs on the offload device. Keep the text encoder out of group offloading. + # `encode_prompt` drives the Qwen3-VL decoder layers directly instead of calling `text_encoder.forward`, and + # pins its inputs to `self.text_encoder.device`. Leaf-level hooks onload each leaf on its own forward while the + # module keeps reporting the offload device, so the inputs are left behind; block-level onloads the whole group + # up front and is unaffected, which is where the text encoder does get covered. group_offloading_leaf_level_exclude_modules = ["text_encoder"] def get_dummy_components(self, num_layers: int = 1): @@ -283,7 +285,8 @@ class TestIdeogram4PipelineMemory(Ideogram4PipelineTesterConfig, MemoryTesterMix pins its inputs to `self.text_encoder.device` so they follow the weights under `enable_model_cpu_offload` (whose `CpuOffload` hook wraps the bypassed `forward` and so never fires). That pinning is wrong for every mechanism that hooks the submodules instead: they onload to the accelerator while the module still reports the - offload device, so the inputs are left behind. Hence the skips below. + offload device, so the inputs are left behind. Hence the skips below, and the text encoder's leaf-level group + offload exclusion on the config class. """ _SUBMODULE_OFFLOAD_SKIP = ( @@ -304,14 +307,3 @@ def test_sequential_cpu_offload_forward_pass(self, base_pipe_output, expected_ma ) def test_sequential_offload_forward_pass_twice(self, expected_max_diff=2e-4): pass - - @pytest.mark.skip( - reason=( - "Block-level group offloading cannot cover `text_encoder`: it leaves ungrouped leaves such as " - "`embed_tokens` to the root module's forward pre-hook, which never fires because `encode_prompt` " - "drives the decoder layers directly. Leaf-level offloading hooks those leaves individually and is " - "bit-exact here; only the block-level half of this test fails." - ) - ) - def test_group_offloading_inference(self): - pass diff --git a/tests/pipelines/ltx2/test_ltx2.py b/tests/pipelines/ltx2/test_ltx2.py index 89b7724b4351..21b0359a94c1 100644 --- a/tests/pipelines/ltx2/test_ltx2.py +++ b/tests/pipelines/ltx2/test_ltx2.py @@ -26,7 +26,7 @@ from diffusers.pipelines.ltx2 import LTX2DurationHead, LTX2TextConnectors from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder -from ...testing_utils import assert_tensors_close, enable_full_determinism, require_torch_accelerator, torch_device +from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device from ..testing_utils import ( BasePipelineTesterConfig, LoraMemoryTesterMixin, @@ -46,6 +46,12 @@ class LTX2PipelineTesterConfig(BasePipelineTesterConfig): ) batch_input_params = frozenset(["prompt", "negative_prompt"]) output_shape = (5, 3, 32, 32) + # `audio_vae` belongs with the other VAEs the group offload tests keep on the accelerator: its decode-time + # convolutions read weights the offload hooks have not onloaded yet. + group_offloading_onload_component_names = [ + *BasePipelineTesterConfig.group_offloading_onload_component_names, + "audio_vae", + ] # LTX2 is a video pipeline (`num_videos_per_prompt`, not `num_images_per_prompt`) and takes a second latent # input for the audio stream. optional_input_params = frozenset( @@ -407,14 +413,6 @@ def test_invalid_duration_bounds_raise(self): class TestLTX2PipelineMemory(LTX2PipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the LTX2 pipeline.""" - @require_torch_accelerator - def test_group_offloading_inference(self): - # The shared helper only offloads a fixed set of component names and leaves LTX2's extra module - # components (`connectors`, `audio_vae`, `vocoder`) on CPU, so the forward pass mixes devices. - # Pipeline-level offloading, which walks every component, is exercised by - # `test_pipeline_level_group_offloading_inference`. - pytest.skip("Using test_pipeline_level_group_offloading_inference instead") - class TestLTX2PipelineLoRA(LTX2PipelineTesterConfig, LoraTesterMixin): """LoRA tests for the LTX2 pipeline.""" diff --git a/tests/pipelines/testing_utils/common.py b/tests/pipelines/testing_utils/common.py index 23db523f1e1d..5bff4d6ac5dd 100644 --- a/tests/pipelines/testing_utils/common.py +++ b/tests/pipelines/testing_utils/common.py @@ -65,14 +65,30 @@ class BasePipelineTesterConfig: ] ) + # The group offload tests derive what they offload: every `torch.nn.Module` component of the pipeline is + # offloaded unless it is named in one of the three lists below, which are kept on the accelerator instead. A + # component that is covered by default is the point — a pipeline that adds a second denoiser or an extra + # encoder gets it exercised without touching this file, and dropping something from the tests takes naming it + # next to a reason. + # Components that cannot be offloaded at leaf level, e.g. a `transformers` model whose attention is a # `torch.nn.MultiheadAttention` (it reads its projection weights directly instead of calling the submodules, so # the leaf-level onload hooks never fire and the weights stay on the offload device). Such a component is often - # fine at block level, hence the level in the name. Listed components are kept on the accelerator by - # `test_pipeline_level_group_offloading_inference` so the remaining ones are still covered, instead of skipping - # the test outright. + # fine at block level, hence the level in the name, and it is still covered by the block-level test. group_offloading_leaf_level_exclude_modules = [] + # Components that cannot be group offloaded at either level. Prefer the leaf-level list above — this one drops + # the component from every group offload test, so state why in a comment next to the name. + group_offloading_exclude_modules = [] + + # Components the component-scoped tests keep on the accelerator rather than offloading. Unlike the two + # exclusion lists above, this one does not reach `test_pipeline_level_group_offloading_inference`, which walks + # the whole pipeline — a component listed here is still leaf offloaded there. The VAE is the reason the list + # exists: some tests enable tiling, and when accelerator streams are used the execution order of a tiled + # forward pass is not traced correctly, which errors out. Group offloading a VAE wants a warmup forward pass + # first (even on dummy inputs). + group_offloading_onload_component_names = ["vae", "vqvae", "image_encoder"] + # ==================== Required interface ==================== @property diff --git a/tests/pipelines/testing_utils/memory.py b/tests/pipelines/testing_utils/memory.py index 6c7986b1bb5c..d30c5d121989 100644 --- a/tests/pipelines/testing_utils/memory.py +++ b/tests/pipelines/testing_utils/memory.py @@ -267,96 +267,100 @@ def test_layerwise_casting_inference(self): class GroupOffloadTesterMixin(BasePipelineOutputMixin): """Block/leaf-level group offload, both component-scoped and pipeline-level orchestration.""" - @require_torch_accelerator - def test_group_offloading_inference(self): - pipe = self.get_pipeline() - for name, component in pipe.components.items(): + def _skip_if_group_offloading_unsupported(self, pipe): + for component in pipe.components.values(): if hasattr(component, "_supports_group_offloading") and not component._supports_group_offloading: pytest.skip(f"{self.pipeline_class.__name__} has a component that does not support group offloading.") - def create_pipe(): - torch.manual_seed(0) - return self.get_pipeline() - - def enable_group_offload_on_component(pipe, group_offloading_kwargs): - # We intentionally don't test VAE's here. This is because some tests enable tiling on the VAE. If - # tiling is enabled and a forward pass is run, when accelerator streams are used, the execution order - # of the layers is not traced correctly. This causes errors. For apply group offloading to VAE, a - # warmup forward pass (even with dummy small inputs) is recommended. - for component_name in [ - "text_encoder", - "text_encoder_2", - "text_encoder_3", - "transformer", - "transformer_2", - "unet", - "controlnet", - "adapter", - ]: - if not hasattr(pipe, component_name): - continue - component = getattr(pipe, component_name) - if component is None: - continue - if not getattr(component, "_supports_group_offloading", True): - continue - if hasattr(component, "enable_group_offload"): - # For diffusers ModelMixin implementations - component.enable_group_offload(torch.device(torch_device), **group_offloading_kwargs) - else: - # For other models not part of diffusers - apply_group_offloading( - component, onload_device=torch.device(torch_device), **group_offloading_kwargs - ) - assert all( - module._diffusers_hook.get_hook("group_offloading") is not None - for module in component.modules() - if hasattr(module, "_diffusers_hook") - ) - for component_name in ["vae", "vqvae", "image_encoder"]: - component = getattr(pipe, component_name, None) - if isinstance(component, torch.nn.Module): - component.to(torch_device) - - def run_forward(pipe): - torch.manual_seed(0) - inputs = self.get_dummy_inputs() - return pipe(**inputs)[0] - - pipe = create_pipe().to(torch_device) - output_without_group_offloading = run_forward(pipe) - - pipe = create_pipe() - enable_group_offload_on_component(pipe, {"offload_type": "block_level", "num_blocks_per_group": 1}) - output_with_group_offloading1 = run_forward(pipe) - - pipe = create_pipe() - enable_group_offload_on_component(pipe, {"offload_type": "leaf_level"}) - output_with_group_offloading2 = run_forward(pipe) + def _group_offload_exclude_modules(self, pipe, offload_type): + """Config-declared components to keep out of group offloading at `offload_type`. + + Every group offload test routes its exclusions through here, so a name that matches no component on the + pipeline is reported as the typo it is rather than silently costing coverage and surfacing later as a + device mismatch. The onload names are not checked: they are a shared default covering several pipelines, + most of which have only some of them. + """ + exclude = set(self.group_offloading_exclude_modules) + if offload_type == "leaf_level": + exclude |= set(self.group_offloading_leaf_level_exclude_modules) + + # Checked against every registered component rather than the module-valued ones, so that excluding an + # optional component a config leaves unset reads as the no-op it is instead of a typo. + unknown = sorted(exclude - set(pipe.components)) + assert not unknown, ( + f"{type(self).__name__} excludes {unknown} from group offloading, but " + f"{self.pipeline_class.__name__} has no such component. Its components are " + f"{sorted(pipe.components)}." + ) + return exclude + + def _split_group_offload_components(self, pipe, offload_type): + """Split the pipeline's module components into the ones to offload and the ones to keep on the accelerator. + + Everything is offloaded unless the config lists it, so a component a pipeline adds under a name this file + has never heard of is covered by default rather than silently left on CPU. See the three list attributes on + `BasePipelineTesterConfig`. + """ + module_names = [name for name, component in pipe.components.items() if isinstance(component, torch.nn.Module)] + onload_names = self._group_offload_exclude_modules(pipe, offload_type) | set( + self.group_offloading_onload_component_names + ) + offload = [name for name in module_names if name not in onload_names] + onload = [name for name in module_names if name in onload_names] + return offload, onload - assert_tensors_close( - output_with_group_offloading1, - output_without_group_offloading, - atol=1e-4, - rtol=1e-5, + def _enable_group_offload_on_components(self, pipe, **group_offloading_kwargs): + offload_names, onload_names = self._split_group_offload_components( + pipe, group_offloading_kwargs["offload_type"] + ) + for component_name in offload_names: + component = getattr(pipe, component_name) + if hasattr(component, "enable_group_offload"): + # For diffusers ModelMixin implementations + component.enable_group_offload(torch.device(torch_device), **group_offloading_kwargs) + else: + # For other models not part of diffusers + apply_group_offloading(component, onload_device=torch.device(torch_device), **group_offloading_kwargs) + assert all( + module._diffusers_hook.get_hook("group_offloading") is not None + for module in component.modules() + if hasattr(module, "_diffusers_hook") + ) + for component_name in onload_names: + getattr(pipe, component_name).to(torch_device) + + def _run_group_offload_inference(self, base_pipe_output, expected_max_difference, msg, **group_offloading_kwargs): + # Build the offload pipeline the same way as `base_pipe_output` so that group offloading is the only + # difference under test. It stays on CPU here — the components are placed as they are hooked. + pipe = self.get_pipeline() + 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) + + @require_torch_accelerator + def test_group_offloading_inference_block_level(self, base_pipe_output, expected_max_difference=1e-4): + self._run_group_offload_inference( + base_pipe_output, + expected_max_difference, msg="block-level group offloading should not affect the inference results", + offload_type="block_level", + num_blocks_per_group=1, ) - assert_tensors_close( - output_with_group_offloading2, - output_without_group_offloading, - atol=1e-4, - rtol=1e-5, + + @require_torch_accelerator + def test_group_offloading_inference_leaf_level(self, base_pipe_output, expected_max_difference=1e-4): + self._run_group_offload_inference( + base_pipe_output, + expected_max_difference, msg="leaf-level group offloading should not affect the inference results", + offload_type="leaf_level", ) @require_torch_accelerator def test_pipeline_level_group_offloading_sanity_checks(self): pipe: DiffusionPipeline = self.get_pipeline() - - for name, component in pipe.components.items(): - if hasattr(component, "_supports_group_offloading"): - if not component._supports_group_offloading: - pytest.skip(f"{self.pipeline_class.__name__} is not suitable for this test.") + self._skip_if_group_offloading_unsupported(pipe) module_names = sorted( [name for name, component in pipe.components.items() if isinstance(component, torch.nn.Module)] @@ -387,18 +391,15 @@ def test_pipeline_level_group_offloading_inference(self, base_pipe_output, expec # Build the offload pipeline the same way as `base_pipe_output` so that group offloading is the only # difference under test. It stays on CPU here — `enable_group_offload` places the components. pipe: DiffusionPipeline = self.get_pipeline() - - for name, component in pipe.components.items(): - if hasattr(component, "_supports_group_offloading"): - if not component._supports_group_offloading: - pytest.skip(f"{self.pipeline_class.__name__} is not suitable for this test.") + self._skip_if_group_offloading_unsupported(pipe) offload_device = "cpu" + offload_type = "leaf_level" pipe.enable_group_offload( onload_device=torch_device, offload_device=offload_device, - offload_type="leaf_level", - exclude_modules=self.group_offloading_leaf_level_exclude_modules, + offload_type=offload_type, + exclude_modules=sorted(self._group_offload_exclude_modules(pipe, offload_type)), ) pipe.set_progress_bar_config(disable=None) inputs = self.get_dummy_inputs() From 9f63110bb39f1a42c4857045e30b218ed5aeb3cf Mon Sep 17 00:00:00 2001 From: Daniel Gu Date: Sat, 29 Aug 2026 08:24:27 +0200 Subject: [PATCH 2/3] style: fix `create_pipe` indentation in the group offload tester The method body was indented one level too deep, which `ruff format` rewrites. Co-Authored-By: Claude Opus 5 (1M context) --- tests/pipelines/testing_utils/memory.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/pipelines/testing_utils/memory.py b/tests/pipelines/testing_utils/memory.py index c2d014418f06..6c0b999f8f1e 100644 --- a/tests/pipelines/testing_utils/memory.py +++ b/tests/pipelines/testing_utils/memory.py @@ -268,8 +268,8 @@ class GroupOffloadTesterMixin(BasePipelineOutputMixin): """Block/leaf-level group offload, both component-scoped and pipeline-level orchestration.""" def create_pipe(self): - torch.manual_seed(0) - return self.get_pipeline() + torch.manual_seed(0) + return self.get_pipeline() def _skip_if_group_offloading_unsupported(self, pipe): for component in pipe.components.values(): From dfdc6b686def0db1cee86c542ebe6191631d9e2a Mon Sep 17 00:00:00 2001 From: Daniel Gu Date: Sat, 29 Aug 2026 08:24:49 +0200 Subject: [PATCH 3/3] tests: give group offloading one exclusion list per level MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `group_offloading_onload_component_names` named an implementation detail — components the tests keep resident — and justified itself with a mechanism that cannot occur: lazy prefetch mis-tracing a tiled VAE decode needs `use_stream=True`, which no test on this mixin sets, and most of the pipelines it covers never enable tiling either. What it actually encodes is a level-specific incapability, the mirror of `group_offloading_leaf_level_exclude_modules`. Leaf-level offloading onloads each leaf on its own `forward`, so it breaks on compute that reads a leaf's `.weight` directly. Block-level onloads a group when the group's leader runs its `forward`, so it breaks on compute that re-enters submodules without going through that leader — which is what a VAE decode path does. Rename it to `group_offloading_block_level_exclude_modules` and scope it to the level it describes. `group_offloading_exclude_modules` is then redundant, since a component that fails at both levels goes in both lists, and it never had a user. `vqvae` leaves the default: no pipeline in the new harness has one. Measured across the suite, offloading the `vae` at block level breaks 28 of 86 test classes, while all 86 offload it at leaf level with no failures -- coverage the component-scoped tests were skipping and the pipeline-level test has always had. Sweep before and after: the same 9 pre-existing failures, 782 -> 810 passing. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/references/testing.md | 4 +-- tests/pipelines/ltx2/test_ltx2.py | 8 ++--- tests/pipelines/testing_utils/common.py | 37 +++++++++++------------ tests/pipelines/testing_utils/memory.py | 40 ++++++++++++++----------- 4 files changed, 45 insertions(+), 44 deletions(-) diff --git a/.ai/references/testing.md b/.ai/references/testing.md index 9a55f3670b4e..cfc471930878 100644 --- a/.ai/references/testing.md +++ b/.ai/references/testing.md @@ -30,8 +30,8 @@ Follow the style introduced in [#14113](https://github.com/huggingface/diffusers - `MemoryTesterMixin` — CPU offload, group offload, layerwise casting. - Cache mixins — `PyramidAttentionBroadcastTesterMixin`, `FasterCacheTesterMixin`, `FirstBlockCacheTesterMixin`, `TaylorSeerCacheTesterMixin`, `MagCacheTesterMixin`. Guidance-distilled models override the cache config (e.g. `FASTER_CACHE_CONFIG = {... "is_guidance_distilled": True}`). Don't introduce caching related tests in the first iteration. These tests are added on a case-by-case basis. - In the first pass, just add tests related to `PipelineTesterMixin` and `MemoryTesterMixin`. -- **Declare a component that can't be offloaded — don't hand-write a skip.** Leaf-level offloading hooks only the supported leaf types (`nn.Linear`, `nn.Conv*`, `nn.Embedding` — see `_GO_LC_SUPPORTED_PYTORCH_LAYERS` in `src/diffusers/hooks/_common.py`) and onloads each on its own `forward`, so any code that reads a leaf's `.weight` instead of calling the leaf bypasses that leaf's hook and computes against offloaded weights. Which fix applies depends on who owns the component. For a diffusers model, set `_supports_group_offloading = False` on the `ModelMixin` subclass (as `HunyuanDiT2DModel` does) — both offload mixins honor the flag and skip themselves, so the gap is declared on the model instead of buried in a test file. For a third-party component you can't annotate, such as a `transformers` encoder, list it in `group_offloading_leaf_level_exclude_modules` on the config class; excluded components are kept on the accelerator, so every other component stays covered. Block-level offloading is usually unaffected, hence the level in the name, and `test_group_offloading_inference_block_level` still covers the component — a component that fails at both levels goes in `group_offloading_exclude_modules` instead, with a comment saying why. - - **Every `nn.Module` component is group offloaded unless a config list names it** — `group_offloading_leaf_level_exclude_modules`, `group_offloading_exclude_modules`, or `group_offloading_onload_component_names` (the VAE and friends, kept on the accelerator because tiling breaks stream tracing). A pipeline that adds a second denoiser or an extra encoder therefore gets it exercised without touching the shared mixin, and losing coverage takes naming the component. A name in an exclusion list that matches no component on the pipeline fails the test as a typo. +- **Declare a component that can't be offloaded — don't hand-write a skip.** Leaf-level offloading hooks only the supported leaf types (`nn.Linear`, `nn.Conv*`, `nn.Embedding` — see `_GO_LC_SUPPORTED_PYTORCH_LAYERS` in `src/diffusers/hooks/_common.py`) and onloads each on its own `forward`, so any code that reads a leaf's `.weight` instead of calling the leaf bypasses that leaf's hook and computes against offloaded weights. Which fix applies depends on who owns the component. For a diffusers model, set `_supports_group_offloading = False` on the `ModelMixin` subclass (as `HunyuanDiT2DModel` does) — both offload mixins honor the flag and skip themselves, so the gap is declared on the model instead of buried in a test file. For a third-party component you can't annotate, such as a `transformers` encoder, list it in `group_offloading_leaf_level_exclude_modules` on the config class; excluded components are kept on the accelerator, so every other component stays covered, and `test_group_offloading_inference_block_level` still exercises it at the other level. + - **Every `nn.Module` component is group offloaded unless the list for that level names it.** The two levels fail on opposite hazards, so each has its own list and a component that fails at both goes in both. `group_offloading_leaf_level_exclude_modules` is for a component whose compute reads a leaf's `.weight` instead of calling the leaf, so the leaf's hook never fires. `group_offloading_block_level_exclude_modules` is for the mirror case: block-level onloads a group when the group's leader runs its `forward`, so a component that re-enters submodules without going through that leader finds its weights offloaded — VAE decode paths are the usual instance, which is why `vae` and `image_encoder` are the default. A pipeline that adds a second denoiser or an extra encoder therefore gets it exercised without touching the shared mixin, and losing coverage takes naming the component. A name in an exclusion list that matches no component on the pipeline fails the test as a typo. - `torch.nn.MultiheadAttention` is the common instance: it passes `self.out_proj.weight` straight to `torch.nn.functional.multi_head_attention_forward` instead of calling `self.out_proj`, so the hook on `out_proj` never fires. `SiglipVisionModel`'s attention pooling head wraps one — see `tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py`, whose `image_encoder` is excluded for this reason. - `HunyuanDiTAttentionPool` (`src/diffusers/models/embeddings.py`) shows the same failure without an MHA module: a plain `nn.Module` that hands its `q_proj` / `k_proj` / `v_proj` / `c_proj` weights to `torch.nn.functional.multi_head_attention_forward`, so all four projections stay offloaded rather than just one. `HunyuanDiT2DModel` opts out of group offloading entirely with `_supports_group_offloading = False`. - Before adding a skip or an exclusion, confirm the failure still reproduces — several existing skips are stale, having outlived the upstream cause. diff --git a/tests/pipelines/ltx2/test_ltx2.py b/tests/pipelines/ltx2/test_ltx2.py index 21b0359a94c1..91dca581a6bf 100644 --- a/tests/pipelines/ltx2/test_ltx2.py +++ b/tests/pipelines/ltx2/test_ltx2.py @@ -46,10 +46,10 @@ class LTX2PipelineTesterConfig(BasePipelineTesterConfig): ) batch_input_params = frozenset(["prompt", "negative_prompt"]) output_shape = (5, 3, 32, 32) - # `audio_vae` belongs with the other VAEs the group offload tests keep on the accelerator: its decode-time - # convolutions read weights the offload hooks have not onloaded yet. - group_offloading_onload_component_names = [ - *BasePipelineTesterConfig.group_offloading_onload_component_names, + # `audio_vae` fails at block level for the same reason the other VAEs do: its decode-time convolutions run + # without the group leader's `forward` having onloaded the group. + group_offloading_block_level_exclude_modules = [ + *BasePipelineTesterConfig.group_offloading_block_level_exclude_modules, "audio_vae", ] # LTX2 is a video pipeline (`num_videos_per_prompt`, not `num_images_per_prompt`) and takes a second latent diff --git a/tests/pipelines/testing_utils/common.py b/tests/pipelines/testing_utils/common.py index 5bff4d6ac5dd..260962d8a884 100644 --- a/tests/pipelines/testing_utils/common.py +++ b/tests/pipelines/testing_utils/common.py @@ -66,28 +66,25 @@ class BasePipelineTesterConfig: ) # The group offload tests derive what they offload: every `torch.nn.Module` component of the pipeline is - # offloaded unless it is named in one of the three lists below, which are kept on the accelerator instead. A - # component that is covered by default is the point — a pipeline that adds a second denoiser or an extra - # encoder gets it exercised without touching this file, and dropping something from the tests takes naming it - # next to a reason. - - # Components that cannot be offloaded at leaf level, e.g. a `transformers` model whose attention is a - # `torch.nn.MultiheadAttention` (it reads its projection weights directly instead of calling the submodules, so - # the leaf-level onload hooks never fire and the weights stay on the offload device). Such a component is often - # fine at block level, hence the level in the name, and it is still covered by the block-level test. + # offloaded unless the list for that level names it, in which case it is kept on the accelerator. A component + # that is covered by default is the point — a pipeline that adds a second denoiser or an extra encoder gets it + # exercised without touching this file, and dropping something from the tests takes naming it next to a reason. + # The two levels fail on opposite hazards, so each gets its own list; a component that fails at both goes in + # both. + + # Components that cannot be offloaded at leaf level. Leaf-level offloading onloads each supported leaf on its + # own `forward`, so any code that reads a leaf's `.weight` instead of calling the leaf bypasses that leaf's + # hook and computes against offloaded weights — `torch.nn.MultiheadAttention` being the usual instance. Such a + # component is normally fine at block level, where the whole group is onloaded at once, and stays covered + # there. group_offloading_leaf_level_exclude_modules = [] - # Components that cannot be group offloaded at either level. Prefer the leaf-level list above — this one drops - # the component from every group offload test, so state why in a comment next to the name. - group_offloading_exclude_modules = [] - - # Components the component-scoped tests keep on the accelerator rather than offloading. Unlike the two - # exclusion lists above, this one does not reach `test_pipeline_level_group_offloading_inference`, which walks - # the whole pipeline — a component listed here is still leaf offloaded there. The VAE is the reason the list - # exists: some tests enable tiling, and when accelerator streams are used the execution order of a tiled - # forward pass is not traced correctly, which errors out. Group offloading a VAE wants a warmup forward pass - # first (even on dummy inputs). - group_offloading_onload_component_names = ["vae", "vqvae", "image_encoder"] + # Components that cannot be offloaded at block level. Block-level offloading onloads a group when the group's + # leader runs its `forward`, so a component whose compute re-enters submodules without going through that + # leader finds its weights still on the offload device. VAE decode paths are the usual instance — measured + # across the suite, offloading a `vae` at block level breaks 28 of 86 pipeline test classes, while every one of + # them offloads it at leaf level without complaint, which is what the pipeline-level test has always done. + group_offloading_block_level_exclude_modules = ["vae", "image_encoder"] # ==================== Required interface ==================== diff --git a/tests/pipelines/testing_utils/memory.py b/tests/pipelines/testing_utils/memory.py index 6c0b999f8f1e..8ebc117611df 100644 --- a/tests/pipelines/testing_utils/memory.py +++ b/tests/pipelines/testing_utils/memory.py @@ -30,7 +30,7 @@ require_torch_accelerator, torch_device, ) -from .common import BasePipelineOutputMixin +from .common import BasePipelineOutputMixin, BasePipelineTesterConfig if is_accelerate_available(): @@ -279,20 +279,27 @@ def _skip_if_group_offloading_unsupported(self, pipe): def _group_offload_exclude_modules(self, pipe, offload_type): """Config-declared components to keep out of group offloading at `offload_type`. - Every group offload test routes its exclusions through here, so a name that matches no component on the - pipeline is reported as the typo it is rather than silently costing coverage and surfacing later as a - device mismatch. The onload names are not checked: they are a shared default covering several pipelines, - most of which have only some of them. + The two levels fail on opposite hazards, so each reads its own list — see the attributes on + `BasePipelineTesterConfig`. Every group offload test routes its exclusions through here, so a name a config + adds that matches no component on the pipeline is reported as the typo it is rather than silently costing + coverage and surfacing later as a device mismatch. """ - exclude = set(self.group_offloading_exclude_modules) if offload_type == "leaf_level": - exclude |= set(self.group_offloading_leaf_level_exclude_modules) - - # Checked against every registered component rather than the module-valued ones, so that excluding an - # optional component a config leaves unset reads as the no-op it is instead of a typo. - unknown = sorted(exclude - set(pipe.components)) + exclude = set(self.group_offloading_leaf_level_exclude_modules) + elif offload_type == "block_level": + exclude = set(self.group_offloading_block_level_exclude_modules) + else: + raise ValueError(f"Unknown `offload_type` {offload_type!r}.") + + # Only names a config adds are checked: the shared defaults cover the whole suite and most pipelines have + # just some of them. The check is against every registered component rather than the module-valued ones, so + # that excluding an optional component a config leaves unset reads as the no-op it is instead of a typo. + shared_defaults = set(BasePipelineTesterConfig.group_offloading_leaf_level_exclude_modules) | set( + BasePipelineTesterConfig.group_offloading_block_level_exclude_modules + ) + unknown = sorted(exclude - set(pipe.components) - shared_defaults) assert not unknown, ( - f"{type(self).__name__} excludes {unknown} from group offloading, but " + f"{type(self).__name__} excludes {unknown} from {offload_type} group offloading, but " f"{self.pipeline_class.__name__} has no such component. Its components are " f"{sorted(pipe.components)}." ) @@ -301,14 +308,11 @@ def _group_offload_exclude_modules(self, pipe, offload_type): def _split_group_offload_components(self, pipe, offload_type): """Split the pipeline's module components into the ones to offload and the ones to keep on the accelerator. - Everything is offloaded unless the config lists it, so a component a pipeline adds under a name this file - has never heard of is covered by default rather than silently left on CPU. See the three list attributes on - `BasePipelineTesterConfig`. + Everything is offloaded unless the list for this level names it, so a component a pipeline adds under a + name this file has never heard of is covered by default rather than silently left on CPU. """ module_names = [name for name, component in pipe.components.items() if isinstance(component, torch.nn.Module)] - onload_names = self._group_offload_exclude_modules(pipe, offload_type) | set( - self.group_offloading_onload_component_names - ) + onload_names = self._group_offload_exclude_modules(pipe, offload_type) offload = [name for name in module_names if name not in onload_names] onload = [name for name in module_names if name in onload_names] return offload, onload