diff --git a/.ai/references/testing.md b/.ai/references/testing.md index ac0fda2d440b..30f2e27af855 100644 --- a/.ai/references/testing.md +++ b/.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 diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py b/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py index 3187aa88ddc4..1bb5e06a632d 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py +++ b/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), diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index cb35fc46ee7d..106ba55cf149 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/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) diff --git a/tests/pipelines/testing_utils/common.py b/tests/pipelines/testing_utils/common.py index 0851c85f2bca..23db523f1e1d 100644 --- a/tests/pipelines/testing_utils/common.py +++ b/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 diff --git a/tests/pipelines/testing_utils/memory.py b/tests/pipelines/testing_utils/memory.py index 77b25c5b2a3e..e74d50537c8d 100644 --- a/tests/pipelines/testing_utils/memory.py +++ b/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()