diffusers
e9161bfc - [tests] Split Pipeline Group Offloading Block-Level and Leaf-Level Tests (#14635)

Commit
26 days ago
[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>
Author
Parents
Loading