diff --git a/.ai/references/testing.md b/.ai/references/testing.md index 62e1ca986a07..9a55f3670b4e 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 eaee8dc2ba9a..21b216b55fe9 100644 --- a/tests/pipelines/ideogram4/test_pipeline_ideogram4.py +++ b/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.""" 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()