[modular] improve auto offload

[yiyixuxu]

What does this PR do?

this PR improves ComponentsManager's auto offloading, main features includes:

• automatic OOM recovery
• keep a record of every offloading decision,
• clean up some unused method in ComponentsManager
• a slow test suite for modular pipeline that include offloading behavior test for different GPUs
• updated docs.

Offloading behavior

• memory_reserve (renamed from memory_reserve_margin, default "3GB", accepts size strings): each offloading decision keeps this much memory free as headroom for activations. The docs explain how to size it for a workload using the record's peak column.
• OOM retry (retry_on_oom=True): memory_reserve is an estimate, so a forward pass can still run out of memory. The manager now recovers on its own, when it runs into a OOM, it offloads the smallest resident model and reruns the forward pass, escalating one model at a time until it fits. If nothing is left to offload, it re-raises with advice to use group offloading. This also covers autoencoder encode/decode entry points (apply_forward_hook now marks its wrappers so the hook machinery can discover them).
• Non-disruptive add/remove: adding a component while offloading is enabled hooks it in without rebuilding existing hooks or moving resident models; removing detaches only that component.

Offload record

Every move is now recorded. print(manager.offload_record) shows one row per decision — what loaded, what was offloaded to make room, and the memory picture the decision was based on (Available free memory and Peak allocated memory), plus a Reason column for moves an onload didn't cause (oom_retry:<model>, offloading_disabled):

# | Onload                                 | Offloaded                              | Available | Peak     | Reason
-------------------------------------------------------------------------------------------------------------------
1 | text_encoder_129824856400544 (7.49 GB) | -                                      | 18.02 GB  | 8.50 KB  |
2 | transformer_129824854851072 (11.46 GB) | text_encoder_129824856400544 (7.49 GB) | 10.38 GB  | 7.77 GB  |
3 | vae_129824905851968 (159.87 MB)        | -                                      | 6.36 GB   | 12.07 GB |

The docs walk through reading it: spotting thrashing, finding the run's memory bottleneck with the peak column, and diagnosing an under-sized memory_reserve from an oom_retry row.

Custom offload strategies

New ComponentsManager.set_offload_strategy(): the decision-maker is any callable (hooks, model_id, model, execution_device) -> hooks to offload. The docs show a one-line sequential strategy (equivalent to enable_model_cpu_offload for standard pipelines) and a policy example that keeps co-scheduled models resident together.

ComponentsManager cleanup

get_components_by_names and get_one's pattern matching are removed (get_one is exact-match), OffloadRecord.summary() and unused record fields are gone, display helpers moved to components_manager_utils.py, and device handling is unified in one normalize_execution_device rule. The zh doc is synced for the removals.

Tests

• New simulate_accelerator_memory testing util: makes the device behave like a card of a given size, this allows use to run integration tests use real checkpoints "on" different consumer-sized cards.
• test_components_manager.py reorganized into focused classes (registry / strategy unit tests / util validation / real-accelerator behavior)
• New slow ModularPipelineIntegrationTesterMixin: subclasses declare the expected offloading behavior per simulated card (offload counts, OOM counts, final devices) and the test verifies the record against the declaration plus output equivalence with a fully-resident baseline. First user: Z-Image (test_modular_pipeline_z_image.py), which also gains basic t2i/i2i slice tests. All expected values were captured on real hardware (H100).

Docs

docs/source/en/modular_diffusers/components_manager.md: the offloading section now covers memory_reserve and OOM retry, reading the record, writing and setting a custom strategy, finding the bottleneck with the peak column, and tuning the reserve from an OOM-retry record. All tables in the doc are genuine Z-Image-Turbo output on simulated cards.

Follow up TODOs

• group offloading + auto offloading [Modular] support group offloading in auto offloading #14346
• use diffusers hooks [modular] Migrate ComponentsManager offload hooks to ModelHook / HookRegistry #14328

[HuggingFaceDocBuilderDev]

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.

[yiyixuxu]

I broke this up into a couple of tests and moved things around a bit, all the tests are still there, just got moved to different places

• TestComponentsManagerRegistry: test stuff like add/remove components
• TestAutoOffloadStrategy: test the strategy itself
• TestAutoOffload: unit test for the enable_auto_cpu_offload feature with dummy models and simulated memory

[yiyixuxu]

@sayakpaul I kept this but I think maybe we can remove this test now that we are going to have an integration test that tests real checkpoints. Let me know

[stevhliu]

thanks!

[stevhliu]

Suggested change

The reserve is an estimate, so a forward pass can still run out of memory. By default the manager recovers on its own: it offloads the smallest model on the device and retries the forward pass, escalating one model at a time until it fits — pass `retry_on_oom=False` to raise the error instead. A model that still doesn't fit once everything else is offloaded is too large for the device on its own; use [group offloading](../optimization/memory#group-offloading) for it.

The reserve is an estimate, so a forward pass can still run out of memory. By default the manager recovers on its own: it offloads the smallest model on the device and retries the forward pass, escalating one model at a time until it fits. OOM retry is intended for inference only because each retry reruns the forward call from its original inputs. Set `retry_on_oom=False` during training or when the forward pass is not safe to repeat. A model that still doesn't fit once everything else is offloaded is too large for the device on its own; use [group offloading](../optimization/memory#group-offloading) for it.

[stevhliu]

Suggested change

	This strategy offloads and onloads models in the sequence they are called (pretty much what [`~DiffusionPipeline.enable_model_cpu_offload`] does for standard pipelines). Use it like this:
	This strategy offloads and onloads models in the sequence they are called (like [`~DiffusionPipeline.enable_model_cpu_offload`] for standard pipelines). Use it like this:

[stevhliu]

I think we can clarify how to interpret Peak and memory_reserve a bit more. For example, Peak is cumulative since process start but here it reads a bit like a per-run activation measurement (same for memory_reserve).

This may also work better as a sequential list instead of a paragraph :)

[stevhliu]

This was a bit long and dense, I think we can let the table show and explain instead of explicitly spelling out whats happening.

Suggested change

Each row is one decision: the model that loaded, what was offloaded to make room for it, and the memory picture (`Available` and `Peak`, read just before the decision's first move) it was based on. Here the transformer found 10.38 GB free — not enough for its 11.46 GB of weights while keeping the 3GB `memory_reserve` — so the text encoder was offloaded first. The VAE then fit into the 6.36 GB left next to the transformer, so it loaded without pushing anything off. The `Reason` column stays empty for these planned moves; it is filled when a model is offloaded for any other reason: `oom_retry:<model>` when a forward pass ran out of memory and offloading was the rescue (covered below), or `offloading_disabled` when [`~ComponentsManager.disable_auto_cpu_offload`] moves everything back to CPU. To watch the moves live as they happen instead, enable info logging with `diffusers.logging.set_verbosity_info()`.

Each row represents one offload decision. `Onload` names the model moved to the device, and `Offloaded` lists the models moved to CPU first. `Available` and `Peak` show the memory readings before the decision. Planned offloads leave `Reason` blank. The field is populated for OOM retries and for moves caused by disabling offloading.

[stevhliu]

Suggested change

A model appearing repeatedly in this table is thrashing — it is being offloaded and re-loaded every step, which costs a PCIe transfer each way. That usually means `memory_reserve` is too large (models are pushed off that would have fit) or too small (each step ends in an OOM retry). It can also mean the default strategy is making the wrong call for your workload: a model that runs again inside the same denoise loop should stay resident even when evicting it looks fine on memory alone. That is a case for a custom strategy, which the next section covers.

A model appearing repeatedly in this table is thrashing — it is being offloaded and reloaded every step, which costs a PCIe transfer each way. That usually means `memory_reserve` is too large (models are pushed off that would have fit) or too small (each step ends in an OOM retry). It can also mean the default strategy is making the wrong call for your workload: a model that runs again inside the same denoise loop should stay resident even when evicting it looks fine on memory alone. That is a case for a custom strategy, which the next section covers.