From 4d4bddd2a622886ef11b5fe259f60d08f4b57743 Mon Sep 17 00:00:00 2001 From: Sayak Mondal Date: Wed, 29 Jul 2026 09:20:56 +0530 Subject: [PATCH 1/4] modular_pipelines/components_manager: cap AutoOffloadStrategy's exhaustive offload search --- .../modular_pipelines/components_manager.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 31ba2c942203..de8bd38b69f9 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,24 @@ 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 take the biggest ones + # first (module_sizes is already sorted that way) until we've freed + # enough. not guaranteed to be the smallest possible set, but it's fast. + 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 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): From b634bc2c82eb9d5a62d1f69bd021d3b7525fa9e2 Mon Sep 17 00:00:00 2001 From: Sayak Mondal Date: Wed, 29 Jul 2026 09:22:41 +0530 Subject: [PATCH 2/4] modular_pipelines/components_manager: cap AutoOffloadStrategy's exhaustive offload search --- .../test_components_manager.py | 29 ++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index 42929b04e3d8..03cfa8339b40 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,29 @@ 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 just grabs the biggest models first. 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. The greedy fallback + # walks `module_sizes` (already largest-first; here all tied at UNIT) and stops + # as soon as it has freed enough, so it should pick exactly 5 of the 28 models. + assert len(selected) == 5 + assert sum(hook_sizes[model_id] for model_id in selected) == 5 * UNIT + assert elapsed < 1.0 + # ------------------------------------------------------------------ # Registry tests (hardware-independent) # ------------------------------------------------------------------ From ea6456f1ed4686dd2ee25198bd3658b82e726793 Mon Sep 17 00:00:00 2001 From: Dev-X25874 Date: Wed, 29 Jul 2026 15:39:11 +0530 Subject: [PATCH 3/4] modular_pipelines/components_manager: cap AutoOffloadStrategy's exhaustive offload search --- .../modular_pipelines/components_manager.py | 25 ++++++++++++++++ .../test_components_manager.py | 29 ++++++++++++++++++- 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index 31ba2c942203..de8bd38b69f9 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,24 @@ 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 take the biggest ones + # first (module_sizes is already sorted that way) until we've freed + # enough. not guaranteed to be the smallest possible set, but it's fast. + 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 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..03cfa8339b40 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,29 @@ 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 just grabs the biggest models first. 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. The greedy fallback + # walks `module_sizes` (already largest-first; here all tied at UNIT) and stops + # as soon as it has freed enough, so it should pick exactly 5 of the 28 models. + assert len(selected) == 5 + assert sum(hook_sizes[model_id] for model_id in selected) == 5 * UNIT + assert elapsed < 1.0 + # ------------------------------------------------------------------ # Registry tests (hardware-independent) # ------------------------------------------------------------------ From 594d2bbc43a92ca2f57c025b416e294ddb344080 Mon Sep 17 00:00:00 2001 From: Dev-X25874 Date: Thu, 30 Jul 2026 08:23:33 +0530 Subject: [PATCH 4/4] modular_pipelines/components_manager: evict smallest models first in greedy offload fallback --- .../modular_pipelines/components_manager.py | 10 +++--- .../test_components_manager.py | 32 ++++++++++++++++--- 2 files changed, 33 insertions(+), 9 deletions(-) diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index de8bd38b69f9..ee1f08bfe05f 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -207,16 +207,18 @@ def search_best_candidate(module_sizes, 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 take the biggest ones - # first (module_sizes is already sorted that way) until we've freed - # enough. not guaranteed to be the smallest possible set, but it's fast. + # 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 model_ids: + for model_id in reversed(model_ids): if greedy_size >= min_memory_offload: break greedy_candidate.append(model_id) diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index 03cfa8339b40..5001e16c186a 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -215,8 +215,8 @@ def test_strategy_memory_reserve_margin_changes_decision(self): def test_strategy_uses_greedy_fallback_above_exhaustive_search_limit(self): # past _MAX_EXHAUSTIVE_OFFLOAD_CANDIDATES, search_best_candidate stops trying - # every combination and just grabs the biggest models first. check it still - # frees enough memory and does it fast - 28 models is 2^28 combos the old way. + # 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)} @@ -229,13 +229,35 @@ def test_strategy_uses_greedy_fallback_above_exhaustive_search_limit(self): ) elapsed = time.perf_counter() - start - # usable = 10 - 0 = 10, incoming needs 15 -> must free 5. The greedy fallback - # walks `module_sizes` (already largest-first; here all tied at UNIT) and stops - # as soon as it has freed enough, so it should pick exactly 5 of the 28 models. + # 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) # ------------------------------------------------------------------