mmap unload fix for lowvram
This commit is contained in:
@@ -180,9 +180,32 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
|
||||
device_to = device_to if comfy.model_management.DISABLE_SMART_MEMORY else None
|
||||
return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights)
|
||||
|
||||
mmap_released = False
|
||||
def load(self, *args, force_patch_weights=False, **kwargs):
|
||||
# always call `patch_weight_to_device` even for lowvram
|
||||
return super().load(*args, force_patch_weights=True, **kwargs)
|
||||
super().load(*args, force_patch_weights=True, **kwargs)
|
||||
|
||||
# make sure nothing stays linked to mmap after first load
|
||||
if not self.mmap_released:
|
||||
linked = []
|
||||
if kwargs.get("lowvram_model_memory", 0) > 0:
|
||||
for n, m in self.model.named_modules():
|
||||
if hasattr(m, "weight"):
|
||||
device = getattr(m.weight, "device", None)
|
||||
if device == self.offload_device:
|
||||
linked.append((n, m))
|
||||
continue
|
||||
if hasattr(m, "bias"):
|
||||
device = getattr(m.bias, "device", None)
|
||||
if device == self.offload_device:
|
||||
linked.append((n, m))
|
||||
continue
|
||||
if linked:
|
||||
print(f"Attempting to release mmap ({len(linked)})")
|
||||
for n, m in linked:
|
||||
# TODO: possible to OOM, find better way to detach
|
||||
m.to(self.load_device).to(self.offload_device)
|
||||
self.mmap_released = True
|
||||
|
||||
def clone(self, *args, **kwargs):
|
||||
n = GGUFModelPatcher(self.model, self.load_device, self.offload_device, self.size, weight_inplace_update=self.weight_inplace_update)
|
||||
|
||||
Reference in New Issue
Block a user