diff --git a/invokeai/app/services/session_processor/session_processor_default.py b/invokeai/app/services/session_processor/session_processor_default.py index 38563f772eb..4ab5e3b0703 100644 --- a/invokeai/app/services/session_processor/session_processor_default.py +++ b/invokeai/app/services/session_processor/session_processor_default.py @@ -42,7 +42,7 @@ from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context from invokeai.app.util.profiler import Profiler from invokeai.backend.util.device_pool import GENERATION_DEVICE_POOL -from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.devices import TorchDevice, disable_conv_benchmark_empty_cache # A failed owner lookup is retried before the item is refused, so that a transient error # — a busy-timeout on the shared SQLite connection under multi-GPU write contention, say — @@ -618,6 +618,12 @@ def start(self, invoker: Invoker) -> None: # encoders on (see offload_text_encoders_to_idle_gpus). None means legacy single-device mode. GENERATION_DEVICE_POOL.set_generation_devices([d for d in devices if d is not None]) + # With more than one CUDA/HIP generation device, torch's post-conv-algorithm-search global + # emptyCache() convoys the peer GPU's in-flight step from C++, where the peer-aware + # empty_cache wrapper cannot intercept it. Trade it for cached workspace blocks instead. + if sum(1 for d in devices if d is not None and d.type == "cuda") > 1: + disable_conv_benchmark_empty_cache() + # If profiling is enabled, create a profiler. The same profiler will be used for all sessions. Internally, # the profiler will create a new profile for each session. Profiling uses a process-global cProfile, which # cannot cleanly attribute work when multiple sessions run concurrently, so it is disabled in multi-GPU mode. diff --git a/invokeai/backend/util/devices.py b/invokeai/backend/util/devices.py index 23e1c336756..3d55c3fe950 100644 --- a/invokeai/backend/util/devices.py +++ b/invokeai/backend/util/devices.py @@ -498,3 +498,34 @@ def peer_aware_empty_cache() -> None: setattr(peer_aware_empty_cache, _PEER_AWARE_SENTINEL, True) torch.cuda.empty_cache = peer_aware_empty_cache + + +def disable_conv_benchmark_empty_cache() -> None: + """Stop torch's conv algorithm search from calling global emptyCache() after each find. + + When torch searches for a convolution algorithm (cuDNN benchmark mode; MIOpen on ROCm on + every algo-cache miss, benchmark flag or not), ``findAlgorithm`` ends the search with a + process-global ``CUDACachingAllocator::emptyCache()`` to release benchmarking workspace + (``aten/src/ATen/native/miopen/Conv_miopen.cpp``, likewise the cuDNN path). That call takes + every CUDA/HIP device's allocator mutex and frees their cached blocks — freezing a peer + GPU's worker mid-step exactly like the Python-level ``torch.cuda.empty_cache`` calls + handled by ``install_peer_aware_empty_cache``, but from C++, out of reach of that wrapper + (observed via py-spy on a dual-GPU ROCm rig: one worker inside + ``chooseAlgorithm -> emptyCache -> hipFree`` waiting out the other worker's 40-100 s + denoise step; each new conv shape in the process re-triggers it). + + The call is gated on ``_cudnn_get_conv_benchmark_empty_cache()``, which torch exposes a + setter for. Disabling it leaves the benchmarking workspace blocks cached in the allocator + for reuse instead of returning them to the driver — the same trade the peer-aware skips + already make everywhere else. Called only on multi-GPU installs (see + ``DefaultSessionProcessor.start``); single-GPU installs keep torch's default behavior. + + No-op on torch builds that lack the flag (e.g. CPU-only builds). + """ + setter = getattr(torch._C, "_cudnn_set_conv_benchmark_empty_cache", None) + if setter is None: + return + setter(False) + InvokeAILogger.get_logger("TorchDevice").debug( + "Disabled torch's post-conv-algorithm-search global emptyCache (multi-GPU install)." + ) diff --git a/tests/backend/util/test_devices.py b/tests/backend/util/test_devices.py index 57ac821f390..ac6a7642466 100644 --- a/tests/backend/util/test_devices.py +++ b/tests/backend/util/test_devices.py @@ -714,3 +714,27 @@ def test_install_peer_aware_empty_cache_wraps_torch_entry_point(monkeypatch): # monkeypatch restores the attribute we set; make sure the true original is back for # other tests regardless of ordering. torch_mod.cuda.empty_cache = original + + +def test_disable_conv_benchmark_empty_cache_flips_torch_flag(): + """The multi-GPU startup path must clear torch's post-conv-find emptyCache flag (and no-op + gracefully on builds that lack it).""" + from invokeai.backend.util.devices import disable_conv_benchmark_empty_cache + + getter = getattr(torch._C, "_cuda_get_conv_benchmark_empty_cache", None) + setter = getattr(torch._C, "_cudnn_set_conv_benchmark_empty_cache", None) + if getter is None or setter is None: + # CPU-only torch builds lack the flag; the function must still be a safe no-op. + disable_conv_benchmark_empty_cache() + return + + original = getter() + try: + setter(True) + disable_conv_benchmark_empty_cache() + assert getter() is False + # Idempotent: a second call keeps it disabled without raising. + disable_conv_benchmark_empty_cache() + assert getter() is False + finally: + setter(original)