Skip to content

Multi-threaded load of models from disk (big load time speedups & Offload to disk) (CORE-43,CORE-152,CORE-164,CORE-165,CORE-117) #13802

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
comfyanonymous merged 27 commits into from
May 21, 2026
Merged

Multi-threaded load of models from disk (big load time speedups & Offload to disk) (CORE-43,CORE-152,CORE-164,CORE-165,CORE-117) #13802

Show file tree
Hide file tree
Changes from 15 commits
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
1fe3a13
model_management: disable non-dynamic smart memory
rattus128 May 7, 2026
157965a
pinned_memory: implement with aimdo growable buffer
rattus128 May 4, 2026
b66b642
mm: use aimdo to do transfer from disk to pin
rattus128 May 4, 2026
8070cb7
Add stream host pin buffer for AIMDO casts
rattus128 May 7, 2026
1795523
remove old pin path
rattus128 May 7, 2026
8187cd7
Implement JIT pinned memory pressure
rattus128 May 7, 2026
2b927e1
LowVRAMPatch: change to two-phase visit
rattus128 May 7, 2026
8e473d7
lora: re-implement as inplace swiss-army-knife operation
rattus128 May 7, 2026
e48dace
prepare for multiple pin sets
rattus128 May 7, 2026
3a3b75a
implement pinned loras
rattus128 May 8, 2026
c395f2d
requirements: comfy-aimdo 0.4.0
rattus128 May 8, 2026
44c0a06
ops: remove unused arg
rattus128 May 11, 2026
ee927aa
ops: sync the CPU with only the offload stream activity
rattus128 May 9, 2026
d61026d
pins: implement freeing intermediate for pinned memory
rattus128 May 12, 2026
3f71781
execution: implement pin eviction on RAM presure
rattus128 May 13, 2026
3115053
implement pin registration swaps
rattus128 May 14, 2026
18a74cb
cli_args/execution: Implement lower background cache-ram threshold
rattus128 May 13, 2026
d8b4427
make default
rattus128 May 13, 2026
55197d8
bump aimdo
rattus128 May 15, 2026
ea5775c
Merge pull request #9 from rattus128/dev/threaded-loader-2-ram-cache
rattus128 May 15, 2026
0242954
model-patcher: force-cast tiny weights
rattus128 May 15, 2026
ed15d62
ops: refactor in prep for chunking
rattus128 May 15, 2026
4386563
mm: delegate pin-on-the-way to aimdo
rattus128 May 15, 2026
52a68b9
bump aimdo
rattus128 May 15, 2026
c451854
pinning updates
rattus128 May 17, 2026
09a98a9
specify hostbuf max allocation size
rattus128 May 18, 2026
c836a5d
tests: update execution tests for caching
rattus128 May 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 13 additions & 6 deletions comfy/lora.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -475,16 +475,23 @@ def calculate_weight(patches, weight, key, intermediate_dtype=torch.float32, ori

return weight

def prefetch_prepared_value(value, allocate_buffer, stream):
def prefetch_prepared_value(value, counter, destination, stream, copy):
if isinstance(value, torch.Tensor):
dest = allocate_buffer(comfy.memory_management.vram_aligned_size(value))
comfy.model_management.cast_to_gathered([value], dest, non_blocking=True, stream=stream)
size = comfy.memory_management.vram_aligned_size(value)
offset = counter[0]
counter[0] += size
if destination is None:
return value

dest = destination[offset:offset + size]
if copy:
comfy.model_management.cast_to_gathered([value], dest, non_blocking=True, stream=stream)
return comfy.memory_management.interpret_gathered_like([value], dest)[0]
elif isinstance(value, weight_adapter.WeightAdapterBase):
return type(value)(value.loaded_keys, prefetch_prepared_value(value.weights, allocate_buffer, stream))
return type(value)(value.loaded_keys, prefetch_prepared_value(value.weights, counter, destination, stream, copy))
elif isinstance(value, tuple):
return tuple(prefetch_prepared_value(item, allocate_buffer, stream) for item in value)
return tuple(prefetch_prepared_value(item, counter, destination, stream, copy) for item in value)
elif isinstance(value, list):
return [prefetch_prepared_value(item, allocate_buffer, stream) for item in value]
return [prefetch_prepared_value(item, counter, destination, stream, copy) for item in value]

return value
10 changes: 8 additions & 2 deletions comfy/memory_management.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,12 @@ def read_tensor_file_slice_into(tensor, destination):
if info.size == 0:
return True

hostbuf = getattr(destination.untyped_storage(), "_comfy_hostbuf", None)
if hostbuf is not None:
hostbuf.read_file_slice(file_obj, info.offset, info.size,
offset=destination.data_ptr() - hostbuf.get_raw_address())
return True

buf_type = ctypes.c_ubyte * info.size
view = memoryview(buf_type.from_address(destination.data_ptr()))

Expand Down Expand Up @@ -151,7 +157,7 @@ def set_ram_cache_release_state(callback, headroom):
extra_ram_release_callback = callback
RAM_CACHE_HEADROOM = max(0, int(headroom))

def extra_ram_release(target):
def extra_ram_release(target, free_active=False):
if extra_ram_release_callback is None:
return 0
return extra_ram_release_callback(target)
return extra_ram_release_callback(target, free_active=free_active)
151 changes: 84 additions & 67 deletions comfy/model_management.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
import comfy.memory_management
import comfy.utils
import comfy.quant_ops
import comfy_aimdo.host_buffer
import comfy_aimdo.vram_buffer

class VRAMState(Enum):
Expand Down Expand Up @@ -495,6 +496,10 @@ def get_torch_device_name(device):

current_loaded_models = []

DIRTY_MMAPS = set()

PIN_PRESSURE_HYSTERESIS = 128 * 1024 * 1024

def module_size(module):
module_mem = 0
sd = module.state_dict()
Expand All @@ -503,27 +508,31 @@ def module_size(module):
module_mem += t.nbytes
return module_mem

def module_mmap_residency(module, free=False):
mmap_touched_mem = 0
module_mem = 0
bounced_mmaps = set()
sd = module.state_dict()
for k in sd:
t = sd[k]
module_mem += t.nbytes
storage = t._qdata.untyped_storage() if isinstance(t, comfy.quant_ops.QuantizedTensor) else t.untyped_storage()
if not getattr(storage, "_comfy_tensor_mmap_touched", False):
continue
mmap_touched_mem += t.nbytes
if not free:
continue
storage._comfy_tensor_mmap_touched = False
mmap_obj = storage._comfy_tensor_mmap_refs[0]
if mmap_obj in bounced_mmaps:
continue
mmap_obj.bounce()
bounced_mmaps.add(mmap_obj)
return mmap_touched_mem, module_mem
def mark_mmap_dirty(storage):
mmap_refs = getattr(storage, "_comfy_tensor_mmap_refs", None)
if mmap_refs is not None:
DIRTY_MMAPS.add(mmap_refs[0])

def free_pins(size, evict_active=False):
if size <= 0:
return

for loaded_model in reversed(current_loaded_models):
model = loaded_model.model
if model is not None and model.is_dynamic() and (evict_active or not model.model.dynamic_pins[model.load_device]["active"]):
size -= model.partially_unload_ram(size)
if size <= 0:
break

@coderabbitai coderabbitai Bot May 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Confirm dynamic_pins is always populated for load_device before any code path that might
# trigger ensure_pin_budget/free_pins.
ast-grep --pattern 'dynamic_pins[$_] = $_'
rg -nP --type=py -C3 '\bdynamic_pins\b'

Repository: Comfy-Org/ComfyUI

Length of output: 3537

🏁 Script executed:

#!/bin/bash
# Understand the class structure where dynamic_pins is initialized
cd comfy
# Find the class containing the constructor with the initialization
rg -B20 'if self.load_device not in self.model.dynamic_pins:' model_patcher.py | head -40

# Check what function/method contains the problematic line 522
rg -B10 'def free_pins' model_management.py

# Find where current_loaded_models is populated
rg 'current_loaded_models.append|current_loaded_models +=' model_management.py

Repository: Comfy-Org/ComfyUI

Length of output: 1511

🏁 Script executed:

#!/bin/bash
# Find is_dynamic() method definition
rg 'def is_dynamic' comfy/

# Check where models are added to current_loaded_models with more context
rg -B5 -A5 'current_loaded_models.append|current_loaded_models +=' comfy/model_management.py | head -50

# Check if ModelPatcherDynamic is the only class with is_dynamic
rg 'class.*ModelPatcher' comfy/model_patcher.py

# Look for any guard/check before accessing dynamic_pins in other places
rg -C2 'is_dynamic\(\)' comfy/model_management.py

Repository: Comfy-Org/ComfyUI

Length of output: 1513

🏁 Script executed:

#!/bin/bash
# Get full definition of is_dynamic methods
rg -A5 'def is_dynamic' comfy/model_patcher.py

# Check the LoadedModel class and when models are added
rg -B10 -A10 'class LoadedModel' comfy/model_management.py

# Find where models are loaded and added to current_loaded_models
rg -B8 -A8 'current_loaded_models.append.*LoadedModel' comfy/model_management.py

Repository: Comfy-Org/ComfyUI

Length of output: 969

🏁 Script executed:

#!/bin/bash
# Find where models get loaded/added to current_loaded_models
rg -B15 'loaded_model = LoadedModel' comfy/model_management.py | head -50

# Check for load_model function and how models are prepared
rg 'def load_model' comfy/model_management.py

# Look for any git history or PR context - check if this is new code
git log --oneline comfy/model_management.py 2>/dev/null | head -5 || echo "Git log not available"
git diff HEAD^ comfy/model_management.py 2>/dev/null | grep -A10 -B10 "free_pins\|dynamic_pins" | head -50 || echo "Git diff not available"

Repository: Comfy-Org/ComfyUI

Length of output: 656

🏁 Script executed:

#!/bin/bash
# Check if there's any code that changes load_device after construction
rg 'load_device\s*=' comfy/model_patcher.py | head -20

# Look for any re-attachment or device change logic
rg 'attach|detach|reload' comfy/model_patcher.py | grep -i 'def '

# Check if dynamic_pins gets updated for new devices anywhere
rg -B5 -A5 'dynamic_pins\[.*\] = {' comfy/

# Look at the full context of line 522 and around it
sed -n '510,535p' comfy/model_management.py

Repository: Comfy-Org/ComfyUI

Length of output: 1454

🏁 Script executed:

#!/bin/bash
# Fix the regex and search for dynamic_pins initialization pattern
rg -C5 'dynamic_pins\[' comfy/model_patcher.py | grep -E '(def |class |dynamic_pins)'

# Look for any code that might change load_device or reuse models across devices
rg 'model\.load_device' comfy/model_management.py | head -20

# Check if there's any scenario where a model might be used with different load_devices
rg -B10 'for loaded_model in.*current_loaded_models' comfy/model_management.py | head -40

# Search for attach/detach patterns that might change device association
rg -B3 -A8 'def.*attach|def.*prepare' comfy/model_patcher.py | head -60

Repository: Comfy-Org/ComfyUI

Length of output: 3858

🏁 Script executed:

#!/bin/bash
# Check if load_device can change or if models can be reused across devices
rg -B5 -A5 'load_device.*=' comfy/model_patcher.py | grep -A8 -B8 '__init__'

# Search for any scenario where a model might be attached/prepared with a new device
rg -C10 'load_device\s*=' comfy/model_management.py

# Check the git log to understand when this code was added
git log --all --full-history -p -- comfy/model_management.py 2>/dev/null | grep -A30 'def free_pins' | head -50 || echo "Checking commit history..."

# Look for any recent changes to dynamic_pins handling
git show HEAD:comfy/model_management.py 2>/dev/null | sed -n '516,525p' || git log -1 --format=%h comfy/model_management.py

Repository: Comfy-Org/ComfyUI

Length of output: 4627

🏁 Script executed:

#!/bin/bash
# Verify if load_device is ever reassigned after __init__
rg 'self\.load_device\s*=' comfy/model_patcher.py

# Check if there's any code path where a model could have load_device changed
git log -1 --format="%H %s" comfy/model_management.py

# Verify the exact changes in this commit
git show --stat comfy/model_management.py 2>/dev/null | head -20

# Check if prepare() method modifies load_device or affects dynamic_pins initialization
rg -B5 -A15 'def prepare\(' comfy/model_patcher.py | head -40

Repository: Comfy-Org/ComfyUI

Length of output: 1494

🏁 Script executed:

#!/bin/bash
# Final check: verify there's no other code path that could create a dynamic model without full initialization
rg 'is_dynamic\s*=\s*True|class.*Dynamic' comfy/ --type py

# Confirm that every place accessing dynamic_pins for write operations has it initialized
rg 'dynamic_pins\[.*\]\s*=' comfy/

Repository: Comfy-Org/ComfyUI

Length of output: 533

Use .get() to safely access dynamic_pins[load_device] in case initialization is incomplete.

While the current initialization logic in ModelPatcherDynamic.__init__() should populate dynamic_pins[load_device] before any dynamic model reaches free_pins(), defensive access is prudent in this RAM-pressure code path. Consider:

if model is not None and model.is_dynamic():
    pin_state = model.model.dynamic_pins.get(model.load_device)
    if pin_state is not None and (evict_active or not pin_state["active"]):
        size -= model.partially_unload_ram(size)
        if size <= 0:
            break

This same pattern should also be applied at lines 1249 and 1251 in the same file for consistency.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@comfy/model_management.py` around lines 516 - 525, In free_pins(), avoid
indexing into model.model.dynamic_pins[model.load_device] directly; instead
safely .get(model.load_device) and check for None before accessing ["active"] so
incomplete initialization won't raise; update the condition in free_pins
(referencing current_loaded_models, loaded_model.model, model.is_dynamic(),
model.partially_unload_ram) and apply the same defensive .get() pattern to the
other two occurrences noted around lines where dynamic pin checks occur (the
other methods that reference model.model.dynamic_pins[...] at the end of file)
to ensure you only call partially_unload_ram when pin_state is present and its
"active" flag is evaluated.


def ensure_pin_budget(size, evict_active=False):
if MAX_PINNED_MEMORY <= 0:
return

shortfall = TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY
if shortfall <= 0:
return

free_pins(shortfall + PIN_PRESSURE_HYSTERESIS, evict_active=evict_active)

class LoadedModel:
def __init__(self, model):
Expand Down Expand Up @@ -553,9 +562,6 @@ def model(self):
def model_memory(self):
return self.model.model_size()

def model_mmap_residency(self, free=False):
return self.model.model_mmap_residency(free=free)

def model_loaded_memory(self):
return self.model.loaded_size()

Expand Down Expand Up @@ -635,15 +641,9 @@ def offloaded_memory(loaded_models, device):

EXTRA_RESERVED_VRAM = 400 * 1024 * 1024
if WINDOWS:
import comfy.windows
EXTRA_RESERVED_VRAM = 600 * 1024 * 1024 #Windows is higher because of the shared vram issue
if total_vram > (15 * 1024): # more extra reserved vram on 16GB+ cards
EXTRA_RESERVED_VRAM += 100 * 1024 * 1024
def get_free_ram():
return comfy.windows.get_free_ram()
else:
def get_free_ram():
return psutil.virtual_memory().available

if args.reserve_vram is not None:
EXTRA_RESERVED_VRAM = args.reserve_vram * 1024 * 1024 * 1024
Expand All @@ -657,7 +657,6 @@ def minimum_inference_memory():

def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins_required=0, ram_required=0):
cleanup_models_gc()
comfy.memory_management.extra_ram_release(max(pins_required, ram_required))
unloaded_model = []
can_unload = []
unloaded_models = []
Expand All @@ -673,30 +672,16 @@ def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins
for x in can_unload_sorted:
i = x[-1]
memory_to_free = 1e32
pins_to_free = 1e32
if not DISABLE_SMART_MEMORY or device is None:
if current_loaded_models[i].model.is_dynamic() and (not DISABLE_SMART_MEMORY or device is None):
memory_to_free = 0 if device is None else memory_required - get_free_memory(device)
pins_to_free = pins_required - get_free_ram()
if current_loaded_models[i].model.is_dynamic() and for_dynamic:
if for_dynamic:
#don't actually unload dynamic models for the sake of other dynamic models
#as that works on-demand.
memory_required -= current_loaded_models[i].model.loaded_size()
memory_to_free = 0
if memory_to_free > 0 and current_loaded_models[i].model_unload(memory_to_free):
logging.debug(f"Unloading {current_loaded_models[i].model.model.__class__.__name__}")
unloaded_model.append(i)
if pins_to_free > 0:
logging.debug(f"PIN Unloading {current_loaded_models[i].model.model.__class__.__name__}")
current_loaded_models[i].model.partially_unload_ram(pins_to_free)

for x in can_unload_sorted:
i = x[-1]
ram_to_free = ram_required - psutil.virtual_memory().available
if ram_to_free <= 0 and i not in unloaded_model:
continue
resident_memory, _ = current_loaded_models[i].model_mmap_residency(free=True)
if resident_memory > 0:
logging.debug(f"RAM Unloading {current_loaded_models[i].model.model.__class__.__name__}")

for i in sorted(unloaded_model, reverse=True):
unloaded_models.append(current_loaded_models.pop(i))
Expand Down Expand Up @@ -762,29 +747,16 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
model_to_unload.model.detach(unpatch_all=False)
model_to_unload.model_finalizer.detach()


total_memory_required = {}
total_pins_required = {}
total_ram_required = {}
for loaded_model in models_to_load:
device = loaded_model.device
total_memory_required[device] = total_memory_required.get(device, 0) + loaded_model.model_memory_required(device)
resident_memory, model_memory = loaded_model.model_mmap_residency()
pinned_memory = loaded_model.model.pinned_memory_size()
#FIXME: This can over-free the pins as it budgets to pin the entire model. We should
#make this JIT to keep as much pinned as possible.
pins_required = model_memory - pinned_memory
ram_required = model_memory - resident_memory
total_pins_required[device] = total_pins_required.get(device, 0) + pins_required
total_ram_required[device] = total_ram_required.get(device, 0) + ram_required

for device in total_memory_required:
if device != torch.device("cpu"):
free_memory(total_memory_required[device] * 1.1 + extra_mem,
device,
for_dynamic=free_for_dynamic,
pins_required=total_pins_required[device],
ram_required=total_ram_required[device])
for_dynamic=free_for_dynamic)

for device in total_memory_required:
if device != torch.device("cpu"):
Expand Down Expand Up @@ -1180,6 +1152,7 @@ def current_stream(device):
LARGEST_CASTED_WEIGHT = (None, 0)
STREAM_AIMDO_CAST_BUFFERS = {}
LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
STREAM_PIN_BUFFERS = {}

DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE = 16 * 1024 ** 3

Expand Down Expand Up @@ -1220,21 +1193,66 @@ def get_aimdo_cast_buffer(offload_stream, device):
if cast_buffer is None:
cast_buffer = comfy_aimdo.vram_buffer.VRAMBuffer(DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE, device.index)
STREAM_AIMDO_CAST_BUFFERS[offload_stream] = cast_buffer

return cast_buffer

def get_pin_buffer(offload_stream):
pin_buffer = STREAM_PIN_BUFFERS.get(offload_stream, None)
if pin_buffer is None:
pin_buffer = comfy_aimdo.host_buffer.HostBuffer(0)
STREAM_PIN_BUFFERS[offload_stream] = pin_buffer
elif offload_stream is not None:
event = getattr(pin_buffer, "_comfy_event", None)
if event is not None:
event.synchronize()
delattr(pin_buffer, "_comfy_event")
return pin_buffer
Comment thread
coderabbitai[bot] marked this conversation as resolved.

def resize_pin_buffer(pin_buffer, size):
global TOTAL_PINNED_MEMORY
old_size = pin_buffer.size
if size <= old_size:
return True
growth = size - old_size
comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM, free_active=True)
ensure_pin_budget(growth, evict_active=True)
try:
pin_buffer.extend(size=size, reallocate=True)
except RuntimeError:
return False
TOTAL_PINNED_MEMORY += pin_buffer.size - old_size
return True

def reset_cast_buffers():
global TOTAL_PINNED_MEMORY
global LARGEST_CASTED_WEIGHT
global LARGEST_AIMDO_CASTED_WEIGHT

LARGEST_CASTED_WEIGHT = (None, 0)
LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
for offload_stream in set(STREAM_CAST_BUFFERS) | set(STREAM_AIMDO_CAST_BUFFERS):
for offload_stream in set(STREAM_CAST_BUFFERS) | set(STREAM_AIMDO_CAST_BUFFERS) | set(STREAM_PIN_BUFFERS):
if offload_stream is not None:
offload_stream.synchronize()
synchronize()

for mmap_obj in DIRTY_MMAPS:
mmap_obj.bounce()
DIRTY_MMAPS.clear()

for pin_buffer in STREAM_PIN_BUFFERS.values():
TOTAL_PINNED_MEMORY -= pin_buffer.size
if TOTAL_PINNED_MEMORY < 0:
TOTAL_PINNED_MEMORY = 0

for loaded_model in current_loaded_models:
model = loaded_model.model
if model is not None and model.is_dynamic():
model.model.dynamic_pins[model.load_device]["active"] = False
model.partially_unload_ram(1e30, subsets=[ "patches" ])
model.model.dynamic_pins[model.load_device]["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024), [])

STREAM_CAST_BUFFERS.clear()
STREAM_AIMDO_CAST_BUFFERS.clear()
STREAM_PIN_BUFFERS.clear()
soft_empty_cache()

def get_offload_stream(device):
Expand Down Expand Up @@ -1296,8 +1314,7 @@ def cast_to_gathered(tensors, r, non_blocking=False, stream=None):
if comfy.memory_management.read_tensor_file_slice_into(tensor, dest_view):
continue
storage = tensor._qdata.untyped_storage() if isinstance(tensor, comfy.quant_ops.QuantizedTensor) else tensor.untyped_storage()
if hasattr(storage, "_comfy_tensor_mmap_touched"):
storage._comfy_tensor_mmap_touched = True
mark_mmap_dirty(storage)
dest_view.copy_(tensor, non_blocking=non_blocking)


Expand Down Expand Up @@ -1378,8 +1395,8 @@ def pin_memory(tensor):
return False

size = tensor.nbytes
if (TOTAL_PINNED_MEMORY + size) > MAX_PINNED_MEMORY:
return False
comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
ensure_pin_budget(size)

ptr = tensor.data_ptr()
if ptr == 0:
Expand Down
Loading
Loading