xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2025-12-12 15:35:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

mm: guard against double pin and unpin explicitly (#10672)

Browse Source 
As commented, if you let cuda be the one to detect double pin/unpinning
it actually creates an asyc GPU error.
This commit is contained in:
rattus 2025-11-07 12:20:48 +10:00 committed by GitHub
parent eb1c42f649
commit cf97b033ee
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 12 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
comfy/model_management.py
View File

@ -1103,6 +1103,12 @@ def pin_memory(tensor):
if not is_device_cpu(tensor.device): if not is_device_cpu(tensor.device):
return False return False
if tensor.is_pinned():
#NOTE: Cuda does detect when a tensor is already pinned and would
#error below, but there are proven cases where this also queues an error
#on the GPU async. So dont trust the CUDA API and guard here
return False
size = tensor.numel() * tensor.element_size() size = tensor.numel() * tensor.element_size()
if (TOTAL_PINNED_MEMORY + size) > MAX_PINNED_MEMORY: if (TOTAL_PINNED_MEMORY + size) > MAX_PINNED_MEMORY:
return False return False
@ -1123,6 +1129,12 @@ def unpin_memory(tensor):
if not is_device_cpu(tensor.device): if not is_device_cpu(tensor.device):
return False return False
if not tensor.is_pinned():
#NOTE: Cuda does detect when a tensor is already pinned and would
#error below, but there are proven cases where this also queues an error
#on the GPU async. So dont trust the CUDA API and guard here
return False
ptr = tensor.data_ptr() ptr = tensor.data_ptr()
if torch.cuda.cudart().cudaHostUnregister(ptr) == 0: if torch.cuda.cudart().cudaHostUnregister(ptr) == 0:
TOTAL_PINNED_MEMORY -= PINNED_MEMORY.pop(ptr) TOTAL_PINNED_MEMORY -= PINNED_MEMORY.pop(ptr)