Add Group Offloading Summary Utility - #14702
Conversation
Group offloading leaves no way to ask what grouping it installed: the structure lives on the `GroupOffloadingHook` in each participating submodule's `_diffusers_hook` registry, several submodules share one `ModuleGroup`, and the groups hold module objects rather than names. So a device-mismatch failure names only the layer that tripped over it, never the group whose weights were never onloaded. Add `_get_group_offload_summary` next to the existing introspection cluster (`_get_top_level_group_offload_hook` and friends). It dedupes the groups and resolves the module objects back to qualified names, reporting each group against the module whose `forward` actually brings its weights over. That is not always the group's own leader: with `use_stream=True`, lazy prefetch sets `onload_self = False` and the group is onloaded by whichever group names it as `next_group`, so the prefetch chain is read off the per-module hooks rather than the shared group. That chain is wired only once a module's `forward` has completed once, which the docstring calls out — a module that had not run yet when the failure hit reports every group as onloading itself, and reading that as "this module is not prefetching" would be wrong. A group can reach modules that are not under the one being walked, so calling this on a submodule of an offloaded model — one component of a pipeline, say — finds the group its siblings belong to and cannot name them. Those are reported as `<outside this module>` rather than raising, since a debugging aid that dies when pointed at the thing being debugged is no use. It stays private. What a component reports would shift the day it gains `_group_offload_block_modules`, and reporting the grouping is all this can soundly do — whether a component is safe to offload at block level depends on which entry points the pipeline drives it through, which the module graph alone does not say. A component can group normally and still break: `PriorTransformer` splits its `transformer_blocks` into per-block groups yet reads `clip_mean` / `clip_std` from the root group in `post_process_latents`, and `AutoencoderKL` groups entirely into its `down_blocks` / `up_blocks` while `decode` bypasses `forward`. Use it in the pipeline group offload tests, where a `RuntimeError` from the offloaded run is now re-raised with the grouping of every offloaded component attached. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`_get_group_offload_summary` is only read once something has already gone wrong, so a regression in it is invisible: a summary is still produced, it is just wrong, and it sends the reader to the wrong module. Nothing else in the suite would notice. Cover what it claims rather than how it formats it — the number of groups, and which module's `forward` brings each one over — asserting against the groups read straight off the hooks. Only "no group offloading applied" and "prefetched by" are matched as text, so the layout stays free to change. The two cases that carry their weight are prefetching, which is reported only once the first forward has wired the chain, and summarizing a submodule whose group reaches outside it. Both were live bugs in the helper: reverting either fix fails exactly its own test and no other. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Self-Review ReportSelf-reviewRan the What this adds
Group offloading installs its grouping on the Each group is reported against the module whose Non-blocking findings1. The test oracle shares the implementation's algorithm. 2. 3. The handler decorates every Fixed during review
What the tests actually pinThe helper only runs once something has already failed, so a regression in it is invisible — a summary is still produced, it is just wrong, and it sends the reader to the wrong module. Nothing else in the suite would notice. Each fix was reverted in turn to confirm the tests catch it:
Each fails exactly one test and leaves the others green. The tests assert against groups read straight off the hooks; only Open questionsShould the Should the leading underscore come off? The original argument for keeping it private was that Should there be a test for the exclusion lists? An earlier draft added one asserting that a config's block-level exclusions are entirely root-gated. I dropped it: the predicate is unsound. Verification
|
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
What does this PR do?
This PR adds a group offloading summary utility
_get_group_offload_summarywhich outputs a string showing the offloading groups of a module. The motivation is to allow users to easily see the created offloading groups and to aid in debugging group offloading bugs. It can also be used for more informative test messages (an example is inGroupOffloadTesterMixin._run_group_offload_inferencein this PR).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