[modular] improve auto offload - #14304
Draft
yiyixuxu wants to merge 14 commits into
Draft
Conversation
…tive add/remove - each offloading decision checks the memory actually available on the device (`mem_get_info` free plus the allocator's reusable cache) and keeps a tunable `memory_reserve` of it free; `memory_reserve_margin` is renamed to `memory_reserve` - if a forward still runs out of device memory, offload the smallest model on the device and retry, escalating one model at a time until it fits (`retry_on_oom=True` by default); when nothing else is resident, point at group offloading - adding or removing a component attaches/detaches a single hook instead of re-running `enable_auto_cpu_offload` and offloading the resident working set - add a `simulate_accelerator_memory` test util that makes a large accelerator behave like a smaller card, so offloading can be exercised with real models Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
github-actions
Bot
added
size/L
|
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. |
…ments `memory_budget` (and with it the budget/dynamic split) is not part of this PR, so there is only one strategy behavior left to describe. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`enable_auto_cpu_offload` makes its decisions silently, so there is no way to tell a well-tuned `memory_reserve` from one that thrashes, and tests have to monkeypatch internals to see the sequence. Add an always-on `OffloadRecord` on the manager: every onload, offload and OOM retry is appended as an `OffloadEvent` carrying the model, its size, the reason (`needed_by:<model>`, `oom_retry:<model>`, `component_added`, `offloading_disabled`), and what the strategy saw when it decided. Printing `manager.offload_record` gives a table plus a summary (bytes moved, OOM retries, peak co-residency); the event log is bounded and notes what it dropped. Add opt-in `measure_activations=True`, which brackets each forward with `reset_peak_memory_stats()` and reports `activation_peak` and `suggested_memory_reserve` — the number `memory_reserve` is supposed to cover, measured on the user's own hardware at their own settings, instead of a calibration run that assumes spare memory. `AutoOffloadStrategy.last_decision` is cleared at the top of `__call__` so an early-returning decision cannot inherit the previous call's readings. The two test spies that monkeypatched `UserCustomOffloadHook.offload` and `AutoOffloadStrategy.__call__` now read the record instead. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
yiyixuxu
marked this pull request as draft
July 28, 2026 22:18
Each printed row is now one decision: the model that loaded, what was evicted to make room for it (folded from its needed_by offload events), and the free memory the decision saw. Offloads an onload did not cause (OOM retry, disabling offloading) keep their own rows, and column widths follow the content. The event log and summary() are unchanged. The doc example is replaced with real output from running Z-Image-Turbo under auto offload on a (simulated) 20GB card - the previous hand-written table showed sizes that contradicted the component listing above it and an Available value on an offload row, which never has one. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…measurement; cover encode/decode in the OOM retry - format_size, a shared format_table renderer, and summarize_dict_by_value_and_parts (simplified; its no-common-prefix branch wrote a stale loop variable) move to components_manager_utils.py. Both table reprs render through format_table, the components table Size column uses format_size, and get_model_info drops a deepcopy of the attention processors. Doc tables regenerated from a real run. - measure_activations is removed: the reserve can be deduced from torch's own peak counter plus the weights the record already shows, so the flag, the per-load peak-stat resets, and OffloadRecord.activations go away; the doc shows the deduction on a real Z-Image run (peak 14.15 GB - 11.6 GB resident weights -> ~2.5 GB decode headroom, covered by the default 3GB reserve). - The OOM retry now wraps encode/decode too: autoencoders enter through apply_forward_hook, which fires pre_forward but routes around forward, so a VAE decode OOM was previously never retried. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ry itself - OffloadEvent drops data nothing read: seconds, the constant reason="forward" on onloads, the OOM message, and memory_reserve (static config). The Reason column now only speaks on moves an onload did not cause (oom_retry:<model>, offloading_disabled); the doc points at set_verbosity_info() for watching moves live. - AutoOffloadStrategy.last_decision is gone: pre_forward reads the available memory itself just before the eviction decision, via a shared available_device_memory() helper (driver-free + reusable allocator cache) used by both the hook and the strategy. Every onload row now records a reading - including the first, which previously showed "-" - and custom strategies get the Available column for free. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
An event is now only ever a move - action "onload" or "offload" - with reason explaining why. The oom pseudo-action is gone (an OOM appears as the eviction it caused, reason oom_retry:<model>, which summary() now counts), and OffloadEvent drops the redundant offloaded/resident_before tuples: each eviction is its own event whose needed_by:<onloader> reason names its cause, and the printed table correlates them back into one row per decision. The dropped counter and truncation notice are removed (bounded deque truncates silently at 10k events). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The record is now just the bounded event deque plus the printed table - summary(), the repr's footer line, and __len__ are gone. Consumers read events directly; the tests replay peak co-residency from the moves with a small helper. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- apply_forward_hook stamps its wrapper with _is_forward_entry_point, so wrap_forward statically discovers which methods besides forward are device entry points (autoencoders' encode/decode) and wraps them all at attach for the OOM retry - no hardcoded name list, first calls included. The retry loop moves to a _with_oom_retry method. - Eviction reason renamed needed_by:<model> -> release_memory_for:<model>. - MAX_RECORDED_OFFLOAD_EVENTS becomes OffloadRecord.MAX_EVENTS. - pre_forward reads available memory unguarded (enable_auto_cpu_offload already rejects devices without mem_get_info); the cpu hook tests keep their fake mem_get_info alive with an autouse fixture instead. - resident_other_hooks() inlined into its only caller. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Checked every method against its real consumers (diffusers itself and the Mellon ModularDiffusers nodes): - search_components and its pattern-matching machinery (wildcards, !, |) are removed: nothing used them - Mellon fetches by component_id, the one in-repo caller was get_one's search path, which now goes through _lookup_ids (exact name/collection/load_id). - get_components_by_names is removed (no callers anywhere); the zh doc's workflow example now builds the dict with a get_one comprehension. - Single-caller helpers inlined: _attach_offload_hook into add(), _detach_offload_hook into remove(), get_ids into (the late) get_components_by_names. Kept with receipts: remove_from_collection (Mellon x3), _lookup_ids (Mellon x3 + 3 internal callers), get_components_by_ids (Mellon x8), get_one/get_model_info/enable+disable_auto_cpu_offload (Mellon + tests). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…arts - OffloadEvent recording moves to the site of each decision; offload() is a pure move, and every event (onload and offload) carries the free-memory reading taken just before its move. The repr's Available column shows the decision's first reading, so the table keeps its meaning. - add()'s offload(reason="component_added") was dead code: accelerate's add_hook_to_module already moves the model to CPU at attach, so the call never moved or recorded anything. Deleted. - One device rule: normalize_execution_device() (accelerators indexed, cpu not - the form tensors report) is shared by enable_auto_cpu_offload and CustomOffloadHook, and every resident check is now a plain ==. The hook defaults via get_device(); the PartialState import is gone. - Removed unused set_strategy; _offload_retry_on_oom is None outside the enabled window; remove() finds the hook by component id; pre_forward got a docstring; the strategy docstring no longer describes hook behavior; stale pre-PR comment deleted; "evict" replaced by "offload" throughout. - Tests: the four record tests that relied on cpu != cpu:0 to observe onloads moved to the accelerator section under simulated pressure, and disable's recorded final move is now asserted; the strategy-record test also checks the offload event's own reading. Verified with a real Z-Image-Turbo run on a simulated 20GB card: same 3-row table, peak 14.15 GB as documented. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Four focused test groups: registry (no hardware), strategy unit tests (cuda descriptor + scripted readings), simulate_accelerator_memory validation, and TestAutoOffload (real device moves), plus the pipeline offload mixin. House-style require_accelerate/require_accelerator decorators per test. - The OOM-retry tests now recover from real torch.OutOfMemoryErrors on a hard-capped simulated card (weights fit, the forward's output does not; one eviction is exactly what makes the retry fit) instead of scripted fake OOMs - the fake-OOM wrapper and its choreography tests are gone. - One patch helper (_patch_memory_stats) replaces the three reading fakes; _simulate_card_with_headroom sizes a simulated card relative to whatever the device currently holds. - Patterns adopted from the group-offloading tests: output equivalence (an OOM-survived run reproduces a plain run's output bit-for-bit), torch.no_grad around measured phases, hooks-installed guards, backend_* memory helpers, and a three-tier peak-memory test (baseline > partial offloading on a 160MB card > fully serialized on an 80MB card). - New coverage: enable rejects backends without mem_get_info; adding or removing models mid-run keeps residents in place and links hooks correctly; per-record asserts walk the offload record event by event. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.