[tests] Add Config to Exclude Modules from Leaf-Level Group Offloading

.ai/references/testing.md

@@ -30,6 +30,10 @@ 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 (the attribute is also on the old-style `PipelineTesterMixin` in `tests/pipelines/test_pipelines_common.py`); `enable_group_offload` keeps excluded components on the accelerator, so every other component is still covered — where a hand-written skip on `test_pipeline_level_group_offloading_inference` would drop offload coverage for the whole pipeline, including the VAE, which the component-scoped `test_group_offloading_inference` deliberately excludes. Block-level offloading is usually unaffected, hence the level in the name — a component that fails at both levels does need a skip.
+  - `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.
 - **IP-Adapter tests** live in their own class decorated with `@is_ip_adapter`, subclassing only the config (not `PipelineTesterMixin`).
 
 ### Modular pipelines

tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py

@@ -74,6 +74,15 @@ class HunyuanVideoFramepackPipelineFastTests(
     test_layerwise_casting = True
     test_group_offloading = True
 
+    # `image_encoder` is a `SiglipVisionModel`, whose attention pooling head
+    # (`SiglipMultiheadAttentionPoolingHead`) wraps a `torch.nn.MultiheadAttention`. That hands
+    # `self.out_proj.weight` to `torch.nn.functional.multi_head_attention_forward` instead of calling
+    # `self.out_proj`, so the leaf-level onload hook on `out_proj` never fires and its weights stay on the offload
+    # device. Same root cause as the sequential CPU offloading skips below. Block-level offloading is unaffected
+    # (the whole head is onloaded as one unmatched module), and every other component offloads fine at leaf level,
+    # so exclude just this one rather than skipping the test.
+    group_offloading_leaf_level_exclude_modules = ["image_encoder"]
+
     faster_cache_config = FasterCacheConfig(
         spatial_attention_block_skip_range=2,
         spatial_attention_timestep_skip_range=(-1, 901),

tests/pipelines/test_pipelines_common.py

@@ -1057,6 +1057,10 @@ class PipelineTesterMixin:
     test_layerwise_casting = False
     test_group_offloading = False
 
+    # Components that cannot be offloaded at leaf level. See `BasePipelineTesterConfig` in
+    # `tests/pipelines/testing_utils/common.py`, which declares the same attribute, for the rationale.
+    group_offloading_leaf_level_exclude_modules = []
+
     def get_generator(self, seed):
         device = torch_device if torch_device != "mps" else "cpu"
         generator = torch.Generator(device).manual_seed(seed)
@@ -2470,6 +2474,7 @@ def test_pipeline_level_group_offloading_inference(self, expected_max_difference
             onload_device=torch_device,
             offload_device=offload_device,
             offload_type="leaf_level",
+            exclude_modules=self.group_offloading_leaf_level_exclude_modules,
         )
         pipe.set_progress_bar_config(disable=None)
         inputs["generator"] = torch.manual_seed(0)

tests/pipelines/testing_utils/common.py

@@ -65,6 +65,14 @@ 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.
+    group_offloading_leaf_level_exclude_modules = []
+
     # ==================== Required interface ====================
 
     @property

tests/pipelines/testing_utils/memory.py

@@ -397,6 +397,7 @@ def test_pipeline_level_group_offloading_inference(self, base_pipe_output, expec
             onload_device=torch_device,
             offload_device=offload_device,
             offload_type="leaf_level",
+            exclude_modules=self.group_offloading_leaf_level_exclude_modules,
         )
         pipe.set_progress_bar_config(disable=None)
         inputs = self.get_dummy_inputs()