[tests] Split Pipeline Group Offloading Block-Level and Leaf-Level Tests (#14635)
* tests: split block/leaf group offloading and derive the offloaded components
`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>
* style: fix `create_pipe` indentation in the group offload tester
The method body was indented one level too deep, which `ruff format` rewrites.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* tests: give group offloading one exclusion list per level
`group_offloading_onload_component_names` named an implementation detail —
components the tests keep resident — and justified itself with a mechanism
that cannot occur: lazy prefetch mis-tracing a tiled VAE decode needs
`use_stream=True`, which no test on this mixin sets, and most of the
pipelines it covers never enable tiling either.
What it actually encodes is a level-specific incapability, the mirror of
`group_offloading_leaf_level_exclude_modules`. Leaf-level offloading onloads
each leaf on its own `forward`, so it breaks on compute that reads a leaf's
`.weight` directly. Block-level onloads a group when the group's leader runs
its `forward`, so it breaks on compute that re-enters submodules without
going through that leader — which is what a VAE decode path does. Rename it
to `group_offloading_block_level_exclude_modules` and scope it to the level
it describes. `group_offloading_exclude_modules` is then redundant, since a
component that fails at both levels goes in both lists, and it never had a
user. `vqvae` leaves the default: no pipeline in the new harness has one.
Measured across the suite, offloading the `vae` at block level breaks 28 of
86 test classes, while all 86 offload it at leaf level with no failures --
coverage the component-scoped tests were skipping and the pipeline-level
test has always had. Sweep before and after: the same 9 pre-existing
failures, 782 -> 810 passing.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Move group offloading changes to shared LTX-2 test mixins
* tests: cover the group offloading stream prefetch path
`use_stream=True` is the only way into `LazyPrefetchGroupOffloadingHook`,
which traces the group execution order on the first forward pass, wires each
group to its successor and flips `non_blocking` for the passes after. Nothing
in the pipeline suite reached that code, which is how a stale comment about
streams mis-tracing a tiled VAE decode survived unexamined for so long.
Add `test_group_offloading_inference_block_level_streaming` and
`..._leaf_level_streaming` alongside the synchronous pair, as separate methods
rather than a parametrized test so a pipeline can skip one level without
losing the other. Both route through `_run_group_offload_inference`, so they
compare against the same `base_pipe_output` and honour the same per-level
exclusion lists.
`apply_group_offloading` raises rather than degrading when there is nowhere to
put a stream, and `@require_torch_accelerator` still passes on MPS, so the two
tests also check for CUDA or XPU before running.
Guard against the flag quietly ceasing to have an effect: after enabling
offloading with `use_stream=True`, assert at least one group across the
pipeline actually carries a stream. The count is pipeline-wide because a
component with no `ModuleList`/`Sequential` children has nothing to prefetch
and legitimately streams nothing, and it asserts rather than skips because a
regression that drops `use_stream` would zero every count and a skip would
turn the whole streaming suite green by omission.
Sweep over `tests/pipelines -k group_offloading`: 31 failed, 1145 passed,
identical to the run without these tests. All 13 streaming failures have a
non-streaming twin that fails the same way, so the prefetch path adds no
failure of its own.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* tests: port the group offload overrides orphaned by the level split
Splitting `test_group_offloading_inference` into a block-level and a
leaf-level test removed the name six other test files override, so their
markers guarded a method the mixin no longer defines while the real tests ran
unprotected. The `xfail(strict=True)` ones hid it: `super()` raised
`AttributeError`, which counts as the expected failure, so they reported green
while testing nothing.
Rechecking each reason rather than renaming mechanically, three of the six
turned out to describe behaviour that no longer exists:
`LTX2DFRPipelineTesterConfig` subclasses `BasePipelineTesterConfig` rather
than `LTX2BaseTesterConfig`, so the family's `audio_vae` block-level exclusion
never reached it and its audio VAE was being offloaded. It needs the
declaration, not a skip, and its group offload tests now pass.
AudioLDM2 and Motif Video only fail at block level. `get_text_features()` and
the vision tower bypass the `forward` that block-level offloading gates the
group on, while leaf level onloads each leaf on its own `forward` and passes,
so the skip and the xfail are scoped to block level and the leaf-level tests
now run.
The Kandinsky prior xfails stay at both levels but their reason was describing
the hardcoded component list this branch removed. The actual cause is that the
pipeline calls `PriorTransformer.post_process_latents()` after the denoising
loop, reading `clip_mean` / `clip_std` that group offloading onloads only for
the duration of `forward`.
Also drop the class-count measurement from the block-level exclusion comment:
it was already stale one merge later. The mechanism it was evidence for —
`vae.decode()` never running `vae.forward()` — is the durable part.
Sweep over `tests/pipelines -k group_offloading`: 31 failed -> 13, all of them
failing identically on `main`, except `kolors`'s pipeline-level test, which
fails in isolation on `main` too and only passes there because of the test
order the split changed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Clarify the Ideogram 4 block-level group offloading module exclusion comment
* tests: parametrize `use_stream` instead of splitting the streaming tests
The streaming tests were separate methods so a pipeline could skip or xfail
them independently of their synchronous twins. Nothing has ever needed that:
across three full sweeps every streaming failure had a synchronous twin at the
same level and vice versa, and none of the sixteen downstream overrides marked
a streaming variant differently from its twin. The two *levels* do get scoped
apart — AudioLDM2 and Motif Video are block-level only — so those keep their
own methods.
Fold `use_stream` into a parameter on each level's test. Four mixin tests
become two and sixteen override methods become eight, with the same coverage:
a full sweep of the parametrized and the separate-method versions on the same
tree gives 15 failed, 1263 passed, 152 skipped, 22 xfailed, 1 xpassed on both,
with identical failure sets once `[stream]` and the `_streaming` suffix are
normalised.
Marks do not carry through a subclass override, so an override that forgets to
re-declare the parametrize collapses to one un-parametrized test that errors on
the missing argument — which an `xfail` then reports as green, the same failure
mode the orphaned overrides had. `_USE_STREAM` carries a note saying so, and the
six overrides re-declare it.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* tests: point the PNDM marker comment at the split group offload tests
`test_group_offloading_inference` no longer exists — it became one test per
offload level — so the comment listing which tests share the class-level xfail
named a test that is not there, and undercounted them by one.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* docs: record the group offload streaming parameter in the testing guide
The shared memory mixin now runs each offload level with and without
`use_stream`, which is the only path that reaches group offloading's lazy
prefetch. Say so next to the exclusion-list guidance, along with the reason it
is a parameter rather than a separate test, and the requirement that a subclass
overriding either test re-declare the `parametrize` — marks do not carry
through an override, and forgetting leaves a test that errors on the missing
argument and is reported as green by an `xfail`.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* tests: inline the single-caller group offload helpers
`_split_group_offload_components` and `_count_streamed_groups` each had one
caller, so per `code_style.md` they read better at the call site. The reasons
they existed survive as comments: that a component is offloaded unless the
level's list names it, and that several submodules share one `ModuleGroup` so
the stream count has to dedupe on group identity.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* docs: condense the group offload guidance to one bullet
The three group offload bullets in the testing guide had grown to explain
the hook mechanism at a level of detail that answers a reviewer's questions
rather than telling a test author what to write. Collapse them into one
bullet keeping the parts that change what a PR does: which remedy to reach
for, which exclusion list matches which hazard, that an unmatched name fails
as a typo, and that an override must re-declare the `use_stream` parametrize.
The worked examples below the bullet already carry the concrete diagnosis,
and the parametrize trap is explained in full at the code site.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Explicitly declares group offloading leaf level excluded components for LTX-2
* Simplify comment regarding _USE_STREAM in GroupOffloadTesterMixin
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>