[tests] Add Config to Exclude Modules from Leaf-Level Group Offloading - #14564
[tests] Add Config to Exclude Modules from Leaf-Level Group Offloading#14564dg845 wants to merge 3 commits into
Conversation
`test_pipeline_level_group_offloading_inference` was skipped outright for HunyuanVideoFramepack because `image_encoder` is a `SiglipVisionModel`, whose attention pooling head 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. Add a `group_offloading_leaf_level_exclude_modules` knob to the old-style `PipelineTesterMixin` and the new-style `BasePipelineTesterConfig` (empty by default, so no behavior change elsewhere), pass it through to `enable_group_offload(exclude_modules=...)` in both implementations of the test, and set it to `["image_encoder"]` for framepack instead of skipping. Block-level offloading is unaffected — the whole head is onloaded as one unmatched module — hence the level in the name. The test now passes and covers leaf-level offloading of the transformer, VAE and both text encoders. The VAE is coverage nothing else provided: `test_group_offloading_inference` deliberately excludes `vae` and `image_encoder`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Records the rule behind `group_offloading_leaf_level_exclude_modules`: leaf-level offloading hooks only the supported leaf types and onloads each on its own forward, so any code that reads a leaf's `.weight` instead of calling the leaf bypasses that hook. Routes the fix by who owns the component. A diffusers model declares the gap with `_supports_group_offloading = False` on the `ModelMixin` subclass, which both offload mixins honor. A third-party component that can't be annotated goes in `group_offloading_leaf_level_exclude_modules`, which keeps offload coverage for every other component — where a hand-written skip would drop it for the whole pipeline, the VAE included, since the component-scoped `test_group_offloading_inference` deliberately excludes it. `torch.nn.MultiheadAttention` is called out as the common instance rather than as the definition, with `HunyuanDiTAttentionPool` as a case that fails the same way with no MHA module involved, so the guidance still applies when a future component fails for a different reason. Also notes that a failure should be reproduced before a skip or exclusion is added: of the five pipelines currently skipping the pipeline-level test, only framepack and motif_video still fail. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Self-Review ReportSelf-review — Diff: 2 commits, 5 files, +27/−8 Blocking issuesNone. Non-blocking issues
Both are house-style calls; neither is obviously correct enough to settle without the reviewer. VerificationCode commit, re-confirmed: attribute resolves for every new-style user (AST scan of Tests: Dead CodeNot applicable — no source changes. SummaryREADY. Nothing left to fix before submitting; the two non-blocking items above are the only things that should reach the reviewer unresolved. |
dg845
changed the title
| - `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. |
There was a problem hiding this comment.
For a third-party component you can't annotate, such as a
transformersencoder, list it ingroup_offloading_leaf_level_exclude_moduleson the config class (the attribute is also on the old-stylePipelineTesterMixinintests/pipelines/test_pipelines_common.py);
I don't think we have to mention the old-style code. Let's also be specific and provide an example file for reference here.
enable_group_offloadkeeps excluded components on the accelerator, so every other component is still covered — where a hand-written skip ontest_pipeline_level_group_offloading_inferencewould drop offload coverage for the whole pipeline, including the VAE, which the component-scopedtest_group_offloading_inferencedeliberately excludes.
This could be further simplified a bit.
| test_layerwise_casting = False | ||
| test_group_offloading = False | ||
|
|
||
| # Components that cannot be offloaded at leaf level. See `BasePipelineTesterConfig` in |
There was a problem hiding this comment.
I think we can remove this? Let's not update stuff we plan on removing.
2 participants
What does this PR do?
This PR adds a
group_offloading_leaf_level_exclude_modulespipeline test config attribute, which allows leaf-level group offloading tests such astest_pipeline_level_group_offloading_inferenceto exclude particular modules which are not leaf-level offloadable while still covering other modules which are. In practice, modules which are not offloadable at the leaf level tend to still be offloadable at the block level (such as thetransformersSiglipVisionModelused by the Hunyuan Video Framepack pipeline which motivated this PR), which is why we focus on leaf-level offloading here.The new
group_offloading_leaf_level_exclude_modulesattribute can potentially be used by the following test files, but is only used for the motivating Hunyuan Video Framepack case in this PR:tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py(this PR)tests/pipelines/glm_image/test_glm_image.pytests/pipelines/joyimage/test_joyimage_edit.pytests/pipelines/joyimage/test_joyimage_edit_plus.pytests/pipelines/motif_video/test_motif_video_image2video.pyCan open up a follow-up PR to address the other tests.
Based on PR #14551.
Before submitting
self-reviewskill on the diff?documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@sayakpaul
@DN6