From 780e112c8f504cd38989d7747fc21e810b587a5d Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 27 Jul 2026 21:26:21 +0000 Subject: [PATCH 01/14] [modular] auto offload: tunable memory reserve, OOM retry, non-disruptive 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 --- .../modular_diffusers/components_manager.md | 10 +- .../modular_pipelines/components_manager.py | 223 +++++++++--- .../test_components_manager.py | 334 ++++++++++++++++-- tests/testing_utils.py | 49 +++ 4 files changed, 540 insertions(+), 76 deletions(-) diff --git a/docs/source/en/modular_diffusers/components_manager.md b/docs/source/en/modular_diffusers/components_manager.md index 426739347f27..2bac75a876c6 100644 --- a/docs/source/en/modular_diffusers/components_manager.md +++ b/docs/source/en/modular_diffusers/components_manager.md @@ -87,7 +87,15 @@ The [`~ComponentsManager.enable_auto_cpu_offload`] method is a global offloading manager.enable_auto_cpu_offload(device="cuda") ``` -All models begin on the CPU and [`ComponentsManager`] moves them to the appropriate device right before they're needed, and moves other models back to the CPU when GPU memory is low. +All models begin on the CPU and [`ComponentsManager`] moves them to the appropriate device right before they're needed, and moves other models back to the CPU when the incoming model wouldn't fit. Each offloading decision checks the memory actually available on the device at that moment and keeps `memory_reserve` (3GB unless you change it) of it free as headroom for activations, which scale with resolution, batch size, and sequence length rather than with the number of models — image models around 1024px are comfortable with a few GB, while video models need more. + +```py +manager.enable_auto_cpu_offload(device="cuda", memory_reserve="1GB") +``` + +Set `memory_reserve=0` to keep as much on the device as possible. If a forward pass runs out of memory anyway, the manager offloads the smallest model on the device and retries it, 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. + +Components added while offloading is enabled join the managed set without disturbing it: the new component starts on CPU and everything already resident stays where it is; removing a component likewise detaches only that component. Call [`~ComponentsManager.disable_auto_cpu_offload`] to disable offloading. diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 31ba2c942203..ab3a080805e4 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -15,6 +15,7 @@ from __future__ import annotations import copy +import functools import time from collections import OrderedDict from itertools import combinations @@ -49,6 +50,9 @@ class CustomOffloadHook(ModelHook): execution_device(`str`, `int` or `torch.device`, *optional*): The device on which the model should be executed. Will default to the MPS device if it's available, then GPU 0 if there is a GPU, and finally to the CPU. + retry_on_oom(`bool`, *optional*, defaults to `True`): + Whether to recover from a forward pass that runs out of device memory by offloading the other models one + at a time and retrying. If `False`, the error is raised to the caller. """ no_grad = False @@ -58,11 +62,14 @@ def __init__( execution_device: str | int | torch.device | None = None, other_hooks: list["UserCustomOffloadHook"] | None = None, offload_strategy: "AutoOffloadStrategy" | None = None, + retry_on_oom: bool = True, ): self.execution_device = execution_device if execution_device is not None else PartialState().default_device self.other_hooks = other_hooks self.offload_strategy = offload_strategy + self.retry_on_oom = retry_on_oom self.model_id = None + self._forward_before_oom_wrap = None def set_strategy(self, offload_strategy: "AutoOffloadStrategy"): self.offload_strategy = offload_strategy @@ -107,6 +114,69 @@ def pre_forward(self, module, *args, **kwargs): module.to(self.execution_device) return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device) + def resident_other_hooks(self): + """ + The other managed models that are currently on the execution device. + """ + return [ + other + for other in (self.other_hooks or []) + # compare with a normalized index: model.device reports no index for the default device + if other.model.device.type == self.execution_device.type + and (other.model.device.index or 0) == (self.execution_device.index or 0) + ] + + # YiYi TODO: diffusers' own `ModelHook` (`diffusers.hooks.hooks`) supports a `new_forward` around-hook, which + # would replace this manual wrapping - we may migrate this class to it in the future. + def wrap_forward(self, module): + # `pre_forward`/`post_forward` cannot see an exception raised by the forward itself, so wrap the (already + # hooked) forward here: if it runs out of device memory, free the smallest resident model and retry, + # escalating until it fits. The memory readings cannot guide this - the failed forward has already unwound, + # so its activations are freed and the device looks free again. Inference only: the retried forward re-runs + # cleanly from its original inputs. + if not self.retry_on_oom: + return + + hooked_forward = module.forward + + @functools.wraps(hooked_forward) + def forward_with_oom_retry(*args, **kwargs): + # each pass offloads one more model, so this terminates once they all have been. Tracked explicitly + # rather than by device, so that a model that did not actually move cannot be picked twice. + offloaded = set() + while True: + try: + return hooked_forward(*args, **kwargs) + except torch.OutOfMemoryError as e: + resident = sorted( + (hook for hook in self.resident_other_hooks() if id(hook) not in offloaded), + key=lambda hook: hook.model.get_memory_footprint(), + ) + if not resident: + raise torch.OutOfMemoryError( + f"{self.model_id} ran out of device memory ({e}) with every other managed model already " + "offloaded, so it does not fit on its own. Consider group offloading " + "(`ModelMixin.enable_group_offload`), which offloads a single model in groups of " + "internal layers." + ) from e + smallest = resident[0] + logger.warning( + f"{self.model_id} ran out of device memory ({e}); offloading {smallest.model_id} and " + "retrying. If this happens repeatedly, set a larger `memory_reserve` in " + "`enable_auto_cpu_offload`." + ) + smallest.offload() + offloaded.add(id(smallest)) + clear_device_cache() + + module.forward = forward_with_oom_retry + self._forward_before_oom_wrap = hooked_forward + + def unwrap_forward(self, module): + if self._forward_before_oom_wrap is not None: + module.forward = self._forward_before_oom_wrap + self._forward_before_oom_wrap = None + class UserCustomOffloadHook: """ @@ -125,8 +195,10 @@ def offload(self): def attach(self): add_hook_to_module(self.model, self.hook) self.hook.model_id = self.model_id + self.hook.wrap_forward(self.model) def remove(self): + self.hook.unwrap_forward(self.model) remove_hook_from_module(self.model) self.hook.model_id = None @@ -139,8 +211,11 @@ def custom_offload_with_hook( model: torch.nn.Module, execution_device: str | int | torch.device = None, offload_strategy: "AutoOffloadStrategy" | None = None, + retry_on_oom: bool = True, ): - hook = CustomOffloadHook(execution_device=execution_device, offload_strategy=offload_strategy) + hook = CustomOffloadHook( + execution_device=execution_device, offload_strategy=offload_strategy, retry_on_oom=retry_on_oom + ) user_hook = UserCustomOffloadHook(model_id=model_id, model=model, hook=hook) user_hook.attach() return user_hook @@ -149,14 +224,17 @@ def custom_offload_with_hook( # this is the class that user can customize to implement their own offload strategy class AutoOffloadStrategy: """ - Offload strategy that should be used with `CustomOffloadHook` to automatically offload models to the CPU based on - the available memory on the device. + Offload strategy that should be used with `CustomOffloadHook` to automatically offload models to the CPU so the + incoming model fits on the device: at each offload decision, check the memory actually available on the device + and keep `memory_reserve` of it free. + + The sizes cover the weights managed by this strategy only — the actual memory requirements will include + activations and any other allocations, so `memory_reserve` covers exactly that headroom; a `memory_reserve` of 0 + packs the device as full as the weights allow, relying on the OOM retry as a backstop. """ - # YiYi TODO: instead of memory_reserve_margin, we should let user set the maximum_total_models_size to keep on device - # the actual memory usage would be higher. But it's simpler this way, and can be tested - def __init__(self, memory_reserve_margin="3GB"): - self.memory_reserve_margin = convert_file_size_to_int(memory_reserve_margin) + def __init__(self, memory_reserve="3GB"): + self.memory_reserve = convert_file_size_to_int(memory_reserve) def __call__(self, hooks, model_id, model, execution_device): if len(hooks) == 0: @@ -167,18 +245,30 @@ def __call__(self, hooks, model_id, model, execution_device): except AttributeError: raise AttributeError(f"Do not know how to compute memory footprint of `{model.__class__.__name__}.") - device_type = execution_device.type - device_module = getattr(torch, device_type, torch.cuda) - try: - mem_on_device = device_module.mem_get_info(execution_device.index)[0] - except AttributeError: - raise AttributeError(f"Do not know how to obtain obtain memory info for {str(device_module)}.") + resident_size = sum(hook.model.get_memory_footprint() for hook in hooks) + + # Check the memory actually available on the device right now. The allocator keeps freed memory in its + # cache (mem_get_info reports it as used), but it is reusable — count it as available. + device_module = getattr(torch, execution_device.type, torch.cuda) + available_memory = device_module.mem_get_info(execution_device.index)[0] + if hasattr(device_module, "memory_reserved") and hasattr(device_module, "memory_allocated"): + available_memory += device_module.memory_reserved(execution_device) - device_module.memory_allocated( + execution_device + ) - mem_on_device = mem_on_device - self.memory_reserve_margin - if current_module_size < mem_on_device: + if current_module_size <= available_memory - self.memory_reserve: return [] - min_memory_offload = current_module_size - mem_on_device + min_memory_offload = current_module_size - (available_memory - self.memory_reserve) + if min_memory_offload >= resident_size: + logger.warning( + f"fitting {model_id} ({current_module_size / 1024**3:.2f} GB) needs " + f"{min_memory_offload / 1024**3:.2f} GB but only {resident_size / 1024**3:.2f} GB of managed " + "weights are resident, offloading all other models. If it still does not fit, consider group " + "offloading (`ModelMixin.enable_group_offload`)." + ) + return hooks + logger.info(f" search for models to offload in order to free up {min_memory_offload / 1024**3:.2f} GB memory") # exlucde models that's not currently loaded on the device @@ -214,16 +304,11 @@ def search_best_candidate(module_sizes, min_memory_offload): return best_candidate + # a combination is guaranteed to exist: offloading everything frees `resident_size`, which is more than + # `min_memory_offload` (the case where it isn't returned early above) best_offload_model_ids = search_best_candidate(module_sizes, min_memory_offload) - if best_offload_model_ids is None: - # if no combination is found, meaning that we cannot meet the memory requirement, offload all models - logger.warning("no combination of models to offload to cpu is found, offloading all models") - hooks_to_offload = hooks - else: - hooks_to_offload = [hook for hook in hooks if hook.model_id in best_offload_model_ids] - - return hooks_to_offload + return [hook for hook in hooks if hook.model_id in best_offload_model_ids] # utils for display component info in a readable format @@ -331,11 +416,12 @@ class ComponentsManager: def __init__(self): self.components = OrderedDict() - # YiYi TODO: can remove once confirm we don't need this in mellon self.added_time = OrderedDict() # Store when components were added self.collections = OrderedDict() # collection_name -> set of component_names self.model_hooks = None self._auto_offload_enabled = False + self._offload_strategy = None + self._offload_retry_on_oom = True def _lookup_ids( self, @@ -450,8 +536,8 @@ def add(self, name: str, component: Any, collection: str | None = None): else: logger.info(f"ComponentsManager: added component '{name}' as '{component_id}'") - if self._auto_offload_enabled and is_new_component: - self.enable_auto_cpu_offload(self._auto_offload_device) + if self._auto_offload_enabled and is_new_component and isinstance(component, torch.nn.Module): + self._attach_offload_hook(component_id, component) return component_id @@ -491,11 +577,10 @@ def remove(self, component_id: str = None): if component_id in self.collections[collection]: self.collections[collection].remove(component_id) - if self._auto_offload_enabled: - self.enable_auto_cpu_offload(self._auto_offload_device) - else: - if isinstance(component, torch.nn.Module): - component.to("cpu") + if isinstance(component, torch.nn.Module): + if self._auto_offload_enabled: + self._detach_offload_hook(component) + component.to("cpu") del component import gc @@ -692,22 +777,37 @@ def matches_pattern(component_id, pattern, exact_match=False): return get_return_dict(matches, return_dict_with_names) - def enable_auto_cpu_offload(self, device: str | int | torch.device = None, memory_reserve_margin="3GB"): + def enable_auto_cpu_offload( + self, + device: str | int | torch.device = None, + memory_reserve: str | int = "3GB", + retry_on_oom: bool = True, + ): """ Enable automatic CPU offloading for all components. The algorithm works as follows: 1. All models start on CPU by default 2. When a model's forward pass is called, it's moved to the execution device - 3. If there's insufficient memory, other models on the device are moved back to CPU + 3. If it doesn't fit into the memory currently available on the device minus `memory_reserve`, other models + on the device are moved back to CPU first 4. The system tries to offload the smallest combination of models that frees enough memory 5. Models stay on the execution device until another model needs memory and forces them off + 6. If a forward pass still runs out of device memory, the smallest model on the device is offloaded and the + forward is retried, escalating one model at a time until it fits (inference only: each retried forward + re-runs from its original inputs) Args: device (str | int | torch.device): The execution device where models are moved for forward passes - memory_reserve_margin (str): The memory reserve margin to use, default is 3GB. This is the amount of - memory to keep free on the device to avoid running out of memory during model - execution (e.g., for intermediate activations, gradients, etc.) + memory_reserve (str | int, *optional*, defaults to `"3GB"`): + The amount of available device memory to keep free when deciding whether an incoming model fits, + checked at each offloading decision — e.g. `"3GB"` or a number of bytes. The reserve is what covers + allocations the offloading cannot see, mainly activations, which scale with resolution / batch size / + sequence length. Set it to `0` to keep as much on the device as possible, relying on the OOM retry. + retry_on_oom (bool, *optional*, defaults to `True`): + Whether to recover from a forward pass that runs out of device memory by offloading the models on the + device one at a time, smallest first, and retrying until it fits. Set it to `False` to raise the error + instead — the forward passes are then left untouched. """ if not is_accelerate_available(): raise ImportError("Make sure to install accelerate to use auto_cpu_offload") @@ -717,27 +817,28 @@ def enable_auto_cpu_offload(self, device: str | int | torch.device = None, memor if not isinstance(device, torch.device): device = torch.device(device) - device_type = device.type - device_module = getattr(torch, device_type, torch.cuda) + if device.index is None: + device = torch.device(f"{device.type}:{0}") + + device_module = getattr(torch, device.type, torch.cuda) if not hasattr(device_module, "mem_get_info"): raise NotImplementedError( - f"`enable_auto_cpu_offload() relies on the `mem_get_info()` method. It's not implemented for {str(device.type)}." + f"Offloading decisions rely on `mem_get_info()`, which is not implemented for {str(device.type)}." ) - if device.index is None: - device = torch.device(f"{device.type}:{0}") - for name, component in self.components.items(): if isinstance(component, torch.nn.Module) and hasattr(component, "_hf_hook"): remove_hook_from_module(component, recurse=True) self.disable_auto_cpu_offload() - offload_strategy = AutoOffloadStrategy(memory_reserve_margin=memory_reserve_margin) + offload_strategy = AutoOffloadStrategy(memory_reserve=memory_reserve) all_hooks = [] for name, component in self.components.items(): if isinstance(component, torch.nn.Module): - hook = custom_offload_with_hook(name, component, device, offload_strategy=offload_strategy) + hook = custom_offload_with_hook( + name, component, device, offload_strategy=offload_strategy, retry_on_oom=retry_on_oom + ) all_hooks.append(hook) for hook in all_hooks: @@ -749,6 +850,36 @@ def enable_auto_cpu_offload(self, device: str | int | torch.device = None, memor self.model_hooks = all_hooks self._auto_offload_enabled = True self._auto_offload_device = device + self._offload_strategy = offload_strategy + self._offload_retry_on_oom = retry_on_oom + + def _attach_offload_hook(self, component_id: str, component: torch.nn.Module): + """ + Attach an offload hook to a newly added component without disturbing the models already managed: the new + component starts on CPU (like every model under auto offload), everything else stays where it is. + """ + hook = custom_offload_with_hook( + component_id, + component, + self._auto_offload_device, + offload_strategy=self._offload_strategy, + retry_on_oom=self._offload_retry_on_oom, + ) + for other_hook in self.model_hooks: + if other_hook.hook.execution_device == hook.hook.execution_device: + hook.add_other_hook(other_hook) + other_hook.add_other_hook(hook) + hook.offload() + self.model_hooks.append(hook) + + def _detach_offload_hook(self, component: torch.nn.Module): + """Detach a removed component's offload hook, leaving all other managed models where they are.""" + hook = next(user_hook for user_hook in self.model_hooks if user_hook.model is component) + hook.remove() + self.model_hooks.remove(hook) + for other_hook in self.model_hooks: + if other_hook.hook.other_hooks and hook in other_hook.hook.other_hooks: + other_hook.hook.other_hooks.remove(hook) def disable_auto_cpu_offload(self): """ @@ -765,6 +896,8 @@ def disable_auto_cpu_offload(self): clear_device_cache() self.model_hooks = None self._auto_offload_enabled = False + self._offload_strategy = None + self._offload_retry_on_oom = True def get_model_info( self, diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index 42929b04e3d8..506eda308654 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import gc from unittest import mock @@ -22,7 +23,13 @@ from diffusers.models import ModelMixin from diffusers.utils import is_accelerate_available -from ..testing_utils import backend_empty_cache, require_accelerate, require_accelerator, torch_device +from ..testing_utils import ( + backend_empty_cache, + require_accelerate, + require_accelerator, + simulate_accelerator_memory, + torch_device, +) if is_accelerate_available(): @@ -36,6 +43,10 @@ # hardware (an 80GB GPU never runs low on a handful of KB-sized models). UNIT = 1024 +# More free memory than any tiny test checkpoint could ever need, so the strategy never +# decides to offload. Used to assert the *negative*: no eviction without memory pressure. +_AMPLE_FREE_BYTES = 1024**4 + class DummyModel(ModelMixin): def __init__(self, footprint_bytes: int = UNIT): @@ -48,6 +59,19 @@ def forward(self, x): return x + self.weight.sum() +class OOMOnceModel(DummyModel): + """Raises a fake device OOM on its first `_oom_calls` forwards (or on every forward with `_repeated_oom`).""" + + _repeated_oom = False + _oom_calls = 1 + + def forward(self, x): + self.forward_calls = getattr(self, "forward_calls", 0) + 1 + if self.forward_calls <= self._oom_calls or self._repeated_oom: + raise torch.OutOfMemoryError("fake OOM") + return super().forward(x) + + class _FakeHook: """ Minimal stand-in for `UserCustomOffloadHook` in strategy-level unit tests. @@ -62,20 +86,33 @@ def __init__(self, model_id: str, model: torch.nn.Module): self.model = model -def _patch_cuda_mem_get_info(free_bytes: int, total_bytes: int = 80 * UNIT): +@contextlib.contextmanager +def _patch_memory_stats(device_module, free_bytes, total_bytes, cached_bytes=0): + # `mem_get_info` returns `(free, total)` and is where the strategy learns how much + # memory is free; the dynamic mode additionally counts the allocator's reusable + # cache (`memory_reserved - memory_allocated`) as available, so those are pinned + # too — `cached_bytes` simulates reusable cache on top of `free_bytes`. + with contextlib.ExitStack() as stack: + stack.enter_context(mock.patch.object(device_module, "mem_get_info", return_value=(free_bytes, total_bytes))) + if hasattr(device_module, "memory_reserved") and hasattr(device_module, "memory_allocated"): + stack.enter_context(mock.patch.object(device_module, "memory_reserved", return_value=cached_bytes)) + stack.enter_context(mock.patch.object(device_module, "memory_allocated", return_value=0)) + yield + + +def _patch_cuda_mem_get_info(free_bytes: int, total_bytes: int = 80 * UNIT, cached_bytes: int = 0): # Strategy unit tests use a `cuda:0` execution-device *descriptor* (which needs no - # real GPU), so they patch `torch.cuda.mem_get_info` directly. - return mock.patch.object(torch.cuda, "mem_get_info", return_value=(free_bytes, total_bytes)) + # real GPU), so they patch the `torch.cuda` memory introspection directly. + return _patch_memory_stats(torch.cuda, free_bytes, total_bytes, cached_bytes=cached_bytes) def _patch_free_memory(free_bytes: int, total_bytes: int = 80 * UNIT): - # Integration tests run on the real `torch_device`; patch `mem_get_info` on - # whichever backend module (cuda/xpu/...) actually backs it. `mem_get_info` returns - # `(free, total)` and is the single point where the strategy learns how much memory - # is available, so patching it simulates arbitrary memory pressure. + # Integration tests run on the real `torch_device`; patch the memory introspection + # on whichever backend module (cuda/xpu/...) actually backs it to simulate + # arbitrary memory pressure. device_type = torch.device(torch_device).type device_module = getattr(torch, device_type, torch.cuda) - return mock.patch.object(device_module, "mem_get_info", return_value=(free_bytes, total_bytes)) + return _patch_memory_stats(device_module, free_bytes, total_bytes) @require_accelerate @@ -85,7 +122,7 @@ class ComponentsManagerTesterMixin: """ # A `cuda:0` device descriptor is enough to drive the strategy's device-type and - # index logic; no real GPU is required because `mem_get_info` is mocked. + # index logic; no real GPU is required because the memory readings are mocked. strategy_execution_device = torch.device("cuda:0") def setup_method(self): @@ -106,11 +143,19 @@ def get_dummy_model(self, footprint_bytes: int = UNIT) -> ModelMixin: # ------------------------------------------------------------------ # AutoOffloadStrategy unit tests (hardware-independent) # ------------------------------------------------------------------ - def _select_offload(self, *, incoming_footprint, free_bytes, hook_sizes, memory_reserve_margin=UNIT): - strategy = AutoOffloadStrategy(memory_reserve_margin=memory_reserve_margin) + def _select_offload( + self, + *, + incoming_footprint, + free_bytes, + hook_sizes, + memory_reserve=UNIT, + cached_bytes=0, + ): + strategy = AutoOffloadStrategy(memory_reserve=memory_reserve) hooks = [_FakeHook(model_id, self.get_dummy_model(fp)) for model_id, fp in hook_sizes.items()] incoming = self.get_dummy_model(incoming_footprint) - with _patch_cuda_mem_get_info(free_bytes): + with _patch_cuda_mem_get_info(free_bytes, cached_bytes=cached_bytes): selected = strategy( hooks=hooks, model_id="incoming", @@ -166,20 +211,20 @@ def test_strategy_no_hooks_returns_empty(self): ) assert selected == [] - def test_strategy_memory_reserve_margin_changes_decision(self): - # Same device free memory and incoming model; only the reserve margin differs. - # A small margin leaves enough room; a large margin forces an offload. We check - # this both with a single resident model and with several, to confirm the margin + def test_strategy_memory_reserve_changes_decision(self): + # Same device free memory and incoming model; only the reserve differs. + # A small reserve leaves enough room; a large reserve forces an offload. We check + # this both with a single resident model and with several, to confirm the reserve # participates in the selection regardless of how many candidates exist. - # Single candidate: free=5, incoming=3. margin 1 -> usable 4 (fits); margin 3 -> + # Single candidate: free=5, incoming=3. reserve 1 -> usable 4 (fits); reserve 3 -> # usable 2, must free 1 -> offload "a". assert ( self._select_offload( incoming_footprint=3 * UNIT, free_bytes=5 * UNIT, hook_sizes={"a": 2 * UNIT}, - memory_reserve_margin=1 * UNIT, + memory_reserve=1 * UNIT, ) == [] ) @@ -187,10 +232,10 @@ def test_strategy_memory_reserve_margin_changes_decision(self): incoming_footprint=3 * UNIT, free_bytes=5 * UNIT, hook_sizes={"a": 2 * UNIT}, - memory_reserve_margin=3 * UNIT, + memory_reserve=3 * UNIT, ) == ["a"] - # Multiple candidates: free=6, incoming=4. margin 1 -> usable 5 (fits); margin 3 + # Multiple candidates: free=6, incoming=4. reserve 1 -> usable 5 (fits); reserve 3 # -> usable 3, must free 1 -> smallest sufficient model "c" (1) is offloaded. multi_hooks = {"a": 3 * UNIT, "b": 2 * UNIT, "c": 1 * UNIT} assert ( @@ -198,7 +243,7 @@ def test_strategy_memory_reserve_margin_changes_decision(self): incoming_footprint=4 * UNIT, free_bytes=6 * UNIT, hook_sizes=multi_hooks, - memory_reserve_margin=1 * UNIT, + memory_reserve=1 * UNIT, ) == [] ) @@ -206,9 +251,43 @@ def test_strategy_memory_reserve_margin_changes_decision(self): incoming_footprint=4 * UNIT, free_bytes=6 * UNIT, hook_sizes=multi_hooks, - memory_reserve_margin=3 * UNIT, + memory_reserve=3 * UNIT, ) == ["c"] + def test_strategy_reserve_zero_packs_tight(self): + # `memory_reserve=0` uses every last byte for weights: incoming exactly equals + # the free memory and still fits without evicting. + selected = self._select_offload( + incoming_footprint=4 * UNIT, + free_bytes=4 * UNIT, + hook_sizes={"a": 4 * UNIT}, + memory_reserve=0, + ) + assert selected == [] + + def test_strategy_allocator_cache_counts_as_available(self): + # `mem_get_info` reports only 2 units free, but 6 units of the allocator's cache + # are reusable: available = 2 + 6 = 8, minus 1 reserve -> the incoming 4 fits + # with nothing evicted. Without the cache add-back this would over-evict. + selected = self._select_offload( + incoming_footprint=4 * UNIT, + free_bytes=2 * UNIT, + hook_sizes={"a": 4 * UNIT}, + cached_bytes=6 * UNIT, + ) + assert selected == [] + + def test_strategy_memory_reserve_accepts_file_size_strings(self): + # "3KiB" = 3 * 1024 bytes = a 3-unit reserve; free 4 - 3 = 1 usable but the + # incoming needs 4, so the resident model must go. + selected = self._select_offload( + incoming_footprint=4 * UNIT, + free_bytes=4 * UNIT, + hook_sizes={"a": 4 * UNIT}, + memory_reserve="3KiB", + ) + assert selected == ["a"] + # ------------------------------------------------------------------ # Registry tests (hardware-independent) # ------------------------------------------------------------------ @@ -249,7 +328,7 @@ def test_auto_offload_starts_with_all_components_on_cpu(self): cm = ComponentsManager() model = self.get_dummy_model(4 * UNIT) cm.add("m1", model) - cm.enable_auto_cpu_offload(device=torch_device, memory_reserve_margin=UNIT) + cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) try: assert next(model.parameters()).device.type == "cpu" finally: @@ -263,7 +342,7 @@ def test_auto_offload_evicts_resident_model_under_memory_pressure(self): m2 = self.get_dummy_model(4 * UNIT) cm.add("m1", m1) cm.add("m2", m2) - cm.enable_auto_cpu_offload(device=torch_device, memory_reserve_margin=UNIT) + cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) try: # Both components start offloaded on the CPU. assert next(m1.parameters()).device.type == "cpu" @@ -294,7 +373,7 @@ def test_auto_offload_keeps_models_resident_when_memory_is_ample(self): m2 = self.get_dummy_model(4 * UNIT) cm.add("m1", m1) cm.add("m2", m2) - cm.enable_auto_cpu_offload(device=torch_device, memory_reserve_margin=UNIT) + cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) try: x = torch.randn(2, 4, device=torch_device) with _patch_free_memory(70 * UNIT): @@ -311,9 +390,126 @@ class TestComponentsManager(ComponentsManagerTesterMixin): pass -# More free memory than any tiny test checkpoint could ever need, so the strategy never -# decides to offload. Used to assert the *negative*: no eviction without memory pressure. -_AMPLE_FREE_BYTES = 1024**4 +@require_accelerate +class TestOffloadHookBehavior: + """CPU-level tests of the hook machinery: OOM retry and non-disruptive add/remove.""" + + def _manager(self, components, **offload_kwargs): + manager = ComponentsManager() + for name, component in components.items(): + manager.add(name, component) + # These tests exercise the hook machinery only: with a cpu execution device the models + # never move, so the strategy is never consulted — but enabling requires the device to + # implement `mem_get_info`, which cpu doesn't, so provide an ample fake one. + with mock.patch.object( + torch.cpu, "mem_get_info", create=True, return_value=(_AMPLE_FREE_BYTES, _AMPLE_FREE_BYTES) + ): + manager.enable_auto_cpu_offload(device="cpu", **offload_kwargs) + return manager + + @staticmethod + def _record_offloads(manager): + """ + Record which models the OOM retry picks, in order. + + The execution device is the cpu here, so the offload itself is a no-op move and leaves nothing to observe — + what these tests are about is *which* model gets selected. + """ + offloaded = [] + + def recorder(user_hook): + # model ids are `_`; record the name the test added it under + return lambda: offloaded.append(user_hook.model_id.rsplit("_", 1)[0]) + + for user_hook in manager.model_hooks: + user_hook.offload = recorder(user_hook) + return offloaded + + def test_oom_retry_offloads_smallest_and_reruns(self): + model = OOMOnceModel() + smallest = DummyModel(UNIT) + larger = DummyModel(4 * UNIT) + manager = self._manager({"model": model, "larger": larger, "smallest": smallest}) + offloaded = self._record_offloads(manager) + + out = model(torch.zeros(2, 4)) + assert out.shape == (2, 4) + assert model.forward_calls == 2 + # one OOM costs the smallest resident model, not everything on the device + assert offloaded == ["smallest"] + + def test_oom_retry_escalates_one_model_at_a_time(self): + model = OOMOnceModel() + model._oom_calls = 2 + smallest = DummyModel(UNIT) + larger = DummyModel(4 * UNIT) + manager = self._manager({"model": model, "larger": larger, "smallest": smallest}) + offloaded = self._record_offloads(manager) + + out = model(torch.zeros(2, 4)) + assert out.shape == (2, 4) + assert model.forward_calls == 3 + assert offloaded == ["smallest", "larger"] + + def test_oom_reraises_when_nothing_to_evict(self): + model = OOMOnceModel() + self._manager({"model": model}) + + with pytest.raises(torch.OutOfMemoryError, match="group offloading"): + model(torch.zeros(2, 4)) + + def test_oom_reraises_once_everything_is_offloaded(self): + model = OOMOnceModel() + model._repeated_oom = True + smallest = DummyModel(UNIT) + larger = DummyModel(4 * UNIT) + manager = self._manager({"model": model, "larger": larger, "smallest": smallest}) + offloaded = self._record_offloads(manager) + + with pytest.raises(torch.OutOfMemoryError, match="group offloading"): + model(torch.zeros(2, 4)) + # one attempt per model offloaded, plus the initial one + assert model.forward_calls == 3 + assert offloaded == ["smallest", "larger"] + + def test_retry_on_oom_false_leaves_the_forward_alone(self): + model = OOMOnceModel() + other = DummyModel() + manager = self._manager({"model": model, "other": other}, retry_on_oom=False) + + with pytest.raises(torch.OutOfMemoryError, match="fake OOM"): + model(torch.zeros(2, 4)) + assert model.forward_calls == 1 + assert all(user_hook.hook._forward_before_oom_wrap is None for user_hook in manager.model_hooks) + + def test_add_does_not_rebuild_existing_hooks(self): + model_a = DummyModel() + manager = self._manager({"a": model_a}) + hooks_before = list(manager.model_hooks) + + model_b = DummyModel() + manager.add("b", model_b) + + assert len(manager.model_hooks) == 2 + assert manager.model_hooks[0] is hooks_before[0], "existing hook was rebuilt" + hook_a, hook_b = manager.model_hooks + assert hook_b.hook.other_hooks == [hook_a] + assert hook_a.hook.other_hooks == [hook_b] + + def test_remove_detaches_only_that_hook(self): + model_a = DummyModel() + model_b = DummyModel() + manager = self._manager({"a": model_a, "b": model_b}) + hook_a = manager.model_hooks[0] + b_id = [cid for cid in manager.components if cid.startswith("b_")][0] + + manager.remove(b_id) + + assert manager.model_hooks == [hook_a] + assert hook_a.hook.other_hooks == [] + assert not hasattr(model_b, "_hf_hook") + # model a still hooked and functional + assert model_a(torch.zeros(2, 4)).shape == (2, 4) class ModularPipelineOffloadTesterMixin: @@ -346,7 +542,7 @@ def _run_offloaded(self, free_bytes): """ cm = ComponentsManager() pipe = self.get_pipeline(components_manager=cm) - cm.enable_auto_cpu_offload(device=torch_device, memory_reserve_margin=0) + cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=0) records = [] original_call = AutoOffloadStrategy.__call__ @@ -444,3 +640,81 @@ def test_auto_cpu_offload_inference_consistent_under_memory_pressure(self, expec assert max_diff < expected_max_diff, f"offloaded output diverged from baseline (max diff {max_diff})" finally: cm.disable_auto_cpu_offload() + + +@require_accelerator +class TestSimulateAcceleratorMemory: + """ + Validates the `simulate_accelerator_memory` testing util itself against the real device: unlike + `_patch_memory_stats` (which pins scripted readings for hardware-independent unit tests), the + util wraps `mem_get_info` so real allocations show up live on a simulated smaller card, and + hard-caps the allocator so overshooting the simulated capacity raises a real OOM. + """ + + MB = 2**20 + + def setup_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + @staticmethod + def _device_module(): + return getattr(torch, torch.device(torch_device).type) + + def test_readings_translate_to_simulated_card(self): + device_module = self._device_module() + free, real_total = device_module.mem_get_info() + # Simulate a card with 64MB of headroom on top of whatever the device currently holds. + sim_total = (real_total - free) + 64 * self.MB + + with simulate_accelerator_memory(sim_total, device=torch_device, hard=False): + free0, total0 = device_module.mem_get_info() + assert total0 == sim_total + assert free0 <= 64 * self.MB + + # A real allocation is visible on the simulated card: free drops by at least its size. + x = torch.zeros(8 * self.MB, dtype=torch.float32, device=torch_device) # 32MB + free1, total1 = device_module.mem_get_info() + assert total1 == sim_total + assert free1 <= free0 - 32 * self.MB + del x + + # Exiting the context restores the real readings. + assert device_module.mem_get_info()[1] == real_total + + def test_run_ooms_when_simulated_card_is_too_small(self): + device_module = self._device_module() + model = DummyModel(footprint_bytes=64 * self.MB) + x = torch.randn(4, device=torch_device) + + # Measure what the run actually needs on the device (weights + forward). + device_module.reset_peak_memory_stats() + model.to(torch_device) + model(x) + requirement = device_module.max_memory_allocated() + model.to("cpu") + backend_empty_cache(torch_device) + + try: + # Control: a simulated card with comfortable headroom runs the model fine. + free, real_total = device_module.mem_get_info() + with simulate_accelerator_memory((real_total - free) + 2 * requirement, device=torch_device): + model.to(torch_device) + model(x) + model.to("cpu") + backend_empty_cache(torch_device) + + # A card with only half the requirement left cannot: the hard cap turns the + # overshoot into a real device OOM instead of using the extra physical memory. + free, real_total = device_module.mem_get_info() + with simulate_accelerator_memory((real_total - free) + requirement // 2, device=torch_device): + with pytest.raises(torch.OutOfMemoryError): + model.to(torch_device) + model(x) + finally: + model.to("cpu") + backend_empty_cache(torch_device) diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 0fb03e92028f..6de43fa49203 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -702,6 +702,55 @@ def require_accelerator(test_case): return pytest.mark.skipif(torch_device == "cpu", reason="test requires a hardware accelerator")(test_case) +@contextmanager +def simulate_accelerator_memory(total: Union[int, str], device=None, hard: bool = True): + """ + Make the accelerator behave like a card with `total` memory (e.g. `"20GB"`) for the duration of the context, so + behavior under consumer-sized memory constraints can be exercised with real models on a large dev GPU. + + Unlike pinning `mem_get_info` to fixed values, this *wraps* it: the free-memory reading keeps tracking real + allocations live, translated onto the simulated capacity. Consumers that make decisions from memory readings + (e.g. the dynamic mode of `ComponentsManager.enable_auto_cpu_offload`) therefore behave exactly as they would on + the smaller card. + + With `hard=True` (the default) the caching allocator is additionally capped, so an allocation that would not fit + on the simulated card raises a real `torch.OutOfMemoryError` instead of silently using the extra physical + memory. Memory the allocator does not manage (the device context, other processes) occupies the simulated card + too — as it would on real hardware — so it is deducted from the allocator's share of the capacity. + + Requires a backend that implements `mem_get_info` (CUDA/XPU). `total` accepts bytes or a size string (size + strings require `accelerate`). + """ + if isinstance(total, str): + from accelerate.utils.modeling import convert_file_size_to_int + + total = convert_file_size_to_int(total) + + device = torch.device(device if device is not None else torch_device) + device_module = getattr(torch, device.type) + real_mem_get_info = device_module.mem_get_info + free_bytes, real_total = real_mem_get_info(device.index) + + def simulated_mem_get_info(dev=None): + free, _ = real_mem_get_info(dev if dev is not None else device.index) + used = real_total - free + return max(total - used, 0), total + + hard = hard and hasattr(device_module, "set_per_process_memory_fraction") + if hard: + non_allocator_overhead = (real_total - free_bytes) - device_module.memory_reserved(device.index) + allocator_share = max(total - non_allocator_overhead, 0) + device_module.set_per_process_memory_fraction(min(allocator_share / real_total, 1.0), device.index) + + device_module.mem_get_info = simulated_mem_get_info + try: + yield + finally: + device_module.mem_get_info = real_mem_get_info + if hard: + device_module.set_per_process_memory_fraction(1.0, device.index) + + def require_torchsde(test_case): """ Decorator marking a test that requires torchsde. These tests are skipped when torchsde isn't installed. From d4b4edd15f958b05792e1fb1372846a1166c8a6f Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 28 Jul 2026 00:17:12 +0000 Subject: [PATCH 02/14] [modular] drop stale "dynamic mode" wording from the offload test comments `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 --- tests/modular_pipelines/test_components_manager.py | 6 +++--- tests/testing_utils.py | 3 +-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index 506eda308654..3c6acb9ca357 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -89,9 +89,9 @@ def __init__(self, model_id: str, model: torch.nn.Module): @contextlib.contextmanager def _patch_memory_stats(device_module, free_bytes, total_bytes, cached_bytes=0): # `mem_get_info` returns `(free, total)` and is where the strategy learns how much - # memory is free; the dynamic mode additionally counts the allocator's reusable - # cache (`memory_reserved - memory_allocated`) as available, so those are pinned - # too — `cached_bytes` simulates reusable cache on top of `free_bytes`. + # memory is free; the strategy additionally counts the allocator's reusable cache + # (`memory_reserved - memory_allocated`) as available, so those are pinned too — + # `cached_bytes` simulates reusable cache on top of `free_bytes`. with contextlib.ExitStack() as stack: stack.enter_context(mock.patch.object(device_module, "mem_get_info", return_value=(free_bytes, total_bytes))) if hasattr(device_module, "memory_reserved") and hasattr(device_module, "memory_allocated"): diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 6de43fa49203..bc38604b688d 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -710,8 +710,7 @@ def simulate_accelerator_memory(total: Union[int, str], device=None, hard: bool Unlike pinning `mem_get_info` to fixed values, this *wraps* it: the free-memory reading keeps tracking real allocations live, translated onto the simulated capacity. Consumers that make decisions from memory readings - (e.g. the dynamic mode of `ComponentsManager.enable_auto_cpu_offload`) therefore behave exactly as they would on - the smaller card. + (e.g. `ComponentsManager.enable_auto_cpu_offload`) therefore behave exactly as they would on the smaller card. With `hard=True` (the default) the caching allocator is additionally capped, so an allocation that would not fit on the simulated card raises a real `torch.OutOfMemoryError` instead of silently using the extra physical From 34036c1fadec4e1fc5fe2c12dabb3ab3075f8660 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 28 Jul 2026 04:24:05 +0000 Subject: [PATCH 03/14] Record what the auto-offloader does MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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:`, `oom_retry:`, `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 --- .../modular_diffusers/components_manager.md | 34 +++ .../modular_pipelines/components_manager.py | 277 ++++++++++++++++-- .../test_components_manager.py | 236 +++++++++++---- 3 files changed, 466 insertions(+), 81 deletions(-) diff --git a/docs/source/en/modular_diffusers/components_manager.md b/docs/source/en/modular_diffusers/components_manager.md index 2bac75a876c6..854f48d8aa68 100644 --- a/docs/source/en/modular_diffusers/components_manager.md +++ b/docs/source/en/modular_diffusers/components_manager.md @@ -97,6 +97,40 @@ Set `memory_reserve=0` to keep as much on the device as possible. If a forward p Components added while offloading is enabled join the managed set without disturbing it: the new component starts on CPU and everything already resident stays where it is; removing a component likewise detaches only that component. +### Inspecting what the offloader did + +Every move is recorded. Print [`ComponentsManager.offload_record`] after a run to see the sequence, what each move cost, and where a forward pass ran out of memory. + +```py +manager.enable_auto_cpu_offload(device="cuda") +pipe(prompt="a cat") +print(manager.offload_record) +``` + +```py +# | Action | Model | Size | Available | Reason +------------------------------------------------------------------------------------ +1 | onload | text_encoder_140458257514 | 9.52 GB | - | forward +2 | offload | text_encoder_140458257514 | 9.52 GB | 4.31 GB | needed_by:transformer_140458257515616 +3 | onload | transformer_1404582575156 | 23.80 GB | 4.31 GB | forward +------------------------------------------------------------------------------------ +2 onloads (33.32 GB) / 1 offloads (9.52 GB), 0 OOM retries, peak co-residency 1 +``` + +A model appearing repeatedly in this table is thrashing — it is being evicted 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). + +To find the right value, let the manager measure it. `measure_activations=True` records how much memory the forward passes actually need on top of the weights — the exact thing `memory_reserve` covers — and reports the reserve that would have covered your run: + +```py +manager.enable_auto_cpu_offload(device="cuda", memory_reserve=0, measure_activations=True) +pipe(prompt="a cat") + +manager.offload_record.activation_peak # what the forward passes needed +manager.offload_record.suggested_memory_reserve # a reserve that would have covered it +``` + +Measure at the resolution, batch size, and sequence length you intend to use — activations scale with all three. Leave `measure_activations` off if you are measuring peak memory yourself, since it resets the device's peak-memory stats around every forward pass. + Call [`~ComponentsManager.disable_auto_cpu_offload`] to disable offloading. ```py diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index ab3a080805e4..3a2328b4f7ef 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -17,7 +17,8 @@ import copy import functools import time -from collections import OrderedDict +from collections import OrderedDict, deque +from dataclasses import dataclass from itertools import combinations from typing import Any @@ -40,6 +41,147 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name +# Default number of events an `OffloadRecord` keeps. A generation step produces at most a handful, so this +# holds a long run while staying bounded for a long-lived manager (e.g. in a server). +MAX_RECORDED_OFFLOAD_EVENTS = 10000 + + +def format_size(num_bytes: int | None) -> str: + """Bytes as a human-readable size, so that both a 20GB transformer and a KB-sized test model read sensibly.""" + if num_bytes is None: + return "-" + for unit in ("B", "KB", "MB"): + if abs(num_bytes) < 1024: + return f"{num_bytes:.0f} {unit}" if unit == "B" else f"{num_bytes:.2f} {unit}" + num_bytes /= 1024 + return f"{num_bytes:.2f} GB" + + +@dataclass +class OffloadEvent: + """ + A single thing the offloader did: a model moved onto the execution device (`"onload"`), a model moved back to the + CPU (`"offload"`), or a forward pass that ran out of device memory (`"oom"`). + """ + + action: str + model_id: str + model_size: int | None = None + # why the model moved: which model needed the room, the OOM retry, attaching a new component, ... + reason: str | None = None + # what the strategy saw when it decided (`None` for a custom strategy that does not report it) + available_memory: int | None = None + memory_reserve: int | None = None + resident_before: tuple[str, ...] = () + offloaded: tuple[str, ...] = () + seconds: float | None = None + + +class OffloadRecord: + """ + What the offloader did, in order. + + Every offloading decision appends to this record, so after a run it holds the full sequence of moves, what each + move cost, and where the forward passes ran out of memory. Printing it shows that sequence as a table followed by + a summary; [`~OffloadRecord.summary`] returns the same numbers as a dict. + + With `measure_activations=True` in [`~ComponentsManager.enable_auto_cpu_offload`], the record additionally + measures how much memory the forward passes need on top of the weights, and reports the `memory_reserve` that + would have covered it — measured on your hardware, at your resolution and batch size. + """ + + def __init__(self, maxlen: int = MAX_RECORDED_OFFLOAD_EVENTS): + self.events: deque[OffloadEvent] = deque(maxlen=maxlen) + self.activations: dict[str, int] = {} + self._recorded = 0 + + def add(self, event: OffloadEvent): + self.events.append(event) + self._recorded += 1 + + def note_activation(self, model_id: str, activation_bytes: int): + """Keep the largest activation cost seen for a model.""" + self.activations[model_id] = max(self.activations.get(model_id, 0), activation_bytes) + + def clear(self): + self.events.clear() + self.activations.clear() + self._recorded = 0 + + @property + def dropped(self) -> int: + """How many events fell out of the (bounded) log.""" + return max(self._recorded - len(self.events), 0) + + @property + def activation_peak(self) -> int | None: + """ + The largest activation cost measured across all forward passes, or `None` if activations were not measured. + """ + return max(self.activations.values()) if self.activations else None + + @property + def suggested_memory_reserve(self) -> int | None: + """ + A `memory_reserve` that would have covered the measured activations, with a margin. `None` if activations were + not measured. Only as representative as the run it was measured on — activations grow with resolution, batch + size and sequence length. + """ + peak = self.activation_peak + return int(peak * 1.2) if peak is not None else None + + def summary(self) -> dict[str, Any]: + onloads = [event for event in self.events if event.action == "onload"] + offloads = [event for event in self.events if event.action == "offload"] + resident, peak_co_residency = set(), 0 + for event in self.events: + if event.action == "onload": + resident.add(event.model_id) + peak_co_residency = max(peak_co_residency, len(resident)) + elif event.action == "offload": + resident.discard(event.model_id) + return { + "onloads": len(onloads), + "offloads": len(offloads), + "oom_retries": sum(event.action == "oom" for event in self.events), + "peak_co_residency": peak_co_residency, + "bytes_onloaded": sum(event.model_size or 0 for event in onloads), + "bytes_offloaded": sum(event.model_size or 0 for event in offloads), + "activation_peak": self.activation_peak, + "suggested_memory_reserve": self.suggested_memory_reserve, + } + + def __len__(self): + return len(self.events) + + def __repr__(self): + if not self.events: + return "Offload record: nothing recorded yet" + + header = f"{'#':<5} | {'Action':<8} | {'Model':<40} | {'Size':<10} | {'Available':<10} | Reason" + lines = [header, "-" * len(header)] + for index, event in enumerate(self.events, start=self.dropped + 1): + lines.append( + f"{index:<5} | {event.action:<8} | {event.model_id:<40} | {format_size(event.model_size):<10} | " + f"{format_size(event.available_memory):<10} | {event.reason or ''}" + ) + if self.dropped: + lines.insert(2, f"... {self.dropped} earlier events dropped from the log ...") + + summary = self.summary() + lines += [ + "-" * len(header), + f"{summary['onloads']} onloads ({format_size(summary['bytes_onloaded'])}) / " + f"{summary['offloads']} offloads ({format_size(summary['bytes_offloaded'])}), " + f"{summary['oom_retries']} OOM retries, peak co-residency {summary['peak_co_residency']}", + ] + if summary["activation_peak"] is not None: + lines.append( + f"activations peaked at {format_size(summary['activation_peak'])} -> a `memory_reserve` of " + f"{format_size(summary['suggested_memory_reserve'])} would have covered this run" + ) + return "\n".join(lines) + class CustomOffloadHook(ModelHook): """ @@ -53,6 +195,11 @@ class CustomOffloadHook(ModelHook): retry_on_oom(`bool`, *optional*, defaults to `True`): Whether to recover from a forward pass that runs out of device memory by offloading the other models one at a time and retrying. If `False`, the error is raised to the caller. + record(`OffloadRecord`, *optional*): + Where to record the moves this hook makes. + measure_activations(`bool`, *optional*, defaults to `False`): + Whether to measure how much device memory each forward pass needs on top of the weights. Owns the + device's peak-memory stats while enabled. """ no_grad = False @@ -63,13 +210,18 @@ def __init__( other_hooks: list["UserCustomOffloadHook"] | None = None, offload_strategy: "AutoOffloadStrategy" | None = None, retry_on_oom: bool = True, + record: OffloadRecord | None = None, + measure_activations: bool = False, ): self.execution_device = execution_device if execution_device is not None else PartialState().default_device self.other_hooks = other_hooks self.offload_strategy = offload_strategy self.retry_on_oom = retry_on_oom + self.record = record if record is not None else OffloadRecord() + self.measure_activations = measure_activations self.model_id = None self._forward_before_oom_wrap = None + self._allocated_before_forward = None def set_strategy(self, offload_strategy: "AutoOffloadStrategy"): self.offload_strategy = offload_strategy @@ -87,8 +239,10 @@ def init_hook(self, module): def pre_forward(self, module, *args, **kwargs): if module.device != self.execution_device: + resident_before, hooks_to_offload, elapsed = (), [], None if self.other_hooks is not None: hooks_to_offload = [hook for hook in self.other_hooks if hook.model.device == self.execution_device] + resident_before = tuple(hook.model_id for hook in hooks_to_offload) # offload all other hooks start_time = time.perf_counter() if self.offload_strategy is not None: @@ -99,19 +253,39 @@ def pre_forward(self, module, *args, **kwargs): execution_device=self.execution_device, ) end_time = time.perf_counter() - logger.info( - f" time taken to apply offload strategy for {self.model_id}: {(end_time - start_time):.2f} seconds" - ) + elapsed = end_time - start_time + logger.info(f" time taken to apply offload strategy for {self.model_id}: {elapsed:.2f} seconds") for hook in hooks_to_offload: logger.info( f"moving {self.model_id} to {self.execution_device}, offloading {hook.model_id} to cpu" ) - hook.offload() + hook.offload(reason=f"needed_by:{self.model_id}") if hooks_to_offload: clear_device_cache() module.to(self.execution_device) + # what the strategy saw, when it is one that reports it + decision = getattr(self.offload_strategy, "last_decision", None) or {} + self.record.add( + OffloadEvent( + action="onload", + model_id=self.model_id, + model_size=module.get_memory_footprint(), + reason="forward", + available_memory=decision.get("available_memory"), + memory_reserve=decision.get("memory_reserve"), + resident_before=resident_before, + offloaded=tuple(hook.model_id for hook in hooks_to_offload), + seconds=elapsed, + ) + ) + if self.measure_activations: + # baseline for the activation measurement: everything allocated once this model is on the device + device_module = getattr(torch, self.execution_device.type, torch.cuda) + if hasattr(device_module, "reset_peak_memory_stats"): + device_module.reset_peak_memory_stats(self.execution_device) + self._allocated_before_forward = device_module.memory_allocated(self.execution_device) return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device) def resident_other_hooks(self): @@ -126,15 +300,25 @@ def resident_other_hooks(self): and (other.model.device.index or 0) == (self.execution_device.index or 0) ] + def _run_forward(self, hooked_forward, args, kwargs): + if not self.measure_activations: + return hooked_forward(*args, **kwargs) + output = hooked_forward(*args, **kwargs) + device_module = getattr(torch, self.execution_device.type, torch.cuda) + if self._allocated_before_forward is not None and hasattr(device_module, "max_memory_allocated"): + peak = device_module.max_memory_allocated(self.execution_device) + self.record.note_activation(self.model_id, max(peak - self._allocated_before_forward, 0)) + return output + # YiYi TODO: diffusers' own `ModelHook` (`diffusers.hooks.hooks`) supports a `new_forward` around-hook, which # would replace this manual wrapping - we may migrate this class to it in the future. def wrap_forward(self, module): - # `pre_forward`/`post_forward` cannot see an exception raised by the forward itself, so wrap the (already - # hooked) forward here: if it runs out of device memory, free the smallest resident model and retry, - # escalating until it fits. The memory readings cannot guide this - the failed forward has already unwound, - # so its activations are freed and the device looks free again. Inference only: the retried forward re-runs - # cleanly from its original inputs. - if not self.retry_on_oom: + # `pre_forward`/`post_forward` cannot see the forward pass itself - neither its memory peak nor an exception + # it raises - so wrap the (already hooked) forward here. On an OOM, free the smallest resident model and + # retry, escalating until it fits. The memory readings cannot guide this: the failed forward has already + # unwound, so its activations are freed and the device looks free again. Inference only: the retried forward + # re-runs cleanly from its original inputs. + if not (self.retry_on_oom or self.measure_activations): return hooked_forward = module.forward @@ -146,8 +330,11 @@ def forward_with_oom_retry(*args, **kwargs): offloaded = set() while True: try: - return hooked_forward(*args, **kwargs) + return self._run_forward(hooked_forward, args, kwargs) except torch.OutOfMemoryError as e: + if not self.retry_on_oom: + raise + self.record.add(OffloadEvent(action="oom", model_id=self.model_id, reason=str(e).split(".")[0])) resident = sorted( (hook for hook in self.resident_other_hooks() if id(hook) not in offloaded), key=lambda hook: hook.model.get_memory_footprint(), @@ -165,7 +352,7 @@ def forward_with_oom_retry(*args, **kwargs): "retrying. If this happens repeatedly, set a larger `memory_reserve` in " "`enable_auto_cpu_offload`." ) - smallest.offload() + smallest.offload(reason=f"oom_retry:{self.model_id}") offloaded.add(id(smallest)) clear_device_cache() @@ -189,8 +376,18 @@ def __init__(self, model_id, model, hook): self.model = model self.hook = hook - def offload(self): + def offload(self, reason: str | None = None): + was_resident = self.model.device.type == self.hook.execution_device.type self.hook.init_hook(self.model) + if was_resident: + self.hook.record.add( + OffloadEvent( + action="offload", + model_id=self.model_id, + model_size=self.model.get_memory_footprint(), + reason=reason, + ) + ) def attach(self): add_hook_to_module(self.model, self.hook) @@ -212,9 +409,15 @@ def custom_offload_with_hook( execution_device: str | int | torch.device = None, offload_strategy: "AutoOffloadStrategy" | None = None, retry_on_oom: bool = True, + record: OffloadRecord | None = None, + measure_activations: bool = False, ): hook = CustomOffloadHook( - execution_device=execution_device, offload_strategy=offload_strategy, retry_on_oom=retry_on_oom + execution_device=execution_device, + offload_strategy=offload_strategy, + retry_on_oom=retry_on_oom, + record=record, + measure_activations=measure_activations, ) user_hook = UserCustomOffloadHook(model_id=model_id, model=model, hook=hook) user_hook.attach() @@ -235,8 +438,13 @@ class AutoOffloadStrategy: def __init__(self, memory_reserve="3GB"): self.memory_reserve = convert_file_size_to_int(memory_reserve) + # what the last decision saw, for the hook to record. A custom strategy that does not set it simply + # leaves those columns empty in the record. + self.last_decision = None def __call__(self, hooks, model_id, model, execution_device): + # cleared first so that a decision that returns early is never recorded with the previous decision's readings + self.last_decision = None if len(hooks) == 0: return [] @@ -255,6 +463,7 @@ def __call__(self, hooks, model_id, model, execution_device): available_memory += device_module.memory_reserved(execution_device) - device_module.memory_allocated( execution_device ) + self.last_decision = {"available_memory": available_memory, "memory_reserve": self.memory_reserve} if current_module_size <= available_memory - self.memory_reserve: return [] @@ -422,6 +631,8 @@ def __init__(self): self._auto_offload_enabled = False self._offload_strategy = None self._offload_retry_on_oom = True + self._offload_measure_activations = False + self._offload_record = OffloadRecord() def _lookup_ids( self, @@ -782,6 +993,7 @@ def enable_auto_cpu_offload( device: str | int | torch.device = None, memory_reserve: str | int = "3GB", retry_on_oom: bool = True, + measure_activations: bool = False, ): """ Enable automatic CPU offloading for all components. @@ -808,6 +1020,13 @@ def enable_auto_cpu_offload( Whether to recover from a forward pass that runs out of device memory by offloading the models on the device one at a time, smallest first, and retrying until it fits. Set it to `False` to raise the error instead — the forward passes are then left untouched. + measure_activations (bool, *optional*, defaults to `False`): + Whether to measure how much device memory each forward pass needs on top of the weights, and report + the `memory_reserve` that would have covered it (see [`~ComponentsManager.offload_record`]). This + resets the device's peak-memory stats around every forward pass, so leave it off if you are measuring + peak memory yourself. + + Every move the offloader makes is recorded in [`~ComponentsManager.offload_record`]. """ if not is_accelerate_available(): raise ImportError("Make sure to install accelerate to use auto_cpu_offload") @@ -832,12 +1051,19 @@ def enable_auto_cpu_offload( self.disable_auto_cpu_offload() offload_strategy = AutoOffloadStrategy(memory_reserve=memory_reserve) + self._offload_record.clear() all_hooks = [] for name, component in self.components.items(): if isinstance(component, torch.nn.Module): hook = custom_offload_with_hook( - name, component, device, offload_strategy=offload_strategy, retry_on_oom=retry_on_oom + name, + component, + device, + offload_strategy=offload_strategy, + retry_on_oom=retry_on_oom, + record=self._offload_record, + measure_activations=measure_activations, ) all_hooks.append(hook) @@ -852,6 +1078,18 @@ def enable_auto_cpu_offload( self._auto_offload_device = device self._offload_strategy = offload_strategy self._offload_retry_on_oom = retry_on_oom + self._offload_measure_activations = measure_activations + + @property + def offload_record(self) -> OffloadRecord: + """ + What the offloader has done so far: every model moved onto or off the device, in order, and where a forward + pass ran out of memory. Print it to see the sequence, or read + [`~OffloadRecord.summary`]/[`~OffloadRecord.suggested_memory_reserve`] for the numbers behind it. Kept across + [`~ComponentsManager.disable_auto_cpu_offload`] (it is the post-mortem of the run), and cleared when + offloading is enabled again. + """ + return self._offload_record def _attach_offload_hook(self, component_id: str, component: torch.nn.Module): """ @@ -864,12 +1102,14 @@ def _attach_offload_hook(self, component_id: str, component: torch.nn.Module): self._auto_offload_device, offload_strategy=self._offload_strategy, retry_on_oom=self._offload_retry_on_oom, + record=self._offload_record, + measure_activations=self._offload_measure_activations, ) for other_hook in self.model_hooks: if other_hook.hook.execution_device == hook.hook.execution_device: hook.add_other_hook(other_hook) other_hook.add_other_hook(hook) - hook.offload() + hook.offload(reason="component_added") self.model_hooks.append(hook) def _detach_offload_hook(self, component: torch.nn.Module): @@ -890,7 +1130,7 @@ def disable_auto_cpu_offload(self): return for hook in self.model_hooks: - hook.offload() + hook.offload(reason="offloading_disabled") hook.remove() if self.model_hooks: clear_device_cache() @@ -898,6 +1138,7 @@ def disable_auto_cpu_offload(self): self._auto_offload_enabled = False self._offload_strategy = None self._offload_retry_on_oom = True + self._offload_measure_activations = False def get_model_info( self, diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index 3c6acb9ca357..85ea3d3e4336 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import collections import contextlib import gc from unittest import mock @@ -72,6 +73,18 @@ def forward(self, x): return super().forward(x) +class ActivationHungryModel(DummyModel): + """Allocates a scratch tensor of a known size during its forward, giving the activation measurement a floor.""" + + def __init__(self, footprint_bytes: int = UNIT, activation_bytes: int = 8 * UNIT): + super().__init__(footprint_bytes=footprint_bytes) + self.activation_bytes = activation_bytes + + def forward(self, x): + scratch = torch.zeros(self.activation_bytes // 4, device=x.device) + return super().forward(x) + scratch[0] + + class _FakeHook: """ Minimal stand-in for `UserCustomOffloadHook` in strategy-level unit tests. @@ -365,6 +378,70 @@ def test_auto_offload_evicts_resident_model_under_memory_pressure(self): finally: cm.disable_auto_cpu_offload() + @require_accelerator + def test_record_captures_what_the_strategy_saw(self): + cm = ComponentsManager() + m1 = self.get_dummy_model(4 * UNIT) + m2 = self.get_dummy_model(4 * UNIT) + cm.add("m1", m1) + cm.add("m2", m2) + cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) + try: + x = torch.randn(2, 4, device=torch_device) + with _patch_free_memory(70 * UNIT): + m1(x) + # m2 does not fit alongside m1 (usable 4 - 1 = 3 < 4), so m1 is evicted for it + with _patch_free_memory(4 * UNIT): + m2(x) + + onloads = [event for event in cm.offload_record.events if event.action == "onload"] + names = [event.model_id.rsplit("_", 1)[0] for event in onloads] + assert names == ["m1", "m2"] + # Nothing was resident when m1 loaded, so the strategy returned without reading memory and + # there is nothing to report for that decision. + assert onloads[0].available_memory is None + # The readings behind a real decision are recorded, so a run can be explained after the fact. + assert onloads[1].available_memory == 4 * UNIT + assert onloads[1].memory_reserve == UNIT + assert onloads[1].resident_before == (cm.model_hooks[0].model_id,) + assert onloads[1].offloaded == (cm.model_hooks[0].model_id,) + assert cm.offload_record.summary()["peak_co_residency"] == 1 + finally: + cm.disable_auto_cpu_offload() + + @require_accelerator + def test_record_measures_activations_and_suggests_a_reserve(self): + # The point of the measurement: a user who cannot spare memory for a calibration run still learns + # what their forward passes need on top of the weights, which is what `memory_reserve` covers. + activation_bytes = 4 * 1024**2 + cm = ComponentsManager() + model = ActivationHungryModel(footprint_bytes=4 * UNIT, activation_bytes=activation_bytes) + cm.add("model", model) + cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT, measure_activations=True) + try: + model(torch.randn(2, 4, device=torch_device)) + + peak = cm.offload_record.activation_peak + assert peak >= activation_bytes, f"expected at least the scratch tensor ({activation_bytes}), saw {peak}" + assert cm.offload_record.suggested_memory_reserve > peak + assert "memory_reserve" in repr(cm.offload_record) + finally: + cm.disable_auto_cpu_offload() + + @require_accelerator + def test_activations_are_not_measured_by_default(self): + cm = ComponentsManager() + model = ActivationHungryModel() + cm.add("model", model) + cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) + try: + model(torch.randn(2, 4, device=torch_device)) + # peak stats are the user's to own unless they opt in + assert cm.offload_record.activation_peak is None + assert cm.offload_record.suggested_memory_reserve is None + finally: + cm.disable_auto_cpu_offload() + @require_accelerator def test_auto_offload_keeps_models_resident_when_memory_is_ample(self): device_type = torch.device(torch_device).type @@ -408,35 +485,24 @@ def _manager(self, components, **offload_kwargs): return manager @staticmethod - def _record_offloads(manager): - """ - Record which models the OOM retry picks, in order. - - The execution device is the cpu here, so the offload itself is a no-op move and leaves nothing to observe — - what these tests are about is *which* model gets selected. - """ - offloaded = [] - - def recorder(user_hook): - # model ids are `_`; record the name the test added it under - return lambda: offloaded.append(user_hook.model_id.rsplit("_", 1)[0]) - - for user_hook in manager.model_hooks: - user_hook.offload = recorder(user_hook) - return offloaded + def _offloaded_names(manager): + """The models the record shows going back to the CPU, in order, by the name they were added under.""" + # model ids are `_` + return [ + event.model_id.rsplit("_", 1)[0] for event in manager.offload_record.events if event.action == "offload" + ] def test_oom_retry_offloads_smallest_and_reruns(self): model = OOMOnceModel() smallest = DummyModel(UNIT) larger = DummyModel(4 * UNIT) manager = self._manager({"model": model, "larger": larger, "smallest": smallest}) - offloaded = self._record_offloads(manager) out = model(torch.zeros(2, 4)) assert out.shape == (2, 4) assert model.forward_calls == 2 # one OOM costs the smallest resident model, not everything on the device - assert offloaded == ["smallest"] + assert self._offloaded_names(manager) == ["smallest"] def test_oom_retry_escalates_one_model_at_a_time(self): model = OOMOnceModel() @@ -444,12 +510,11 @@ def test_oom_retry_escalates_one_model_at_a_time(self): smallest = DummyModel(UNIT) larger = DummyModel(4 * UNIT) manager = self._manager({"model": model, "larger": larger, "smallest": smallest}) - offloaded = self._record_offloads(manager) out = model(torch.zeros(2, 4)) assert out.shape == (2, 4) assert model.forward_calls == 3 - assert offloaded == ["smallest", "larger"] + assert self._offloaded_names(manager) == ["smallest", "larger"] def test_oom_reraises_when_nothing_to_evict(self): model = OOMOnceModel() @@ -464,13 +529,12 @@ def test_oom_reraises_once_everything_is_offloaded(self): smallest = DummyModel(UNIT) larger = DummyModel(4 * UNIT) manager = self._manager({"model": model, "larger": larger, "smallest": smallest}) - offloaded = self._record_offloads(manager) with pytest.raises(torch.OutOfMemoryError, match="group offloading"): model(torch.zeros(2, 4)) # one attempt per model offloaded, plus the initial one assert model.forward_calls == 3 - assert offloaded == ["smallest", "larger"] + assert self._offloaded_names(manager) == ["smallest", "larger"] def test_retry_on_oom_false_leaves_the_forward_alone(self): model = OOMOnceModel() @@ -482,6 +546,80 @@ def test_retry_on_oom_false_leaves_the_forward_alone(self): assert model.forward_calls == 1 assert all(user_hook.hook._forward_before_oom_wrap is None for user_hook in manager.model_hooks) + def test_record_captures_the_onload_sequence(self): + model_a = DummyModel(UNIT) + model_b = DummyModel(4 * UNIT) + manager = self._manager({"a": model_a, "b": model_b}) + + model_a(torch.zeros(2, 4)) + model_b(torch.zeros(2, 4)) + + onloads = [event for event in manager.offload_record.events if event.action == "onload"] + assert [event.model_id.rsplit("_", 1)[0] for event in onloads] == ["a", "b"] + assert [event.model_size for event in onloads] == [UNIT, 4 * UNIT] + summary = manager.offload_record.summary() + assert summary["onloads"] == 2 + assert summary["bytes_onloaded"] == 5 * UNIT + assert summary["oom_retries"] == 0 + + def test_record_captures_oom_and_the_escalation(self): + model = OOMOnceModel() + smallest = DummyModel(UNIT) + manager = self._manager({"model": model, "smallest": smallest}) + + model(torch.zeros(2, 4)) + + actions = [(event.action, event.model_id.rsplit("_", 1)[0]) for event in manager.offload_record.events] + assert ("oom", "model") in actions + assert ("offload", "smallest") in actions + summary = manager.offload_record.summary() + assert summary["oom_retries"] == 1 + assert summary["offloads"] == 1 + assert summary["bytes_offloaded"] == UNIT + # the eviction is attributed to the model that ran out of memory + offload_event = next(event for event in manager.offload_record.events if event.action == "offload") + assert offload_event.reason.startswith("oom_retry:") + + def test_record_tracks_peak_co_residency(self): + model_a = DummyModel(UNIT) + model_b = DummyModel(UNIT) + manager = self._manager({"a": model_a, "b": model_b}) + + model_a(torch.zeros(2, 4)) + model_b(torch.zeros(2, 4)) + assert manager.offload_record.summary()["peak_co_residency"] == 2 + + # an offload brings residency back down, so a later onload does not raise the peak + manager.model_hooks[0].offload(reason="test") + model_a(torch.zeros(2, 4)) + assert manager.offload_record.summary()["peak_co_residency"] == 2 + + def test_record_is_bounded_and_clearable(self): + model = DummyModel() + manager = self._manager({"model": model}) + manager.offload_record.events = collections.deque(maxlen=2) + + for _ in range(4): + model(torch.zeros(2, 4)) + + assert len(manager.offload_record.events) == 2 + assert manager.offload_record.dropped == 2 + assert "earlier events dropped" in repr(manager.offload_record) + + manager.offload_record.clear() + assert len(manager.offload_record) == 0 + assert manager.offload_record.dropped == 0 + assert "nothing recorded yet" in repr(manager.offload_record) + + def test_record_survives_disable(self): + model = DummyModel() + manager = self._manager({"model": model}) + model(torch.zeros(2, 4)) + + manager.disable_auto_cpu_offload() + # the record is the post-mortem of the run that just ended + assert manager.offload_record.summary()["onloads"] == 1 + def test_add_does_not_rebuild_existing_hooks(self): model_a = DummyModel() manager = self._manager({"a": model_a}) @@ -536,62 +674,34 @@ def _run_offloaded(self, free_bytes): Run the pipeline with auto offload on and `free_bytes` of *simulated* device memory, recording every offload decision the strategy makes. - Each record is `{"incoming", "resident_before", "offloaded"}` (lists of model - ids), captured by spying on `AutoOffloadStrategy.__call__`, which the hooks call - each time a model is about to be moved onto the device. + Each decision is an `"onload"` event in `cm.offload_record`, holding the incoming + model, what was resident just before, and what the strategy evicted to make room. """ cm = ComponentsManager() pipe = self.get_pipeline(components_manager=cm) cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=0) - records = [] - original_call = AutoOffloadStrategy.__call__ - - def spy_call(strategy, hooks, model_id, model, execution_device): - selected = original_call( - strategy, hooks=hooks, model_id=model_id, model=model, execution_device=execution_device - ) - records.append( - { - "incoming": model_id, - "resident_before": [hook.model_id for hook in hooks], - "offloaded": [hook.model_id for hook in selected], - } - ) - return selected - - with _patch_free_memory(free_bytes), mock.patch.object(AutoOffloadStrategy, "__call__", spy_call): + with _patch_free_memory(free_bytes): output = pipe(**self.get_dummy_inputs(), output=self.output_name) - return cm, records, output - - @staticmethod - def _peak_co_residency(records): - """ - Largest number of models simultaneously on the device, reconstructed from the - strategy's view of residency just before each load. - """ - peak = 0 - for record in records: - resident = (set(record["resident_before"]) - set(record["offloaded"])) | {record["incoming"]} - peak = max(peak, len(resident)) - return peak + onloads = [event for event in cm.offload_record.events if event.action == "onload"] + return cm, onloads, output @require_accelerate @require_accelerator def test_auto_cpu_offload_serializes_models_under_memory_pressure(self): # Zero simulated free memory: every model that runs must first evict whatever is # currently resident (comfy-style serialized execution). - cm, records, _ = self._run_offloaded(free_bytes=0) + cm, onloads, _ = self._run_offloaded(free_bytes=0) try: - distinct_models = {record["incoming"] for record in records} + distinct_models = {event.model_id for event in onloads} if len(distinct_models) < 2: pytest.skip("pipeline has fewer than two offloadable model components") # Offloading actually fired (at least one eviction happened). - assert any(record["offloaded"] for record in records), "expected at least one eviction" + assert any(event.offloaded for event in onloads), "expected at least one eviction" # Sequencing: models run one at a time, never two co-resident on the device. - peak = self._peak_co_residency(records) + peak = cm.offload_record.summary()["peak_co_residency"] assert peak == 1, f"expected serialized execution under pressure, saw {peak} models co-resident" # Device placement after the run: at most the last-run model stays on the @@ -608,17 +718,17 @@ def test_auto_cpu_offload_serializes_models_under_memory_pressure(self): def test_auto_cpu_offload_keeps_models_resident_without_memory_pressure(self): # Negative case: with ample simulated memory the strategy is still consulted on # every load, but it must never decide to evict anything. - cm, records, _ = self._run_offloaded(free_bytes=_AMPLE_FREE_BYTES) + cm, onloads, _ = self._run_offloaded(free_bytes=_AMPLE_FREE_BYTES) try: - distinct_models = {record["incoming"] for record in records} + distinct_models = {event.model_id for event in onloads} if len(distinct_models) < 2: pytest.skip("pipeline has fewer than two offloadable model components") # Nothing was ever offloaded... - assert all(record["offloaded"] == [] for record in records), "no model should be evicted" + assert all(event.offloaded == () for event in onloads), "no model should be evicted" # ...and models accumulate on the device instead of being serialized. - peak = self._peak_co_residency(records) + peak = cm.offload_record.summary()["peak_co_residency"] assert peak >= 2, f"expected models to co-reside without pressure, saw peak {peak}" models = self._managed_models(cm) From 24460ae67aeb98c64dfd1e542cb34add05168308 Mon Sep 17 00:00:00 2001 From: "yiyi@huggingface.co" Date: Wed, 29 Jul 2026 00:16:54 +0000 Subject: [PATCH 04/14] Group the offload record table by decision 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 --- .../modular_diffusers/components_manager.md | 18 +++++++----- .../modular_pipelines/components_manager.py | 28 ++++++++++++++++--- .../test_components_manager.py | 3 ++ 3 files changed, 38 insertions(+), 11 deletions(-) diff --git a/docs/source/en/modular_diffusers/components_manager.md b/docs/source/en/modular_diffusers/components_manager.md index 854f48d8aa68..6224712e8dd9 100644 --- a/docs/source/en/modular_diffusers/components_manager.md +++ b/docs/source/en/modular_diffusers/components_manager.md @@ -107,16 +107,20 @@ pipe(prompt="a cat") print(manager.offload_record) ``` +On a 20GB card, the Z-Image pipeline from the example above records this: + ```py -# | Action | Model | Size | Available | Reason ------------------------------------------------------------------------------------- -1 | onload | text_encoder_140458257514 | 9.52 GB | - | forward -2 | offload | text_encoder_140458257514 | 9.52 GB | 4.31 GB | needed_by:transformer_140458257515616 -3 | onload | transformer_1404582575156 | 23.80 GB | 4.31 GB | forward ------------------------------------------------------------------------------------- -2 onloads (33.32 GB) / 1 offloads (9.52 GB), 0 OOM retries, peak co-residency 1 +# | Onload | Offloaded | Available | Reason +------------------------------------------------------------------------------------------------------------- +1 | text_encoder_140458257514752 (7.49 GB) | - | - | forward +2 | transformer_140458257515616 (11.46 GB) | text_encoder_140458257514752 (7.49 GB) | 10.38 GB | forward +3 | vae_140458257515376 (159.87 MB) | - | 6.36 GB | forward +------------------------------------------------------------------------------------------------------------- +3 onloads (19.11 GB) / 1 offloads (7.49 GB), 0 OOM retries, peak co-residency 2 ``` +Each row is one decision: the model that loaded, what was evicted to make room for it, and the free device memory (`Available`) the decision saw at that moment. 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 evicted first. The VAE then fit into the 6.36 GB left next to the transformer, so it loaded without evicting anything (peak co-residency 2). `Available` reads `-` when nothing was resident, so the decision needed no memory reading (row 1). Offloads an onload did not cause — an OOM retry, disabling offloading — appear as their own rows, explained by their `Reason`. + A model appearing repeatedly in this table is thrashing — it is being evicted 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). To find the right value, let the manager measure it. `measure_activations=True` records how much memory the forward passes actually need on top of the weights — the exact thing `memory_reserve` covers — and reports the reserve that would have covered your run: diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 3a2328b4f7ef..174637d01d23 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -158,12 +158,32 @@ def __repr__(self): if not self.events: return "Offload record: nothing recorded yet" - header = f"{'#':<5} | {'Action':<8} | {'Model':<40} | {'Size':<10} | {'Available':<10} | Reason" + # One row per decision: an onload and the evictions it caused (its "needed_by" offload events) share a + # row. Offloads with other causes (OOM retry, offloading disabled) and OOMs get their own rows. + sizes = {event.model_id: event.model_size for event in self.events if event.model_size is not None} + + def label(model_id): + return f"{model_id} ({format_size(sizes.get(model_id))})" + + rows = [] + for event in self.events: + if event.action == "onload": + offloaded = ", ".join(label(model_id) for model_id in event.offloaded) or "-" + rows.append((label(event.model_id), offloaded, format_size(event.available_memory), event.reason)) + elif event.action == "offload": + if event.reason is not None and event.reason.startswith("needed_by:"): + continue # shown on the row of the onload that caused it + rows.append(("-", label(event.model_id), "-", event.reason or "")) + else: # oom + rows.append(("-", "-", "-", f"oom:{event.model_id}")) + + onload_width = max(len("Onload"), *(len(row[0]) for row in rows)) + offload_width = max(len("Offloaded"), *(len(row[1]) for row in rows)) + header = f"{'#':<5} | {'Onload':<{onload_width}} | {'Offloaded':<{offload_width}} | {'Available':<10} | Reason" lines = [header, "-" * len(header)] - for index, event in enumerate(self.events, start=self.dropped + 1): + for index, row in enumerate(rows, start=1): lines.append( - f"{index:<5} | {event.action:<8} | {event.model_id:<40} | {format_size(event.model_size):<10} | " - f"{format_size(event.available_memory):<10} | {event.reason or ''}" + f"{index:<5} | {row[0]:<{onload_width}} | {row[1]:<{offload_width}} | {row[2]:<10} | {row[3]}" ) if self.dropped: lines.insert(2, f"... {self.dropped} earlier events dropped from the log ...") diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index 85ea3d3e4336..3f473f725849 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -406,6 +406,9 @@ def test_record_captures_what_the_strategy_saw(self): assert onloads[1].resident_before == (cm.model_hooks[0].model_id,) assert onloads[1].offloaded == (cm.model_hooks[0].model_id,) assert cm.offload_record.summary()["peak_co_residency"] == 1 + # the printed table shows one row per decision: m2's row carries the m1 eviction it caused + m2_row = next(line for line in repr(cm.offload_record).splitlines() if line.startswith("2 ")) + assert cm.model_hooks[1].model_id in m2_row and cm.model_hooks[0].model_id in m2_row finally: cm.disable_auto_cpu_offload() From b755e08df3be0b637b3260fdabfd1604d6993743 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 29 Jul 2026 04:19:07 +0000 Subject: [PATCH 05/14] Split display helpers into components_manager_utils; drop activation 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 --- .../modular_diffusers/components_manager.md | 49 +- .../modular_pipelines/components_manager.py | 523 +++++------------- .../components_manager_utils.py | 75 +++ .../test_components_manager.py | 66 +-- 4 files changed, 267 insertions(+), 446 deletions(-) create mode 100644 src/diffusers/modular_pipelines/components_manager_utils.py diff --git a/docs/source/en/modular_diffusers/components_manager.md b/docs/source/en/modular_diffusers/components_manager.md index 6224712e8dd9..dc53033ff1bc 100644 --- a/docs/source/en/modular_diffusers/components_manager.md +++ b/docs/source/en/modular_diffusers/components_manager.md @@ -58,23 +58,23 @@ The output below corresponds to the `from_pretrained` example above. ```py Components: -============================================================================================================================= +======================================================================================================================================================================= Models: ------------------------------------------------------------------------------------------------------------------------------ -Name_ID | Class | Device: act(exec) | Dtype | Size (GB) | Load ID ------------------------------------------------------------------------------------------------------------------------------ -text_encoder_140458257514752 | Qwen3Model | cpu | torch.bfloat16 | 7.49 | Tongyi-MAI/Z-Image-Turbo|text_encoder|null|null -vae_140458257515376 | AutoencoderKL | cpu | torch.bfloat16 | 0.16 | Tongyi-MAI/Z-Image-Turbo|vae|null|null -transformer_140458257515616 | ZImageTransformer2DModel | cpu | torch.bfloat16 | 11.46 | Tongyi-MAI/Z-Image-Turbo|transformer|null|null ------------------------------------------------------------------------------------------------------------------------------ +----------------------------------------------------------------------------------------------------------------------------------------------------------------------- +Name_ID | Class | Device: act(exec) | Dtype | Size | Load ID | Collection +----------------------------------------------------------------------------------------------------------------------------------------------------------------------- +vae_140458257515376 | AutoencoderKL | cpu | torch.bfloat16 | 159.87 MB | Tongyi-MAI/Z-Image-Turbo|vae|null|null | N/A +text_encoder_140458257514752 | Qwen3Model | cpu | torch.bfloat16 | 7.49 GB | Tongyi-MAI/Z-Image-Turbo|text_encoder|null|null | N/A +transformer_140458257515616 | ZImageTransformer2DModel | cpu | torch.bfloat16 | 11.46 GB | Tongyi-MAI/Z-Image-Turbo|transformer|null|null | N/A +----------------------------------------------------------------------------------------------------------------------------------------------------------------------- Other Components: ------------------------------------------------------------------------------------------------------------------------------ -ID | Class | Collection ------------------------------------------------------------------------------------------------------------------------------ -scheduler_140461023555264 | FlowMatchEulerDiscreteScheduler | N/A -tokenizer_140458256346432 | Qwen2Tokenizer | N/A ------------------------------------------------------------------------------------------------------------------------------ +------------------------------------------------------------------------ +ID | Class | Collection +------------------------------------------------------------------------ +scheduler_140461023555264 | FlowMatchEulerDiscreteScheduler | N/A +tokenizer_140458256346432 | Qwen2Tokenizer | N/A +------------------------------------------------------------------------ ``` The table shows models (with device, dtype, and memory info) separately from other components like schedulers and tokenizers. If any models have LoRA adapters, IP-Adapters, or quantization applied, that information is displayed in an additional section at the bottom. @@ -110,12 +110,12 @@ print(manager.offload_record) On a 20GB card, the Z-Image pipeline from the example above records this: ```py -# | Onload | Offloaded | Available | Reason -------------------------------------------------------------------------------------------------------------- -1 | text_encoder_140458257514752 (7.49 GB) | - | - | forward -2 | transformer_140458257515616 (11.46 GB) | text_encoder_140458257514752 (7.49 GB) | 10.38 GB | forward -3 | vae_140458257515376 (159.87 MB) | - | 6.36 GB | forward -------------------------------------------------------------------------------------------------------------- +# | Onload | Offloaded | Available | Reason +-------------------------------------------------------------------------------------------------------- +1 | text_encoder_140458257514752 (7.49 GB) | - | - | forward +2 | transformer_140458257515616 (11.46 GB) | text_encoder_140458257514752 (7.49 GB) | 10.38 GB | forward +3 | vae_140458257515376 (159.87 MB) | - | 6.36 GB | forward +-------------------------------------------------------------------------------------------------------- 3 onloads (19.11 GB) / 1 offloads (7.49 GB), 0 OOM retries, peak co-residency 2 ``` @@ -123,17 +123,16 @@ Each row is one decision: the model that loaded, what was evicted to make room f A model appearing repeatedly in this table is thrashing — it is being evicted 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). -To find the right value, let the manager measure it. `measure_activations=True` records how much memory the forward passes actually need on top of the weights — the exact thing `memory_reserve` covers — and reports the reserve that would have covered your run: +To find the right value, measure your workflow once: run it at the resolution, batch size, and sequence length you intend to use (activations scale with all three), and read the device's peak memory afterwards: ```py -manager.enable_auto_cpu_offload(device="cuda", memory_reserve=0, measure_activations=True) +manager.enable_auto_cpu_offload(device="cuda") pipe(prompt="a cat") -manager.offload_record.activation_peak # what the forward passes needed -manager.offload_record.suggested_memory_reserve # a reserve that would have covered it +torch.cuda.max_memory_allocated() # peak of the run: resident weights + activations ``` -Measure at the resolution, batch size, and sequence length you intend to use — activations scale with all three. Leave `measure_activations` off if you are measuring peak memory yourself, since it resets the device's peak-memory stats around every forward pass. +The peak is the heaviest moment of the run: the weights resident at that point plus the running model's activations. The record shows the weights side — which models were on the device together and their sizes — so subtracting them from the peak gives the activation headroom, and rounding that up generously gives the `memory_reserve`. In the 20GB run above the peak reads 14.15 GB and falls in the decode stage, where the transformer (11.46 GB) and the VAE (159.87 MB) are both resident; the ~2.5 GB left over is what the VAE's decode needed on top of the weights, so the default 3GB reserve covers this workload. Call [`~ComponentsManager.disable_auto_cpu_offload`] to disable offloading. diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 174637d01d23..35318a957fa7 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -14,7 +14,6 @@ from __future__ import annotations -import copy import functools import time from collections import OrderedDict, deque @@ -30,6 +29,7 @@ logging, ) from ..utils.torch_utils import get_device +from .components_manager_utils import format_size, format_table, summarize_dict_by_value_and_parts if is_accelerate_available(): @@ -46,17 +46,6 @@ MAX_RECORDED_OFFLOAD_EVENTS = 10000 -def format_size(num_bytes: int | None) -> str: - """Bytes as a human-readable size, so that both a 20GB transformer and a KB-sized test model read sensibly.""" - if num_bytes is None: - return "-" - for unit in ("B", "KB", "MB"): - if abs(num_bytes) < 1024: - return f"{num_bytes:.0f} {unit}" if unit == "B" else f"{num_bytes:.2f} {unit}" - num_bytes /= 1024 - return f"{num_bytes:.2f} GB" - - @dataclass class OffloadEvent: """ @@ -84,28 +73,18 @@ class OffloadRecord: Every offloading decision appends to this record, so after a run it holds the full sequence of moves, what each move cost, and where the forward passes ran out of memory. Printing it shows that sequence as a table followed by a summary; [`~OffloadRecord.summary`] returns the same numbers as a dict. - - With `measure_activations=True` in [`~ComponentsManager.enable_auto_cpu_offload`], the record additionally - measures how much memory the forward passes need on top of the weights, and reports the `memory_reserve` that - would have covered it — measured on your hardware, at your resolution and batch size. """ def __init__(self, maxlen: int = MAX_RECORDED_OFFLOAD_EVENTS): self.events: deque[OffloadEvent] = deque(maxlen=maxlen) - self.activations: dict[str, int] = {} self._recorded = 0 def add(self, event: OffloadEvent): self.events.append(event) self._recorded += 1 - def note_activation(self, model_id: str, activation_bytes: int): - """Keep the largest activation cost seen for a model.""" - self.activations[model_id] = max(self.activations.get(model_id, 0), activation_bytes) - def clear(self): self.events.clear() - self.activations.clear() self._recorded = 0 @property @@ -113,23 +92,6 @@ def dropped(self) -> int: """How many events fell out of the (bounded) log.""" return max(self._recorded - len(self.events), 0) - @property - def activation_peak(self) -> int | None: - """ - The largest activation cost measured across all forward passes, or `None` if activations were not measured. - """ - return max(self.activations.values()) if self.activations else None - - @property - def suggested_memory_reserve(self) -> int | None: - """ - A `memory_reserve` that would have covered the measured activations, with a margin. `None` if activations were - not measured. Only as representative as the run it was measured on — activations grow with resolution, batch - size and sequence length. - """ - peak = self.activation_peak - return int(peak * 1.2) if peak is not None else None - def summary(self) -> dict[str, Any]: onloads = [event for event in self.events if event.action == "onload"] offloads = [event for event in self.events if event.action == "offload"] @@ -147,8 +109,6 @@ def summary(self) -> dict[str, Any]: "peak_co_residency": peak_co_residency, "bytes_onloaded": sum(event.model_size or 0 for event in onloads), "bytes_offloaded": sum(event.model_size or 0 for event in offloads), - "activation_peak": self.activation_peak, - "suggested_memory_reserve": self.suggested_memory_reserve, } def __len__(self): @@ -169,37 +129,29 @@ def label(model_id): for event in self.events: if event.action == "onload": offloaded = ", ".join(label(model_id) for model_id in event.offloaded) or "-" - rows.append((label(event.model_id), offloaded, format_size(event.available_memory), event.reason)) + rows.append([label(event.model_id), offloaded, format_size(event.available_memory), event.reason]) elif event.action == "offload": if event.reason is not None and event.reason.startswith("needed_by:"): continue # shown on the row of the onload that caused it - rows.append(("-", label(event.model_id), "-", event.reason or "")) + rows.append(["-", label(event.model_id), "-", event.reason or ""]) else: # oom - rows.append(("-", "-", "-", f"oom:{event.model_id}")) - - onload_width = max(len("Onload"), *(len(row[0]) for row in rows)) - offload_width = max(len("Offloaded"), *(len(row[1]) for row in rows)) - header = f"{'#':<5} | {'Onload':<{onload_width}} | {'Offloaded':<{offload_width}} | {'Available':<10} | Reason" - lines = [header, "-" * len(header)] - for index, row in enumerate(rows, start=1): - lines.append( - f"{index:<5} | {row[0]:<{onload_width}} | {row[1]:<{offload_width}} | {row[2]:<10} | {row[3]}" - ) + rows.append(["-", "-", "-", f"oom:{event.model_id}"]) + + table = format_table( + ["#", "Onload", "Offloaded", "Available", "Reason"], + [[str(index), *row] for index, row in enumerate(rows, start=1)], + ) + lines = [table[0], "-" * len(table[0]), *table[1:]] if self.dropped: lines.insert(2, f"... {self.dropped} earlier events dropped from the log ...") summary = self.summary() lines += [ - "-" * len(header), + "-" * len(table[0]), f"{summary['onloads']} onloads ({format_size(summary['bytes_onloaded'])}) / " f"{summary['offloads']} offloads ({format_size(summary['bytes_offloaded'])}), " f"{summary['oom_retries']} OOM retries, peak co-residency {summary['peak_co_residency']}", ] - if summary["activation_peak"] is not None: - lines.append( - f"activations peaked at {format_size(summary['activation_peak'])} -> a `memory_reserve` of " - f"{format_size(summary['suggested_memory_reserve'])} would have covered this run" - ) return "\n".join(lines) @@ -217,9 +169,6 @@ class CustomOffloadHook(ModelHook): at a time and retrying. If `False`, the error is raised to the caller. record(`OffloadRecord`, *optional*): Where to record the moves this hook makes. - measure_activations(`bool`, *optional*, defaults to `False`): - Whether to measure how much device memory each forward pass needs on top of the weights. Owns the - device's peak-memory stats while enabled. """ no_grad = False @@ -231,17 +180,14 @@ def __init__( offload_strategy: "AutoOffloadStrategy" | None = None, retry_on_oom: bool = True, record: OffloadRecord | None = None, - measure_activations: bool = False, ): self.execution_device = execution_device if execution_device is not None else PartialState().default_device self.other_hooks = other_hooks self.offload_strategy = offload_strategy self.retry_on_oom = retry_on_oom self.record = record if record is not None else OffloadRecord() - self.measure_activations = measure_activations self.model_id = None - self._forward_before_oom_wrap = None - self._allocated_before_forward = None + self._wrapped_methods = {} def set_strategy(self, offload_strategy: "AutoOffloadStrategy"): self.offload_strategy = offload_strategy @@ -300,12 +246,6 @@ def pre_forward(self, module, *args, **kwargs): seconds=elapsed, ) ) - if self.measure_activations: - # baseline for the activation measurement: everything allocated once this model is on the device - device_module = getattr(torch, self.execution_device.type, torch.cuda) - if hasattr(device_module, "reset_peak_memory_stats"): - device_module.reset_peak_memory_stats(self.execution_device) - self._allocated_before_forward = device_module.memory_allocated(self.execution_device) return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device) def resident_other_hooks(self): @@ -320,69 +260,66 @@ def resident_other_hooks(self): and (other.model.device.index or 0) == (self.execution_device.index or 0) ] - def _run_forward(self, hooked_forward, args, kwargs): - if not self.measure_activations: - return hooked_forward(*args, **kwargs) - output = hooked_forward(*args, **kwargs) - device_module = getattr(torch, self.execution_device.type, torch.cuda) - if self._allocated_before_forward is not None and hasattr(device_module, "max_memory_allocated"): - peak = device_module.max_memory_allocated(self.execution_device) - self.record.note_activation(self.model_id, max(peak - self._allocated_before_forward, 0)) - return output - # YiYi TODO: diffusers' own `ModelHook` (`diffusers.hooks.hooks`) supports a `new_forward` around-hook, which # would replace this manual wrapping - we may migrate this class to it in the future. def wrap_forward(self, module): - # `pre_forward`/`post_forward` cannot see the forward pass itself - neither its memory peak nor an exception - # it raises - so wrap the (already hooked) forward here. On an OOM, free the smallest resident model and - # retry, escalating until it fits. The memory readings cannot guide this: the failed forward has already - # unwound, so its activations are freed and the device looks free again. Inference only: the retried forward - # re-runs cleanly from its original inputs. - if not (self.retry_on_oom or self.measure_activations): + # `pre_forward`/`post_forward` cannot see an exception the forward pass raises, so wrap the (already hooked) + # entry points here. `forward` is how most models run; autoencoders enter through `encode`/`decode` instead + # (decorated with `apply_forward_hook`, which fires `pre_forward` but routes around `forward`). On an OOM, + # free the smallest resident model and retry, escalating until it fits. The memory readings cannot guide + # this: the failed forward has already unwound, so its activations are freed and the device looks free again. + # Inference only: the retried forward re-runs cleanly from its original inputs. + if not self.retry_on_oom: return - hooked_forward = module.forward - - @functools.wraps(hooked_forward) - def forward_with_oom_retry(*args, **kwargs): - # each pass offloads one more model, so this terminates once they all have been. Tracked explicitly - # rather than by device, so that a model that did not actually move cannot be picked twice. - offloaded = set() - while True: - try: - return self._run_forward(hooked_forward, args, kwargs) - except torch.OutOfMemoryError as e: - if not self.retry_on_oom: - raise - self.record.add(OffloadEvent(action="oom", model_id=self.model_id, reason=str(e).split(".")[0])) - resident = sorted( - (hook for hook in self.resident_other_hooks() if id(hook) not in offloaded), - key=lambda hook: hook.model.get_memory_footprint(), - ) - if not resident: - raise torch.OutOfMemoryError( - f"{self.model_id} ran out of device memory ({e}) with every other managed model already " - "offloaded, so it does not fit on its own. Consider group offloading " - "(`ModelMixin.enable_group_offload`), which offloads a single model in groups of " - "internal layers." - ) from e - smallest = resident[0] - logger.warning( - f"{self.model_id} ran out of device memory ({e}); offloading {smallest.model_id} and " - "retrying. If this happens repeatedly, set a larger `memory_reserve` in " - "`enable_auto_cpu_offload`." - ) - smallest.offload(reason=f"oom_retry:{self.model_id}") - offloaded.add(id(smallest)) - clear_device_cache() + def with_oom_retry(entry_point): + @functools.wraps(entry_point) + def run_with_oom_retry(*args, **kwargs): + # each pass offloads one more model, so this terminates once they all have been. Tracked explicitly + # rather than by device, so that a model that did not actually move cannot be picked twice. + offloaded = set() + while True: + try: + return entry_point(*args, **kwargs) + except torch.OutOfMemoryError as e: + if not self.retry_on_oom: + raise + self.record.add( + OffloadEvent(action="oom", model_id=self.model_id, reason=str(e).split(".")[0]) + ) + resident = sorted( + (hook for hook in self.resident_other_hooks() if id(hook) not in offloaded), + key=lambda hook: hook.model.get_memory_footprint(), + ) + if not resident: + raise torch.OutOfMemoryError( + f"{self.model_id} ran out of device memory ({e}) with every other managed model " + "already offloaded, so it does not fit on its own. Consider group offloading " + "(`ModelMixin.enable_group_offload`), which offloads a single model in groups of " + "internal layers." + ) from e + smallest = resident[0] + logger.warning( + f"{self.model_id} ran out of device memory ({e}); offloading {smallest.model_id} and " + "retrying. If this happens repeatedly, set a larger `memory_reserve` in " + "`enable_auto_cpu_offload`." + ) + smallest.offload(reason=f"oom_retry:{self.model_id}") + offloaded.add(id(smallest)) + clear_device_cache() - module.forward = forward_with_oom_retry - self._forward_before_oom_wrap = hooked_forward + return run_with_oom_retry + + for name in ("forward", "encode", "decode"): + entry_point = getattr(module, name, None) + if callable(entry_point): + self._wrapped_methods[name] = entry_point + setattr(module, name, with_oom_retry(entry_point)) def unwrap_forward(self, module): - if self._forward_before_oom_wrap is not None: - module.forward = self._forward_before_oom_wrap - self._forward_before_oom_wrap = None + for name, entry_point in self._wrapped_methods.items(): + setattr(module, name, entry_point) + self._wrapped_methods.clear() class UserCustomOffloadHook: @@ -430,14 +367,12 @@ def custom_offload_with_hook( offload_strategy: "AutoOffloadStrategy" | None = None, retry_on_oom: bool = True, record: OffloadRecord | None = None, - measure_activations: bool = False, ): hook = CustomOffloadHook( execution_device=execution_device, offload_strategy=offload_strategy, retry_on_oom=retry_on_oom, record=record, - measure_activations=measure_activations, ) user_hook = UserCustomOffloadHook(model_id=model_id, model=model, hook=hook) user_hook.attach() @@ -540,67 +475,6 @@ def search_best_candidate(module_sizes, min_memory_offload): return [hook for hook in hooks if hook.model_id in best_offload_model_ids] -# utils for display component info in a readable format -# TODO: move to a different file -def summarize_dict_by_value_and_parts(d: dict[str, Any]) -> dict[str, Any]: - """Summarizes a dictionary by finding common prefixes that share the same value. - - For a dictionary with dot-separated keys like: { - 'down_blocks.1.attentions.1.transformer_blocks.0.attn2.processor': [0.6], - 'down_blocks.1.attentions.1.transformer_blocks.1.attn2.processor': [0.6], - 'up_blocks.1.attentions.0.transformer_blocks.0.attn2.processor': [0.3], - } - - Returns a dictionary where keys are the shortest common prefixes and values are their shared values: { - 'down_blocks': [0.6], 'up_blocks': [0.3] - } - """ - # First group by values - convert lists to tuples to make them hashable - value_to_keys = {} - for key, value in d.items(): - value_tuple = tuple(value) if isinstance(value, list) else value - if value_tuple not in value_to_keys: - value_to_keys[value_tuple] = [] - value_to_keys[value_tuple].append(key) - - def find_common_prefix(keys: list[str]) -> str: - """Find the shortest common prefix among a list of dot-separated keys.""" - if not keys: - return "" - if len(keys) == 1: - return keys[0] - - # Split all keys into parts - key_parts = [k.split(".") for k in keys] - - # Find how many initial parts are common - common_length = 0 - for parts in zip(*key_parts): - if len(set(parts)) == 1: # All parts at this position are the same - common_length += 1 - else: - break - - if common_length == 0: - return "" - - # Return the common prefix - return ".".join(key_parts[0][:common_length]) - - # Create summary by finding common prefixes for each value group - summary = {} - for value_tuple, keys in value_to_keys.items(): - prefix = find_common_prefix(keys) - if prefix: # Only add if we found a common prefix - # Convert tuple back to list if it was originally a list - value = list(value_tuple) if isinstance(d[keys[0]], list) else value_tuple - summary[prefix] = value - else: - summary[""] = value # Use empty string if no common prefix - - return summary - - class ComponentsManager: """ A central registry and management system for model components across multiple pipelines. @@ -651,7 +525,6 @@ def __init__(self): self._auto_offload_enabled = False self._offload_strategy = None self._offload_retry_on_oom = True - self._offload_measure_activations = False self._offload_record = OffloadRecord() def _lookup_ids( @@ -1013,7 +886,6 @@ def enable_auto_cpu_offload( device: str | int | torch.device = None, memory_reserve: str | int = "3GB", retry_on_oom: bool = True, - measure_activations: bool = False, ): """ Enable automatic CPU offloading for all components. @@ -1040,11 +912,6 @@ def enable_auto_cpu_offload( Whether to recover from a forward pass that runs out of device memory by offloading the models on the device one at a time, smallest first, and retrying until it fits. Set it to `False` to raise the error instead — the forward passes are then left untouched. - measure_activations (bool, *optional*, defaults to `False`): - Whether to measure how much device memory each forward pass needs on top of the weights, and report - the `memory_reserve` that would have covered it (see [`~ComponentsManager.offload_record`]). This - resets the device's peak-memory stats around every forward pass, so leave it off if you are measuring - peak memory yourself. Every move the offloader makes is recorded in [`~ComponentsManager.offload_record`]. """ @@ -1083,7 +950,6 @@ def enable_auto_cpu_offload( offload_strategy=offload_strategy, retry_on_oom=retry_on_oom, record=self._offload_record, - measure_activations=measure_activations, ) all_hooks.append(hook) @@ -1098,16 +964,14 @@ def enable_auto_cpu_offload( self._auto_offload_device = device self._offload_strategy = offload_strategy self._offload_retry_on_oom = retry_on_oom - self._offload_measure_activations = measure_activations @property def offload_record(self) -> OffloadRecord: """ What the offloader has done so far: every model moved onto or off the device, in order, and where a forward - pass ran out of memory. Print it to see the sequence, or read - [`~OffloadRecord.summary`]/[`~OffloadRecord.suggested_memory_reserve`] for the numbers behind it. Kept across - [`~ComponentsManager.disable_auto_cpu_offload`] (it is the post-mortem of the run), and cleared when - offloading is enabled again. + pass ran out of memory. Print it to see the sequence, or read [`~OffloadRecord.summary`] for the numbers + behind it. Kept across [`~ComponentsManager.disable_auto_cpu_offload`] (it is the post-mortem of the run), + and cleared when offloading is enabled again. """ return self._offload_record @@ -1123,7 +987,6 @@ def _attach_offload_hook(self, component_id: str, component: torch.nn.Module): offload_strategy=self._offload_strategy, retry_on_oom=self._offload_retry_on_oom, record=self._offload_record, - measure_activations=self._offload_measure_activations, ) for other_hook in self.model_hooks: if other_hook.hook.execution_device == hook.hook.execution_device: @@ -1158,227 +1021,131 @@ def disable_auto_cpu_offload(self): self._auto_offload_enabled = False self._offload_strategy = None self._offload_retry_on_oom = True - self._offload_measure_activations = False def get_model_info( self, component_id: str, fields: str | list[str] | None = None, - ) -> dict[str, Any] | None: + ) -> dict[str, Any]: """Get comprehensive information about a component. Args: component_id (str): Name of the component to get info for - fields (str | list[str] | None): - Field(s) to return. Can be a string for single field or list of fields. If None, uses the - available_info_fields setting. + fields (str | list[str] | None): Field(s) to return, all fields if `None`. Returns: - Dictionary containing requested component metadata. If fields is specified, returns only those fields. - Otherwise, returns all fields. + Dictionary containing the requested component metadata. """ if component_id not in self.components: raise ValueError(f"Component '{component_id}' not found in ComponentsManager") - component = self.components[component_id] - # Validate fields if specified - if fields is not None: - if isinstance(fields, str): - fields = [fields] - for field in fields: - if field not in self._available_info_fields: - raise ValueError(f"Field '{field}' not found in available_info_fields") - - # Build complete info dict first info = { "model_id": component_id, "added_time": self.added_time[component_id], - "collection": ", ".join([coll for coll, comps in self.collections.items() if component_id in comps]) - or None, + "collection": ", ".join(coll for coll, comps in self.collections.items() if component_id in comps) or None, } - # Additional info for torch.nn.Module components if isinstance(component, torch.nn.Module): - # Check for hook information - has_hook = hasattr(component, "_hf_hook") - execution_device = None - if has_hook and hasattr(component._hf_hook, "execution_device"): - execution_device = component._hf_hook.execution_device - + hook = getattr(component, "_hf_hook", None) info.update( { "class_name": component.__class__.__name__, - "size_gb": component.get_memory_footprint() / (1024**3), - "adapters": None, # Default to None - "has_hook": has_hook, - "execution_device": execution_device, + "size_gb": component.get_memory_footprint() / 1024**3, + "adapters": list(component.peft_config.keys()) if hasattr(component, "peft_config") else None, + "has_hook": hook is not None, + "execution_device": getattr(hook, "execution_device", None), } ) - # Get adapters if applicable - if hasattr(component, "peft_config"): - info["adapters"] = list(component.peft_config.keys()) - - # Check for IP-Adapter scales + # IP-Adapter attention processor scales, summarized by shared layer prefix if hasattr(component, "_load_ip_adapter_weights") and hasattr(component, "attn_processors"): - processors = copy.deepcopy(component.attn_processors) - # First check if any processor is an IP-Adapter - processor_types = [v.__class__.__name__ for v in processors.values()] - if any("IPAdapter" in ptype for ptype in processor_types): - # Then get scales only from IP-Adapter processors - scales = { - k: v.scale - for k, v in processors.items() - if hasattr(v, "scale") and "IPAdapter" in v.__class__.__name__ - } - if scales: - info["ip_adapter"] = summarize_dict_by_value_and_parts(scales) - - # Check for quantization + scales = { + name: processor.scale + for name, processor in component.attn_processors.items() + if "IPAdapter" in processor.__class__.__name__ and hasattr(processor, "scale") + } + if scales: + info["ip_adapter"] = summarize_dict_by_value_and_parts(scales) + hf_quantizer = getattr(component, "hf_quantizer", None) - if hf_quantizer is not None: - quant_config = hf_quantizer.quantization_config - if hasattr(quant_config, "to_diff_dict"): - info["quantization"] = quant_config.to_diff_dict() - else: - info["quantization"] = quant_config.to_dict() - else: + if hf_quantizer is None: info["quantization"] = None + else: + quant_config = hf_quantizer.quantization_config + info["quantization"] = ( + quant_config.to_diff_dict() if hasattr(quant_config, "to_diff_dict") else quant_config.to_dict() + ) - # If fields specified, filter info - if fields is not None: - return {k: v for k, v in info.items() if k in fields} - else: + if fields is None: return info - - # YiYi TODO: (1) add display fields, allow user to set which fields to display in the comnponents table + if isinstance(fields, str): + fields = [fields] + for field in fields: + if field not in self._available_info_fields: + raise ValueError(f"Field '{field}' not found in available_info_fields") + return {k: v for k, v in info.items() if k in fields} + + # YiYi TODO: (1) add display fields, allow user to set which fields to display in the components table def __repr__(self): - # Handle empty components case if not self.components: return "Components:\n" + "=" * 50 + "\nNo components registered.\n" + "=" * 50 - # Extract load_id if available - def get_load_id(component): - if hasattr(component, "_diffusers_load_id"): - return component._diffusers_load_id - return "N/A" - - # Format device info compactly - def format_device(component, info): - if not info["has_hook"]: - return str(getattr(component, "device", "N/A")) - else: - device = str(getattr(component, "device", "N/A")) - exec_device = str(info["execution_device"] or "N/A") - return f"{device}({exec_device})" - - # Get max length of load_ids for models - load_ids = [ - get_load_id(component) - for component in self.components.values() - if isinstance(component, torch.nn.Module) and hasattr(component, "_diffusers_load_id") - ] - max_load_id_len = max([15] + [len(str(lid)) for lid in load_ids]) if load_ids else 15 - - # Get all collections for each component - component_collections = {} - for name in self.components.keys(): - component_collections[name] = [] - for coll, comps in self.collections.items(): - if name in comps: - component_collections[name].append(coll) - if not component_collections[name]: - component_collections[name] = ["N/A"] - - # Find the maximum collection name length - all_collections = [coll for colls in component_collections.values() for coll in colls] - max_collection_len = max(10, max(len(str(c)) for c in all_collections)) if all_collections else 10 - - col_widths = { - "id": max(15, max(len(name) for name in self.components.keys())), - "class": max(25, max(len(component.__class__.__name__) for component in self.components.values())), - "device": 20, - "dtype": 15, - "size": 10, - "load_id": max_load_id_len, - "collection": max_collection_len, + infos = {name: self.get_model_info(name) for name in self.components} + # every collection a component belongs to; the first goes on the component's own row, the rest on + # continuation rows below it + component_collections = { + name: [coll for coll, comps in self.collections.items() if name in comps] or ["N/A"] + for name in self.components } - # Create the header lines - sep_line = "=" * (sum(col_widths.values()) + len(col_widths) * 3 - 1) + "\n" - dash_line = "-" * (sum(col_widths.values()) + len(col_widths) * 3 - 1) + "\n" + def rows_with_collections(name: str, cells: list[str]) -> list[list[str]]: + first, *rest = component_collections[name] + return [[*cells, first]] + [[""] * len(cells) + [coll] for coll in rest] - output = "Components:\n" + sep_line + models = {name: c for name, c in self.components.items() if isinstance(c, torch.nn.Module)} + others = {name: c for name, c in self.components.items() if not isinstance(c, torch.nn.Module)} - # Separate components into models and others - models = {k: v for k, v in self.components.items() if isinstance(v, torch.nn.Module)} - others = {k: v for k, v in self.components.items() if not isinstance(v, torch.nn.Module)} - - # Models section + sections = [] if models: - output += "Models:\n" + dash_line - # Column headers - output += f"{'Name_ID':<{col_widths['id']}} | {'Class':<{col_widths['class']}} | " - output += f"{'Device: act(exec)':<{col_widths['device']}} | {'Dtype':<{col_widths['dtype']}} | " - output += f"{'Size (GB)':<{col_widths['size']}} | {'Load ID':<{col_widths['load_id']}} | Collection\n" - output += dash_line - - # Model entries + rows = [] for name, component in models.items(): - info = self.get_model_info(name) - device_str = format_device(component, info) - dtype = str(component.dtype) if hasattr(component, "dtype") else "N/A" - load_id = get_load_id(component) - - # Print first collection on the main line - first_collection = component_collections[name][0] if component_collections[name] else "N/A" - - output += f"{name:<{col_widths['id']}} | {info['class_name']:<{col_widths['class']}} | " - output += f"{device_str:<{col_widths['device']}} | {dtype:<{col_widths['dtype']}} | " - output += f"{info['size_gb']:<{col_widths['size']}.2f} | {load_id:<{col_widths['load_id']}} | {first_collection}\n" - - # Print additional collections on separate lines if they exist - for i in range(1, len(component_collections[name])): - collection = component_collections[name][i] - output += f"{'':<{col_widths['id']}} | {'':<{col_widths['class']}} | " - output += f"{'':<{col_widths['device']}} | {'':<{col_widths['dtype']}} | " - output += f"{'':<{col_widths['size']}} | {'':<{col_widths['load_id']}} | {collection}\n" - - output += dash_line - - # Other components section + info = infos[name] + device = str(getattr(component, "device", "N/A")) + if info["has_hook"]: + device = f"{device}({info['execution_device'] or 'N/A'})" + rows += rows_with_collections( + name, + [ + name, + info["class_name"], + device, + str(component.dtype) if hasattr(component, "dtype") else "N/A", + format_size(component.get_memory_footprint()), + str(getattr(component, "_diffusers_load_id", "N/A")), + ], + ) + headers = ["Name_ID", "Class", "Device: act(exec)", "Dtype", "Size", "Load ID", "Collection"] + sections.append(("Models:", format_table(headers, rows))) if others: - if models: # Add extra newline if we had models section + rows = [ + row + for name, component in others.items() + for row in rows_with_collections(name, [name, component.__class__.__name__]) + ] + sections.append(("Other Components:", format_table(["ID", "Class", "Collection"], rows))) + + output = "Components:\n" + "=" * max(len(table[0]) for _, table in sections) + "\n" + for index, (title, table) in enumerate(sections): + dash_line = "-" * len(table[0]) + "\n" + if index: output += "\n" - output += "Other Components:\n" + dash_line - # Column headers for other components - output += f"{'ID':<{col_widths['id']}} | {'Class':<{col_widths['class']}} | Collection\n" - output += dash_line - - # Other component entries - for name, component in others.items(): - info = self.get_model_info(name) - - # Print first collection on the main line - first_collection = component_collections[name][0] if component_collections[name] else "N/A" - - output += f"{name:<{col_widths['id']}} | {component.__class__.__name__:<{col_widths['class']}} | {first_collection}\n" - - # Print additional collections on separate lines if they exist - for i in range(1, len(component_collections[name])): - collection = component_collections[name][i] - output += f"{'':<{col_widths['id']}} | {'':<{col_widths['class']}} | {collection}\n" - - output += dash_line + output += title + "\n" + dash_line + table[0] + "\n" + dash_line + output += "\n".join(table[1:]) + "\n" + dash_line - # Add additional component info output += "\nAdditional Component Info:\n" + "=" * 50 + "\n" - for name in self.components: - info = self.get_model_info(name) - if info is not None and ( - info.get("adapters") is not None or info.get("ip_adapter") or info.get("quantization") - ): + for name, info in infos.items(): + if info.get("adapters") is not None or info.get("ip_adapter") or info.get("quantization"): output += f"\n{name}:\n" if info.get("adapters") is not None: output += f" Adapters: {info['adapters']}\n" diff --git a/src/diffusers/modular_pipelines/components_manager_utils.py b/src/diffusers/modular_pipelines/components_manager_utils.py new file mode 100644 index 000000000000..84c4e7c83982 --- /dev/null +++ b/src/diffusers/modular_pipelines/components_manager_utils.py @@ -0,0 +1,75 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Display helpers for `ComponentsManager`: human-readable sizes, text tables, per-layer value summaries.""" + +from __future__ import annotations + +from typing import Any + + +def format_size(num_bytes: int | None) -> str: + """Bytes as a human-readable size, so that both a 20GB transformer and a KB-sized test model read sensibly.""" + if num_bytes is None: + return "-" + for unit in ("B", "KB", "MB"): + if abs(num_bytes) < 1024: + return f"{num_bytes:.0f} {unit}" if unit == "B" else f"{num_bytes:.2f} {unit}" + num_bytes /= 1024 + return f"{num_bytes:.2f} GB" + + +def format_table(headers: list[str], rows: list[list[str]]) -> list[str]: + """ + Align `headers` and `rows` into text-table lines (header first), each column as wide as its longest cell and + cells joined by " | ". The last column is not padded, so a long final cell (a reason, a load id) runs free + without stretching the table. Separator lines are the caller's to add — `len(lines[0])` is the table width. + """ + widths = [len(header) for header in headers] + for row in rows: + for index, cell in enumerate(row[:-1]): + widths[index] = max(widths[index], len(cell)) + + def line(cells: list[str]) -> str: + return " | ".join([cell.ljust(width) for cell, width in zip(cells[:-1], widths)] + [cells[-1]]) + + return [line(headers), *(line(row) for row in rows)] + + +def summarize_dict_by_value_and_parts(d: dict[str, Any]) -> dict[str, Any]: + """ + Summarize a dict with dot-separated keys by grouping keys that share a value under their longest common prefix. + + For example IP-Adapter scales per attention processor: { + 'down_blocks.1.attentions.1.transformer_blocks.0.attn2.processor': [0.6], + 'down_blocks.1.attentions.1.transformer_blocks.1.attn2.processor': [0.6], + 'up_blocks.1.attentions.0.transformer_blocks.0.attn2.processor': [0.3], + } becomes {'down_blocks.1.attentions.1.transformer_blocks': [0.6], 'up_blocks': [0.3]}. + """ + value_to_keys: dict[Any, list[str]] = {} + for key, value in d.items(): + hashable = tuple(value) if isinstance(value, list) else value + value_to_keys.setdefault(hashable, []).append(key) + + summary = {} + for keys in value_to_keys.values(): + split_keys = [key.split(".") for key in keys] + common_parts = [] + for parts in zip(*split_keys): + if len(set(parts)) != 1: + break + common_parts.append(parts[0]) + value = d[keys[0]] + summary[".".join(common_parts)] = list(value) if isinstance(value, list) else value + return summary diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index 3f473f725849..c5f525a4235a 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -23,6 +23,7 @@ from diffusers import ComponentsManager from diffusers.models import ModelMixin from diffusers.utils import is_accelerate_available +from diffusers.utils.accelerate_utils import apply_forward_hook from ..testing_utils import ( backend_empty_cache, @@ -73,16 +74,16 @@ def forward(self, x): return super().forward(x) -class ActivationHungryModel(DummyModel): - """Allocates a scratch tensor of a known size during its forward, giving the activation measurement a floor.""" +class OOMOnceDecodeModel(DummyModel): + """Autoencoder-style model: runs through a `decode` entry point (via `apply_forward_hook`, like the VAEs) + instead of `forward`, and raises a fake device OOM on its first decode.""" - def __init__(self, footprint_bytes: int = UNIT, activation_bytes: int = 8 * UNIT): - super().__init__(footprint_bytes=footprint_bytes) - self.activation_bytes = activation_bytes - - def forward(self, x): - scratch = torch.zeros(self.activation_bytes // 4, device=x.device) - return super().forward(x) + scratch[0] + @apply_forward_hook + def decode(self, x): + self.decode_calls = getattr(self, "decode_calls", 0) + 1 + if self.decode_calls == 1: + raise torch.OutOfMemoryError("fake OOM") + return x + self.weight.sum() class _FakeHook: @@ -412,39 +413,6 @@ def test_record_captures_what_the_strategy_saw(self): finally: cm.disable_auto_cpu_offload() - @require_accelerator - def test_record_measures_activations_and_suggests_a_reserve(self): - # The point of the measurement: a user who cannot spare memory for a calibration run still learns - # what their forward passes need on top of the weights, which is what `memory_reserve` covers. - activation_bytes = 4 * 1024**2 - cm = ComponentsManager() - model = ActivationHungryModel(footprint_bytes=4 * UNIT, activation_bytes=activation_bytes) - cm.add("model", model) - cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT, measure_activations=True) - try: - model(torch.randn(2, 4, device=torch_device)) - - peak = cm.offload_record.activation_peak - assert peak >= activation_bytes, f"expected at least the scratch tensor ({activation_bytes}), saw {peak}" - assert cm.offload_record.suggested_memory_reserve > peak - assert "memory_reserve" in repr(cm.offload_record) - finally: - cm.disable_auto_cpu_offload() - - @require_accelerator - def test_activations_are_not_measured_by_default(self): - cm = ComponentsManager() - model = ActivationHungryModel() - cm.add("model", model) - cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) - try: - model(torch.randn(2, 4, device=torch_device)) - # peak stats are the user's to own unless they opt in - assert cm.offload_record.activation_peak is None - assert cm.offload_record.suggested_memory_reserve is None - finally: - cm.disable_auto_cpu_offload() - @require_accelerator def test_auto_offload_keeps_models_resident_when_memory_is_ample(self): device_type = torch.device(torch_device).type @@ -547,7 +515,19 @@ def test_retry_on_oom_false_leaves_the_forward_alone(self): with pytest.raises(torch.OutOfMemoryError, match="fake OOM"): model(torch.zeros(2, 4)) assert model.forward_calls == 1 - assert all(user_hook.hook._forward_before_oom_wrap is None for user_hook in manager.model_hooks) + assert all(not user_hook.hook._wrapped_methods for user_hook in manager.model_hooks) + + def test_oom_retry_covers_apply_forward_hook_entry_points(self): + # Autoencoders run through `encode`/`decode` (via `apply_forward_hook`), not `forward` - an OOM + # there must be retried just like one raised inside `forward`. + model = OOMOnceDecodeModel() + smallest = DummyModel(UNIT) + manager = self._manager({"model": model, "smallest": smallest}) + + out = model.decode(torch.zeros(2, 4)) + assert out.shape == (2, 4) + assert model.decode_calls == 2 + assert self._offloaded_names(manager) == ["smallest"] def test_record_captures_the_onload_sequence(self): model_a = DummyModel(UNIT) From d2eb134a5cb2e082898311c68992b38eb8d13528 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 29 Jul 2026 08:12:29 +0000 Subject: [PATCH 06/14] Trim the offload record to consumed fields; hook reads available memory 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:, 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 --- .../modular_diffusers/components_manager.md | 8 +-- .../modular_pipelines/components_manager.py | 59 +++++++++---------- .../components_manager_utils.py | 3 +- .../test_components_manager.py | 7 +-- 4 files changed, 36 insertions(+), 41 deletions(-) diff --git a/docs/source/en/modular_diffusers/components_manager.md b/docs/source/en/modular_diffusers/components_manager.md index dc53033ff1bc..faa5440a46b8 100644 --- a/docs/source/en/modular_diffusers/components_manager.md +++ b/docs/source/en/modular_diffusers/components_manager.md @@ -112,14 +112,14 @@ On a 20GB card, the Z-Image pipeline from the example above records this: ```py # | Onload | Offloaded | Available | Reason -------------------------------------------------------------------------------------------------------- -1 | text_encoder_140458257514752 (7.49 GB) | - | - | forward -2 | transformer_140458257515616 (11.46 GB) | text_encoder_140458257514752 (7.49 GB) | 10.38 GB | forward -3 | vae_140458257515376 (159.87 MB) | - | 6.36 GB | forward +1 | text_encoder_140458257514752 (7.49 GB) | - | 18.02 GB | +2 | transformer_140458257515616 (11.46 GB) | text_encoder_140458257514752 (7.49 GB) | 10.38 GB | +3 | vae_140458257515376 (159.87 MB) | - | 6.36 GB | -------------------------------------------------------------------------------------------------------- 3 onloads (19.11 GB) / 1 offloads (7.49 GB), 0 OOM retries, peak co-residency 2 ``` -Each row is one decision: the model that loaded, what was evicted to make room for it, and the free device memory (`Available`) the decision saw at that moment. 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 evicted first. The VAE then fit into the 6.36 GB left next to the transformer, so it loaded without evicting anything (peak co-residency 2). `Available` reads `-` when nothing was resident, so the decision needed no memory reading (row 1). Offloads an onload did not cause — an OOM retry, disabling offloading — appear as their own rows, explained by their `Reason`. +Each row is one decision: the model that loaded, what was evicted to make room for it, and the free device memory (`Available`, read just before each load) the decision 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 evicted first. The VAE then fit into the 6.36 GB left next to the transformer, so it loaded without evicting anything (peak co-residency 2). The `Reason` column only speaks on moves an onload did not cause — an OOM retry (`oom_retry:`), disabling offloading — which appear as their own rows. To watch the moves live as they happen instead, enable info logging with `diffusers.logging.set_verbosity_info()`. A model appearing repeatedly in this table is thrashing — it is being evicted 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). diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 35318a957fa7..55f48622e2f4 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -46,6 +46,20 @@ MAX_RECORDED_OFFLOAD_EVENTS = 10000 +def available_device_memory(execution_device: torch.device) -> int: + """ + The device memory available for new weights right now: what the driver reports free, plus the allocator's + reusable cache (`mem_get_info` counts the cache as used, but freed tensors in it can be reallocated). + """ + device_module = getattr(torch, execution_device.type, torch.cuda) + available_memory = device_module.mem_get_info(execution_device.index)[0] + if hasattr(device_module, "memory_reserved") and hasattr(device_module, "memory_allocated"): + available_memory += device_module.memory_reserved(execution_device) - device_module.memory_allocated( + execution_device + ) + return available_memory + + @dataclass class OffloadEvent: """ @@ -58,12 +72,10 @@ class OffloadEvent: model_size: int | None = None # why the model moved: which model needed the room, the OOM retry, attaching a new component, ... reason: str | None = None - # what the strategy saw when it decided (`None` for a custom strategy that does not report it) + # free device memory read just before the onload's eviction decision (`None` if the device cannot report it) available_memory: int | None = None - memory_reserve: int | None = None resident_before: tuple[str, ...] = () offloaded: tuple[str, ...] = () - seconds: float | None = None class OffloadRecord: @@ -129,7 +141,9 @@ def label(model_id): for event in self.events: if event.action == "onload": offloaded = ", ".join(label(model_id) for model_id in event.offloaded) or "-" - rows.append([label(event.model_id), offloaded, format_size(event.available_memory), event.reason]) + rows.append( + [label(event.model_id), offloaded, format_size(event.available_memory), event.reason or ""] + ) elif event.action == "offload": if event.reason is not None and event.reason.startswith("needed_by:"): continue # shown on the row of the onload that caused it @@ -205,7 +219,12 @@ def init_hook(self, module): def pre_forward(self, module, *args, **kwargs): if module.device != self.execution_device: - resident_before, hooks_to_offload, elapsed = (), [], None + # read before any eviction: this is the number the eviction decision is based on + device_module = getattr(torch, self.execution_device.type, torch.cuda) + available_memory = ( + available_device_memory(self.execution_device) if hasattr(device_module, "mem_get_info") else None + ) + resident_before, hooks_to_offload = (), [] if self.other_hooks is not None: hooks_to_offload = [hook for hook in self.other_hooks if hook.model.device == self.execution_device] resident_before = tuple(hook.model_id for hook in hooks_to_offload) @@ -218,8 +237,7 @@ def pre_forward(self, module, *args, **kwargs): model=module, execution_device=self.execution_device, ) - end_time = time.perf_counter() - elapsed = end_time - start_time + elapsed = time.perf_counter() - start_time logger.info(f" time taken to apply offload strategy for {self.model_id}: {elapsed:.2f} seconds") for hook in hooks_to_offload: @@ -231,19 +249,14 @@ def pre_forward(self, module, *args, **kwargs): if hooks_to_offload: clear_device_cache() module.to(self.execution_device) - # what the strategy saw, when it is one that reports it - decision = getattr(self.offload_strategy, "last_decision", None) or {} self.record.add( OffloadEvent( action="onload", model_id=self.model_id, model_size=module.get_memory_footprint(), - reason="forward", - available_memory=decision.get("available_memory"), - memory_reserve=decision.get("memory_reserve"), + available_memory=available_memory, resident_before=resident_before, offloaded=tuple(hook.model_id for hook in hooks_to_offload), - seconds=elapsed, ) ) return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device) @@ -284,9 +297,7 @@ def run_with_oom_retry(*args, **kwargs): except torch.OutOfMemoryError as e: if not self.retry_on_oom: raise - self.record.add( - OffloadEvent(action="oom", model_id=self.model_id, reason=str(e).split(".")[0]) - ) + self.record.add(OffloadEvent(action="oom", model_id=self.model_id)) resident = sorted( (hook for hook in self.resident_other_hooks() if id(hook) not in offloaded), key=lambda hook: hook.model.get_memory_footprint(), @@ -393,13 +404,8 @@ class AutoOffloadStrategy: def __init__(self, memory_reserve="3GB"): self.memory_reserve = convert_file_size_to_int(memory_reserve) - # what the last decision saw, for the hook to record. A custom strategy that does not set it simply - # leaves those columns empty in the record. - self.last_decision = None def __call__(self, hooks, model_id, model, execution_device): - # cleared first so that a decision that returns early is never recorded with the previous decision's readings - self.last_decision = None if len(hooks) == 0: return [] @@ -410,16 +416,7 @@ def __call__(self, hooks, model_id, model, execution_device): resident_size = sum(hook.model.get_memory_footprint() for hook in hooks) - # Check the memory actually available on the device right now. The allocator keeps freed memory in its - # cache (mem_get_info reports it as used), but it is reusable — count it as available. - device_module = getattr(torch, execution_device.type, torch.cuda) - available_memory = device_module.mem_get_info(execution_device.index)[0] - if hasattr(device_module, "memory_reserved") and hasattr(device_module, "memory_allocated"): - available_memory += device_module.memory_reserved(execution_device) - device_module.memory_allocated( - execution_device - ) - self.last_decision = {"available_memory": available_memory, "memory_reserve": self.memory_reserve} - + available_memory = available_device_memory(execution_device) if current_module_size <= available_memory - self.memory_reserve: return [] diff --git a/src/diffusers/modular_pipelines/components_manager_utils.py b/src/diffusers/modular_pipelines/components_manager_utils.py index 84c4e7c83982..05d542c9e934 100644 --- a/src/diffusers/modular_pipelines/components_manager_utils.py +++ b/src/diffusers/modular_pipelines/components_manager_utils.py @@ -42,7 +42,8 @@ def format_table(headers: list[str], rows: list[list[str]]) -> list[str]: widths[index] = max(widths[index], len(cell)) def line(cells: list[str]) -> str: - return " | ".join([cell.ljust(width) for cell, width in zip(cells[:-1], widths)] + [cells[-1]]) + # rstrip so a row whose last cell is empty does not end in "| " + trailing space + return " | ".join([cell.ljust(width) for cell, width in zip(cells[:-1], widths)] + [cells[-1]]).rstrip() return [line(headers), *(line(row) for row in rows)] diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index c5f525a4235a..308f78a0ffa6 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -398,12 +398,9 @@ def test_record_captures_what_the_strategy_saw(self): onloads = [event for event in cm.offload_record.events if event.action == "onload"] names = [event.model_id.rsplit("_", 1)[0] for event in onloads] assert names == ["m1", "m2"] - # Nothing was resident when m1 loaded, so the strategy returned without reading memory and - # there is nothing to report for that decision. - assert onloads[0].available_memory is None - # The readings behind a real decision are recorded, so a run can be explained after the fact. + # The reading behind each decision is recorded, so a run can be explained after the fact. + assert onloads[0].available_memory == 70 * UNIT assert onloads[1].available_memory == 4 * UNIT - assert onloads[1].memory_reserve == UNIT assert onloads[1].resident_before == (cm.model_hooks[0].model_id,) assert onloads[1].offloaded == (cm.model_hooks[0].model_id,) assert cm.offload_record.summary()["peak_co_residency"] == 1 From 09d5ec7c20bca556acc8d10df1b26fde38390cea Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 29 Jul 2026 08:33:54 +0000 Subject: [PATCH 07/14] Make offload record events pure moves 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:, which summary() now counts), and OffloadEvent drops the redundant offloaded/resident_before tuples: each eviction is its own event whose needed_by: 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 --- .../modular_pipelines/components_manager.py | 57 ++++++++----------- .../test_components_manager.py | 38 ++++++------- 2 files changed, 39 insertions(+), 56 deletions(-) diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 55f48622e2f4..6153e50ee61d 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -63,46 +63,38 @@ def available_device_memory(execution_device: torch.device) -> int: @dataclass class OffloadEvent: """ - A single thing the offloader did: a model moved onto the execution device (`"onload"`), a model moved back to the - CPU (`"offload"`), or a forward pass that ran out of device memory (`"oom"`). + A single move the offloader made: a model moved onto the execution device (`"onload"`) or back to the CPU + (`"offload"`), with `reason` explaining why. """ action: str + # the model this event moved model_id: str model_size: int | None = None - # why the model moved: which model needed the room, the OOM retry, attaching a new component, ... + # why the model moved: which model needed the room ("needed_by:"), an OOM retry + # ("oom_retry:"), attaching a new component, disabling offloading, ... reason: str | None = None # free device memory read just before the onload's eviction decision (`None` if the device cannot report it) available_memory: int | None = None - resident_before: tuple[str, ...] = () - offloaded: tuple[str, ...] = () class OffloadRecord: """ What the offloader did, in order. - Every offloading decision appends to this record, so after a run it holds the full sequence of moves, what each - move cost, and where the forward passes ran out of memory. Printing it shows that sequence as a table followed by - a summary; [`~OffloadRecord.summary`] returns the same numbers as a dict. + Every move appends to this record, so after a run it holds the full sequence, what each move cost, and where the + forward passes ran out of memory. Printing it shows that sequence as a table followed by a summary; + [`~OffloadRecord.summary`] returns the same numbers as a dict. """ def __init__(self, maxlen: int = MAX_RECORDED_OFFLOAD_EVENTS): self.events: deque[OffloadEvent] = deque(maxlen=maxlen) - self._recorded = 0 def add(self, event: OffloadEvent): self.events.append(event) - self._recorded += 1 def clear(self): self.events.clear() - self._recorded = 0 - - @property - def dropped(self) -> int: - """How many events fell out of the (bounded) log.""" - return max(self._recorded - len(self.events), 0) def summary(self) -> dict[str, Any]: onloads = [event for event in self.events if event.action == "onload"] @@ -112,12 +104,12 @@ def summary(self) -> dict[str, Any]: if event.action == "onload": resident.add(event.model_id) peak_co_residency = max(peak_co_residency, len(resident)) - elif event.action == "offload": + else: resident.discard(event.model_id) return { "onloads": len(onloads), "offloads": len(offloads), - "oom_retries": sum(event.action == "oom" for event in self.events), + "oom_retries": sum(event.reason.startswith("oom_retry:") for event in offloads if event.reason), "peak_co_residency": peak_co_residency, "bytes_onloaded": sum(event.model_size or 0 for event in onloads), "bytes_offloaded": sum(event.model_size or 0 for event in offloads), @@ -130,34 +122,36 @@ def __repr__(self): if not self.events: return "Offload record: nothing recorded yet" - # One row per decision: an onload and the evictions it caused (its "needed_by" offload events) share a - # row. Offloads with other causes (OOM retry, offloading disabled) and OOMs get their own rows. + # One row per decision: an onload and the evictions it caused share a row, correlated through the + # eviction's "needed_by:" reason (evictions precede their onload in the event sequence). + # Offloads with other causes (OOM retry, offloading disabled) get their own rows. sizes = {event.model_id: event.model_size for event in self.events if event.model_size is not None} def label(model_id): return f"{model_id} ({format_size(sizes.get(model_id))})" - rows = [] + rows, pending_evictions = [], {} for event in self.events: if event.action == "onload": - offloaded = ", ".join(label(model_id) for model_id in event.offloaded) or "-" + evicted = pending_evictions.pop(event.model_id, []) + offloaded = ", ".join(label(model_id) for model_id in evicted) or "-" rows.append( [label(event.model_id), offloaded, format_size(event.available_memory), event.reason or ""] ) - elif event.action == "offload": - if event.reason is not None and event.reason.startswith("needed_by:"): - continue # shown on the row of the onload that caused it + elif event.reason is not None and event.reason.startswith("needed_by:"): + pending_evictions.setdefault(event.reason.removeprefix("needed_by:"), []).append(event.model_id) + else: rows.append(["-", label(event.model_id), "-", event.reason or ""]) - else: # oom - rows.append(["-", "-", "-", f"oom:{event.model_id}"]) + # evictions whose onload never happened (its forward raised) still deserve a row + for reason_target, evicted in pending_evictions.items(): + for model_id in evicted: + rows.append(["-", label(model_id), "-", f"needed_by:{reason_target}"]) table = format_table( ["#", "Onload", "Offloaded", "Available", "Reason"], [[str(index), *row] for index, row in enumerate(rows, start=1)], ) lines = [table[0], "-" * len(table[0]), *table[1:]] - if self.dropped: - lines.insert(2, f"... {self.dropped} earlier events dropped from the log ...") summary = self.summary() lines += [ @@ -224,10 +218,8 @@ def pre_forward(self, module, *args, **kwargs): available_memory = ( available_device_memory(self.execution_device) if hasattr(device_module, "mem_get_info") else None ) - resident_before, hooks_to_offload = (), [] if self.other_hooks is not None: hooks_to_offload = [hook for hook in self.other_hooks if hook.model.device == self.execution_device] - resident_before = tuple(hook.model_id for hook in hooks_to_offload) # offload all other hooks start_time = time.perf_counter() if self.offload_strategy is not None: @@ -255,8 +247,6 @@ def pre_forward(self, module, *args, **kwargs): model_id=self.model_id, model_size=module.get_memory_footprint(), available_memory=available_memory, - resident_before=resident_before, - offloaded=tuple(hook.model_id for hook in hooks_to_offload), ) ) return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device) @@ -297,7 +287,6 @@ def run_with_oom_retry(*args, **kwargs): except torch.OutOfMemoryError as e: if not self.retry_on_oom: raise - self.record.add(OffloadEvent(action="oom", model_id=self.model_id)) resident = sorted( (hook for hook in self.resident_other_hooks() if id(hook) not in offloaded), key=lambda hook: hook.model.get_memory_footprint(), diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index 308f78a0ffa6..4b22a7ba328c 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -401,8 +401,10 @@ def test_record_captures_what_the_strategy_saw(self): # The reading behind each decision is recorded, so a run can be explained after the fact. assert onloads[0].available_memory == 70 * UNIT assert onloads[1].available_memory == 4 * UNIT - assert onloads[1].resident_before == (cm.model_hooks[0].model_id,) - assert onloads[1].offloaded == (cm.model_hooks[0].model_id,) + # the eviction is its own event, attributed to the onload that caused it + eviction = next(event for event in cm.offload_record.events if event.action == "offload") + assert eviction.model_id == cm.model_hooks[0].model_id + assert eviction.reason == f"needed_by:{cm.model_hooks[1].model_id}" assert cm.offload_record.summary()["peak_co_residency"] == 1 # the printed table shows one row per decision: m2's row carries the m1 eviction it caused m2_row = next(line for line in repr(cm.offload_record).splitlines() if line.startswith("2 ")) @@ -549,16 +551,14 @@ def test_record_captures_oom_and_the_escalation(self): model(torch.zeros(2, 4)) - actions = [(event.action, event.model_id.rsplit("_", 1)[0]) for event in manager.offload_record.events] - assert ("oom", "model") in actions - assert ("offload", "smallest") in actions summary = manager.offload_record.summary() assert summary["oom_retries"] == 1 assert summary["offloads"] == 1 assert summary["bytes_offloaded"] == UNIT - # the eviction is attributed to the model that ran out of memory + # the OOM shows up as the eviction it caused, attributed to the model that ran out of memory offload_event = next(event for event in manager.offload_record.events if event.action == "offload") - assert offload_event.reason.startswith("oom_retry:") + assert offload_event.model_id.rsplit("_", 1)[0] == "smallest" + assert offload_event.reason == f"oom_retry:{manager.model_hooks[0].model_id}" def test_record_tracks_peak_co_residency(self): model_a = DummyModel(UNIT) @@ -583,12 +583,9 @@ def test_record_is_bounded_and_clearable(self): model(torch.zeros(2, 4)) assert len(manager.offload_record.events) == 2 - assert manager.offload_record.dropped == 2 - assert "earlier events dropped" in repr(manager.offload_record) manager.offload_record.clear() assert len(manager.offload_record) == 0 - assert manager.offload_record.dropped == 0 assert "nothing recorded yet" in repr(manager.offload_record) def test_record_survives_disable(self): @@ -652,10 +649,8 @@ def _is_resident(model): def _run_offloaded(self, free_bytes): """ Run the pipeline with auto offload on and `free_bytes` of *simulated* device - memory, recording every offload decision the strategy makes. - - Each decision is an `"onload"` event in `cm.offload_record`, holding the incoming - model, what was resident just before, and what the strategy evicted to make room. + memory, recording every move in `cm.offload_record`: an `"onload"` event per model + that ran, and an `"offload"` event (`reason="needed_by:"`) per eviction. """ cm = ComponentsManager() pipe = self.get_pipeline(components_manager=cm) @@ -663,22 +658,21 @@ def _run_offloaded(self, free_bytes): with _patch_free_memory(free_bytes): output = pipe(**self.get_dummy_inputs(), output=self.output_name) - onloads = [event for event in cm.offload_record.events if event.action == "onload"] - return cm, onloads, output + return cm, list(cm.offload_record.events), output @require_accelerate @require_accelerator def test_auto_cpu_offload_serializes_models_under_memory_pressure(self): # Zero simulated free memory: every model that runs must first evict whatever is # currently resident (comfy-style serialized execution). - cm, onloads, _ = self._run_offloaded(free_bytes=0) + cm, events, _ = self._run_offloaded(free_bytes=0) try: - distinct_models = {event.model_id for event in onloads} + distinct_models = {event.model_id for event in events if event.action == "onload"} if len(distinct_models) < 2: pytest.skip("pipeline has fewer than two offloadable model components") # Offloading actually fired (at least one eviction happened). - assert any(event.offloaded for event in onloads), "expected at least one eviction" + assert any(event.action == "offload" for event in events), "expected at least one eviction" # Sequencing: models run one at a time, never two co-resident on the device. peak = cm.offload_record.summary()["peak_co_residency"] @@ -698,14 +692,14 @@ def test_auto_cpu_offload_serializes_models_under_memory_pressure(self): def test_auto_cpu_offload_keeps_models_resident_without_memory_pressure(self): # Negative case: with ample simulated memory the strategy is still consulted on # every load, but it must never decide to evict anything. - cm, onloads, _ = self._run_offloaded(free_bytes=_AMPLE_FREE_BYTES) + cm, events, _ = self._run_offloaded(free_bytes=_AMPLE_FREE_BYTES) try: - distinct_models = {event.model_id for event in onloads} + distinct_models = {event.model_id for event in events if event.action == "onload"} if len(distinct_models) < 2: pytest.skip("pipeline has fewer than two offloadable model components") # Nothing was ever offloaded... - assert all(event.offloaded == () for event in onloads), "no model should be evicted" + assert all(event.action == "onload" for event in events), "no model should be evicted" # ...and models accumulate on the device instead of being serialized. peak = cm.offload_record.summary()["peak_co_residency"] From d3c1e9d7968a7f654b6388aa3b6ee3600e75a605 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 29 Jul 2026 08:52:56 +0000 Subject: [PATCH 08/14] Drop OffloadRecord.summary(): the table is the record 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 --- .../modular_diffusers/components_manager.md | 3 +- .../modular_pipelines/components_manager.py | 42 +++---------------- .../test_components_manager.py | 40 +++++++++++------- 3 files changed, 30 insertions(+), 55 deletions(-) diff --git a/docs/source/en/modular_diffusers/components_manager.md b/docs/source/en/modular_diffusers/components_manager.md index faa5440a46b8..bf0b4ad3997a 100644 --- a/docs/source/en/modular_diffusers/components_manager.md +++ b/docs/source/en/modular_diffusers/components_manager.md @@ -116,10 +116,9 @@ On a 20GB card, the Z-Image pipeline from the example above records this: 2 | transformer_140458257515616 (11.46 GB) | text_encoder_140458257514752 (7.49 GB) | 10.38 GB | 3 | vae_140458257515376 (159.87 MB) | - | 6.36 GB | -------------------------------------------------------------------------------------------------------- -3 onloads (19.11 GB) / 1 offloads (7.49 GB), 0 OOM retries, peak co-residency 2 ``` -Each row is one decision: the model that loaded, what was evicted to make room for it, and the free device memory (`Available`, read just before each load) the decision 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 evicted first. The VAE then fit into the 6.36 GB left next to the transformer, so it loaded without evicting anything (peak co-residency 2). The `Reason` column only speaks on moves an onload did not cause — an OOM retry (`oom_retry:`), disabling offloading — which appear as their own rows. To watch the moves live as they happen instead, enable info logging with `diffusers.logging.set_verbosity_info()`. +Each row is one decision: the model that loaded, what was evicted to make room for it, and the free device memory (`Available`, read just before each load) the decision 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 evicted first. The VAE then fit into the 6.36 GB left next to the transformer, so it loaded without evicting anything. The `Reason` column only speaks on moves an onload did not cause — an OOM retry (`oom_retry:`), disabling offloading — which appear as their own rows. To watch the moves live as they happen instead, enable info logging with `diffusers.logging.set_verbosity_info()`. A model appearing repeatedly in this table is thrashing — it is being evicted 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). diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 6153e50ee61d..b02dc6e6fdab 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -83,8 +83,7 @@ class OffloadRecord: What the offloader did, in order. Every move appends to this record, so after a run it holds the full sequence, what each move cost, and where the - forward passes ran out of memory. Printing it shows that sequence as a table followed by a summary; - [`~OffloadRecord.summary`] returns the same numbers as a dict. + forward passes ran out of memory. Printing it shows that sequence as a table, one row per decision. """ def __init__(self, maxlen: int = MAX_RECORDED_OFFLOAD_EVENTS): @@ -96,28 +95,6 @@ def add(self, event: OffloadEvent): def clear(self): self.events.clear() - def summary(self) -> dict[str, Any]: - onloads = [event for event in self.events if event.action == "onload"] - offloads = [event for event in self.events if event.action == "offload"] - resident, peak_co_residency = set(), 0 - for event in self.events: - if event.action == "onload": - resident.add(event.model_id) - peak_co_residency = max(peak_co_residency, len(resident)) - else: - resident.discard(event.model_id) - return { - "onloads": len(onloads), - "offloads": len(offloads), - "oom_retries": sum(event.reason.startswith("oom_retry:") for event in offloads if event.reason), - "peak_co_residency": peak_co_residency, - "bytes_onloaded": sum(event.model_size or 0 for event in onloads), - "bytes_offloaded": sum(event.model_size or 0 for event in offloads), - } - - def __len__(self): - return len(self.events) - def __repr__(self): if not self.events: return "Offload record: nothing recorded yet" @@ -151,16 +128,7 @@ def label(model_id): ["#", "Onload", "Offloaded", "Available", "Reason"], [[str(index), *row] for index, row in enumerate(rows, start=1)], ) - lines = [table[0], "-" * len(table[0]), *table[1:]] - - summary = self.summary() - lines += [ - "-" * len(table[0]), - f"{summary['onloads']} onloads ({format_size(summary['bytes_onloaded'])}) / " - f"{summary['offloads']} offloads ({format_size(summary['bytes_offloaded'])}), " - f"{summary['oom_retries']} OOM retries, peak co-residency {summary['peak_co_residency']}", - ] - return "\n".join(lines) + return "\n".join([table[0], "-" * len(table[0]), *table[1:]]) class CustomOffloadHook(ModelHook): @@ -955,9 +923,9 @@ def enable_auto_cpu_offload( def offload_record(self) -> OffloadRecord: """ What the offloader has done so far: every model moved onto or off the device, in order, and where a forward - pass ran out of memory. Print it to see the sequence, or read [`~OffloadRecord.summary`] for the numbers - behind it. Kept across [`~ComponentsManager.disable_auto_cpu_offload`] (it is the post-mortem of the run), - and cleared when offloading is enabled again. + pass ran out of memory. Print it to see the sequence, or read `events` for the data behind it. Kept across + [`~ComponentsManager.disable_auto_cpu_offload`] (it is the post-mortem of the run), and cleared when + offloading is enabled again. """ return self._offload_record diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index 4b22a7ba328c..fa60c7464612 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -50,6 +50,18 @@ _AMPLE_FREE_BYTES = 1024**4 +def _peak_co_residency(events): + """Largest number of models simultaneously on the device, replayed from the recorded moves.""" + resident, peak = set(), 0 + for event in events: + if event.action == "onload": + resident.add(event.model_id) + peak = max(peak, len(resident)) + else: + resident.discard(event.model_id) + return peak + + class DummyModel(ModelMixin): def __init__(self, footprint_bytes: int = UNIT): super().__init__() @@ -405,7 +417,7 @@ def test_record_captures_what_the_strategy_saw(self): eviction = next(event for event in cm.offload_record.events if event.action == "offload") assert eviction.model_id == cm.model_hooks[0].model_id assert eviction.reason == f"needed_by:{cm.model_hooks[1].model_id}" - assert cm.offload_record.summary()["peak_co_residency"] == 1 + assert _peak_co_residency(cm.offload_record.events) == 1 # the printed table shows one row per decision: m2's row carries the m1 eviction it caused m2_row = next(line for line in repr(cm.offload_record).splitlines() if line.startswith("2 ")) assert cm.model_hooks[1].model_id in m2_row and cm.model_hooks[0].model_id in m2_row @@ -539,10 +551,7 @@ def test_record_captures_the_onload_sequence(self): onloads = [event for event in manager.offload_record.events if event.action == "onload"] assert [event.model_id.rsplit("_", 1)[0] for event in onloads] == ["a", "b"] assert [event.model_size for event in onloads] == [UNIT, 4 * UNIT] - summary = manager.offload_record.summary() - assert summary["onloads"] == 2 - assert summary["bytes_onloaded"] == 5 * UNIT - assert summary["oom_retries"] == 0 + assert [event.action for event in manager.offload_record.events] == ["onload", "onload"] def test_record_captures_oom_and_the_escalation(self): model = OOMOnceModel() @@ -551,12 +560,11 @@ def test_record_captures_oom_and_the_escalation(self): model(torch.zeros(2, 4)) - summary = manager.offload_record.summary() - assert summary["oom_retries"] == 1 - assert summary["offloads"] == 1 - assert summary["bytes_offloaded"] == UNIT # the OOM shows up as the eviction it caused, attributed to the model that ran out of memory - offload_event = next(event for event in manager.offload_record.events if event.action == "offload") + offloads = [event for event in manager.offload_record.events if event.action == "offload"] + assert len(offloads) == 1 + offload_event = offloads[0] + assert offload_event.model_size == UNIT assert offload_event.model_id.rsplit("_", 1)[0] == "smallest" assert offload_event.reason == f"oom_retry:{manager.model_hooks[0].model_id}" @@ -567,12 +575,12 @@ def test_record_tracks_peak_co_residency(self): model_a(torch.zeros(2, 4)) model_b(torch.zeros(2, 4)) - assert manager.offload_record.summary()["peak_co_residency"] == 2 + assert _peak_co_residency(manager.offload_record.events) == 2 # an offload brings residency back down, so a later onload does not raise the peak manager.model_hooks[0].offload(reason="test") model_a(torch.zeros(2, 4)) - assert manager.offload_record.summary()["peak_co_residency"] == 2 + assert _peak_co_residency(manager.offload_record.events) == 2 def test_record_is_bounded_and_clearable(self): model = DummyModel() @@ -585,7 +593,7 @@ def test_record_is_bounded_and_clearable(self): assert len(manager.offload_record.events) == 2 manager.offload_record.clear() - assert len(manager.offload_record) == 0 + assert len(manager.offload_record.events) == 0 assert "nothing recorded yet" in repr(manager.offload_record) def test_record_survives_disable(self): @@ -595,7 +603,7 @@ def test_record_survives_disable(self): manager.disable_auto_cpu_offload() # the record is the post-mortem of the run that just ended - assert manager.offload_record.summary()["onloads"] == 1 + assert sum(event.action == "onload" for event in manager.offload_record.events) == 1 def test_add_does_not_rebuild_existing_hooks(self): model_a = DummyModel() @@ -675,7 +683,7 @@ def test_auto_cpu_offload_serializes_models_under_memory_pressure(self): assert any(event.action == "offload" for event in events), "expected at least one eviction" # Sequencing: models run one at a time, never two co-resident on the device. - peak = cm.offload_record.summary()["peak_co_residency"] + peak = _peak_co_residency(events) assert peak == 1, f"expected serialized execution under pressure, saw {peak} models co-resident" # Device placement after the run: at most the last-run model stays on the @@ -702,7 +710,7 @@ def test_auto_cpu_offload_keeps_models_resident_without_memory_pressure(self): assert all(event.action == "onload" for event in events), "no model should be evicted" # ...and models accumulate on the device instead of being serialized. - peak = cm.offload_record.summary()["peak_co_residency"] + peak = _peak_co_residency(events) assert peak >= 2, f"expected models to co-reside without pressure, saw peak {peak}" models = self._managed_models(cm) From b3604203159b5b9030707fd20fd8ab5b92e18434 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 29 Jul 2026 10:15:21 +0000 Subject: [PATCH 09/14] Address review: discover entry points via an apply_forward_hook marker - 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: -> release_memory_for:. - 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 --- .../modular_pipelines/components_manager.py | 145 +++++++++--------- src/diffusers/utils/accelerate_utils.py | 1 + .../test_components_manager.py | 36 +++-- 3 files changed, 96 insertions(+), 86 deletions(-) diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index b02dc6e6fdab..c3895aab5bf0 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -41,10 +41,6 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -# Default number of events an `OffloadRecord` keeps. A generation step produces at most a handful, so this -# holds a long run while staying bounded for a long-lived manager (e.g. in a server). -MAX_RECORDED_OFFLOAD_EVENTS = 10000 - def available_device_memory(execution_device: torch.device) -> int: """ @@ -71,10 +67,10 @@ class OffloadEvent: # the model this event moved model_id: str model_size: int | None = None - # why the model moved: which model needed the room ("needed_by:"), an OOM retry + # why the model moved: which model needed the room ("release_memory_for:"), an OOM retry # ("oom_retry:"), attaching a new component, disabling offloading, ... reason: str | None = None - # free device memory read just before the onload's eviction decision (`None` if the device cannot report it) + # free device memory read just before the onload's eviction decision available_memory: int | None = None @@ -86,7 +82,11 @@ class OffloadRecord: forward passes ran out of memory. Printing it shows that sequence as a table, one row per decision. """ - def __init__(self, maxlen: int = MAX_RECORDED_OFFLOAD_EVENTS): + # Default number of events kept. A generation step produces at most a handful, so this holds a long run + # while staying bounded for a long-lived manager (e.g. in a server). + MAX_EVENTS = 10000 + + def __init__(self, maxlen: int = MAX_EVENTS): self.events: deque[OffloadEvent] = deque(maxlen=maxlen) def add(self, event: OffloadEvent): @@ -100,8 +100,8 @@ def __repr__(self): return "Offload record: nothing recorded yet" # One row per decision: an onload and the evictions it caused share a row, correlated through the - # eviction's "needed_by:" reason (evictions precede their onload in the event sequence). - # Offloads with other causes (OOM retry, offloading disabled) get their own rows. + # eviction's "release_memory_for:" reason (evictions precede their onload in the event + # sequence). Offloads with other causes (OOM retry, offloading disabled) get their own rows. sizes = {event.model_id: event.model_size for event in self.events if event.model_size is not None} def label(model_id): @@ -115,14 +115,16 @@ def label(model_id): rows.append( [label(event.model_id), offloaded, format_size(event.available_memory), event.reason or ""] ) - elif event.reason is not None and event.reason.startswith("needed_by:"): - pending_evictions.setdefault(event.reason.removeprefix("needed_by:"), []).append(event.model_id) + elif event.reason is not None and event.reason.startswith("release_memory_for:"): + pending_evictions.setdefault(event.reason.removeprefix("release_memory_for:"), []).append( + event.model_id + ) else: rows.append(["-", label(event.model_id), "-", event.reason or ""]) # evictions whose onload never happened (its forward raised) still deserve a row for reason_target, evicted in pending_evictions.items(): for model_id in evicted: - rows.append(["-", label(model_id), "-", f"needed_by:{reason_target}"]) + rows.append(["-", label(model_id), "-", f"release_memory_for:{reason_target}"]) table = format_table( ["#", "Onload", "Offloaded", "Available", "Reason"], @@ -182,10 +184,7 @@ def init_hook(self, module): def pre_forward(self, module, *args, **kwargs): if module.device != self.execution_device: # read before any eviction: this is the number the eviction decision is based on - device_module = getattr(torch, self.execution_device.type, torch.cuda) - available_memory = ( - available_device_memory(self.execution_device) if hasattr(device_module, "mem_get_info") else None - ) + available_memory = available_device_memory(self.execution_device) if self.other_hooks is not None: hooks_to_offload = [hook for hook in self.other_hooks if hook.model.device == self.execution_device] # offload all other hooks @@ -204,7 +203,7 @@ def pre_forward(self, module, *args, **kwargs): logger.info( f"moving {self.model_id} to {self.execution_device}, offloading {hook.model_id} to cpu" ) - hook.offload(reason=f"needed_by:{self.model_id}") + hook.offload(reason=f"release_memory_for:{self.model_id}") if hooks_to_offload: clear_device_cache() @@ -219,70 +218,68 @@ def pre_forward(self, module, *args, **kwargs): ) return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device) - def resident_other_hooks(self): - """ - The other managed models that are currently on the execution device. - """ - return [ - other - for other in (self.other_hooks or []) - # compare with a normalized index: model.device reports no index for the default device - if other.model.device.type == self.execution_device.type - and (other.model.device.index or 0) == (self.execution_device.index or 0) - ] + def _with_oom_retry(self, entry_point): + # On an OOM, free the smallest resident model and retry, escalating until it fits. The memory readings + # cannot guide this: the failed forward has already unwound, so its activations are freed and the device + # looks free again. Inference only: the retried forward re-runs cleanly from its original inputs. + @functools.wraps(entry_point) + def run_with_oom_retry(*args, **kwargs): + # each pass offloads one more model, so this terminates once they all have been. Tracked explicitly + # rather than by device, so that a model that did not actually move cannot be picked twice. + offloaded = set() + while True: + try: + return entry_point(*args, **kwargs) + except torch.OutOfMemoryError as e: + resident = sorted( + ( + hook + for hook in (self.other_hooks or []) + # compare with a normalized index: model.device reports no index for the + # default device + if hook.model.device.type == self.execution_device.type + and (hook.model.device.index or 0) == (self.execution_device.index or 0) + and id(hook) not in offloaded + ), + key=lambda hook: hook.model.get_memory_footprint(), + ) + if not resident: + raise torch.OutOfMemoryError( + f"{self.model_id} ran out of device memory ({e}) with every other managed model " + "already offloaded, so it does not fit on its own. Consider group offloading " + "(`ModelMixin.enable_group_offload`), which offloads a single model in groups of " + "internal layers." + ) from e + smallest = resident[0] + logger.warning( + f"{self.model_id} ran out of device memory ({e}); offloading {smallest.model_id} and " + "retrying. If this happens repeatedly, set a larger `memory_reserve` in " + "`enable_auto_cpu_offload`." + ) + smallest.offload(reason=f"oom_retry:{self.model_id}") + offloaded.add(id(smallest)) + clear_device_cache() + + return run_with_oom_retry # YiYi TODO: diffusers' own `ModelHook` (`diffusers.hooks.hooks`) supports a `new_forward` around-hook, which # would replace this manual wrapping - we may migrate this class to it in the future. def wrap_forward(self, module): - # `pre_forward`/`post_forward` cannot see an exception the forward pass raises, so wrap the (already hooked) - # entry points here. `forward` is how most models run; autoencoders enter through `encode`/`decode` instead - # (decorated with `apply_forward_hook`, which fires `pre_forward` but routes around `forward`). On an OOM, - # free the smallest resident model and retry, escalating until it fits. The memory readings cannot guide - # this: the failed forward has already unwound, so its activations are freed and the device looks free again. - # Inference only: the retried forward re-runs cleanly from its original inputs. + # `pre_forward`/`post_forward` cannot see an exception the forward pass raises, so wrap the (already + # hooked) entry points here. `forward` is how most models run; autoencoders enter through the methods + # `apply_forward_hook` marks as device entry points (their encode/decode), which fire `pre_forward` but + # route around `forward`. if not self.retry_on_oom: return - def with_oom_retry(entry_point): - @functools.wraps(entry_point) - def run_with_oom_retry(*args, **kwargs): - # each pass offloads one more model, so this terminates once they all have been. Tracked explicitly - # rather than by device, so that a model that did not actually move cannot be picked twice. - offloaded = set() - while True: - try: - return entry_point(*args, **kwargs) - except torch.OutOfMemoryError as e: - if not self.retry_on_oom: - raise - resident = sorted( - (hook for hook in self.resident_other_hooks() if id(hook) not in offloaded), - key=lambda hook: hook.model.get_memory_footprint(), - ) - if not resident: - raise torch.OutOfMemoryError( - f"{self.model_id} ran out of device memory ({e}) with every other managed model " - "already offloaded, so it does not fit on its own. Consider group offloading " - "(`ModelMixin.enable_group_offload`), which offloads a single model in groups of " - "internal layers." - ) from e - smallest = resident[0] - logger.warning( - f"{self.model_id} ran out of device memory ({e}); offloading {smallest.model_id} and " - "retrying. If this happens repeatedly, set a larger `memory_reserve` in " - "`enable_auto_cpu_offload`." - ) - smallest.offload(reason=f"oom_retry:{self.model_id}") - offloaded.add(id(smallest)) - clear_device_cache() - - return run_with_oom_retry - - for name in ("forward", "encode", "decode"): - entry_point = getattr(module, name, None) - if callable(entry_point): - self._wrapped_methods[name] = entry_point - setattr(module, name, with_oom_retry(entry_point)) + entry_point_names = {"forward"} + for klass in type(module).__mro__: + for name, attr in vars(klass).items(): + if getattr(attr, "_is_forward_entry_point", False): + entry_point_names.add(name) + for name in sorted(entry_point_names): + self._wrapped_methods[name] = getattr(module, name) + setattr(module, name, self._with_oom_retry(self._wrapped_methods[name])) def unwrap_forward(self, module): for name, entry_point in self._wrapped_methods.items(): diff --git a/src/diffusers/utils/accelerate_utils.py b/src/diffusers/utils/accelerate_utils.py index ae6a8ca747ac..83182a04cc0d 100644 --- a/src/diffusers/utils/accelerate_utils.py +++ b/src/diffusers/utils/accelerate_utils.py @@ -45,4 +45,5 @@ def wrapper(self, *args, **kwargs): self._hf_hook.pre_forward(self) return method(self, *args, **kwargs) + wrapper._is_forward_entry_point = True # lets hook machinery statically discover decorated entry points return wrapper diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index fa60c7464612..ad77f8dadd68 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -86,9 +86,17 @@ def forward(self, x): return super().forward(x) -class OOMOnceDecodeModel(DummyModel): +class DecodeModel(DummyModel): """Autoencoder-style model: runs through a `decode` entry point (via `apply_forward_hook`, like the VAEs) - instead of `forward`, and raises a fake device OOM on its first decode.""" + instead of `forward`.""" + + @apply_forward_hook + def decode(self, x): + return x + self.weight.sum() + + +class OOMOnceDecodeModel(DecodeModel): + """`DecodeModel` that raises a fake device OOM on its first decode.""" @apply_forward_hook def decode(self, x): @@ -416,7 +424,7 @@ def test_record_captures_what_the_strategy_saw(self): # the eviction is its own event, attributed to the onload that caused it eviction = next(event for event in cm.offload_record.events if event.action == "offload") assert eviction.model_id == cm.model_hooks[0].model_id - assert eviction.reason == f"needed_by:{cm.model_hooks[1].model_id}" + assert eviction.reason == f"release_memory_for:{cm.model_hooks[1].model_id}" assert _peak_co_residency(cm.offload_record.events) == 1 # the printed table shows one row per decision: m2's row carries the m1 eviction it caused m2_row = next(line for line in repr(cm.offload_record).splitlines() if line.startswith("2 ")) @@ -453,17 +461,21 @@ class TestComponentsManager(ComponentsManagerTesterMixin): class TestOffloadHookBehavior: """CPU-level tests of the hook machinery: OOM retry and non-disruptive add/remove.""" + @pytest.fixture(autouse=True) + def _fake_cpu_mem_get_info(self): + # These tests exercise the hook machinery only, on a cpu execution device — but both enabling + # offloading and every onload read `mem_get_info`, which cpu doesn't implement, so provide an + # ample fake one for the whole test. + with mock.patch.object( + torch.cpu, "mem_get_info", create=True, return_value=(_AMPLE_FREE_BYTES, _AMPLE_FREE_BYTES) + ): + yield + def _manager(self, components, **offload_kwargs): manager = ComponentsManager() for name, component in components.items(): manager.add(name, component) - # These tests exercise the hook machinery only: with a cpu execution device the models - # never move, so the strategy is never consulted — but enabling requires the device to - # implement `mem_get_info`, which cpu doesn't, so provide an ample fake one. - with mock.patch.object( - torch.cpu, "mem_get_info", create=True, return_value=(_AMPLE_FREE_BYTES, _AMPLE_FREE_BYTES) - ): - manager.enable_auto_cpu_offload(device="cpu", **offload_kwargs) + manager.enable_auto_cpu_offload(device="cpu", **offload_kwargs) return manager @staticmethod @@ -530,7 +542,7 @@ def test_retry_on_oom_false_leaves_the_forward_alone(self): def test_oom_retry_covers_apply_forward_hook_entry_points(self): # Autoencoders run through `encode`/`decode` (via `apply_forward_hook`), not `forward` - an OOM - # there must be retried just like one raised inside `forward`. + # there must be retried just like one raised inside `forward`, first call included. model = OOMOnceDecodeModel() smallest = DummyModel(UNIT) manager = self._manager({"model": model, "smallest": smallest}) @@ -658,7 +670,7 @@ def _run_offloaded(self, free_bytes): """ Run the pipeline with auto offload on and `free_bytes` of *simulated* device memory, recording every move in `cm.offload_record`: an `"onload"` event per model - that ran, and an `"offload"` event (`reason="needed_by:"`) per eviction. + that ran, and an `"offload"` event (`reason="release_memory_for:"`) per eviction. """ cm = ComponentsManager() pipe = self.get_pipeline(components_manager=cm) From b313f79eb303bff8e0be83bdb50205725e04e4b0 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 29 Jul 2026 10:20:39 +0000 Subject: [PATCH 10/14] Inline the OOM-retry helper into wrap_forward, its only caller Co-Authored-By: Claude Fable 5 --- .../modular_pipelines/components_manager.py | 92 +++++++++---------- 1 file changed, 46 insertions(+), 46 deletions(-) diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index c3895aab5bf0..941e3abd0373 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -218,60 +218,60 @@ def pre_forward(self, module, *args, **kwargs): ) return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device) - def _with_oom_retry(self, entry_point): - # On an OOM, free the smallest resident model and retry, escalating until it fits. The memory readings - # cannot guide this: the failed forward has already unwound, so its activations are freed and the device - # looks free again. Inference only: the retried forward re-runs cleanly from its original inputs. - @functools.wraps(entry_point) - def run_with_oom_retry(*args, **kwargs): - # each pass offloads one more model, so this terminates once they all have been. Tracked explicitly - # rather than by device, so that a model that did not actually move cannot be picked twice. - offloaded = set() - while True: - try: - return entry_point(*args, **kwargs) - except torch.OutOfMemoryError as e: - resident = sorted( - ( - hook - for hook in (self.other_hooks or []) - # compare with a normalized index: model.device reports no index for the - # default device - if hook.model.device.type == self.execution_device.type - and (hook.model.device.index or 0) == (self.execution_device.index or 0) - and id(hook) not in offloaded - ), - key=lambda hook: hook.model.get_memory_footprint(), - ) - if not resident: - raise torch.OutOfMemoryError( - f"{self.model_id} ran out of device memory ({e}) with every other managed model " - "already offloaded, so it does not fit on its own. Consider group offloading " - "(`ModelMixin.enable_group_offload`), which offloads a single model in groups of " - "internal layers." - ) from e - smallest = resident[0] - logger.warning( - f"{self.model_id} ran out of device memory ({e}); offloading {smallest.model_id} and " - "retrying. If this happens repeatedly, set a larger `memory_reserve` in " - "`enable_auto_cpu_offload`." - ) - smallest.offload(reason=f"oom_retry:{self.model_id}") - offloaded.add(id(smallest)) - clear_device_cache() - - return run_with_oom_retry - # YiYi TODO: diffusers' own `ModelHook` (`diffusers.hooks.hooks`) supports a `new_forward` around-hook, which # would replace this manual wrapping - we may migrate this class to it in the future. def wrap_forward(self, module): # `pre_forward`/`post_forward` cannot see an exception the forward pass raises, so wrap the (already # hooked) entry points here. `forward` is how most models run; autoencoders enter through the methods # `apply_forward_hook` marks as device entry points (their encode/decode), which fire `pre_forward` but - # route around `forward`. + # route around `forward`. On an OOM, free the smallest resident model and retry, escalating until it + # fits. The memory readings cannot guide this: the failed forward has already unwound, so its activations + # are freed and the device looks free again. Inference only: the retried forward re-runs cleanly from its + # original inputs. if not self.retry_on_oom: return + def with_oom_retry(entry_point): + @functools.wraps(entry_point) + def run_with_oom_retry(*args, **kwargs): + # each pass offloads one more model, so this terminates once they all have been. Tracked explicitly + # rather than by device, so that a model that did not actually move cannot be picked twice. + offloaded = set() + while True: + try: + return entry_point(*args, **kwargs) + except torch.OutOfMemoryError as e: + resident = sorted( + ( + hook + for hook in (self.other_hooks or []) + # compare with a normalized index: model.device reports no index for the + # default device + if hook.model.device.type == self.execution_device.type + and (hook.model.device.index or 0) == (self.execution_device.index or 0) + and id(hook) not in offloaded + ), + key=lambda hook: hook.model.get_memory_footprint(), + ) + if not resident: + raise torch.OutOfMemoryError( + f"{self.model_id} ran out of device memory ({e}) with every other managed model " + "already offloaded, so it does not fit on its own. Consider group offloading " + "(`ModelMixin.enable_group_offload`), which offloads a single model in groups of " + "internal layers." + ) from e + smallest = resident[0] + logger.warning( + f"{self.model_id} ran out of device memory ({e}); offloading {smallest.model_id} and " + "retrying. If this happens repeatedly, set a larger `memory_reserve` in " + "`enable_auto_cpu_offload`." + ) + smallest.offload(reason=f"oom_retry:{self.model_id}") + offloaded.add(id(smallest)) + clear_device_cache() + + return run_with_oom_retry + entry_point_names = {"forward"} for klass in type(module).__mro__: for name, attr in vars(klass).items(): @@ -279,7 +279,7 @@ def wrap_forward(self, module): entry_point_names.add(name) for name in sorted(entry_point_names): self._wrapped_methods[name] = getattr(module, name) - setattr(module, name, self._with_oom_retry(self._wrapped_methods[name])) + setattr(module, name, with_oom_retry(self._wrapped_methods[name])) def unwrap_forward(self, module): for name, entry_point in self._wrapped_methods.items(): From 0e512ddf6b39e28b1a8631cb3f158fdeaac3aaf1 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 29 Jul 2026 10:36:53 +0000 Subject: [PATCH 11/14] Shrink ComponentsManager to its used surface 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 --- .../modular_diffusers/components_manager.md | 22 +- .../modular_pipelines/components_manager.py | 286 ++---------------- 2 files changed, 31 insertions(+), 277 deletions(-) diff --git a/docs/source/zh/modular_diffusers/components_manager.md b/docs/source/zh/modular_diffusers/components_manager.md index 39fef0651dd8..70eee420d709 100644 --- a/docs/source/zh/modular_diffusers/components_manager.md +++ b/docs/source/zh/modular_diffusers/components_manager.md @@ -55,13 +55,13 @@ pipe.load_components() pipe2 = ModularPipeline.from_pretrained("YiYiXu/modular-demo-auto", components_manager=comp, collection="test2") ``` -使用 [`~ModularPipeline.null_component_names`] 属性来识别需要加载的任何组件，使用 [`~ComponentsManager.get_components_by_names`] 检索它们，然后调用 [`~ModularPipeline.update_components`] 来添加缺失的组件。 +使用 [`~ModularPipeline.null_component_names`] 属性来识别需要加载的任何组件，使用 [`~ComponentsManager.get_one`] 检索它们，然后调用 [`~ModularPipeline.update_components`] 来添加缺失的组件。 ```py pipe2.null_component_names ['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'image_encoder', 'unet', 'vae', 'scheduler', 'controlnet'] -comp_dict = comp.get_components_by_names(names=pipe2.null_component_names) +comp_dict = {name: comp.get_one(name=name) for name in pipe2.null_component_names} pipe2.update_components(**comp_dict) ``` @@ -87,14 +87,7 @@ comp.remove("text_encoder_139917733042864") ### get_one -[`~ComponentsManager.get_one`] 方法返回单个组件，并支持对 `name` 参数进行模式匹配。如果多个组件匹配，[`~ComponentsManager.get_one`] 会返回错误。 - -| 模式 | 示例 | 描述 | -|-------------|----------------------------------|-------------------------------------------| -| exact | `comp.get_one(name="unet")` | 精确名称匹配 | -| wildcard | `comp.get_one(name="unet*")` | 名称以 "unet" 开头 | -| exclusion | `comp.get_one(name="!unet")` | 排除名为 "unet" 的组件 | -| or | `comp.get_one(name="unet|vae")` | 名称为 "unet" 或 "vae" | +[`~ComponentsManager.get_one`] 方法按 `name` 精确匹配并返回单个组件。如果多个组件匹配，[`~ComponentsManager.get_one`] 会返回错误。 [`~ComponentsManager.get_one`] 还通过 `collection` 参数或 `load_id` 参数过滤组件。 @@ -102,15 +95,6 @@ comp.remove("text_encoder_139917733042864") comp.get_one(name="unet", collection="sdxl") ``` -### get_components_by_names - -[`~ComponentsManager.get_components_by_names`] 方法接受一个名称列表，并返回一个将名称映射到组件的字典。这在 [`ModularPipeline`] 中特别有用，因为它们提供了所需组件名称的列表，并且返回的字典可以直接传递给 [`~ModularPipeline.update_components`]。 - -```py -component_dict = comp.get_components_by_names(names=["text_encoder", "unet", "vae"]) -{"text_encoder": component1, "unet": component2, "vae": component3} -``` - ## 重复检测 建议使用 [`ComponentSpec`] 加载模型组件，以分配具有唯一 id 的组件，该 id 编码了它们的加载参数。这允许 [`ComponentsManager`] 自动检测并防止重复的模型实例，即使不同的对象代表相同的底层检查点。 diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 941e3abd0373..912a3851f3be 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -592,7 +592,22 @@ def add(self, name: str, component: Any, collection: str | None = None): logger.info(f"ComponentsManager: added component '{name}' as '{component_id}'") if self._auto_offload_enabled and is_new_component and isinstance(component, torch.nn.Module): - self._attach_offload_hook(component_id, component) + # attach an offload hook without disturbing the models already managed: the new component starts + # on CPU (like every model under auto offload), everything else stays where it is + hook = custom_offload_with_hook( + component_id, + component, + self._auto_offload_device, + offload_strategy=self._offload_strategy, + retry_on_oom=self._offload_retry_on_oom, + record=self._offload_record, + ) + for other_hook in self.model_hooks: + if other_hook.hook.execution_device == hook.hook.execution_device: + hook.add_other_hook(other_hook) + other_hook.add_other_hook(hook) + hook.offload(reason="component_added") + self.model_hooks.append(hook) return component_id @@ -634,7 +649,13 @@ def remove(self, component_id: str = None): if isinstance(component, torch.nn.Module): if self._auto_offload_enabled: - self._detach_offload_hook(component) + # detach only this component's offload hook, leaving all other managed models where they are + hook = next(user_hook for user_hook in self.model_hooks if user_hook.model is component) + hook.remove() + self.model_hooks.remove(hook) + for other_hook in self.model_hooks: + if other_hook.hook.other_hooks and hook in other_hook.hook.other_hooks: + other_hook.hook.other_hooks.remove(hook) component.to("cpu") del component import gc @@ -645,193 +666,6 @@ def remove(self, component_id: str = None): if torch.xpu.is_available(): torch.xpu.empty_cache() - # YiYi TODO: rename to search_components for now, may remove this method - def search_components( - self, - names: str | None = None, - collection: str | None = None, - load_id: str | None = None, - return_dict_with_names: bool = True, - ): - """ - Search components by name with simple pattern matching. Optionally filter by collection or load_id. - - Args: - names: Component name(s) or pattern(s) - Patterns: - - "unet" : match any component with base name "unet" (e.g., unet_123abc) - - "!unet" : everything except components with base name "unet" - - "unet*" : anything with base name starting with "unet" - - "!unet*" : anything with base name NOT starting with "unet" - - "*unet*" : anything with base name containing "unet" - - "!*unet*" : anything with base name NOT containing "unet" - - "refiner|vae|unet" : anything with base name exactly matching "refiner", "vae", or "unet" - - "!refiner|vae|unet" : anything with base name NOT exactly matching "refiner", "vae", or "unet" - - "unet*|vae*" : anything with base name starting with "unet" OR starting with "vae" - collection: Optional collection to filter by - load_id: Optional load_id to filter by - return_dict_with_names: - If True, returns a dictionary with component names as keys, throw an error if - multiple components with the same name are found If False, returns a dictionary - with component IDs as keys - - Returns: - Dictionary mapping component names to components if return_dict_with_names=True, or a dictionary mapping - component IDs to components if return_dict_with_names=False - """ - - # select components based on collection and load_id filters - selected_ids = self._lookup_ids(collection=collection, load_id=load_id) - components = {k: self.components[k] for k in selected_ids} - - def get_return_dict(components, return_dict_with_names): - """ - Create a dictionary mapping component names to components if return_dict_with_names=True, or a dictionary - mapping component IDs to components if return_dict_with_names=False, throw an error if duplicate component - names are found when return_dict_with_names=True - """ - if return_dict_with_names: - dict_to_return = {} - for comp_id, comp in components.items(): - comp_name = self._id_to_name(comp_id) - if comp_name in dict_to_return: - raise ValueError( - f"Duplicate component names found in the search results: {comp_name}, please set `return_dict_with_names=False` to return a dictionary with component IDs as keys" - ) - dict_to_return[comp_name] = comp - return dict_to_return - else: - return components - - # if no names are provided, return the filtered components as it is - if names is None: - return get_return_dict(components, return_dict_with_names) - - # if names is not a string, raise an error - elif not isinstance(names, str): - raise ValueError(f"Invalid type for `names: {type(names)}, only support string") - - # Create mapping from component_id to base_name for components to be used for pattern matching - base_names = {comp_id: self._id_to_name(comp_id) for comp_id in components.keys()} - - # Helper function to check if a component matches a pattern based on its base name - def matches_pattern(component_id, pattern, exact_match=False): - """ - Helper function to check if a component matches a pattern based on its base name. - - Args: - component_id: The component ID to check - pattern: The pattern to match against - exact_match: If True, only exact matches to base_name are considered - """ - base_name = base_names[component_id] - - # Exact match with base name - if exact_match: - return pattern == base_name - - # Prefix match (ends with *) - elif pattern.endswith("*"): - prefix = pattern[:-1] - return base_name.startswith(prefix) - - # Contains match (starts with *) - elif pattern.startswith("*"): - search = pattern[1:-1] if pattern.endswith("*") else pattern[1:] - return search in base_name - - # Exact match (no wildcards) - else: - return pattern == base_name - - # Check if this is a "not" pattern - is_not_pattern = names.startswith("!") - if is_not_pattern: - names = names[1:] # Remove the ! prefix - - # Handle OR patterns (containing |) - if "|" in names: - terms = names.split("|") - matches = {} - - for comp_id, comp in components.items(): - # For OR patterns with exact names (no wildcards), we do exact matching on base names - exact_match = all(not (term.startswith("*") or term.endswith("*")) for term in terms) - - # Check if any of the terms match this component - should_include = any(matches_pattern(comp_id, term, exact_match) for term in terms) - - # Flip the decision if this is a NOT pattern - if is_not_pattern: - should_include = not should_include - - if should_include: - matches[comp_id] = comp - - log_msg = "NOT " if is_not_pattern else "" - match_type = "exactly matching" if exact_match else "matching any of patterns" - logger.info(f"Getting components {log_msg}{match_type} {terms}: {list(matches.keys())}") - - # Try exact match with a base name - elif any(names == base_name for base_name in base_names.values()): - # Find all components with this base name - matches = { - comp_id: comp - for comp_id, comp in components.items() - if (base_names[comp_id] == names) != is_not_pattern - } - - if is_not_pattern: - logger.info(f"Getting all components except those with base name '{names}': {list(matches.keys())}") - else: - logger.info(f"Getting components with base name '{names}': {list(matches.keys())}") - - # Prefix match (ends with *) - elif names.endswith("*"): - prefix = names[:-1] - matches = { - comp_id: comp - for comp_id, comp in components.items() - if base_names[comp_id].startswith(prefix) != is_not_pattern - } - if is_not_pattern: - logger.info(f"Getting components NOT starting with '{prefix}': {list(matches.keys())}") - else: - logger.info(f"Getting components starting with '{prefix}': {list(matches.keys())}") - - # Contains match (starts with *) - elif names.startswith("*"): - search = names[1:-1] if names.endswith("*") else names[1:] - matches = { - comp_id: comp - for comp_id, comp in components.items() - if (search in base_names[comp_id]) != is_not_pattern - } - if is_not_pattern: - logger.info(f"Getting components NOT containing '{search}': {list(matches.keys())}") - else: - logger.info(f"Getting components containing '{search}': {list(matches.keys())}") - - # Substring match (no wildcards, but not an exact component name) - elif any(names in base_name for base_name in base_names.values()): - matches = { - comp_id: comp - for comp_id, comp in components.items() - if (names in base_names[comp_id]) != is_not_pattern - } - if is_not_pattern: - logger.info(f"Getting components NOT containing '{names}': {list(matches.keys())}") - else: - logger.info(f"Getting components containing '{names}': {list(matches.keys())}") - - else: - raise ValueError(f"Component or pattern '{names}' not found in ComponentsManager") - - if not matches: - raise ValueError(f"No components found matching pattern '{names}'") - - return get_return_dict(matches, return_dict_with_names) - def enable_auto_cpu_offload( self, device: str | int | torch.device = None, @@ -926,35 +760,6 @@ def offload_record(self) -> OffloadRecord: """ return self._offload_record - def _attach_offload_hook(self, component_id: str, component: torch.nn.Module): - """ - Attach an offload hook to a newly added component without disturbing the models already managed: the new - component starts on CPU (like every model under auto offload), everything else stays where it is. - """ - hook = custom_offload_with_hook( - component_id, - component, - self._auto_offload_device, - offload_strategy=self._offload_strategy, - retry_on_oom=self._offload_retry_on_oom, - record=self._offload_record, - ) - for other_hook in self.model_hooks: - if other_hook.hook.execution_device == hook.hook.execution_device: - hook.add_other_hook(other_hook) - other_hook.add_other_hook(hook) - hook.offload(reason="component_added") - self.model_hooks.append(hook) - - def _detach_offload_hook(self, component: torch.nn.Module): - """Detach a removed component's offload hook, leaving all other managed models where they are.""" - hook = next(user_hook for user_hook in self.model_hooks if user_hook.model is component) - hook.remove() - self.model_hooks.remove(hook) - for other_hook in self.model_hooks: - if other_hook.hook.other_hooks and hook in other_hook.hook.other_hooks: - other_hook.hook.other_hooks.remove(hook) - def disable_auto_cpu_offload(self): """ Disable automatic CPU offloading for all components. @@ -1116,13 +921,13 @@ def get_one( ) -> Any: """ Get a single component by either: - - searching name (pattern matching), collection, or load_id. + - searching name, collection, or load_id. - passing in a component_id Raises an error if multiple components match or none are found. Args: component_id (str | None): Optional component ID to get - name (str | None): Component name or pattern + name (str | None): Component name collection (str | None): Optional collection to filter by load_id (str | None): Optional load_id to filter by @@ -1142,33 +947,15 @@ def get_one( raise ValueError(f"Component '{component_id}' not found in ComponentsManager") return self.components[component_id] # search with name/collection/load_id - results = self.search_components(name, collection, load_id) + results = self._lookup_ids(name=name, collection=collection, load_id=load_id) if not results: raise ValueError(f"No components found matching '{name}'") if len(results) > 1: - raise ValueError(f"Multiple components found matching '{name}': {list(results.keys())}") + raise ValueError(f"Multiple components found matching '{name}': {sorted(results)}") - return next(iter(results.values())) - - def get_ids(self, names: str | list[str] = None, collection: str | None = None): - """ - Get component IDs by a list of names, optionally filtered by collection. - - Args: - names (str | list[str]): list of component names - collection (str | None): Optional collection to filter by - - Returns: - list[str]: list of component IDs - """ - ids = set() - if not isinstance(names, list): - names = [names] - for name in names: - ids.update(self._lookup_ids(name=name, collection=collection)) - return list(ids) + return self.components[next(iter(results))] def get_components_by_ids(self, ids: list[str], return_dict_with_names: bool | None = True): """ @@ -1202,20 +989,3 @@ def get_components_by_ids(self, ids: list[str], return_dict_with_names: bool | N return dict_to_return else: return components - - def get_components_by_names(self, names: list[str], collection: str | None = None): - """ - Get components by a list of names, optionally filtered by collection. - - Args: - names (list[str]): list of component names - collection (str | None): Optional collection to filter by - - Returns: - dict[str, Any]: Dictionary of components with component names as keys - - Raises: - ValueError: If duplicate component names are found in the search results - """ - ids = self.get_ids(names, collection) - return self.get_components_by_ids(ids) From 77615c07af141a3b850a84d6513511af1f3601e7 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 29 Jul 2026 21:02:37 +0000 Subject: [PATCH 12/14] link issue to TODO --- src/diffusers/modular_pipelines/components_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 912a3851f3be..15f5d7e90881 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -219,7 +219,7 @@ def pre_forward(self, module, *args, **kwargs): return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device) # YiYi TODO: diffusers' own `ModelHook` (`diffusers.hooks.hooks`) supports a `new_forward` around-hook, which - # would replace this manual wrapping - we may migrate this class to it in the future. + # would replace this manual wrapping - see issue tracked in https://github.com/huggingface/diffusers/issues/14328. def wrap_forward(self, module): # `pre_forward`/`post_forward` cannot see an exception the forward pass raises, so wrap the (already # hooked) entry points here. `forward` is how most models run; autoencoders enter through the methods From ba14c2a4184c12907156e98a1fd7b92eba2bb3ce Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Thu, 30 Jul 2026 01:28:47 +0000 Subject: [PATCH 13/14] Address review: record at call sites, one device rule, fewer moving parts - 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 --- .../modular_diffusers/components_manager.md | 4 +- .../modular_pipelines/components_manager.py | 151 ++++++++------ .../test_components_manager.py | 192 +++++++++++------- 3 files changed, 209 insertions(+), 138 deletions(-) diff --git a/docs/source/en/modular_diffusers/components_manager.md b/docs/source/en/modular_diffusers/components_manager.md index bf0b4ad3997a..d6d4a912b56d 100644 --- a/docs/source/en/modular_diffusers/components_manager.md +++ b/docs/source/en/modular_diffusers/components_manager.md @@ -118,9 +118,9 @@ On a 20GB card, the Z-Image pipeline from the example above records this: -------------------------------------------------------------------------------------------------------- ``` -Each row is one decision: the model that loaded, what was evicted to make room for it, and the free device memory (`Available`, read just before each load) the decision 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 evicted first. The VAE then fit into the 6.36 GB left next to the transformer, so it loaded without evicting anything. The `Reason` column only speaks on moves an onload did not cause — an OOM retry (`oom_retry:`), disabling offloading — which appear as their own rows. To watch the moves live as they happen instead, enable info logging with `diffusers.logging.set_verbosity_info()`. +Each row is one decision: the model that loaded, what was offloaded to make room for it, and the free device memory (`Available`, 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 only speaks on moves an onload did not cause — an OOM retry (`oom_retry:`), disabling offloading — which appear as their own rows. To watch the moves live as they happen instead, enable info logging with `diffusers.logging.set_verbosity_info()`. -A model appearing repeatedly in this table is thrashing — it is being evicted 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). +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). To find the right value, measure your workflow once: run it at the resolution, batch size, and sequence length you intend to use (activations scale with all three), and read the device's peak memory afterwards: diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 15f5d7e90881..c02f8aa825e4 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -34,7 +34,6 @@ if is_accelerate_available(): from accelerate.hooks import add_hook_to_module, remove_hook_from_module - from accelerate.state import PartialState from accelerate.utils import send_to_device from accelerate.utils.memory import clear_device_cache from accelerate.utils.modeling import convert_file_size_to_int @@ -42,6 +41,18 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name +def normalize_execution_device(device: str | int | torch.device) -> torch.device: + """ + The device in the form tensors report theirs in: accelerator devices always carry an explicit index + (`cuda:0`), cpu never does. Normalizing to it lets a model's device be compared with the execution device + directly. + """ + device = torch.device(device) + if device.type != "cpu" and device.index is None: + device = torch.device(f"{device.type}:0") + return device + + def available_device_memory(execution_device: torch.device) -> int: """ The device memory available for new weights right now: what the driver reports free, plus the allocator's @@ -68,9 +79,9 @@ class OffloadEvent: model_id: str model_size: int | None = None # why the model moved: which model needed the room ("release_memory_for:"), an OOM retry - # ("oom_retry:"), attaching a new component, disabling offloading, ... + # ("oom_retry:"), offloading being disabled, ... reason: str | None = None - # free device memory read just before the onload's eviction decision + # free device memory read just before this move available_memory: int | None = None @@ -99,32 +110,39 @@ def __repr__(self): if not self.events: return "Offload record: nothing recorded yet" - # One row per decision: an onload and the evictions it caused share a row, correlated through the - # eviction's "release_memory_for:" reason (evictions precede their onload in the event + # One row per decision: an onload and the offloads it caused share a row, correlated through the + # offload's "release_memory_for:" reason (those offloads precede their onload in the event # sequence). Offloads with other causes (OOM retry, offloading disabled) get their own rows. sizes = {event.model_id: event.model_size for event in self.events if event.model_size is not None} def label(model_id): return f"{model_id} ({format_size(sizes.get(model_id))})" - rows, pending_evictions = [], {} + rows, pending_offloads = [], {} for event in self.events: if event.action == "onload": - evicted = pending_evictions.pop(event.model_id, []) - offloaded = ", ".join(label(model_id) for model_id in evicted) or "-" + caused_offloads = pending_offloads.pop(event.model_id, []) + offloaded = ", ".join(label(offload.model_id) for offload in caused_offloads) or "-" + # the decision's memory picture is the reading of its first move + first_event = caused_offloads[0] if caused_offloads else event rows.append( - [label(event.model_id), offloaded, format_size(event.available_memory), event.reason or ""] + [label(event.model_id), offloaded, format_size(first_event.available_memory), event.reason or ""] ) elif event.reason is not None and event.reason.startswith("release_memory_for:"): - pending_evictions.setdefault(event.reason.removeprefix("release_memory_for:"), []).append( - event.model_id - ) + pending_offloads.setdefault(event.reason.removeprefix("release_memory_for:"), []).append(event) else: - rows.append(["-", label(event.model_id), "-", event.reason or ""]) - # evictions whose onload never happened (its forward raised) still deserve a row - for reason_target, evicted in pending_evictions.items(): - for model_id in evicted: - rows.append(["-", label(model_id), "-", f"release_memory_for:{reason_target}"]) + rows.append(["-", label(event.model_id), format_size(event.available_memory), event.reason or ""]) + # offloads whose matching onload never made it into the record still deserve a row + for reason_target, caused_offloads in pending_offloads.items(): + for offload in caused_offloads: + rows.append( + [ + "-", + label(offload.model_id), + format_size(offload.available_memory), + f"release_memory_for:{reason_target}", + ] + ) table = format_table( ["#", "Onload", "Offloaded", "Available", "Reason"], @@ -146,7 +164,8 @@ class CustomOffloadHook(ModelHook): Whether to recover from a forward pass that runs out of device memory by offloading the other models one at a time and retrying. If `False`, the error is raised to the caller. record(`OffloadRecord`, *optional*): - Where to record the moves this hook makes. + Where to record the moves this hook makes. Defaults to a private record; pass a shared one to collect + the moves of several hooks in a single sequence (as `ComponentsManager` does). """ no_grad = False @@ -159,7 +178,9 @@ def __init__( retry_on_oom: bool = True, record: OffloadRecord | None = None, ): - self.execution_device = execution_device if execution_device is not None else PartialState().default_device + self.execution_device = normalize_execution_device( + execution_device if execution_device is not None else get_device() + ) self.other_hooks = other_hooks self.offload_strategy = offload_strategy self.retry_on_oom = retry_on_oom @@ -167,9 +188,6 @@ def __init__( self.model_id = None self._wrapped_methods = {} - def set_strategy(self, offload_strategy: "AutoOffloadStrategy"): - self.offload_strategy = offload_strategy - def add_other_hook(self, hook: "UserCustomOffloadHook"): """ Add a hook to the list of hooks to consider for offloading. @@ -182,12 +200,15 @@ def init_hook(self, module): return module.to("cpu") def pre_forward(self, module, *args, **kwargs): + """ + Runs before every hooked forward pass: if the model is not on the execution device, move it there — first + offloading whatever resident models the strategy picks to make room, recording every move — and move the + forward's args/kwargs to the device. + """ if module.device != self.execution_device: - # read before any eviction: this is the number the eviction decision is based on - available_memory = available_device_memory(self.execution_device) if self.other_hooks is not None: hooks_to_offload = [hook for hook in self.other_hooks if hook.model.device == self.execution_device] - # offload all other hooks + # ask the strategy which of the resident models to offload to make room start_time = time.perf_counter() if self.offload_strategy is not None: hooks_to_offload = self.offload_strategy( @@ -203,19 +224,28 @@ def pre_forward(self, module, *args, **kwargs): logger.info( f"moving {self.model_id} to {self.execution_device}, offloading {hook.model_id} to cpu" ) - hook.offload(reason=f"release_memory_for:{self.model_id}") + self.record.add( + OffloadEvent( + action="offload", + model_id=hook.model_id, + model_size=hook.model.get_memory_footprint(), + reason=f"release_memory_for:{self.model_id}", + available_memory=available_device_memory(self.execution_device), + ) + ) + hook.offload() if hooks_to_offload: clear_device_cache() - module.to(self.execution_device) self.record.add( OffloadEvent( action="onload", model_id=self.model_id, model_size=module.get_memory_footprint(), - available_memory=available_memory, + available_memory=available_device_memory(self.execution_device), ) ) + module.to(self.execution_device) return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device) # YiYi TODO: diffusers' own `ModelHook` (`diffusers.hooks.hooks`) supports a `new_forward` around-hook, which @@ -225,17 +255,14 @@ def wrap_forward(self, module): # hooked) entry points here. `forward` is how most models run; autoencoders enter through the methods # `apply_forward_hook` marks as device entry points (their encode/decode), which fire `pre_forward` but # route around `forward`. On an OOM, free the smallest resident model and retry, escalating until it - # fits. The memory readings cannot guide this: the failed forward has already unwound, so its activations - # are freed and the device looks free again. Inference only: the retried forward re-runs cleanly from its - # original inputs. + # fits. if not self.retry_on_oom: return def with_oom_retry(entry_point): @functools.wraps(entry_point) def run_with_oom_retry(*args, **kwargs): - # each pass offloads one more model, so this terminates once they all have been. Tracked explicitly - # rather than by device, so that a model that did not actually move cannot be picked twice. + # each pass offloads one more model, so this terminates once they all have been offloaded = set() while True: try: @@ -245,11 +272,7 @@ def run_with_oom_retry(*args, **kwargs): ( hook for hook in (self.other_hooks or []) - # compare with a normalized index: model.device reports no index for the - # default device - if hook.model.device.type == self.execution_device.type - and (hook.model.device.index or 0) == (self.execution_device.index or 0) - and id(hook) not in offloaded + if hook.model.device == self.execution_device and id(hook) not in offloaded ), key=lambda hook: hook.model.get_memory_footprint(), ) @@ -266,7 +289,16 @@ def run_with_oom_retry(*args, **kwargs): "retrying. If this happens repeatedly, set a larger `memory_reserve` in " "`enable_auto_cpu_offload`." ) - smallest.offload(reason=f"oom_retry:{self.model_id}") + self.record.add( + OffloadEvent( + action="offload", + model_id=smallest.model_id, + model_size=smallest.model.get_memory_footprint(), + reason=f"oom_retry:{self.model_id}", + available_memory=available_device_memory(self.execution_device), + ) + ) + smallest.offload() offloaded.add(id(smallest)) clear_device_cache() @@ -298,18 +330,8 @@ def __init__(self, model_id, model, hook): self.model = model self.hook = hook - def offload(self, reason: str | None = None): - was_resident = self.model.device.type == self.hook.execution_device.type + def offload(self): self.hook.init_hook(self.model) - if was_resident: - self.hook.record.add( - OffloadEvent( - action="offload", - model_id=self.model_id, - model_size=self.model.get_memory_footprint(), - reason=reason, - ) - ) def attach(self): add_hook_to_module(self.model, self.hook) @@ -353,7 +375,7 @@ class AutoOffloadStrategy: The sizes cover the weights managed by this strategy only — the actual memory requirements will include activations and any other allocations, so `memory_reserve` covers exactly that headroom; a `memory_reserve` of 0 - packs the device as full as the weights allow, relying on the OOM retry as a backstop. + packs the device as full as the weights allow. """ def __init__(self, memory_reserve="3GB"): @@ -386,7 +408,6 @@ def __call__(self, hooks, model_id, model, execution_device): logger.info(f" search for models to offload in order to free up {min_memory_offload / 1024**3:.2f} GB memory") - # exlucde models that's not currently loaded on the device module_sizes = dict( sorted( {hook.model_id: hook.model.get_memory_footprint() for hook in hooks}.items(), @@ -475,7 +496,7 @@ def __init__(self): self.model_hooks = None self._auto_offload_enabled = False self._offload_strategy = None - self._offload_retry_on_oom = True + self._offload_retry_on_oom = None self._offload_record = OffloadRecord() def _lookup_ids( @@ -606,7 +627,6 @@ def add(self, name: str, component: Any, collection: str | None = None): if other_hook.hook.execution_device == hook.hook.execution_device: hook.add_other_hook(other_hook) other_hook.add_other_hook(hook) - hook.offload(reason="component_added") self.model_hooks.append(hook) return component_id @@ -650,7 +670,7 @@ def remove(self, component_id: str = None): if isinstance(component, torch.nn.Module): if self._auto_offload_enabled: # detach only this component's offload hook, leaving all other managed models where they are - hook = next(user_hook for user_hook in self.model_hooks if user_hook.model is component) + hook = next(user_hook for user_hook in self.model_hooks if user_hook.model_id == component_id) hook.remove() self.model_hooks.remove(hook) for other_hook in self.model_hooks: @@ -705,11 +725,7 @@ def enable_auto_cpu_offload( if device is None: device = get_device() - if not isinstance(device, torch.device): - device = torch.device(device) - - if device.index is None: - device = torch.device(f"{device.type}:{0}") + device = normalize_execution_device(device) device_module = getattr(torch, device.type, torch.cuda) if not hasattr(device_module, "mem_get_info"): @@ -769,14 +785,25 @@ def disable_auto_cpu_offload(self): return for hook in self.model_hooks: - hook.offload(reason="offloading_disabled") + # only models actually on the device get a final recorded move; the rest are already on CPU + if hook.model.device == hook.hook.execution_device: + self._offload_record.add( + OffloadEvent( + action="offload", + model_id=hook.model_id, + model_size=hook.model.get_memory_footprint(), + reason="offloading_disabled", + available_memory=available_device_memory(hook.hook.execution_device), + ) + ) + hook.offload() hook.remove() if self.model_hooks: clear_device_cache() self.model_hooks = None self._auto_offload_enabled = False self._offload_strategy = None - self._offload_retry_on_oom = True + self._offload_retry_on_oom = None def get_model_info( self, diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index ad77f8dadd68..c0938fb8c009 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -46,7 +46,7 @@ UNIT = 1024 # More free memory than any tiny test checkpoint could ever need, so the strategy never -# decides to offload. Used to assert the *negative*: no eviction without memory pressure. +# decides to offload. Used to assert the *negative*: no offloading without memory pressure. _AMPLE_FREE_BYTES = 1024**4 @@ -290,7 +290,7 @@ def test_strategy_memory_reserve_changes_decision(self): def test_strategy_reserve_zero_packs_tight(self): # `memory_reserve=0` uses every last byte for weights: incoming exactly equals - # the free memory and still fits without evicting. + # the free memory and still fits without offloading anything. selected = self._select_offload( incoming_footprint=4 * UNIT, free_bytes=4 * UNIT, @@ -302,7 +302,7 @@ def test_strategy_reserve_zero_packs_tight(self): def test_strategy_allocator_cache_counts_as_available(self): # `mem_get_info` reports only 2 units free, but 6 units of the allocator's cache # are reusable: available = 2 + 6 = 8, minus 1 reserve -> the incoming 4 fits - # with nothing evicted. Without the cache add-back this would over-evict. + # with nothing offloaded. Without the cache add-back this would offload needlessly. selected = self._select_offload( incoming_footprint=4 * UNIT, free_bytes=2 * UNIT, @@ -369,7 +369,7 @@ def test_auto_offload_starts_with_all_components_on_cpu(self): cm.disable_auto_cpu_offload() @require_accelerator - def test_auto_offload_evicts_resident_model_under_memory_pressure(self): + def test_auto_offload_offloads_resident_model_under_memory_pressure(self): device_type = torch.device(torch_device).type cm = ComponentsManager() m1 = self.get_dummy_model(4 * UNIT) @@ -384,14 +384,14 @@ def test_auto_offload_evicts_resident_model_under_memory_pressure(self): x = torch.randn(2, 4, device=torch_device) - # Ample free memory: running m1 just moves it onto the device, evicting + # Ample free memory: running m1 just moves it onto the device, offloading # nothing (m2 is not resident, so it is not even a candidate). with _patch_free_memory(70 * UNIT): m1(x) assert next(m1.parameters()).device.type == device_type # Memory pressure: usable = 4 - 1 = 3 but m2 needs 4, so the only resident - # model (m1) must be evicted back to the CPU to make room for m2. + # model (m1) must be offloaded back to the CPU to make room for m2. with _patch_free_memory(4 * UNIT): m2(x) assert next(m2.parameters()).device.type == device_type @@ -411,7 +411,7 @@ def test_record_captures_what_the_strategy_saw(self): x = torch.randn(2, 4, device=torch_device) with _patch_free_memory(70 * UNIT): m1(x) - # m2 does not fit alongside m1 (usable 4 - 1 = 3 < 4), so m1 is evicted for it + # m2 does not fit alongside m1 (usable 4 - 1 = 3 < 4), so m1 is offloaded for it with _patch_free_memory(4 * UNIT): m2(x) @@ -421,12 +421,13 @@ def test_record_captures_what_the_strategy_saw(self): # The reading behind each decision is recorded, so a run can be explained after the fact. assert onloads[0].available_memory == 70 * UNIT assert onloads[1].available_memory == 4 * UNIT - # the eviction is its own event, attributed to the onload that caused it - eviction = next(event for event in cm.offload_record.events if event.action == "offload") - assert eviction.model_id == cm.model_hooks[0].model_id - assert eviction.reason == f"release_memory_for:{cm.model_hooks[1].model_id}" + # the offload is its own event, attributed to the onload that caused it, with its own reading + offload_event = next(event for event in cm.offload_record.events if event.action == "offload") + assert offload_event.model_id == cm.model_hooks[0].model_id + assert offload_event.reason == f"release_memory_for:{cm.model_hooks[1].model_id}" + assert offload_event.available_memory == 4 * UNIT assert _peak_co_residency(cm.offload_record.events) == 1 - # the printed table shows one row per decision: m2's row carries the m1 eviction it caused + # the printed table shows one row per decision: m2's row carries the m1 offload it caused m2_row = next(line for line in repr(cm.offload_record).splitlines() if line.startswith("2 ")) assert cm.model_hooks[1].model_id in m2_row and cm.model_hooks[0].model_id in m2_row finally: @@ -446,12 +447,104 @@ def test_auto_offload_keeps_models_resident_when_memory_is_ample(self): with _patch_free_memory(70 * UNIT): m1(x) m2(x) - # Both fit comfortably, so neither gets evicted. + # Both fit comfortably, so neither gets offloaded. assert next(m1.parameters()).device.type == device_type assert next(m2.parameters()).device.type == device_type finally: cm.disable_auto_cpu_offload() + @require_accelerator + def test_record_captures_the_onload_sequence(self): + cm = ComponentsManager() + model_a = self.get_dummy_model(UNIT) + model_b = self.get_dummy_model(4 * UNIT) + cm.add("a", model_a) + cm.add("b", model_b) + cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) + try: + x = torch.randn(2, 4, device=torch_device) + with _patch_free_memory(70 * UNIT): + model_a(x) + model_b(x) + + onloads = [event for event in cm.offload_record.events if event.action == "onload"] + assert [event.model_id.rsplit("_", 1)[0] for event in onloads] == ["a", "b"] + assert [event.model_size for event in onloads] == [UNIT, 4 * UNIT] + assert [event.action for event in cm.offload_record.events] == ["onload", "onload"] + finally: + cm.disable_auto_cpu_offload() + + @require_accelerator + def test_record_tracks_peak_co_residency(self): + cm = ComponentsManager() + m1 = self.get_dummy_model(4 * UNIT) + m2 = self.get_dummy_model(4 * UNIT) + cm.add("m1", m1) + cm.add("m2", m2) + cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) + try: + x = torch.randn(2, 4, device=torch_device) + with _patch_free_memory(70 * UNIT): + m1(x) + m2(x) + assert _peak_co_residency(cm.offload_record.events) == 2 + + # an offload brings residency back down, so a later onload does not raise the peak: + # m3 needs 4 but only 4 - 1 = 3 is usable, costing one of the resident models its spot + m3 = self.get_dummy_model(4 * UNIT) + cm.add("m3", m3) + with _patch_free_memory(4 * UNIT): + m3(x) + assert _peak_co_residency(cm.offload_record.events) == 2 + finally: + cm.disable_auto_cpu_offload() + + @require_accelerator + def test_record_is_bounded_and_clearable(self): + cm = ComponentsManager() + m1 = self.get_dummy_model(4 * UNIT) + m2 = self.get_dummy_model(4 * UNIT) + cm.add("m1", m1) + cm.add("m2", m2) + cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) + cm.offload_record.events = collections.deque(maxlen=2) + try: + x = torch.randn(2, 4, device=torch_device) + # under pressure the two models trade places on every call, producing a steady + # stream of events; the record keeps only the last `maxlen` + with _patch_free_memory(4 * UNIT): + m1(x) + m2(x) + m1(x) + + assert len(cm.offload_record.events) == 2 + + cm.offload_record.clear() + assert len(cm.offload_record.events) == 0 + assert "nothing recorded yet" in repr(cm.offload_record) + finally: + cm.disable_auto_cpu_offload() + + @require_accelerator + def test_record_survives_disable(self): + cm = ComponentsManager() + m1 = self.get_dummy_model(4 * UNIT) + cm.add("m1", m1) + cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) + x = torch.randn(2, 4, device=torch_device) + with _patch_free_memory(70 * UNIT): + m1(x) + + cm.disable_auto_cpu_offload() + + # the record is the post-mortem of the run that just ended + events = list(cm.offload_record.events) + assert sum(event.action == "onload" for event in events) == 1 + # disabling is recorded too: the resident model's final move back to the CPU + assert events[-1].action == "offload" + assert events[-1].reason == "offloading_disabled" + assert events[-1].available_memory is not None + class TestComponentsManager(ComponentsManagerTesterMixin): pass @@ -463,9 +556,9 @@ class TestOffloadHookBehavior: @pytest.fixture(autouse=True) def _fake_cpu_mem_get_info(self): - # These tests exercise the hook machinery only, on a cpu execution device — but both enabling - # offloading and every onload read `mem_get_info`, which cpu doesn't implement, so provide an - # ample fake one for the whole test. + # These tests exercise the hook machinery only, on a cpu execution device — but enabling + # offloading requires `mem_get_info`, and the OOM-retry records read it, neither of which cpu + # implements, so provide an ample fake one for the whole test. with mock.patch.object( torch.cpu, "mem_get_info", create=True, return_value=(_AMPLE_FREE_BYTES, _AMPLE_FREE_BYTES) ): @@ -510,7 +603,7 @@ def test_oom_retry_escalates_one_model_at_a_time(self): assert model.forward_calls == 3 assert self._offloaded_names(manager) == ["smallest", "larger"] - def test_oom_reraises_when_nothing_to_evict(self): + def test_oom_reraises_when_nothing_to_offload(self): model = OOMOnceModel() self._manager({"model": model}) @@ -552,19 +645,6 @@ def test_oom_retry_covers_apply_forward_hook_entry_points(self): assert model.decode_calls == 2 assert self._offloaded_names(manager) == ["smallest"] - def test_record_captures_the_onload_sequence(self): - model_a = DummyModel(UNIT) - model_b = DummyModel(4 * UNIT) - manager = self._manager({"a": model_a, "b": model_b}) - - model_a(torch.zeros(2, 4)) - model_b(torch.zeros(2, 4)) - - onloads = [event for event in manager.offload_record.events if event.action == "onload"] - assert [event.model_id.rsplit("_", 1)[0] for event in onloads] == ["a", "b"] - assert [event.model_size for event in onloads] == [UNIT, 4 * UNIT] - assert [event.action for event in manager.offload_record.events] == ["onload", "onload"] - def test_record_captures_oom_and_the_escalation(self): model = OOMOnceModel() smallest = DummyModel(UNIT) @@ -572,7 +652,7 @@ def test_record_captures_oom_and_the_escalation(self): model(torch.zeros(2, 4)) - # the OOM shows up as the eviction it caused, attributed to the model that ran out of memory + # the OOM shows up as the offload it caused, attributed to the model that ran out of memory offloads = [event for event in manager.offload_record.events if event.action == "offload"] assert len(offloads) == 1 offload_event = offloads[0] @@ -580,43 +660,6 @@ def test_record_captures_oom_and_the_escalation(self): assert offload_event.model_id.rsplit("_", 1)[0] == "smallest" assert offload_event.reason == f"oom_retry:{manager.model_hooks[0].model_id}" - def test_record_tracks_peak_co_residency(self): - model_a = DummyModel(UNIT) - model_b = DummyModel(UNIT) - manager = self._manager({"a": model_a, "b": model_b}) - - model_a(torch.zeros(2, 4)) - model_b(torch.zeros(2, 4)) - assert _peak_co_residency(manager.offload_record.events) == 2 - - # an offload brings residency back down, so a later onload does not raise the peak - manager.model_hooks[0].offload(reason="test") - model_a(torch.zeros(2, 4)) - assert _peak_co_residency(manager.offload_record.events) == 2 - - def test_record_is_bounded_and_clearable(self): - model = DummyModel() - manager = self._manager({"model": model}) - manager.offload_record.events = collections.deque(maxlen=2) - - for _ in range(4): - model(torch.zeros(2, 4)) - - assert len(manager.offload_record.events) == 2 - - manager.offload_record.clear() - assert len(manager.offload_record.events) == 0 - assert "nothing recorded yet" in repr(manager.offload_record) - - def test_record_survives_disable(self): - model = DummyModel() - manager = self._manager({"model": model}) - model(torch.zeros(2, 4)) - - manager.disable_auto_cpu_offload() - # the record is the post-mortem of the run that just ended - assert sum(event.action == "onload" for event in manager.offload_record.events) == 1 - def test_add_does_not_rebuild_existing_hooks(self): model_a = DummyModel() manager = self._manager({"a": model_a}) @@ -670,7 +713,8 @@ def _run_offloaded(self, free_bytes): """ Run the pipeline with auto offload on and `free_bytes` of *simulated* device memory, recording every move in `cm.offload_record`: an `"onload"` event per model - that ran, and an `"offload"` event (`reason="release_memory_for:"`) per eviction. + that ran, and an `"offload"` event (`reason="release_memory_for:"`) per model + offloaded to make room. """ cm = ComponentsManager() pipe = self.get_pipeline(components_manager=cm) @@ -683,7 +727,7 @@ def _run_offloaded(self, free_bytes): @require_accelerate @require_accelerator def test_auto_cpu_offload_serializes_models_under_memory_pressure(self): - # Zero simulated free memory: every model that runs must first evict whatever is + # Zero simulated free memory: every model that runs must first offload whatever is # currently resident (comfy-style serialized execution). cm, events, _ = self._run_offloaded(free_bytes=0) try: @@ -691,8 +735,8 @@ def test_auto_cpu_offload_serializes_models_under_memory_pressure(self): if len(distinct_models) < 2: pytest.skip("pipeline has fewer than two offloadable model components") - # Offloading actually fired (at least one eviction happened). - assert any(event.action == "offload" for event in events), "expected at least one eviction" + # Offloading actually fired (at least one model was pushed off the device). + assert any(event.action == "offload" for event in events), "expected at least one offload" # Sequencing: models run one at a time, never two co-resident on the device. peak = _peak_co_residency(events) @@ -711,7 +755,7 @@ def test_auto_cpu_offload_serializes_models_under_memory_pressure(self): @require_accelerator def test_auto_cpu_offload_keeps_models_resident_without_memory_pressure(self): # Negative case: with ample simulated memory the strategy is still consulted on - # every load, but it must never decide to evict anything. + # every load, but it must never decide to offload anything. cm, events, _ = self._run_offloaded(free_bytes=_AMPLE_FREE_BYTES) try: distinct_models = {event.model_id for event in events if event.action == "onload"} @@ -719,7 +763,7 @@ def test_auto_cpu_offload_keeps_models_resident_without_memory_pressure(self): pytest.skip("pipeline has fewer than two offloadable model components") # Nothing was ever offloaded... - assert all(event.action == "onload" for event in events), "no model should be evicted" + assert all(event.action == "onload" for event in events), "no model should be offloaded" # ...and models accumulate on the device instead of being serialized. peak = _peak_co_residency(events) From 3f561491fe77b32eb40d76c9a31f4ab26d96f612 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Thu, 30 Jul 2026 09:07:58 +0000 Subject: [PATCH 14/14] Overhaul the ComponentsManager test suite around real behavior - 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 --- .../test_components_manager.py | 970 +++++++++--------- 1 file changed, 509 insertions(+), 461 deletions(-) diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index c0938fb8c009..ec8f2aca976f 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import collections import contextlib import gc from unittest import mock @@ -27,6 +26,8 @@ from ..testing_utils import ( backend_empty_cache, + backend_max_memory_allocated, + backend_reset_peak_memory_stats, require_accelerate, require_accelerator, simulate_accelerator_memory, @@ -50,18 +51,6 @@ _AMPLE_FREE_BYTES = 1024**4 -def _peak_co_residency(events): - """Largest number of models simultaneously on the device, replayed from the recorded moves.""" - resident, peak = set(), 0 - for event in events: - if event.action == "onload": - resident.add(event.model_id) - peak = max(peak, len(resident)) - else: - resident.discard(event.model_id) - return peak - - class DummyModel(ModelMixin): def __init__(self, footprint_bytes: int = UNIT): super().__init__() @@ -73,20 +62,7 @@ def forward(self, x): return x + self.weight.sum() -class OOMOnceModel(DummyModel): - """Raises a fake device OOM on its first `_oom_calls` forwards (or on every forward with `_repeated_oom`).""" - - _repeated_oom = False - _oom_calls = 1 - - def forward(self, x): - self.forward_calls = getattr(self, "forward_calls", 0) + 1 - if self.forward_calls <= self._oom_calls or self._repeated_oom: - raise torch.OutOfMemoryError("fake OOM") - return super().forward(x) - - -class DecodeModel(DummyModel): +class DummyDecodeModel(DummyModel): """Autoencoder-style model: runs through a `decode` entry point (via `apply_forward_hook`, like the VAEs) instead of `forward`.""" @@ -95,17 +71,6 @@ def decode(self, x): return x + self.weight.sum() -class OOMOnceDecodeModel(DecodeModel): - """`DecodeModel` that raises a fake device OOM on its first decode.""" - - @apply_forward_hook - def decode(self, x): - self.decode_calls = getattr(self, "decode_calls", 0) + 1 - if self.decode_calls == 1: - raise torch.OutOfMemoryError("fake OOM") - return x + self.weight.sum() - - class _FakeHook: """ Minimal stand-in for `UserCustomOffloadHook` in strategy-level unit tests. @@ -121,75 +86,80 @@ def __init__(self, model_id: str, model: torch.nn.Module): @contextlib.contextmanager -def _patch_memory_stats(device_module, free_bytes, total_bytes, cached_bytes=0): - # `mem_get_info` returns `(free, total)` and is where the strategy learns how much - # memory is free; the strategy additionally counts the allocator's reusable cache - # (`memory_reserved - memory_allocated`) as available, so those are pinned too — - # `cached_bytes` simulates reusable cache on top of `free_bytes`. +def _patch_memory_stats(free_bytes, total_bytes=80 * UNIT, device_module=None): + # `mem_get_info` returns `(free, total)` and is where the offloader learns how much memory is + # free; it additionally counts the allocator's reusable cache (`memory_reserved - + # memory_allocated`) as available, so those are pinned to zero — `free_bytes` is the whole + # story. `device_module=None` targets the backend of the real `torch_device` (integration + # tests); the strategy unit tests pass `torch.cuda` to match their `cuda:0` device descriptor. + if device_module is None: + device_module = getattr(torch, torch.device(torch_device).type, torch.cuda) with contextlib.ExitStack() as stack: - stack.enter_context(mock.patch.object(device_module, "mem_get_info", return_value=(free_bytes, total_bytes))) + # `create=True` lets this fake `mem_get_info` onto backends that don't implement it (cpu) + stack.enter_context( + mock.patch.object(device_module, "mem_get_info", create=True, return_value=(free_bytes, total_bytes)) + ) if hasattr(device_module, "memory_reserved") and hasattr(device_module, "memory_allocated"): - stack.enter_context(mock.patch.object(device_module, "memory_reserved", return_value=cached_bytes)) + stack.enter_context(mock.patch.object(device_module, "memory_reserved", return_value=0)) stack.enter_context(mock.patch.object(device_module, "memory_allocated", return_value=0)) yield -def _patch_cuda_mem_get_info(free_bytes: int, total_bytes: int = 80 * UNIT, cached_bytes: int = 0): - # Strategy unit tests use a `cuda:0` execution-device *descriptor* (which needs no - # real GPU), so they patch the `torch.cuda` memory introspection directly. - return _patch_memory_stats(torch.cuda, free_bytes, total_bytes, cached_bytes=cached_bytes) +def _simulate_card_with_headroom(headroom_bytes): + """ + A `simulate_accelerator_memory` card sized to whatever the device currently holds plus + `headroom_bytes` — so tests can reason in absolute headroom, independent of the node's state. + """ + device_module = getattr(torch, torch.device(torch_device).type) + free, real_total = device_module.mem_get_info() + return simulate_accelerator_memory((real_total - free) + headroom_bytes, device=torch_device) -def _patch_free_memory(free_bytes: int, total_bytes: int = 80 * UNIT): - # Integration tests run on the real `torch_device`; patch the memory introspection - # on whichever backend module (cuda/xpu/...) actually backs it to simulate - # arbitrary memory pressure. - device_type = torch.device(torch_device).type - device_module = getattr(torch, device_type, torch.cuda) - return _patch_memory_stats(device_module, free_bytes, total_bytes) +class TestComponentsManagerRegistry: + """Registry behavior only: add / look up / remove components. No offloading, no hardware.""" + def test_add_and_get_one(self): + cm = ComponentsManager() + model = DummyModel(UNIT) + component_id = cm.add("unet", model) + assert component_id in cm.components + assert cm.get_one(name="unet") is model + assert cm.get_one(component_id=component_id) is model -@require_accelerate -class ComponentsManagerTesterMixin: - """ - Common tests for `ComponentsManager` and its auto-offload strategy. - """ + def test_add_same_component_twice_reuses_id(self): + cm = ComponentsManager() + model = DummyModel(UNIT) + first_id = cm.add("unet", model) + second_id = cm.add("unet", model) + assert first_id == second_id + assert len(cm.components) == 1 - # A `cuda:0` device descriptor is enough to drive the strategy's device-type and - # index logic; no real GPU is required because the memory readings are mocked. - strategy_execution_device = torch.device("cuda:0") + def test_remove(self): + cm = ComponentsManager() + component_id = cm.add("unet", DummyModel(UNIT)) + cm.remove(component_id) + assert component_id not in cm.components - def setup_method(self): - # Mirror `ModularPipelineTesterMixin` cleanup so this mixin stays interchangeable - # in the MRO when stacked into a pipeline test class. - torch.compiler.reset() - gc.collect() - backend_empty_cache(torch_device) + def test_get_model_info_reports_size(self): + cm = ComponentsManager() + model = DummyModel(footprint_bytes=2 * UNIT) + component_id = cm.add("unet", model) + info = cm.get_model_info(component_id, fields="size_gb") + assert info["size_gb"] == model.get_memory_footprint() / (1024**3) - def teardown_method(self): - torch.compiler.reset() - gc.collect() - backend_empty_cache(torch_device) - def get_dummy_model(self, footprint_bytes: int = UNIT) -> ModelMixin: - return DummyModel(footprint_bytes=footprint_bytes) +@require_accelerate +class TestAutoOffloadStrategy: + """AutoOffloadStrategy unit tests: no real GPU — a `cuda:0` device *descriptor* plus scripted + memory readings are enough to drive the selection logic.""" - # ------------------------------------------------------------------ - # AutoOffloadStrategy unit tests (hardware-independent) - # ------------------------------------------------------------------ - def _select_offload( - self, - *, - incoming_footprint, - free_bytes, - hook_sizes, - memory_reserve=UNIT, - cached_bytes=0, - ): + strategy_execution_device = torch.device("cuda:0") + + def _select_offload(self, *, incoming_footprint, free_bytes, hook_sizes, memory_reserve=UNIT): strategy = AutoOffloadStrategy(memory_reserve=memory_reserve) - hooks = [_FakeHook(model_id, self.get_dummy_model(fp)) for model_id, fp in hook_sizes.items()] - incoming = self.get_dummy_model(incoming_footprint) - with _patch_cuda_mem_get_info(free_bytes, cached_bytes=cached_bytes): + hooks = [_FakeHook(model_id, DummyModel(fp)) for model_id, fp in hook_sizes.items()] + incoming = DummyModel(incoming_footprint) + with _patch_memory_stats(free_bytes, device_module=torch.cuda): selected = strategy( hooks=hooks, model_id="incoming", @@ -288,29 +258,6 @@ def test_strategy_memory_reserve_changes_decision(self): memory_reserve=3 * UNIT, ) == ["c"] - def test_strategy_reserve_zero_packs_tight(self): - # `memory_reserve=0` uses every last byte for weights: incoming exactly equals - # the free memory and still fits without offloading anything. - selected = self._select_offload( - incoming_footprint=4 * UNIT, - free_bytes=4 * UNIT, - hook_sizes={"a": 4 * UNIT}, - memory_reserve=0, - ) - assert selected == [] - - def test_strategy_allocator_cache_counts_as_available(self): - # `mem_get_info` reports only 2 units free, but 6 units of the allocator's cache - # are reusable: available = 2 + 6 = 8, minus 1 reserve -> the incoming 4 fits - # with nothing offloaded. Without the cache add-back this would offload needlessly. - selected = self._select_offload( - incoming_footprint=4 * UNIT, - free_bytes=2 * UNIT, - hook_sizes={"a": 4 * UNIT}, - cached_bytes=6 * UNIT, - ) - assert selected == [] - def test_strategy_memory_reserve_accepts_file_size_strings(self): # "3KiB" = 3 * 1024 bytes = a 3-unit reserve; free 4 - 3 = 1 usable but the # incoming needs 4, so the resident model must go. @@ -322,45 +269,328 @@ def test_strategy_memory_reserve_accepts_file_size_strings(self): ) assert selected == ["a"] + +class TestSimulateAcceleratorMemory: + """ + Validates the `simulate_accelerator_memory` testing util itself against the real device: unlike + `_patch_memory_stats` (which pins scripted readings for hardware-independent unit tests), the + util wraps `mem_get_info` so real allocations show up live on a simulated smaller card, and + hard-caps the allocator so overshooting the simulated capacity raises a real OOM. + """ + + MB = 2**20 + + def setup_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + @require_accelerator + def test_readings_translate_to_simulated_card(self): + # this test validates the simulated readings themselves, so it reads `mem_get_info` raw + device_module = getattr(torch, torch.device(torch_device).type) + free, real_total = device_module.mem_get_info() + # Simulate a card with 64MB of headroom on top of whatever the device currently holds. + sim_total = (real_total - free) + 64 * self.MB + + with simulate_accelerator_memory(sim_total, device=torch_device, hard=False): + free0, total0 = device_module.mem_get_info() + assert total0 == sim_total + assert free0 <= 64 * self.MB + + # A real allocation is visible on the simulated card: free drops by at least its size. + # (8M float32 elements x 4 bytes each = 32MB; "at least" because the allocator grabs + # memory from the driver in larger blocks) + x = torch.zeros(8 * self.MB, dtype=torch.float32, device=torch_device) + free1, total1 = device_module.mem_get_info() + assert total1 == sim_total + assert free1 <= free0 - 32 * self.MB + del x + + # Exiting the context restores the real readings. + assert device_module.mem_get_info()[1] == real_total + + @require_accelerator + def test_run_ooms_when_simulated_card_is_too_small(self): + model = DummyModel(footprint_bytes=64 * self.MB) + x = torch.randn(4, device=torch_device) + + # Measure what the run actually needs on the device (weights + forward). + backend_reset_peak_memory_stats(torch_device) + model.to(torch_device) + model(x) + requirement = backend_max_memory_allocated(torch_device) + model.to("cpu") + backend_empty_cache(torch_device) + + try: + # Control: a simulated card with comfortable headroom runs the model fine. + with _simulate_card_with_headroom(2 * requirement): + model.to(torch_device) + model(x) + model.to("cpu") + backend_empty_cache(torch_device) + + # A card with only half the requirement left cannot: the hard cap turns the + # overshoot into a real device OOM instead of using the extra physical memory. + with _simulate_card_with_headroom(requirement // 2): + with pytest.raises(torch.OutOfMemoryError): + model.to(torch_device) + model(x) + finally: + model.to("cpu") + backend_empty_cache(torch_device) + + +class TestAutoOffload: + """Auto-offload behavior on a real accelerator: genuine device moves. Memory pressure is + simulated two ways — `_patch_memory_stats` scripts the *readings* the strategy decides on, and + `simulate_accelerator_memory` (hard cap) makes the device behave like a smaller card, so the + OOM-retry tests recover from real `torch.OutOfMemoryError`s.""" + + MB = 2**20 + + def setup_method(self): + torch.compiler.reset() + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + torch.compiler.reset() + gc.collect() + backend_empty_cache(torch_device) + # ------------------------------------------------------------------ - # Registry tests (hardware-independent) + # OOM retry (real OOMs on a simulated small card) # ------------------------------------------------------------------ - def test_add_and_get_one(self): - cm = ComponentsManager() - model = self.get_dummy_model() - component_id = cm.add("unet", model) - assert component_id in cm.components - assert cm.get_one(name="unet") is model - assert cm.get_one(component_id=component_id) is model - def test_add_same_component_twice_reuses_id(self): - cm = ComponentsManager() - model = self.get_dummy_model() - first_id = cm.add("unet", model) - second_id = cm.add("unet", model) - assert first_id == second_id - assert len(cm.components) == 1 + @require_accelerate + @require_accelerator + def test_oom_retry_offloads_a_resident_model_and_reruns(self): + device_type = torch.device(torch_device).type - def test_remove(self): - cm = ComponentsManager() - component_id = cm.add("unet", self.get_dummy_model()) - cm.remove(component_id) - assert component_id not in cm.components + # the model that OOMs mid-forward: its 32MB of weights fit the card easily + model = DummyModel(32 * self.MB) + torch.nn.init.normal_(model.weight) # non-zero weights, so the equivalence check below has teeth + # the bystander whose eviction is what saves the retry: offloading it frees 64MB + resident = DummyModel(64 * self.MB) + + x = torch.zeros(10 * self.MB, device=torch_device) # the forward's output weighs 40MB + x_tiny = torch.zeros(4, device=torch_device) + + # baseline output from a plain, unhooked run - the retried run must reproduce it exactly. + # (empty the cache afterwards: blocks this run leaves behind would otherwise be reusable + # free headroom the simulated card below doesn't know about) + with torch.no_grad(): + model.to(torch_device) + baseline = model(x) + model.to("cpu") + backend_empty_cache(torch_device) - def test_get_model_info_reports_size(self): + manager = ComponentsManager() + manager.add("model", model) + manager.add("resident", resident) + # reserve 0 lets the weights pack tight, so it is the forward's own output that overshoots + manager.enable_auto_cpu_offload(device=torch_device, memory_reserve=0) + try: + # both models are actually under the offloader before anything is measured + assert all(hasattr(m, "_hf_hook") for m in (model, resident)) + + # A card where both models' weights fit (64 + 32 = 96 < 120) but the forward's output + # does not fit on top (96 + 40 = 136 > 120): the run really OOMs, and offloading the + # resident model (frees 64) is really what lets the retry succeed (32 + 40 = 72 < 120). + with _simulate_card_with_headroom(120 * self.MB), torch.no_grad(): + # resident onloads (64MB, plenty of room) and runs; it stays on the device + resident(x_tiny) + # model onloads (96MB resident now, still fits) -> its forward OOMs allocating the + # 40MB output -> the retry offloads `resident`, freeing 64MB -> the forward re-runs + # from the original input and fits + out = model(x) + + # surviving an OOM must not change the result + assert torch.allclose(out, baseline, atol=1e-5) + # the OOM cost the resident model its spot; the running model stayed + assert next(resident.parameters()).device.type == "cpu" + assert next(model.parameters()).device.type == device_type + # The retry's offload is the last event: the OOM struck the forward's output allocation, + # after the model's weights had already onloaded - so the re-run needs no new onload, + # only the resident model's eviction. + last_record = list(manager.offload_record.events)[-1] + assert last_record.action == "offload" + assert last_record.model_id == manager.model_hooks[1].model_id + assert last_record.reason == f"oom_retry:{manager.model_hooks[0].model_id}" + finally: + manager.disable_auto_cpu_offload() + + @require_accelerate + @require_accelerator + def test_oom_reraises_when_nothing_to_offload(self): + # 32MB of weights fit the 64MB card on their own; only the forward's output overshoots + model = DummyModel(32 * self.MB) + manager = ComponentsManager() + manager.add("model", model) + manager.enable_auto_cpu_offload(device=torch_device, memory_reserve=0) + try: + x = torch.zeros(12 * self.MB, device=torch_device) # the forward's output weighs 48MB + # The weights (32) fit the 64MB card, weights + output (32 + 48 = 80) never can - and + # with no other managed model to offload, the retry gives up with advice. + with _simulate_card_with_headroom(64 * self.MB), torch.no_grad(): + with pytest.raises(torch.OutOfMemoryError, match="group offloading"): + model(x) + finally: + manager.disable_auto_cpu_offload() + + @require_accelerate + @require_accelerator + def test_oom_retry_covers_apply_forward_hook_entry_points(self): + # Autoencoders run through `encode`/`decode` (via `apply_forward_hook`), which fire the + # offload hook but route around `forward` - an OOM there must be retried just like one + # raised inside `forward`, first call included. + device_type = torch.device(torch_device).type + # the decoder that OOMs: 32MB of weights, called through `decode` instead of `forward` + model = DummyDecodeModel(32 * self.MB) + # the bystander whose 64MB eviction is what lets the retried decode fit + resident = DummyModel(64 * self.MB) + manager = ComponentsManager() + manager.add("model", model) + manager.add("resident", resident) + manager.enable_auto_cpu_offload(device=torch_device, memory_reserve=0) + try: + x = torch.zeros(10 * self.MB, device=torch_device) # the decode's output weighs 40MB + x_tiny = torch.zeros(4, device=torch_device) + # Same card arithmetic as the forward-based retry test above: the weights fit + # (64 + 32 = 96 < 120), the decode's 40MB output on top does not (136 > 120). + with _simulate_card_with_headroom(120 * self.MB), torch.no_grad(): + # resident onloads and runs (recorded: onload); it stays on the device + resident(x_tiny) + # `decode` fires `pre_forward` via `apply_forward_hook`: model onloads (recorded: + # onload) -> decode OOMs allocating its output -> the retry offloads `resident` + # (recorded: offload, reason "oom_retry:") -> decode re-runs and fits + out = model.decode(x) + + assert out.shape == x.shape + assert next(resident.parameters()).device.type == "cpu" + assert next(model.parameters()).device.type == device_type + finally: + manager.disable_auto_cpu_offload() + + @require_accelerate + def test_enable_rejects_devices_without_memory_introspection(self): + # every offloading decision starts from `mem_get_info`, which cpu doesn't implement - + # enabling on such a backend must fail loudly, not silently misbehave cm = ComponentsManager() - model = self.get_dummy_model(footprint_bytes=2 * UNIT) - component_id = cm.add("unet", model) - info = cm.get_model_info(component_id, fields="size_gb") - assert info["size_gb"] == model.get_memory_footprint() / (1024**3) + cm.add("m1", DummyModel(UNIT)) + with pytest.raises(NotImplementedError, match="mem_get_info"): + cm.enable_auto_cpu_offload(device="cpu") + + @require_accelerate + @require_accelerator + def test_offloading_reduces_peak_memory(self): + # The feature's headline claim, measured: the tighter the card, the more the offloader + # serializes the models, and the lower the real peak device memory. + m1 = DummyModel(64 * self.MB) + m2 = DummyModel(64 * self.MB) + m3 = DummyModel(64 * self.MB) + models = (m1, m2, m3) + x = torch.randn(4, device=torch_device) + + # baseline: no offloading, all three resident together -> the peak covers 192MB of weights + backend_reset_peak_memory_stats(torch_device) + with torch.no_grad(): + for model in models: + model.to(torch_device) + model(x) + baseline_peak = backend_max_memory_allocated(torch_device) + for model in models: + model.to("cpu") + + manager = ComponentsManager() + manager.add("m1", m1) + manager.add("m2", m2) + manager.add("m3", m3) + + def peak_on_card(headroom_bytes): + # a fresh enable puts every model back on the CPU, so phases don't leak into each other + manager.enable_auto_cpu_offload(device=torch_device, memory_reserve=0) + backend_empty_cache(torch_device) + backend_reset_peak_memory_stats(torch_device) + with _simulate_card_with_headroom(headroom_bytes), torch.no_grad(): + for model in models: + model(x) + peak = backend_max_memory_allocated(torch_device) + manager.disable_auto_cpu_offload() + return peak + + try: + # a 160MB card fits two models (128) but not three (192): the third load offloads one + # resident, capping the peak at two models + partial_peak = peak_on_card(160 * self.MB) + # an 80MB card fits only one model at a time: fully serialized execution + serialized_peak = peak_on_card(80 * self.MB) + + # the hard cap doubles as a silent assert: neither run OOMed, so the offloader really + # kept each within its card + assert serialized_peak < partial_peak < baseline_peak + finally: + manager.disable_auto_cpu_offload() + + @require_accelerate + @require_accelerator + def test_add_does_not_rebuild_existing_hooks(self): + model_a = DummyModel(UNIT) + manager = ComponentsManager() + manager.add("a", model_a) + manager.enable_auto_cpu_offload(device=torch_device) + try: + hooks_before = list(manager.model_hooks) + + model_b = DummyModel(UNIT) + manager.add("b", model_b) + + assert len(manager.model_hooks) == 2 + assert manager.model_hooks[0] is hooks_before[0], "existing hook was rebuilt" + hook_a, hook_b = manager.model_hooks + assert hook_b.hook.other_hooks == [hook_a] + assert hook_a.hook.other_hooks == [hook_b] + finally: + manager.disable_auto_cpu_offload() + + @require_accelerate + @require_accelerator + def test_remove_detaches_only_that_hook(self): + model_a = DummyModel(UNIT) + model_b = DummyModel(UNIT) + manager = ComponentsManager() + manager.add("a", model_a) + b_id = manager.add("b", model_b) + manager.enable_auto_cpu_offload(device=torch_device) + try: + hook_a = manager.model_hooks[0] + + manager.remove(b_id) + + assert manager.model_hooks == [hook_a] + assert hook_a.hook.other_hooks == [] + assert not hasattr(model_b, "_hf_hook") + # model a keeps its hook: running it still moves it to the device + assert hasattr(model_a, "_hf_hook") + out = model_a(torch.zeros(2, 4, device=torch_device)) + assert out.shape == (2, 4) + assert next(model_a.parameters()).device.type == torch.device(torch_device).type + finally: + manager.disable_auto_cpu_offload() # ------------------------------------------------------------------ - # Auto-offload integration tests (require an accelerator) + # Integration (real accelerator, simulated memory pressure) # ------------------------------------------------------------------ + @require_accelerate @require_accelerator def test_auto_offload_starts_with_all_components_on_cpu(self): cm = ComponentsManager() - model = self.get_dummy_model(4 * UNIT) + model = DummyModel(4 * UNIT) cm.add("m1", model) cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) try: @@ -368,12 +598,13 @@ def test_auto_offload_starts_with_all_components_on_cpu(self): finally: cm.disable_auto_cpu_offload() + @require_accelerate @require_accelerator def test_auto_offload_offloads_resident_model_under_memory_pressure(self): device_type = torch.device(torch_device).type cm = ComponentsManager() - m1 = self.get_dummy_model(4 * UNIT) - m2 = self.get_dummy_model(4 * UNIT) + m1 = DummyModel(4 * UNIT) + m2 = DummyModel(4 * UNIT) cm.add("m1", m1) cm.add("m2", m2) cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) @@ -384,317 +615,212 @@ def test_auto_offload_offloads_resident_model_under_memory_pressure(self): x = torch.randn(2, 4, device=torch_device) - # Ample free memory: running m1 just moves it onto the device, offloading - # nothing (m2 is not resident, so it is not even a candidate). - with _patch_free_memory(70 * UNIT): + # Ample free memory: running m1 just moves it onto the device. + with _patch_memory_stats(70 * UNIT): m1(x) assert next(m1.parameters()).device.type == device_type - # Memory pressure: usable = 4 - 1 = 3 but m2 needs 4, so the only resident - # model (m1) must be offloaded back to the CPU to make room for m2. - with _patch_free_memory(4 * UNIT): + # Memory pressure: usable = 4 - 1 = 3 but m2 needs 4, so the only resident model (m1) + # must be offloaded back to the CPU to make room. + with _patch_memory_stats(4 * UNIT): m2(x) assert next(m2.parameters()).device.type == device_type assert next(m1.parameters()).device.type == "cpu" + + # Ample again: m1 re-loads next to m2 without pushing it off - models co-reside + # whenever memory allows. + with _patch_memory_stats(70 * UNIT): + m1(x) + assert next(m1.parameters()).device.type == device_type + assert next(m2.parameters()).device.type == device_type finally: cm.disable_auto_cpu_offload() + # Disabling moves the residents back to the CPU. + assert next(m1.parameters()).device.type == "cpu" + assert next(m2.parameters()).device.type == "cpu" + + @require_accelerate @require_accelerator def test_record_captures_what_the_strategy_saw(self): cm = ComponentsManager() - m1 = self.get_dummy_model(4 * UNIT) - m2 = self.get_dummy_model(4 * UNIT) + m1 = DummyModel(4 * UNIT) + m2 = DummyModel(4 * UNIT) cm.add("m1", m1) cm.add("m2", m2) cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) try: + m1_id = cm.model_hooks[0].model_id + m2_id = cm.model_hooks[1].model_id + x = torch.randn(2, 4, device=torch_device) - with _patch_free_memory(70 * UNIT): - m1(x) + with _patch_memory_stats(70 * UNIT): + m1(x) # -> records[0]: m1's onload # m2 does not fit alongside m1 (usable 4 - 1 = 3 < 4), so m1 is offloaded for it - with _patch_free_memory(4 * UNIT): - m2(x) + with _patch_memory_stats(4 * UNIT): + m2(x) # -> records[1] + records[2]: m1's offload, then m2's onload + + records = list(cm.offload_record.events) + assert len(records) == 3 + + # records[0]: m1 onloads into ample memory - the 70-unit reading it decided on is kept + assert records[0].action == "onload" + assert records[0].model_id == m1_id + assert records[0].model_size == 4 * UNIT + assert records[0].available_memory == 70 * UNIT + assert records[0].reason is None + + # records[1]: m1 offloads to make room, attributed to m2 with the 4-unit reading that + # forced the decision (usable 4 - 1 = 3 < m2's 4) + assert records[1].action == "offload" + assert records[1].model_id == m1_id + assert records[1].reason == f"release_memory_for:{m2_id}" + assert records[1].available_memory == 4 * UNIT + + # records[2]: m2 onloads into the freed memory + assert records[2].action == "onload" + assert records[2].model_id == m2_id + assert records[2].model_size == 4 * UNIT + assert records[2].available_memory == 4 * UNIT + + # records[3]: disabling offloading is recorded too - m2, the one model still resident, + # moves back to the CPU + cm.disable_auto_cpu_offload() + records = list(cm.offload_record.events) + assert len(records) == 4 + assert records[3].action == "offload" + assert records[3].model_id == m2_id + assert records[3].reason == "offloading_disabled" + assert records[3].available_memory is not None - onloads = [event for event in cm.offload_record.events if event.action == "onload"] - names = [event.model_id.rsplit("_", 1)[0] for event in onloads] - assert names == ["m1", "m2"] - # The reading behind each decision is recorded, so a run can be explained after the fact. - assert onloads[0].available_memory == 70 * UNIT - assert onloads[1].available_memory == 4 * UNIT - # the offload is its own event, attributed to the onload that caused it, with its own reading - offload_event = next(event for event in cm.offload_record.events if event.action == "offload") - assert offload_event.model_id == cm.model_hooks[0].model_id - assert offload_event.reason == f"release_memory_for:{cm.model_hooks[1].model_id}" - assert offload_event.available_memory == 4 * UNIT - assert _peak_co_residency(cm.offload_record.events) == 1 # the printed table shows one row per decision: m2's row carries the m1 offload it caused m2_row = next(line for line in repr(cm.offload_record).splitlines() if line.startswith("2 ")) - assert cm.model_hooks[1].model_id in m2_row and cm.model_hooks[0].model_id in m2_row + assert m2_id in m2_row and m1_id in m2_row finally: cm.disable_auto_cpu_offload() + @require_accelerate @require_accelerator - def test_auto_offload_keeps_models_resident_when_memory_is_ample(self): + def test_add_model_after_a_run_keeps_residents_in_place(self): device_type = torch.device(torch_device).type cm = ComponentsManager() - m1 = self.get_dummy_model(4 * UNIT) - m2 = self.get_dummy_model(4 * UNIT) + m1 = DummyModel(4 * UNIT) + m2 = DummyModel(4 * UNIT) cm.add("m1", m1) cm.add("m2", m2) cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) try: x = torch.randn(2, 4, device=torch_device) - with _patch_free_memory(70 * UNIT): + with _patch_memory_stats(70 * UNIT): m1(x) m2(x) - # Both fit comfortably, so neither gets offloaded. assert next(m1.parameters()).device.type == device_type assert next(m2.parameters()).device.type == device_type - finally: - cm.disable_auto_cpu_offload() - @require_accelerator - def test_record_captures_the_onload_sequence(self): - cm = ComponentsManager() - model_a = self.get_dummy_model(UNIT) - model_b = self.get_dummy_model(4 * UNIT) - cm.add("a", model_a) - cm.add("b", model_b) - cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) - try: - x = torch.randn(2, 4, device=torch_device) - with _patch_free_memory(70 * UNIT): - model_a(x) - model_b(x) - - onloads = [event for event in cm.offload_record.events if event.action == "onload"] - assert [event.model_id.rsplit("_", 1)[0] for event in onloads] == ["a", "b"] - assert [event.model_size for event in onloads] == [UNIT, 4 * UNIT] - assert [event.action for event in cm.offload_record.events] == ["onload", "onload"] + # Adding a model mid-run hooks it into the managed set without disturbing it: the + # residents stay where they are and the newcomer starts on the CPU. + m3 = DummyModel(4 * UNIT) + cm.add("m3", m3) + hook1, hook2, hook3 = cm.model_hooks + assert hook3.hook.other_hooks == [hook1, hook2] + assert hook1.hook.other_hooks == [hook2, hook3] + assert hook2.hook.other_hooks == [hook1, hook3] + assert next(m1.parameters()).device.type == device_type + assert next(m2.parameters()).device.type == device_type + assert next(m3.parameters()).device.type == "cpu" + + # With enough memory the newcomer loads next to the residents. + with _patch_memory_stats(70 * UNIT): + m3(x) + assert next(m1.parameters()).device.type == device_type + assert next(m2.parameters()).device.type == device_type + assert next(m3.parameters()).device.type == device_type finally: cm.disable_auto_cpu_offload() + @require_accelerate @require_accelerator - def test_record_tracks_peak_co_residency(self): + def test_added_model_can_displace_residents_under_pressure(self): + device_type = torch.device(torch_device).type cm = ComponentsManager() - m1 = self.get_dummy_model(4 * UNIT) - m2 = self.get_dummy_model(4 * UNIT) + m1 = DummyModel(4 * UNIT) + m2 = DummyModel(4 * UNIT) cm.add("m1", m1) cm.add("m2", m2) cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) try: x = torch.randn(2, 4, device=torch_device) - with _patch_free_memory(70 * UNIT): + with _patch_memory_stats(70 * UNIT): m1(x) m2(x) - assert _peak_co_residency(cm.offload_record.events) == 2 - # an offload brings residency back down, so a later onload does not raise the peak: - # m3 needs 4 but only 4 - 1 = 3 is usable, costing one of the resident models its spot - m3 = self.get_dummy_model(4 * UNIT) + m3 = DummyModel(4 * UNIT) cm.add("m3", m3) - with _patch_free_memory(4 * UNIT): + # Insufficient memory: m3 needs 4 but usable is 4 - 1 = 3, so the strategy frees one + # resident (the first-added, with all sizes equal) to make room. + with _patch_memory_stats(4 * UNIT): m3(x) - assert _peak_co_residency(cm.offload_record.events) == 2 + assert next(m3.parameters()).device.type == device_type + assert next(m1.parameters()).device.type == "cpu" + assert next(m2.parameters()).device.type == device_type finally: cm.disable_auto_cpu_offload() + @require_accelerate @require_accelerator - def test_record_is_bounded_and_clearable(self): + def test_remove_model_after_a_run_leaves_others_resident(self): + device_type = torch.device(torch_device).type cm = ComponentsManager() - m1 = self.get_dummy_model(4 * UNIT) - m2 = self.get_dummy_model(4 * UNIT) + m1 = DummyModel(4 * UNIT) + m2 = DummyModel(4 * UNIT) + m3 = DummyModel(4 * UNIT) cm.add("m1", m1) - cm.add("m2", m2) + m2_id = cm.add("m2", m2) + cm.add("m3", m3) cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) - cm.offload_record.events = collections.deque(maxlen=2) try: x = torch.randn(2, 4, device=torch_device) - # under pressure the two models trade places on every call, producing a steady - # stream of events; the record keeps only the last `maxlen` - with _patch_free_memory(4 * UNIT): + with _patch_memory_stats(70 * UNIT): m1(x) m2(x) - m1(x) + m3(x) - assert len(cm.offload_record.events) == 2 + # Removing a component detaches only that component: it lands on the CPU unhooked, + # the others keep their hooks (now linked to each other only) and stay resident. + cm.remove(m2_id) + assert not hasattr(m2, "_hf_hook") + assert next(m2.parameters()).device.type == "cpu" + hook1, hook3 = cm.model_hooks + assert hook1.hook.other_hooks == [hook3] + assert hook3.hook.other_hooks == [hook1] + assert next(m1.parameters()).device.type == device_type + assert next(m3.parameters()).device.type == device_type - cm.offload_record.clear() - assert len(cm.offload_record.events) == 0 - assert "nothing recorded yet" in repr(cm.offload_record) + # ...and the remaining models still run. + with _patch_memory_stats(70 * UNIT): + m1(x) + assert next(m1.parameters()).device.type == device_type finally: cm.disable_auto_cpu_offload() - @require_accelerator - def test_record_survives_disable(self): - cm = ComponentsManager() - m1 = self.get_dummy_model(4 * UNIT) - cm.add("m1", m1) - cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=UNIT) - x = torch.randn(2, 4, device=torch_device) - with _patch_free_memory(70 * UNIT): - m1(x) - - cm.disable_auto_cpu_offload() - - # the record is the post-mortem of the run that just ended - events = list(cm.offload_record.events) - assert sum(event.action == "onload" for event in events) == 1 - # disabling is recorded too: the resident model's final move back to the CPU - assert events[-1].action == "offload" - assert events[-1].reason == "offloading_disabled" - assert events[-1].available_memory is not None - - -class TestComponentsManager(ComponentsManagerTesterMixin): - pass - - -@require_accelerate -class TestOffloadHookBehavior: - """CPU-level tests of the hook machinery: OOM retry and non-disruptive add/remove.""" - - @pytest.fixture(autouse=True) - def _fake_cpu_mem_get_info(self): - # These tests exercise the hook machinery only, on a cpu execution device — but enabling - # offloading requires `mem_get_info`, and the OOM-retry records read it, neither of which cpu - # implements, so provide an ample fake one for the whole test. - with mock.patch.object( - torch.cpu, "mem_get_info", create=True, return_value=(_AMPLE_FREE_BYTES, _AMPLE_FREE_BYTES) - ): - yield - - def _manager(self, components, **offload_kwargs): - manager = ComponentsManager() - for name, component in components.items(): - manager.add(name, component) - manager.enable_auto_cpu_offload(device="cpu", **offload_kwargs) - return manager - - @staticmethod - def _offloaded_names(manager): - """The models the record shows going back to the CPU, in order, by the name they were added under.""" - # model ids are `_` - return [ - event.model_id.rsplit("_", 1)[0] for event in manager.offload_record.events if event.action == "offload" - ] - - def test_oom_retry_offloads_smallest_and_reruns(self): - model = OOMOnceModel() - smallest = DummyModel(UNIT) - larger = DummyModel(4 * UNIT) - manager = self._manager({"model": model, "larger": larger, "smallest": smallest}) - - out = model(torch.zeros(2, 4)) - assert out.shape == (2, 4) - assert model.forward_calls == 2 - # one OOM costs the smallest resident model, not everything on the device - assert self._offloaded_names(manager) == ["smallest"] - - def test_oom_retry_escalates_one_model_at_a_time(self): - model = OOMOnceModel() - model._oom_calls = 2 - smallest = DummyModel(UNIT) - larger = DummyModel(4 * UNIT) - manager = self._manager({"model": model, "larger": larger, "smallest": smallest}) - - out = model(torch.zeros(2, 4)) - assert out.shape == (2, 4) - assert model.forward_calls == 3 - assert self._offloaded_names(manager) == ["smallest", "larger"] - - def test_oom_reraises_when_nothing_to_offload(self): - model = OOMOnceModel() - self._manager({"model": model}) - - with pytest.raises(torch.OutOfMemoryError, match="group offloading"): - model(torch.zeros(2, 4)) - - def test_oom_reraises_once_everything_is_offloaded(self): - model = OOMOnceModel() - model._repeated_oom = True - smallest = DummyModel(UNIT) - larger = DummyModel(4 * UNIT) - manager = self._manager({"model": model, "larger": larger, "smallest": smallest}) - - with pytest.raises(torch.OutOfMemoryError, match="group offloading"): - model(torch.zeros(2, 4)) - # one attempt per model offloaded, plus the initial one - assert model.forward_calls == 3 - assert self._offloaded_names(manager) == ["smallest", "larger"] - - def test_retry_on_oom_false_leaves_the_forward_alone(self): - model = OOMOnceModel() - other = DummyModel() - manager = self._manager({"model": model, "other": other}, retry_on_oom=False) - - with pytest.raises(torch.OutOfMemoryError, match="fake OOM"): - model(torch.zeros(2, 4)) - assert model.forward_calls == 1 - assert all(not user_hook.hook._wrapped_methods for user_hook in manager.model_hooks) - - def test_oom_retry_covers_apply_forward_hook_entry_points(self): - # Autoencoders run through `encode`/`decode` (via `apply_forward_hook`), not `forward` - an OOM - # there must be retried just like one raised inside `forward`, first call included. - model = OOMOnceDecodeModel() - smallest = DummyModel(UNIT) - manager = self._manager({"model": model, "smallest": smallest}) - - out = model.decode(torch.zeros(2, 4)) - assert out.shape == (2, 4) - assert model.decode_calls == 2 - assert self._offloaded_names(manager) == ["smallest"] - - def test_record_captures_oom_and_the_escalation(self): - model = OOMOnceModel() - smallest = DummyModel(UNIT) - manager = self._manager({"model": model, "smallest": smallest}) - - model(torch.zeros(2, 4)) - - # the OOM shows up as the offload it caused, attributed to the model that ran out of memory - offloads = [event for event in manager.offload_record.events if event.action == "offload"] - assert len(offloads) == 1 - offload_event = offloads[0] - assert offload_event.model_size == UNIT - assert offload_event.model_id.rsplit("_", 1)[0] == "smallest" - assert offload_event.reason == f"oom_retry:{manager.model_hooks[0].model_id}" - - def test_add_does_not_rebuild_existing_hooks(self): - model_a = DummyModel() - manager = self._manager({"a": model_a}) - hooks_before = list(manager.model_hooks) - - model_b = DummyModel() - manager.add("b", model_b) - - assert len(manager.model_hooks) == 2 - assert manager.model_hooks[0] is hooks_before[0], "existing hook was rebuilt" - hook_a, hook_b = manager.model_hooks - assert hook_b.hook.other_hooks == [hook_a] - assert hook_a.hook.other_hooks == [hook_b] - - def test_remove_detaches_only_that_hook(self): - model_a = DummyModel() - model_b = DummyModel() - manager = self._manager({"a": model_a, "b": model_b}) - hook_a = manager.model_hooks[0] - b_id = [cid for cid in manager.components if cid.startswith("b_")][0] - - manager.remove(b_id) - - assert manager.model_hooks == [hook_a] - assert hook_a.hook.other_hooks == [] - assert not hasattr(model_b, "_hf_hook") - # model a still hooked and functional - assert model_a(torch.zeros(2, 4)).shape == (2, 4) - class ModularPipelineOffloadTesterMixin: """ Auto-CPU-offload tests for a modular pipeline's components. """ + @staticmethod + def _peak_co_residency(events): + """Largest number of models simultaneously on the device, replayed from the recorded moves.""" + resident, peak = set(), 0 + for event in events: + if event.action == "onload": + resident.add(event.model_id) + peak = max(peak, len(resident)) + else: + resident.discard(event.model_id) + return peak + @staticmethod def _managed_models(cm): """The registered components that the offloader actually manages (parameterized @@ -720,7 +846,7 @@ def _run_offloaded(self, free_bytes): pipe = self.get_pipeline(components_manager=cm) cm.enable_auto_cpu_offload(device=torch_device, memory_reserve=0) - with _patch_free_memory(free_bytes): + with _patch_memory_stats(free_bytes): output = pipe(**self.get_dummy_inputs(), output=self.output_name) return cm, list(cm.offload_record.events), output @@ -739,7 +865,7 @@ def test_auto_cpu_offload_serializes_models_under_memory_pressure(self): assert any(event.action == "offload" for event in events), "expected at least one offload" # Sequencing: models run one at a time, never two co-resident on the device. - peak = _peak_co_residency(events) + peak = self._peak_co_residency(events) assert peak == 1, f"expected serialized execution under pressure, saw {peak} models co-resident" # Device placement after the run: at most the last-run model stays on the @@ -766,7 +892,7 @@ def test_auto_cpu_offload_keeps_models_resident_without_memory_pressure(self): assert all(event.action == "onload" for event in events), "no model should be offloaded" # ...and models accumulate on the device instead of being serialized. - peak = _peak_co_residency(events) + peak = self._peak_co_residency(events) assert peak >= 2, f"expected models to co-reside without pressure, saw peak {peak}" models = self._managed_models(cm) @@ -788,81 +914,3 @@ def test_auto_cpu_offload_inference_consistent_under_memory_pressure(self, expec assert max_diff < expected_max_diff, f"offloaded output diverged from baseline (max diff {max_diff})" finally: cm.disable_auto_cpu_offload() - - -@require_accelerator -class TestSimulateAcceleratorMemory: - """ - Validates the `simulate_accelerator_memory` testing util itself against the real device: unlike - `_patch_memory_stats` (which pins scripted readings for hardware-independent unit tests), the - util wraps `mem_get_info` so real allocations show up live on a simulated smaller card, and - hard-caps the allocator so overshooting the simulated capacity raises a real OOM. - """ - - MB = 2**20 - - def setup_method(self): - gc.collect() - backend_empty_cache(torch_device) - - def teardown_method(self): - gc.collect() - backend_empty_cache(torch_device) - - @staticmethod - def _device_module(): - return getattr(torch, torch.device(torch_device).type) - - def test_readings_translate_to_simulated_card(self): - device_module = self._device_module() - free, real_total = device_module.mem_get_info() - # Simulate a card with 64MB of headroom on top of whatever the device currently holds. - sim_total = (real_total - free) + 64 * self.MB - - with simulate_accelerator_memory(sim_total, device=torch_device, hard=False): - free0, total0 = device_module.mem_get_info() - assert total0 == sim_total - assert free0 <= 64 * self.MB - - # A real allocation is visible on the simulated card: free drops by at least its size. - x = torch.zeros(8 * self.MB, dtype=torch.float32, device=torch_device) # 32MB - free1, total1 = device_module.mem_get_info() - assert total1 == sim_total - assert free1 <= free0 - 32 * self.MB - del x - - # Exiting the context restores the real readings. - assert device_module.mem_get_info()[1] == real_total - - def test_run_ooms_when_simulated_card_is_too_small(self): - device_module = self._device_module() - model = DummyModel(footprint_bytes=64 * self.MB) - x = torch.randn(4, device=torch_device) - - # Measure what the run actually needs on the device (weights + forward). - device_module.reset_peak_memory_stats() - model.to(torch_device) - model(x) - requirement = device_module.max_memory_allocated() - model.to("cpu") - backend_empty_cache(torch_device) - - try: - # Control: a simulated card with comfortable headroom runs the model fine. - free, real_total = device_module.mem_get_info() - with simulate_accelerator_memory((real_total - free) + 2 * requirement, device=torch_device): - model.to(torch_device) - model(x) - model.to("cpu") - backend_empty_cache(torch_device) - - # A card with only half the requirement left cannot: the hard cap turns the - # overshoot into a real device OOM instead of using the extra physical memory. - free, real_total = device_module.mem_get_info() - with simulate_accelerator_memory((real_total - free) + requirement // 2, device=torch_device): - with pytest.raises(torch.OutOfMemoryError): - model.to(torch_device) - model(x) - finally: - model.to("cpu") - backend_empty_cache(torch_device)