xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-07-30 09:30:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

ops: put all post async transfer compute on the main stream

Browse Source 
Some models have massive weights that need either complex
dequantization or lora patching. Don't do these patchings on the offload
stream, instead do them on the main stream to syncrhonize the
potentially large vram spikes for these compute processes. This avoids
having to assume a worst case scenario of multiple offload streams
all spiking VRAM is parallel with whatever the main stream is doing.
This commit is contained in:
Rattus 2025-12-03 12:39:30 +10:00
parent e0595fc83e
commit 32d8c7374f
1 changed files with 22 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
comfy/ops.py
View File

@ -111,22 +111,24 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of
if s.bias is not None: if s.bias is not None:
bias = comfy.model_management.cast_to(s.bias, bias_dtype, device, non_blocking=non_blocking, copy=bias_has_function, stream=offload_stream) bias = comfy.model_management.cast_to(s.bias, bias_dtype, device, non_blocking=non_blocking, copy=bias_has_function, stream=offload_stream)
if bias_has_function: comfy.model_management.sync_stream(device, offload_stream)
with wf_context:
for f in s.bias_function: bias_a = bias
bias = f(bias) weight_a = weight
if s.bias is not None:
for f in s.bias_function:
bias = f(bias)
if weight_has_function or weight.dtype != dtype: if weight_has_function or weight.dtype != dtype:
with wf_context: weight = weight.to(dtype=dtype)
weight = weight.to(dtype=dtype) if isinstance(weight, QuantizedTensor):
if isinstance(weight, QuantizedTensor): weight = weight.dequantize()
weight = weight.dequantize() for f in s.weight_function:
for f in s.weight_function: weight = f(weight)
weight = f(weight)
comfy.model_management.sync_stream(device, offload_stream)
if offloadable: if offloadable:
return weight, bias, offload_stream return weight, bias, (offload_stream, weight_a, bias_a)
else: else:
#Legacy function signature #Legacy function signature
return weight, bias return weight, bias
@ -135,13 +137,16 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of
def uncast_bias_weight(s, weight, bias, offload_stream): def uncast_bias_weight(s, weight, bias, offload_stream):
if offload_stream is None: if offload_stream is None:
return return
if weight is not None: os, weight_a, bias_a = offload_stream
device = weight.device if os is None:
return
if weight_a is not None:
device = weight_a.device
else: else:
if bias is None: if bias_a is None:
return return
device = bias.device device = bias_a.device
offload_stream.wait_stream(comfy.model_management.current_stream(device)) os.wait_stream(comfy.model_management.current_stream(device))
class CastWeightBiasOp: class CastWeightBiasOp: