[tests] Split Pipeline Group Offloading Block-Level and Leaf-Level Tests

.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, 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.

tests/pipelines/ideogram4/test_pipeline_ideogram4.py

@@ -50,8 +50,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):
@@ -285,7 +287,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 = (
@@ -307,17 +310,6 @@ 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
-
 
 class TestIdeogram4PipelineLoRA(Ideogram4PipelineTesterConfig, LoraTesterMixin):
     """LoRA tests for the Ideogram4 pipeline."""

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` 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
     # 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."""

tests/pipelines/testing_utils/common.py

@@ -65,14 +65,27 @@ class BasePipelineTesterConfig:
         ]
     )
 
-    # 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.
+    # The group offload tests derive what they offload: every `torch.nn.Module` component of the pipeline is
+    # 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 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 ====================
 
     @property