Cap AutoOffloadStrategy's exhaustive offload search for large candidate counts

src/diffusers/modular_pipelines/components_manager.py

@@ -158,6 +158,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:
  """
@@ -210,6 +217,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):

tests/modular_pipelines/test_components_manager.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import gc
+import time
 from unittest import mock
 
 import torch
@@ -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
@@ -199,6 +203,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)
  # ------------------------------------------------------------------