From bc170db86439d18d7fcf66b7ed0d1a492f2c78ca Mon Sep 17 00:00:00 2001 From: John Pollock Date: Fri, 8 May 2026 17:11:31 +0000 Subject: [PATCH 1/5] Fix aimdo device fallback for MultiGPU --- README.md | 2 + __init__.py | 107 ++++++++++++++++++++++++++++++++++++++++++++++++- pyproject.toml | 2 +- 3 files changed, 109 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index a336155..5ae7357 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,8 @@ What is DisTorch? Standing for "distributed torch", the DisTorch nodes in this c ## 🚀 Compatibility Works with all .safetensors and GGUF-quantized models. +On current ComfyUI builds with DynamicVRAM/comfy-aimdo enabled, MultiGPU keeps DynamicVRAM active on CUDA devices that comfy-aimdo has initialized and falls back to legacy model patching for off-grid MultiGPU CUDA devices. This preserves MultiGPU placement for devices such as `cuda:1` even when comfy-aimdo only initialized the primary device. + ⚙️ Expert users: Like .gguf or exl2/3 LLM loaders, use the expert_mode_alloaction for exact allocations of model shards on as many devices as your setup has!

diff --git a/__init__.py b/__init__.py index e536933..46ee1bc 100644 --- a/__init__.py +++ b/__init__.py @@ -176,6 +176,7 @@ current_unet_offload_device = mm.unet_offload_device() _aimdo_initialized_devices = set() if isinstance(current_device, torch.device) and current_device.type == "cuda" and current_device.index is not None: _aimdo_initialized_devices.add(current_device.index) +_aimdo_legacy_fallback_devices = set() def set_current_device(device): """Set the current device context for MultiGPU operations.""" @@ -330,8 +331,83 @@ def _patch_model_management_current_stream(): logger.info("[MultiGPU] Patched comfy.model_management.current_stream to honor CUDA device arguments") return True +def _aimdo_device_ready(device): + """Return whether comfy-aimdo has a device context for a CUDA device.""" + if not getattr(comfy.memory_management, "aimdo_enabled", False): + return False + + target_device = _coerce_torch_device(device) + if target_device is None or target_device.type != "cuda" or target_device.index is None: + return False + + try: + from comfy_aimdo import control as aimdo_control + get_devctx = getattr(aimdo_control, "get_devctx", None) + if not callable(get_devctx): + return False + get_devctx(target_device.index) + return True + except Exception: + return False + +def _patch_dynamic_model_patcher_for_aimdo_devices(): + """Route DynamicVRAM only to CUDA devices aimdo actually initialized.""" + dynamic_patcher = getattr(comfy.model_patcher, "ModelPatcherDynamic", None) + if dynamic_patcher is None: + return False + dynamic_new = getattr(dynamic_patcher, "__new__", None) + if not callable(dynamic_new): + return False + if getattr(dynamic_new, "_multigpu_aimdo_device_guard", False): + return False + + def dynamic_patcher_new_with_aimdo_device_guard( + cls, + model=None, + load_device=None, + offload_device=None, + size=0, + weight_inplace_update=False, + ): + target_device = _coerce_torch_device(load_device) + if ( + getattr(comfy.memory_management, "aimdo_enabled", False) + and target_device is not None + and target_device.type == "cuda" + and target_device.index is not None + and not _aimdo_device_ready(target_device) + ): + if target_device.index not in _aimdo_legacy_fallback_devices: + logger.warning( + f"[MultiGPU] comfy_aimdo has no context for {target_device}; " + "using legacy ModelPatcher for this device" + ) + _aimdo_legacy_fallback_devices.add(target_device.index) + return comfy.model_patcher.ModelPatcher( + model, + load_device, + offload_device, + size, + weight_inplace_update, + ) + + return dynamic_new( + cls, + model, + load_device, + offload_device, + size, + weight_inplace_update, + ) + + dynamic_patcher_new_with_aimdo_device_guard._multigpu_aimdo_device_guard = True + dynamic_patcher_new_with_aimdo_device_guard._multigpu_original = dynamic_new + dynamic_patcher.__new__ = staticmethod(dynamic_patcher_new_with_aimdo_device_guard) + logger.info("[MultiGPU] Patched ModelPatcherDynamic to guard aimdo device coverage") + return True + def _initialize_aimdo_visible_cuda_devices(): - """Ensure DynamicVRAM initializes every visible CUDA device once when enabled.""" + """Configure MultiGPU behavior around comfy-aimdo's initialized CUDA devices.""" if not getattr(comfy.memory_management, "aimdo_enabled", False): logger.info("[MultiGPU] DynamicVRAM not enabled; skipping multi-device aimdo initialization") return False @@ -345,6 +421,35 @@ def _initialize_aimdo_visible_cuda_devices(): logger.warning("[MultiGPU] comfy_aimdo unavailable during multi-device initialization") return False + device_ids = list(range(torch.cuda.device_count())) + ready_device_ids = [] + for device_id in device_ids: + if _aimdo_device_ready(torch.device("cuda", device_id)): + ready_device_ids.append(device_id) + + _aimdo_initialized_devices.clear() + _aimdo_initialized_devices.update(ready_device_ids) + + if len(ready_device_ids) == len(device_ids): + _aimdo_initialized_devices.update(device_ids) + logger.info(f"[MultiGPU] comfy_aimdo already initialized for CUDA devices {device_ids}") + return False + + if ready_device_ids: + logger.warning( + "[MultiGPU] comfy_aimdo initialized CUDA devices " + f"{ready_device_ids}, not every visible device {device_ids}; " + "leaving initialized devices on DynamicVRAM and falling back per-device elsewhere" + ) + else: + logger.warning( + "[MultiGPU] comfy_aimdo has no initialized CUDA devices; " + "falling back to legacy ModelPatcher per device" + ) + _patch_dynamic_model_patcher_for_aimdo_devices() + if len(device_ids) > 1: + return False + init_device = getattr(aimdo_control, "init_device", None) if not callable(init_device): logger.warning("[MultiGPU] comfy_aimdo.control.init_device missing; skipping multi-device initialization") diff --git a/pyproject.toml b/pyproject.toml index 8a2f7f9..211c5a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-multigpu" description = "Provides a suite of custom nodes to manage multiple GPUs for ComfyUI, including advanced model offloading for both GGUF and Safetensor formats with DisTorch, and bespoke MultiGPU support for WanVideoWrapper and other custom nodes." -version = "2.6.3" +version = "2.6.4" license = {file = "LICENSE"} [project.urls] From f5d9276ae3ed477a4f910cd546cc87fb64141b34 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Fri, 8 May 2026 18:06:55 +0000 Subject: [PATCH 2/5] Address PR review aimdo guard feedback --- __init__.py | 63 +++++++++++++++++++++++++++-------------------------- 1 file changed, 32 insertions(+), 31 deletions(-) diff --git a/__init__.py b/__init__.py index 46ee1bc..6a4537a 100644 --- a/__init__.py +++ b/__init__.py @@ -332,7 +332,7 @@ def _patch_model_management_current_stream(): return True def _aimdo_device_ready(device): - """Return whether comfy-aimdo has a device context for a CUDA device.""" + """Return True/False for known aimdo readiness, or None when it cannot be probed.""" if not getattr(comfy.memory_management, "aimdo_enabled", False): return False @@ -344,11 +344,29 @@ def _aimdo_device_ready(device): from comfy_aimdo import control as aimdo_control get_devctx = getattr(aimdo_control, "get_devctx", None) if not callable(get_devctx): - return False + logger.warning("[MultiGPU] comfy_aimdo.control.get_devctx missing; leaving DynamicVRAM unguarded") + return None get_devctx(target_device.index) return True - except Exception: - return False + except RuntimeError as exc: + if "not initialized" in str(exc).lower(): + return False + logger.warning(f"[MultiGPU] Unexpected comfy_aimdo readiness failure for {target_device}: {exc}") + return None + except Exception as exc: + logger.warning(f"[MultiGPU] Unexpected comfy_aimdo readiness failure for {target_device}: {exc}") + return None + +def _extract_model_patcher_load_device(args, kwargs): + """Resolve the effective ModelPatcher load_device from current and future call shapes.""" + if "load_device" in kwargs: + return kwargs["load_device"] + if len(args) >= 2: + return args[1] + get_torch_device = getattr(mm, "get_torch_device", None) + if callable(get_torch_device): + return get_torch_device() + return None def _patch_dynamic_model_patcher_for_aimdo_devices(): """Route DynamicVRAM only to CUDA devices aimdo actually initialized.""" @@ -361,21 +379,15 @@ def _patch_dynamic_model_patcher_for_aimdo_devices(): if getattr(dynamic_new, "_multigpu_aimdo_device_guard", False): return False - def dynamic_patcher_new_with_aimdo_device_guard( - cls, - model=None, - load_device=None, - offload_device=None, - size=0, - weight_inplace_update=False, - ): - target_device = _coerce_torch_device(load_device) + def dynamic_patcher_new_with_aimdo_device_guard(cls, *args, **kwargs): + target_device = _coerce_torch_device(_extract_model_patcher_load_device(args, kwargs)) + aimdo_ready = _aimdo_device_ready(target_device) if ( getattr(comfy.memory_management, "aimdo_enabled", False) and target_device is not None and target_device.type == "cuda" and target_device.index is not None - and not _aimdo_device_ready(target_device) + and aimdo_ready is False ): if target_device.index not in _aimdo_legacy_fallback_devices: logger.warning( @@ -383,22 +395,9 @@ def _patch_dynamic_model_patcher_for_aimdo_devices(): "using legacy ModelPatcher for this device" ) _aimdo_legacy_fallback_devices.add(target_device.index) - return comfy.model_patcher.ModelPatcher( - model, - load_device, - offload_device, - size, - weight_inplace_update, - ) + return comfy.model_patcher.ModelPatcher(*args, **kwargs) - return dynamic_new( - cls, - model, - load_device, - offload_device, - size, - weight_inplace_update, - ) + return dynamic_new(cls, *args, **kwargs) dynamic_patcher_new_with_aimdo_device_guard._multigpu_aimdo_device_guard = True dynamic_patcher_new_with_aimdo_device_guard._multigpu_original = dynamic_new @@ -424,14 +423,16 @@ def _initialize_aimdo_visible_cuda_devices(): device_ids = list(range(torch.cuda.device_count())) ready_device_ids = [] for device_id in device_ids: - if _aimdo_device_ready(torch.device("cuda", device_id)): + ready = _aimdo_device_ready(torch.device("cuda", device_id)) + if ready is None: + return False + if ready: ready_device_ids.append(device_id) _aimdo_initialized_devices.clear() _aimdo_initialized_devices.update(ready_device_ids) if len(ready_device_ids) == len(device_ids): - _aimdo_initialized_devices.update(device_ids) logger.info(f"[MultiGPU] comfy_aimdo already initialized for CUDA devices {device_ids}") return False From 013d1113a3129d49991dd01499952a243ffb10eb Mon Sep 17 00:00:00 2001 From: John Pollock Date: Fri, 8 May 2026 18:38:27 +0000 Subject: [PATCH 3/5] Address aimdo guard follow-up --- __init__.py | 29 +++++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/__init__.py b/__init__.py index 6a4537a..d043552 100644 --- a/__init__.py +++ b/__init__.py @@ -176,6 +176,7 @@ current_unet_offload_device = mm.unet_offload_device() _aimdo_initialized_devices = set() if isinstance(current_device, torch.device) and current_device.type == "cuda" and current_device.index is not None: _aimdo_initialized_devices.add(current_device.index) +_aimdo_readiness_cache = {} _aimdo_legacy_fallback_devices = set() def set_current_device(device): @@ -339,22 +340,29 @@ def _aimdo_device_ready(device): target_device = _coerce_torch_device(device) if target_device is None or target_device.type != "cuda" or target_device.index is None: return False + if target_device.index in _aimdo_readiness_cache: + return _aimdo_readiness_cache[target_device.index] try: from comfy_aimdo import control as aimdo_control get_devctx = getattr(aimdo_control, "get_devctx", None) if not callable(get_devctx): - logger.warning("[MultiGPU] comfy_aimdo.control.get_devctx missing; leaving DynamicVRAM unguarded") + logger.warning("[MultiGPU] comfy_aimdo.control.get_devctx missing; device readiness is unverified") + _aimdo_readiness_cache[target_device.index] = None return None get_devctx(target_device.index) + _aimdo_readiness_cache[target_device.index] = True return True except RuntimeError as exc: if "not initialized" in str(exc).lower(): + _aimdo_readiness_cache[target_device.index] = False return False logger.warning(f"[MultiGPU] Unexpected comfy_aimdo readiness failure for {target_device}: {exc}") + _aimdo_readiness_cache[target_device.index] = None return None except Exception as exc: logger.warning(f"[MultiGPU] Unexpected comfy_aimdo readiness failure for {target_device}: {exc}") + _aimdo_readiness_cache[target_device.index] = None return None def _extract_model_patcher_load_device(args, kwargs): @@ -387,13 +395,15 @@ def _patch_dynamic_model_patcher_for_aimdo_devices(): and target_device is not None and target_device.type == "cuda" and target_device.index is not None - and aimdo_ready is False + and aimdo_ready is not True ): if target_device.index not in _aimdo_legacy_fallback_devices: - logger.warning( - f"[MultiGPU] comfy_aimdo has no context for {target_device}; " - "using legacy ModelPatcher for this device" + reason = ( + "no context" + if aimdo_ready is False + else "unverified context" ) + logger.warning(f"[MultiGPU] comfy_aimdo has {reason} for {target_device}; using legacy ModelPatcher") _aimdo_legacy_fallback_devices.add(target_device.index) return comfy.model_patcher.ModelPatcher(*args, **kwargs) @@ -422,10 +432,12 @@ def _initialize_aimdo_visible_cuda_devices(): device_ids = list(range(torch.cuda.device_count())) ready_device_ids = [] + unverified_device_ids = [] for device_id in device_ids: ready = _aimdo_device_ready(torch.device("cuda", device_id)) if ready is None: - return False + unverified_device_ids.append(device_id) + continue if ready: ready_device_ids.append(device_id) @@ -442,6 +454,11 @@ def _initialize_aimdo_visible_cuda_devices(): f"{ready_device_ids}, not every visible device {device_ids}; " "leaving initialized devices on DynamicVRAM and falling back per-device elsewhere" ) + elif unverified_device_ids: + logger.warning( + "[MultiGPU] comfy_aimdo readiness could not be verified for CUDA devices " + f"{unverified_device_ids}; falling back to legacy ModelPatcher per device" + ) else: logger.warning( "[MultiGPU] comfy_aimdo has no initialized CUDA devices; " From 2ad5af11200a43c25f594e1ddfd28ddda6ab555b Mon Sep 17 00:00:00 2001 From: John Pollock Date: Fri, 8 May 2026 18:45:08 +0000 Subject: [PATCH 4/5] Refresh aimdo readiness after init --- __init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/__init__.py b/__init__.py index d043552..1ad2674 100644 --- a/__init__.py +++ b/__init__.py @@ -480,8 +480,10 @@ def _initialize_aimdo_visible_cuda_devices(): logger.info(f"[MultiGPU] Initializing comfy_aimdo for CUDA device {device_index}") initialized = bool(init_device(device_index)) logger.info(f"[MultiGPU] comfy_aimdo init_device({device_index}) -> {initialized}") + _aimdo_readiness_cache[device_index] = initialized if initialized: _aimdo_initialized_devices.add(device_index) + _aimdo_legacy_fallback_devices.discard(device_index) initialized_any = True return initialized_any From f6521622fba055163b25163c68af9cd61edc9cbc Mon Sep 17 00:00:00 2001 From: John Pollock Date: Fri, 8 May 2026 18:50:23 +0000 Subject: [PATCH 5/5] Avoid stale aimdo negative readiness cache --- __init__.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/__init__.py b/__init__.py index 1ad2674..2d00d58 100644 --- a/__init__.py +++ b/__init__.py @@ -177,6 +177,7 @@ _aimdo_initialized_devices = set() if isinstance(current_device, torch.device) and current_device.type == "cuda" and current_device.index is not None: _aimdo_initialized_devices.add(current_device.index) _aimdo_readiness_cache = {} +_aimdo_readiness_warning_keys = set() _aimdo_legacy_fallback_devices = set() def set_current_device(device): @@ -340,14 +341,17 @@ def _aimdo_device_ready(device): target_device = _coerce_torch_device(device) if target_device is None or target_device.type != "cuda" or target_device.index is None: return False - if target_device.index in _aimdo_readiness_cache: - return _aimdo_readiness_cache[target_device.index] + if _aimdo_readiness_cache.get(target_device.index) is True: + return True try: from comfy_aimdo import control as aimdo_control get_devctx = getattr(aimdo_control, "get_devctx", None) if not callable(get_devctx): - logger.warning("[MultiGPU] comfy_aimdo.control.get_devctx missing; device readiness is unverified") + warning_key = (target_device.index, "missing_get_devctx") + if warning_key not in _aimdo_readiness_warning_keys: + logger.warning("[MultiGPU] comfy_aimdo.control.get_devctx missing; device readiness is unverified") + _aimdo_readiness_warning_keys.add(warning_key) _aimdo_readiness_cache[target_device.index] = None return None get_devctx(target_device.index) @@ -357,11 +361,17 @@ def _aimdo_device_ready(device): if "not initialized" in str(exc).lower(): _aimdo_readiness_cache[target_device.index] = False return False - logger.warning(f"[MultiGPU] Unexpected comfy_aimdo readiness failure for {target_device}: {exc}") + warning_key = (target_device.index, str(exc)) + if warning_key not in _aimdo_readiness_warning_keys: + logger.warning(f"[MultiGPU] Unexpected comfy_aimdo readiness failure for {target_device}: {exc}") + _aimdo_readiness_warning_keys.add(warning_key) _aimdo_readiness_cache[target_device.index] = None return None except Exception as exc: - logger.warning(f"[MultiGPU] Unexpected comfy_aimdo readiness failure for {target_device}: {exc}") + warning_key = (target_device.index, str(exc)) + if warning_key not in _aimdo_readiness_warning_keys: + logger.warning(f"[MultiGPU] Unexpected comfy_aimdo readiness failure for {target_device}: {exc}") + _aimdo_readiness_warning_keys.add(warning_key) _aimdo_readiness_cache[target_device.index] = None return None