[tests] Split Pipeline Group Offloading Block-Level and Leaf-Level Tests - #14635
[tests] Split Pipeline Group Offloading Block-Level and Leaf-Level Tests#14635dg845 wants to merge 2 commits into
Conversation
…ponents `GroupOffloadTesterMixin.test_group_offloading_inference` ran both offload levels in one test body, so a pipeline that failed at one level had to skip both. Split it into `test_group_offloading_inference_block_level` and `test_group_offloading_inference_leaf_level`, sharing the helpers the test body used to define inline, and compare against the class-scoped `base_pipe_output` rather than rebuilding a baseline per level. The set of components to offload was a hardcoded list of eight names, so a pipeline with a component under any other name had it silently left on CPU for the forward pass to trip over. Derive the set instead: every `nn.Module` component is offloaded unless the config lists it in `group_offloading_leaf_level_exclude_modules`, the new `group_offloading_exclude_modules`, or `group_offloading_onload_component_names`. A name in either exclusion list that matches no component on the pipeline fails as the typo it is. Ideogram4's `unconditional_transformer` was one of the silently dropped components, which is why its group offload test was skipped; it now needs no declaration at all, and `text_encoder` picks up block-level coverage it never had. LTX2's skip goes the same way, with `audio_vae` declared alongside the other VAEs the tests keep on the accelerator. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
github-actions
Bot
added
fixes-issue
size/L
|
Do we want to hold off a bit until all the pipelines are migrated to use the new mixins? This way, it's easier to do this all at once? |
Self-Review ReportSelf-review — tests/group-offloading-separate-levels (round 3)
Blocking
A test-merge ( +<<<<<<< .our The other four files merge clean. Resolution is mechanical — keep main's two new LoRA classes, drop the skip block — but the branch needs a rebase onto Non-blocking
SummaryThe change itself is in good shape — rounds 1 and 2 findings are all fixed and re-verified, the accounting closes with no unexplained test-state changes, and the nine sweep failures are confirmed pre-existing against a clean tree. What changed since round 2 is entirely external: main moved, and it moved into the same file. The finding that matters is (1), and it's a rebase rather than a redesign — but every number in the table above predates it, so they're evidence about the old base, not a green light for the merged result. Verdict: NEEDS CHANGES — solely on (1).
|
My current plan is to just have the core changes + Ideogram 4 changes (as proof of concept) + LTX-2 changes (since the CI will be red otherwise), and hold off on changes to other pipelines for now. Would that work? |
But this should be propagated to most pipelines, right? |
|
Right, I will follow up for the other pipelines either on this PR or a follow-up PR after the mixin migration is finished (or sufficiently stable). |
| - 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. |
There was a problem hiding this comment.
Let's provide an example here for easier navigation.
| 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 |
There was a problem hiding this comment.
What is driving here? I think it's better to clarify that.
| group_offloading_onload_component_names = [ | ||
| *BasePipelineTesterConfig.group_offloading_onload_component_names, | ||
| "audio_vae", | ||
| ] |
There was a problem hiding this comment.
I would prefer explicitly defining them.
| # 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"] |
There was a problem hiding this comment.
We should keep them empty IMO to have users / agents explicitly set them.
| 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) |
There was a problem hiding this comment.
The scope of self.group_offloading_exclude_modules and this method are kind of intertwined. Can we split the method into two?
- One that prepares
exclude - One that takes
excludeandpipeand raises if needed.
2 participants
What does this PR do?
This PR splits the
GroupOffloadTesterMixin.test_group_offloading_inferencetest into separate block-level and leaf-level tests:test_group_offloading_inference_block_leveltest_group_offloading_inference_leaf_levelwith the motivation that some components tend be block-level offloadable while not being leaf-level offloadable. To help support this, the PR makes the current hardcoded component list used by that test configurable via
BasePipelineTesterConfig:group_offloading_leaf_level_exclude_modules: excludes modules from leaf-level offloading tests, added previously in [tests] Add Config to Exclude Modules from Leaf-Level Group Offloading #14564group_offloading_exclude_modules: excludes modules from both leaf-level and block-level offloading testsgroup_offloading_onload_component_names: includes modules which should be kept on the accelerator rather than offloaded; the motivating case isvaecomponents, which do not offload correctly if tiling is enabledThe new changes are then applied to the Ideogram 4 tests (the motivating example) and the LTX-2 tests (which overrides
test_group_offloading_inference, so needs to be changed due to that test name no longer existing on the PR).Fixes #14618
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