Skip to content

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

Open
dg845 wants to merge 3 commits into
mainfrom
tests/group-offloading-leaf-level-exclude-modules
Open

[tests] Add Config to Exclude Modules from Leaf-Level Group Offloading#14564
dg845 wants to merge 3 commits into
mainfrom
tests/group-offloading-leaf-level-exclude-modules

Conversation

@dg845

@dg845 dg845 commented Aug 22, 2026

Copy link
Copy Markdown
Collaborator

What does this PR do?

This PR adds a group_offloading_leaf_level_exclude_modules pipeline test config attribute, which allows leaf-level group offloading tests such as test_pipeline_level_group_offloading_inference to 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 the transformers SiglipVisionModel used 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_modules attribute 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.py
  • tests/pipelines/joyimage/test_joyimage_edit.py
  • tests/pipelines/joyimage/test_joyimage_edit_plus.py
  • tests/pipelines/motif_video/test_motif_video_image2video.py

Can open up a follow-up PR to address the other tests.

Based on PR #14551.

Before submitting

  • Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
    • Did you read the Coding with AI agents guide?
    • Did you run the self-review skill on the diff?
    • Did you share the final self-review notes in the PR description or a comment?
  • Did you read the contributor guideline?
  • Did you read our philosophy doc? (important for complex PRs)
  • Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?
  • Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

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

sayakpaul and others added 3 commits August 21, 2026 05:20
@sayakpaul
skip hunyuan framepack group offloading tests
6e70f19
@dg845 @claude
exclude only the SigLIP image encoder from leaf-level group offloading
319e1e0 
`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>
@dg845 @claude
docs: note how to handle a component that breaks leaf-level offloading
e87a5c3 
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>
@github-actions github-actions Bot added tests size/S PR with diff < 50 LOC labels Aug 22, 2026
@dg845

dg845 commented Aug 22, 2026

Copy link
Copy Markdown
Collaborator Author
Self-Review Report

Self-review — tests/group-offloading-leaf-level-exclude-modules vs skip-go-hunyuan-framepack

Diff: 2 commits, 5 files, +27/−8

Blocking issues

None.

Non-blocking issues

  1. The note's home is a partial fit — .ai/references/testing.md:33. The bullet sits under "### Standard pipelines", but since the redraft it also prescribes a change to a model class (_supports_group_offloading = False), and the file has its own "## Model-level tests" section. Still the best single location — the trigger is a failing pipeline test, and splitting the guidance would break the routing logic in half — but a reviewer may want a cross-reference from the model section.
  2. Mutable class-level default — test_pipelines_common.py:1060, testing_utils/common.py:68. Both declarations use = [] while neighbouring contract attributes are frozenset(...). Nothing mutates it, but a subclass calling .append() would write through to the shared base list. A tuple works at runtime but conflicts with enable_group_offload's exclude_modules: str | list[str] annotation (pipeline_utils.py:1391), so it's keep-the-list or widen-the-annotation.

Both are house-style calls; neither is obviously correct enough to settle without the reviewer.

Verification

Code commit, re-confirmed: attribute resolves for every new-style user (AST scan of tests/pipelines/**/test_*.py — no memory-mixin user lacks a *TesterConfig base); default [] is behavior-preserving (pipeline_utils.py:1471-1505); the leaf-level call sites left untouched genuinely don't need the knob (sanity checks run no forward pass, lora.py:1054 offloads only the denoiser).

Tests: framepack -k offload 5 passed / 2 skipped (just re-run); 30 group-offload tests green across flux, cogvideox, qwenimage, flux2-klein-inpaint, hunyuan_video. ruff check / ruff format --check clean; utils/check_ai.py passes.

Dead Code

Not applicable — no source changes.

Summary

READY. Nothing left to fix before submitting; the two non-blocking items above are the only things that should reach the reviewer unresolved.

@dg845 dg845 changed the title Tests/group offloading leaf level exclude modules [tests] Add Config to Exclude Modules from Leaf-Level Group Offloading Aug 22, 2026
@dg845 dg845 mentioned this pull request Aug 22, 2026
Merged
Base automatically changed from skip-go-hunyuan-framepack to main August 22, 2026 07:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

size/S PR with diff < 50 LOC tests

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@dg845 @sayakpaul