diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 31ba2c942203..ee1f08bfe05f 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -146,6 +146,13 @@ def custom_offload_with_hook( return user_hook +# search_best_candidate tries every combination of models to find the smallest one +# that frees enough memory, so it's O(2^n). fine for a few models, but ComponentsManager +# is meant to hold components from multiple pipelines at once so n can get big. past +# this many candidates just do a greedy pick instead of grinding through all of them. +_MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES = 20 + + # this is the class that user can customize to implement their own offload strategy class AutoOffloadStrategy: """ @@ -198,6 +205,26 @@ def search_best_candidate(module_sizes, min_memory_offload): larger than `min_memory_offload` """ model_ids = list(module_sizes.keys()) + + if len(model_ids) > _MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES: + # too many to brute-force all combinations, so fall back to a greedy + # pick instead. module_sizes is sorted largest-to-smallest (for the + # exhaustive path below), but per the TODO above we want to keep + # bigger models on GPU when we can, so walk it smallest-to-largest + # here and evict small models first. + logger.warning( + f"{len(model_ids)} candidate models to offload exceeds " + f"{_MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES}, using a greedy search instead of an exhaustive one" + ) + greedy_candidate = [] + greedy_size = 0 + for model_id in reversed(model_ids): + if greedy_size >= min_memory_offload: + break + greedy_candidate.append(model_id) + greedy_size += module_sizes[model_id] + return tuple(greedy_candidate) if greedy_size >= min_memory_offload else None + best_candidate = None best_size = float("inf") for r in range(1, len(model_ids) + 1): diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index 42929b04e3d8..5001e16c186a 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -13,6 +13,7 @@ # limitations under the License. import gc +import time from unittest import mock import pytest @@ -26,7 +27,10 @@ if is_accelerate_available(): - from diffusers.modular_pipelines.components_manager import AutoOffloadStrategy + from diffusers.modular_pipelines.components_manager import ( + _MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES, + AutoOffloadStrategy, + ) # The offload logic deals in bytes. We keep the test models tiny (a few KB of real @@ -209,6 +213,51 @@ def test_strategy_memory_reserve_margin_changes_decision(self): memory_reserve_margin=3 * UNIT, ) == ["c"] + def test_strategy_uses_greedy_fallback_above_exhaustive_search_limit(self): + # past _MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES, search_best_candidate stops trying + # every combination and grabs models greedily instead. check it still frees + # enough memory and does it fast - 28 models is 2^28 combos the old way. + num_hooks = _MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES + 4 + hook_sizes = {f"model_{i}": UNIT for i in range(num_hooks)} + + start = time.perf_counter() + selected = self._select_offload( + incoming_footprint=15 * UNIT, + free_bytes=10 * UNIT, + hook_sizes=hook_sizes, + memory_reserve_margin=0, + ) + elapsed = time.perf_counter() - start + + # usable = 10 - 0 = 10, incoming needs 15 -> must free 5. All 28 models are the + # same size here so direction doesn't matter for this assertion, just that it + # picks the right count and stays fast - see the next test for direction. + assert len(selected) == 5 + assert sum(hook_sizes[model_id] for model_id in selected) == 5 * UNIT + assert elapsed < 1.0 + + def test_strategy_greedy_fallback_evicts_smallest_models_first(self): + # the greedy fallback should keep bigger models resident and evict smaller + # ones first - matches the YiYi/Dhruv TODO above search_best_candidate ("we + # would tend to keep the larger models on GPU more often"). Use one big model + # plus enough small ones to push past the exhaustive-search cap. + num_small = _MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES + 3 + hook_sizes = {"big": 5 * UNIT} + hook_sizes.update({f"small_{i}": UNIT for i in range(num_small)}) + + selected = self._select_offload( + incoming_footprint=8 * UNIT, + free_bytes=5 * UNIT, + hook_sizes=hook_sizes, + memory_reserve_margin=0, + ) + + # usable = 5, incoming needs 8 -> must free 3. Three small models cover that + # on their own, so "big" should be left alone. + assert "big" not in selected + assert len(selected) == 3 + assert all(model_id.startswith("small_") for model_id in selected) + # ------------------------------------------------------------------ # Registry tests (hardware-independent) # ------------------------------------------------------------------