Refactored into init.py and distorch.py
This commit is contained in:
John Pollock
+34
-255
@@ -1,3 +1,4 @@
|
||||
# __init__.py
|
||||
import copy
|
||||
import torch
|
||||
import sys
|
||||
@@ -8,6 +9,8 @@ import logging
|
||||
import folder_paths
|
||||
from collections import defaultdict
|
||||
|
||||
from .distorch import register_patched_ggufmodelpatcher, analyze_ggml_loading, override_class_with_distorch
|
||||
|
||||
current_device = comfy.model_management.get_torch_device()
|
||||
current_offload_device = comfy.model_management.get_torch_device()
|
||||
|
||||
@@ -48,181 +51,6 @@ comfy.model_management.unet_offload_device = unet_offload_device_patched
|
||||
comfy.model_management.text_encoder_device = text_encoder_device_patched
|
||||
comfy.model_management.text_encoder_offload_device = text_encoder_offload_device_patched
|
||||
|
||||
def register_patched_ggufmodelpatcher(node_instance):
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["UnetLoaderGGUF"]
|
||||
module = sys.modules[original_loader.__module__]
|
||||
|
||||
if not hasattr(module.GGUFModelPatcher, '_patched'):
|
||||
original_load = module.GGUFModelPatcher.load
|
||||
logging.info("MultiGPU: GGUFDisTorch - GGUF ModelPatcher not yet patched, applying patch")
|
||||
|
||||
def new_load(self, *args, force_patch_weights=False, **kwargs):
|
||||
|
||||
super(module.GGUFModelPatcher, self).load(*args, force_patch_weights=True, **kwargs)
|
||||
linked = []
|
||||
module_count = 0
|
||||
for n, m in self.model.named_modules():
|
||||
module_count += 1
|
||||
if hasattr(m, "weight"):
|
||||
device = getattr(m.weight, "device", None)
|
||||
logging.info(f"MultiGPU: GGUFDisTorch - Weight Module {n} on device {device}, offload_device is {self.offload_device}")
|
||||
if device is not None:
|
||||
linked.append((n, m))
|
||||
continue
|
||||
if hasattr(m, "bias"):
|
||||
device = getattr(m.bias, "device", None)
|
||||
#logging.info(f"MultiGPU: GGUFDisTorch - Bias Module {n} on device {device}, offload_device is {self.offload_device}")
|
||||
if device is not None:
|
||||
linked.append((n, m))
|
||||
continue
|
||||
logging.info(f"MultiGPU: GGUFDisTorch - Found {len(linked)} linked modules out of {module_count} total modules")
|
||||
if linked:
|
||||
logging.info(f"MultiGPU: GGUFDisTorch - Found {len(linked)} linked modules, computing reallocation")
|
||||
device_assignments = analyze_ggml_loading(self.model, node_instance.distorch_allocations)['device_assignments']
|
||||
for device, layers in device_assignments.items():
|
||||
logging.info(f"MultiGPU: GGUFDisTorch - Moving {len(layers)} layers to {device}")
|
||||
target_device = torch.device(device)
|
||||
#logging.info(f"MultiGPU: GGUFDisTorch - Moving {len(layers)} layers to {device}")
|
||||
for n, m, _ in layers:
|
||||
m.to(self.load_device).to(target_device)
|
||||
|
||||
self.mmap_released = True
|
||||
logging.info("MultiGPU: GGUFDisTorch - self.mmap_released = True")
|
||||
|
||||
|
||||
module.GGUFModelPatcher.load = new_load
|
||||
module.GGUFModelPatcher._patched = True
|
||||
logging.info("MultiGPU: GGUFDisTorch - Successfully patched GGUF ModelPatcher")
|
||||
else:
|
||||
logging.info("MultiGPU: GGUFDisTorch - GGUF ModelPatcher already patched")
|
||||
|
||||
def analyze_ggml_loading(model, distorch_allocations):
|
||||
|
||||
DEVICE_RATIOS_DISTORCH = {}
|
||||
device_table = {}
|
||||
primary_dev_name = distorch_allocations.get("compute_device")
|
||||
primary_total_mem_bytes = comfy.model_management.get_total_memory(torch.device(primary_dev_name))
|
||||
primary_fraction = distorch_allocations.get("compute_device_alloc", 0.0)
|
||||
primary_alloc_gb = (primary_total_mem_bytes * primary_fraction) / (1024**3)
|
||||
DEVICE_RATIOS_DISTORCH[primary_dev_name] = primary_alloc_gb
|
||||
device_table[primary_dev_name] = {"fraction": primary_fraction,"total_gb": primary_total_mem_bytes / (1024**3),"alloc_gb": primary_alloc_gb}
|
||||
|
||||
i = 1
|
||||
while f"distorch{i}_device" in distorch_allocations:
|
||||
dev_key = f"distorch{i}_device"
|
||||
alloc_key = f"distorch{i}_alloc"
|
||||
dev_name = distorch_allocations[dev_key]
|
||||
dev_total_mem_bytes = comfy.model_management.get_total_memory(torch.device(dev_name))
|
||||
dev_fraction = distorch_allocations.get(alloc_key, 0.0)
|
||||
dev_alloc_gb = (dev_total_mem_bytes * dev_fraction) / (1024**3)
|
||||
DEVICE_RATIOS_DISTORCH[dev_name] = dev_alloc_gb
|
||||
device_table[dev_name] = {"fraction": dev_fraction,"total_gb": dev_total_mem_bytes / (1024**3),"alloc_gb": dev_alloc_gb}
|
||||
i += 1
|
||||
|
||||
cpu_dev_name = distorch_allocations.get("distorch_cpu", "cpu")
|
||||
cpu_total_mem_bytes = comfy.model_management.get_total_memory(torch.device(cpu_dev_name))
|
||||
cpu_fraction = distorch_allocations.get("distorch_cpu_alloc", 0.0)
|
||||
cpu_alloc_gb = (cpu_total_mem_bytes * cpu_fraction) / (1024**3)
|
||||
DEVICE_RATIOS_DISTORCH[cpu_dev_name] = cpu_alloc_gb
|
||||
device_table[cpu_dev_name] = {"fraction": cpu_fraction,"total_gb": cpu_total_mem_bytes / (1024**3),"alloc_gb": cpu_alloc_gb}
|
||||
|
||||
eq_line = "=" * 47
|
||||
dash_line = "-" * 47
|
||||
fmt_alloc = "{:<12}{:>10}{:>14}{:>10}"
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
logging.info(eq_line)
|
||||
logging.info(" DisTorch Analysis")
|
||||
logging.info(eq_line)
|
||||
logging.info(dash_line)
|
||||
logging.info(" DisTorch Device Allocations")
|
||||
logging.info(dash_line)
|
||||
logging.info(fmt_alloc.format("Device", "Alloc %", "Total (GB)", " Alloc (GB)"))
|
||||
logging.info(dash_line)
|
||||
|
||||
sorted_devices = sorted(device_table.keys(), key=lambda d: (d == "cpu", d))
|
||||
|
||||
for dev in sorted_devices:
|
||||
frac = device_table[dev]["fraction"]
|
||||
tot_gb = device_table[dev]["total_gb"]
|
||||
alloc_gb = device_table[dev]["alloc_gb"]
|
||||
logging.info(fmt_alloc.format(dev,f"{int(frac * 100)}%",f"{tot_gb:.2f}",f"{alloc_gb:.2f}"))
|
||||
|
||||
logging.info(dash_line)
|
||||
|
||||
layer_summary = {}
|
||||
layer_list = []
|
||||
memory_by_type = defaultdict(int)
|
||||
total_memory = 0
|
||||
|
||||
for name, module in model.named_modules():
|
||||
if hasattr(module, "weight"):
|
||||
layer_type = type(module).__name__
|
||||
layer_summary[layer_type] = layer_summary.get(layer_type, 0) + 1
|
||||
layer_list.append((name, module, layer_type))
|
||||
layer_memory = 0
|
||||
if module.weight is not None:
|
||||
layer_memory += module.weight.numel() * module.weight.element_size()
|
||||
if hasattr(module, "bias") and module.bias is not None:
|
||||
layer_memory += module.bias.numel() * module.bias.element_size()
|
||||
memory_by_type[layer_type] += layer_memory
|
||||
total_memory += layer_memory
|
||||
|
||||
logging.info(" DisTorch GGML Layer Distribution")
|
||||
logging.info(dash_line)
|
||||
fmt_layer = "{:<12}{:>10}{:>14}{:>10}"
|
||||
logging.info(fmt_layer.format("Layer Type", "Layers", "Memory (MB)", "% Total"))
|
||||
logging.info(dash_line)
|
||||
for layer_type, count in layer_summary.items():
|
||||
mem_mb = memory_by_type[layer_type] / (1024 * 1024)
|
||||
mem_percent = (memory_by_type[layer_type] / total_memory) * 100 if total_memory > 0 else 0
|
||||
logging.info(fmt_layer.format(layer_type,str(count),f"{mem_mb:.2f}",f"{mem_percent:.1f}%"))
|
||||
logging.info(dash_line)
|
||||
|
||||
nonzero_devices = [d for d, r in DEVICE_RATIOS_DISTORCH.items() if r > 0]
|
||||
nonzero_total_ratio = sum(DEVICE_RATIOS_DISTORCH[d] for d in nonzero_devices)
|
||||
device_assignments = {device: [] for device in DEVICE_RATIOS_DISTORCH.keys()}
|
||||
total_layers = len(layer_list)
|
||||
current_layer = 0
|
||||
|
||||
for idx, device in enumerate(nonzero_devices):
|
||||
ratio = DEVICE_RATIOS_DISTORCH[device]
|
||||
if idx == len(nonzero_devices) - 1:
|
||||
device_layer_count = total_layers - current_layer
|
||||
else:
|
||||
device_layer_count = int((ratio / nonzero_total_ratio) * total_layers)
|
||||
start_idx = current_layer
|
||||
end_idx = current_layer + device_layer_count
|
||||
device_assignments[device] = layer_list[start_idx:end_idx]
|
||||
current_layer += device_layer_count
|
||||
|
||||
logging.info(" DisTorch Final Device/Layer Assignments")
|
||||
logging.info(dash_line)
|
||||
fmt_assign = "{:<12}{:>10}{:>14}{:>10}"
|
||||
logging.info(fmt_assign.format("Device", "Layers", "Memory (MB)", "% Total"))
|
||||
logging.info(dash_line)
|
||||
total_assigned_memory = 0
|
||||
device_memories = {}
|
||||
for device, layers in device_assignments.items():
|
||||
device_memory = 0
|
||||
for layer_type in layer_summary:
|
||||
type_layers = sum(1 for _, _, lt in layers if lt == layer_type)
|
||||
if layer_summary[layer_type] > 0:
|
||||
mem_per_layer = memory_by_type[layer_type] / layer_summary[layer_type]
|
||||
device_memory += mem_per_layer * type_layers
|
||||
device_memories[device] = device_memory
|
||||
total_assigned_memory += device_memory
|
||||
|
||||
sorted_assignments = sorted(device_assignments.keys(), key=lambda d: (d == "cpu", d))
|
||||
|
||||
for dev in sorted_assignments:
|
||||
layers = device_assignments[dev]
|
||||
mem_mb = device_memories[dev] / (1024 * 1024)
|
||||
mem_percent = (device_memories[dev] / total_memory) * 100 if total_memory > 0 else 0
|
||||
logging.info(fmt_assign.format(dev,str(len(layers)),f"{mem_mb:.2f}",f"{mem_percent:.1f}%"))
|
||||
logging.info(dash_line)
|
||||
|
||||
return {"device_assignments": device_assignments}
|
||||
|
||||
def get_device_list():
|
||||
import torch
|
||||
@@ -296,65 +124,17 @@ def override_class_with_offload(cls):
|
||||
|
||||
return NodeOverrideDiffSynth
|
||||
|
||||
def override_class_with_distorch(cls):
|
||||
class NodeOverrideDisTorch(cls):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.distorch_allocations = {}
|
||||
self.distorch_compute_device = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
inputs = copy.deepcopy(cls.INPUT_TYPES())
|
||||
devices = [d for d in get_device_list() if d != "cpu"]
|
||||
inputs["optional"] = inputs.get("optional", {})
|
||||
|
||||
inputs["required"]["compute_device"] = (devices, {"default": devices[0], "tooltip": "Device model will use for computation"})
|
||||
inputs["required"]["compute_device_alloc"] = ("FLOAT", {"default": 0.15, "step": 0.01, "tooltip": "Fraction of memory NOT allocated to active latent space computation, recommended <= 15%"})
|
||||
|
||||
for i in range(len(devices) - 1):
|
||||
inputs["optional"][f"distorch{i+1}_device"] = (devices, {"default": devices[i+1], "tooltip": f"Device for distorch{i+1} model layer VRAM allocation"})
|
||||
inputs["optional"][f"distorch{i+1}_alloc"] = ("FLOAT", {"default": 0.9, "step": 0.01, "tooltip": f"Fraction of memory allocated to distorch{i+1} model layer, recommended >= 90%"})
|
||||
|
||||
inputs["optional"]["distorch_cpu"] = (["cpu"], {"default": "cpu", "tooltip": "Device for distorch CPU memory allocation"})
|
||||
inputs["optional"]["distorch_cpu_alloc"] = ("FLOAT", {"default": 0.0, "step": 0.01, "tooltip": "Fraction of memory allocated to distorch CPU memory (potentially slower than cuda)"})
|
||||
|
||||
return inputs
|
||||
|
||||
CATEGORY = "multigpu"
|
||||
FUNCTION = "override"
|
||||
|
||||
|
||||
def override(self, *args, **kwargs):
|
||||
self.distorch_allocations = {}
|
||||
self.distorch_compute_device = kwargs.get("compute_device", None)
|
||||
if self.distorch_compute_device is not None:
|
||||
global current_device
|
||||
current_device = self.distorch_compute_device
|
||||
|
||||
register_patched_ggufmodelpatcher(self)
|
||||
|
||||
for key, value in list(kwargs.items()):
|
||||
if key not in {"unet_name", "clip_name1", "clip_name2", "clip_name2", "type"}:
|
||||
logging.info(f"MultiGPU: Removing {key} from kwargs")
|
||||
logging.info(f"MultiGPU: Value: {value}")
|
||||
self.distorch_allocations[key] = kwargs.pop(key)
|
||||
|
||||
fn = getattr(super(), cls.FUNCTION)
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return NodeOverrideDisTorch
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"DeviceSelectorMultiGPU": DeviceSelectorMultiGPU}
|
||||
|
||||
def check_module_exists(module_path):
|
||||
full_path = os.path.join("custom_nodes", module_path)
|
||||
logging.info(f"MultiGPU: Checking for module at {full_path}")
|
||||
|
||||
|
||||
if not os.path.exists(full_path):
|
||||
logging.info(f"MultiGPU: Module {module_path} not found - skipping")
|
||||
return False
|
||||
|
||||
|
||||
logging.info(f"MultiGPU: Found {module_path}, creating compatible MultiGPU nodes")
|
||||
return True
|
||||
|
||||
@@ -529,7 +309,7 @@ def register_CLIPLoaderGGUFMultiGPU():
|
||||
|
||||
def register_LTXVLoaderMultiGPU():
|
||||
global NODE_CLASS_MAPPINGS
|
||||
|
||||
|
||||
class LTXVLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -540,14 +320,14 @@ def register_LTXVLoaderMultiGPU():
|
||||
"dtype": (["bfloat16", "float32"], {"default": "bfloat16"})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
RETURN_TYPES = ("MODEL", "VAE")
|
||||
RETURN_NAMES = ("model", "vae")
|
||||
FUNCTION = "load"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Loader"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
|
||||
def load(self, ckpt_name, dtype):
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["LTXVLoader"]()
|
||||
@@ -566,12 +346,12 @@ def register_LTXVLoaderMultiGPU():
|
||||
|
||||
def register_Florence2ModelLoaderMultiGPU():
|
||||
global NODE_CLASS_MAPPINGS
|
||||
|
||||
|
||||
class Florence2ModelLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"model": ([item.name for item in Path(folder_paths.models_dir, "LLM").iterdir() if item.is_dir()],
|
||||
"model": ([item.name for item in Path(folder_paths.models_dir, "LLM").iterdir() if item.is_dir()],
|
||||
{"tooltip": "models are expected to be in Comfyui/models/LLM folder"}),
|
||||
"precision": (['fp16','bf16','fp32'],),
|
||||
"attention": (['flash_attention_2', 'sdpa', 'eager'], {"default": 'sdpa'}),
|
||||
@@ -579,20 +359,20 @@ def register_Florence2ModelLoaderMultiGPU():
|
||||
"optional": {
|
||||
"lora": ("PEFTLORA",),
|
||||
}}
|
||||
|
||||
|
||||
RETURN_TYPES = ("FL2MODEL",)
|
||||
RETURN_NAMES = ("florence2_model",)
|
||||
FUNCTION = "loadmodel"
|
||||
CATEGORY = "Florence2"
|
||||
|
||||
|
||||
def loadmodel(self, model, precision, attention, lora=None):
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["Florence2ModelLoader"]()
|
||||
return original_loader.loadmodel(model, precision, attention, lora)
|
||||
NODE_CLASS_MAPPINGS["Florence2ModelLoaderMultiGPU"] = override_class(Florence2ModelLoader)
|
||||
logging.info(f"MultiGPU: Registered Florence2ModelLoaderMultiGPU")
|
||||
|
||||
def register_DownloadAndLoadFlorence2ModelMultiGPU():
|
||||
logging.info(f"MultiGPU: Registered Florence2ModelLoaderMultiGPU")
|
||||
|
||||
def register_DownloadAndLoadFlorence2ModelMultiGPU():
|
||||
global NODE_CLASS_MAPPINGS
|
||||
|
||||
class DownloadAndLoadFlorence2Model:
|
||||
@@ -620,12 +400,12 @@ def register_DownloadAndLoadFlorence2ModelMultiGPU():
|
||||
"optional": {
|
||||
"lora": ("PEFTLORA",),
|
||||
}}
|
||||
|
||||
|
||||
RETURN_TYPES = ("FL2MODEL",)
|
||||
RETURN_NAMES = ("florence2_model",)
|
||||
FUNCTION = "loadmodel"
|
||||
CATEGORY = "Florence2"
|
||||
|
||||
|
||||
def loadmodel(self, model, precision, attention, lora=None):
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2Model"]()
|
||||
@@ -651,7 +431,7 @@ def register_CheckpointLoaderNF4():
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["CheckpointLoaderNF4"]()
|
||||
return original_loader.load_checkpoint(ckpt_name)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["CheckpointLoaderNF4MultiGPU"] = override_class(CheckpointLoaderNF4)
|
||||
logging.info(f"MultiGPU: Registered CheckpointLoaderNF4MultiGPU")
|
||||
|
||||
@@ -674,7 +454,7 @@ def register_LoadFluxControlNetMultiGPU():
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["LoadFluxControlNet"]()
|
||||
return original_loader.loadmodel(model_name, controlnet_path)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["LoadFluxControlNetMultiGPU"] = override_class(LoadFluxControlNet)
|
||||
logging.info(f"MultiGPU: Registered LoadFluxControlNetMultiGPU")
|
||||
|
||||
@@ -688,7 +468,7 @@ def register_MMAudioModelLoaderMultiGPU():
|
||||
return {
|
||||
"required": {
|
||||
"mmaudio_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from the 'ComfyUI/models/mmaudio' -folder",}),
|
||||
|
||||
|
||||
"base_precision": (["fp16", "fp32", "bf16"], {"default": "fp16"}),
|
||||
},
|
||||
}
|
||||
@@ -702,7 +482,7 @@ def register_MMAudioModelLoaderMultiGPU():
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["MMAudioModelLoader"]()
|
||||
return original_loader.loadmodel(mmaudio_model, base_precision)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MMAudioModelLoaderMultiGPU"] = override_class(MMAudioModelLoader)
|
||||
logging.info(f"MultiGPU: Registered MMAudioModelLoaderMultiGPU")
|
||||
|
||||
@@ -738,7 +518,7 @@ def register_MMAudioFeatureUtilsLoaderMultiGPU():
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["MMAudioFeatureUtilsLoader"]()
|
||||
return original_loader.loadmodel(vae_model, precision, synchformer_model, clip_model, mode, bigvgan_vocoder_model)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MMAudioFeatureUtilsLoaderMultiGPU"] = override_class(MMAudioFeatureUtilsLoader)
|
||||
logging.info(f"MultiGPU: Registered MMAudioFeatureUtilsLoaderMultiGPU")
|
||||
|
||||
@@ -776,10 +556,10 @@ def register_MMAudioSamplerMultiGPU():
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["MMAudioSampler"]()
|
||||
return original_loader.sample(mmaudio_model, seed, feature_utils, duration, steps, cfg, prompt, negative_prompt, mask_away_clip, force_offload, images)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MMAudioSamplerMultiGPU"] = override_class(MMAudioSampler)
|
||||
logging.info(f"MultiGPU: Registered MMAudioSamplerMultiGPU")
|
||||
|
||||
|
||||
def register_PulidModelLoader():
|
||||
|
||||
global NODE_CLASS_MAPPINGS
|
||||
@@ -797,7 +577,7 @@ def register_PulidModelLoader():
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["PulidModelLoader"]()
|
||||
return original_loader.load_model(pulid_file)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["PulidModelLoaderMultiGPU"] = override_class(PulidModelLoader)
|
||||
logging.info(f"MultiGPU: Registered PulidModelLoaderMultiGPU")
|
||||
|
||||
@@ -822,7 +602,7 @@ def register_PulidInsightFaceLoader():
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["PulidInsightFaceLoader"]()
|
||||
return original_loader.load_insightface(provider)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["PulidInsightFaceLoaderMultiGPU"] = override_class(PulidInsightFaceLoader)
|
||||
logging.info(f"MultiGPU: Registered PulidInsightFaceLoaderMultiGPU")
|
||||
|
||||
@@ -845,7 +625,7 @@ def register_PulidEvaClipLoader():
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["PulidEvaClipLoader"]()
|
||||
return original_loader.load_eva_clip()
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["PulidEvaClipLoaderMultiGPU"] = override_class(PulidEvaClipLoader)
|
||||
logging.info(f"MultiGPU: Registered PulidEvaClipLoaderMultiGPU")
|
||||
|
||||
@@ -895,7 +675,7 @@ def register_HyVideoModelLoader():
|
||||
"required": {
|
||||
"model": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder",}),
|
||||
"base_precision": (["fp32", "bf16"], {"default": "bf16"}),
|
||||
"quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_scaled', 'torchao_fp8dq', "torchao_fp8dqrow", "torchao_int8dq", "torchao_fp6", "torchao_int4", "torchao_int8"],
|
||||
"quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_scaled', 'torchao_fp8dq', "torchao_fp8dqrow", "torchao_int8dq", "torchao_fp6", "torchao_int4", "torchao_int8"],
|
||||
{"default": 'disabled', "tooltip": "optional quantization method"}),
|
||||
},
|
||||
"optional": {
|
||||
@@ -920,14 +700,14 @@ def register_HyVideoModelLoader():
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["HyVideoModelLoader"]()
|
||||
# Use DiffSynth's auto offloading approach
|
||||
return original_loader.loadmodel(model, base_precision, "main_device", quantization,
|
||||
compile_args, attention_mode, block_swap_args, lora,
|
||||
return original_loader.loadmodel(model, base_precision, "main_device", quantization,
|
||||
compile_args, attention_mode, block_swap_args, lora,
|
||||
auto_cpu_offload=True)
|
||||
|
||||
# Register both with MultiGPU wrapper
|
||||
NODE_CLASS_MAPPINGS["HyVideoModelLoaderMultiGPU"] = override_class(HyVideoModelLoader)
|
||||
NODE_CLASS_MAPPINGS["HyVideoModelLoaderDiffSynthMultiGPU"] = override_class_with_offload(HyVideoModelLoaderDiffSynth)
|
||||
|
||||
|
||||
logging.info(f"MultiGPU: Registered HyVideoModelLoader nodes")
|
||||
|
||||
def register_HyVideoVAELoader():
|
||||
@@ -959,7 +739,7 @@ def register_HyVideoVAELoader():
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["HyVideoVAELoader"]()
|
||||
return original_loader.loadmodel(model_name, precision, compile_args)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["HyVideoVAELoaderMultiGPU"] = override_class(HyVideoVAELoader)
|
||||
logging.info(f"MultiGPU: Registered HyVideoVAELoaderMultiGPU")
|
||||
|
||||
@@ -995,7 +775,7 @@ def register_DownloadAndLoadHyVideoTextEncoder():
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["DownloadAndLoadHyVideoTextEncoder"]()
|
||||
return original_loader.loadmodel(llm_model, clip_model, precision, apply_final_norm, hidden_state_skip_layer, quantization)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["DownloadAndLoadHyVideoTextEncoderMultiGPU"] = override_class(DownloadAndLoadHyVideoTextEncoder)
|
||||
logging.info(f"MultiGPU: Registered DownloadAndLoadHyVideoTextEncoderMultiGPU")
|
||||
|
||||
@@ -1016,7 +796,6 @@ if check_module_exists("ComfyUI-MMAudio"):
|
||||
register_MMAudioSamplerMultiGPU()
|
||||
if check_module_exists("ComfyUI-GGUF"):
|
||||
register_UnetLoaderGGUFMultiGPU()
|
||||
register_UnetLoaderGGUFDisTorchMultiGPU()
|
||||
register_CLIPLoaderGGUFMultiGPU()
|
||||
if check_module_exists("PuLID_ComfyUI"):
|
||||
register_PulidModelLoader()
|
||||
@@ -1027,4 +806,4 @@ if check_module_exists("ComfyUI-HunyuanVideoWrapper"):
|
||||
register_HyVideoVAELoader()
|
||||
register_DownloadAndLoadHyVideoTextEncoder()
|
||||
|
||||
logging.info(f"MultiGPU: Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}")
|
||||
logging.info(f"MultiGPU: Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}")
|
||||
+237
@@ -0,0 +1,237 @@
|
||||
import sys
|
||||
import logging
|
||||
import torch
|
||||
from collections import defaultdict
|
||||
import comfy.model_management
|
||||
import copy
|
||||
|
||||
def register_patched_ggufmodelpatcher(node_instance):
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
original_loader = NODE_CLASS_MAPPINGS["UnetLoaderGGUF"]
|
||||
module = sys.modules[original_loader.__module__]
|
||||
|
||||
if not hasattr(module.GGUFModelPatcher, '_patched'):
|
||||
original_load = module.GGUFModelPatcher.load
|
||||
logging.info("MultiGPU: GGUFDisTorch - GGUF ModelPatcher not yet patched, applying patch")
|
||||
|
||||
def new_load(self, *args, force_patch_weights=False, **kwargs):
|
||||
|
||||
super(module.GGUFModelPatcher, self).load(*args, force_patch_weights=True, **kwargs)
|
||||
linked = []
|
||||
module_count = 0
|
||||
for n, m in self.model.named_modules():
|
||||
module_count += 1
|
||||
if hasattr(m, "weight"):
|
||||
device = getattr(m.weight, "device", None)
|
||||
logging.info(f"MultiGPU: GGUFDisTorch - Weight Module {n} on device {device}, offload_device is {self.offload_device}")
|
||||
if device is not None:
|
||||
linked.append((n, m))
|
||||
continue
|
||||
if hasattr(m, "bias"):
|
||||
device = getattr(m.bias, "device", None)
|
||||
#logging.info(f"MultiGPU: GGUFDisTorch - Bias Module {n} on device {device}, offload_device is {self.offload_device}")
|
||||
if device is not None:
|
||||
linked.append((n, m))
|
||||
continue
|
||||
logging.info(f"MultiGPU: GGUFDisTorch - Found {len(linked)} linked modules out of {module_count} total modules")
|
||||
if linked:
|
||||
logging.info(f"MultiGPU: GGUFDisTorch - Found {len(linked)} linked modules, computing reallocation")
|
||||
device_assignments = analyze_ggml_loading(self.model, node_instance.distorch_allocations)['device_assignments']
|
||||
for device, layers in device_assignments.items():
|
||||
logging.info(f"MultiGPU: GGUFDisTorch - Moving {len(layers)} layers to {device}")
|
||||
target_device = torch.device(device)
|
||||
#logging.info(f"MultiGPU: GGUFDisTorch - Moving {len(layers)} layers to {device}")
|
||||
for n, m, _ in layers:
|
||||
m.to(self.load_device).to(target_device)
|
||||
|
||||
self.mmap_released = True
|
||||
logging.info("MultiGPU: GGUFDisTorch - self.mmap_released = True")
|
||||
|
||||
|
||||
module.GGUFModelPatcher.load = new_load
|
||||
module.GGUFModelPatcher._patched = True
|
||||
logging.info("MultiGPU: GGUFDisTorch - Successfully patched GGUF ModelPatcher")
|
||||
else:
|
||||
logging.info("MultiGPU: GGUFDisTorch - GGUF ModelPatcher already patched")
|
||||
|
||||
def analyze_ggml_loading(model, distorch_allocations):
|
||||
|
||||
DEVICE_RATIOS_DISTORCH = {}
|
||||
device_table = {}
|
||||
primary_dev_name = distorch_allocations.get("compute_device")
|
||||
primary_total_mem_bytes = comfy.model_management.get_total_memory(torch.device(primary_dev_name))
|
||||
primary_fraction = distorch_allocations.get("compute_device_alloc", 0.0)
|
||||
primary_alloc_gb = (primary_total_mem_bytes * primary_fraction) / (1024**3)
|
||||
DEVICE_RATIOS_DISTORCH[primary_dev_name] = primary_alloc_gb
|
||||
device_table[primary_dev_name] = {"fraction": primary_fraction,"total_gb": primary_total_mem_bytes / (1024**3),"alloc_gb": primary_alloc_gb}
|
||||
|
||||
i = 1
|
||||
while f"distorch{i}_device" in distorch_allocations:
|
||||
dev_key = f"distorch{i}_device"
|
||||
alloc_key = f"distorch{i}_alloc"
|
||||
dev_name = distorch_allocations[dev_key]
|
||||
dev_total_mem_bytes = comfy.model_management.get_total_memory(torch.device(dev_name))
|
||||
dev_fraction = distorch_allocations.get(alloc_key, 0.0)
|
||||
dev_alloc_gb = (dev_total_mem_bytes * dev_fraction) / (1024**3)
|
||||
DEVICE_RATIOS_DISTORCH[dev_name] = dev_alloc_gb
|
||||
device_table[dev_name] = {"fraction": dev_fraction,"total_gb": dev_total_mem_bytes / (1024**3),"alloc_gb": dev_alloc_gb}
|
||||
i += 1
|
||||
|
||||
cpu_dev_name = distorch_allocations.get("distorch_cpu", "cpu")
|
||||
cpu_total_mem_bytes = comfy.model_management.get_total_memory(torch.device(cpu_dev_name))
|
||||
cpu_fraction = distorch_allocations.get("distorch_cpu_alloc", 0.0)
|
||||
cpu_alloc_gb = (cpu_total_mem_bytes * cpu_fraction) / (1024**3)
|
||||
DEVICE_RATIOS_DISTORCH[cpu_dev_name] = cpu_alloc_gb
|
||||
device_table[cpu_dev_name] = {"fraction": cpu_fraction,"total_gb": cpu_total_mem_bytes / (1024**3),"alloc_gb": cpu_alloc_gb}
|
||||
|
||||
eq_line = "=" * 47
|
||||
dash_line = "-" * 47
|
||||
fmt_alloc = "{:<12}{:>10}{:>14}{:>10}"
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
logging.info(eq_line)
|
||||
logging.info(" DisTorch Analysis")
|
||||
logging.info(eq_line)
|
||||
logging.info(dash_line)
|
||||
logging.info(" DisTorch Device Allocations")
|
||||
logging.info(dash_line)
|
||||
logging.info(fmt_alloc.format("Device", "Alloc %", "Total (GB)", " Alloc (GB)"))
|
||||
logging.info(dash_line)
|
||||
|
||||
sorted_devices = sorted(device_table.keys(), key=lambda d: (d == "cpu", d))
|
||||
|
||||
for dev in sorted_devices:
|
||||
frac = device_table[dev]["fraction"]
|
||||
tot_gb = device_table[dev]["total_gb"]
|
||||
alloc_gb = device_table[dev]["alloc_gb"]
|
||||
logging.info(fmt_alloc.format(dev,f"{int(frac * 100)}%",f"{tot_gb:.2f}",f"{alloc_gb:.2f}"))
|
||||
|
||||
logging.info(dash_line)
|
||||
|
||||
layer_summary = {}
|
||||
layer_list = []
|
||||
memory_by_type = defaultdict(int)
|
||||
total_memory = 0
|
||||
|
||||
for name, module in model.named_modules():
|
||||
if hasattr(module, "weight"):
|
||||
layer_type = type(module).__name__
|
||||
layer_summary[layer_type] = layer_summary.get(layer_type, 0) + 1
|
||||
layer_list.append((name, module, layer_type))
|
||||
layer_memory = 0
|
||||
if module.weight is not None:
|
||||
layer_memory += module.weight.numel() * module.weight.element_size()
|
||||
if hasattr(module, "bias") and module.bias is not None:
|
||||
layer_memory += module.bias.numel() * module.bias.element_size()
|
||||
memory_by_type[layer_type] += layer_memory
|
||||
total_memory += layer_memory
|
||||
|
||||
logging.info(" DisTorch GGML Layer Distribution")
|
||||
logging.info(dash_line)
|
||||
fmt_layer = "{:<12}{:>10}{:>14}{:>10}"
|
||||
logging.info(fmt_layer.format("Layer Type", "Layers", "Memory (MB)", "% Total"))
|
||||
logging.info(dash_line)
|
||||
for layer_type, count in layer_summary.items():
|
||||
mem_mb = memory_by_type[layer_type] / (1024 * 1024)
|
||||
mem_percent = (memory_by_type[layer_type] / total_memory) * 100 if total_memory > 0 else 0
|
||||
logging.info(fmt_layer.format(layer_type,str(count),f"{mem_mb:.2f}",f"{mem_percent:.1f}%"))
|
||||
logging.info(dash_line)
|
||||
|
||||
nonzero_devices = [d for d, r in DEVICE_RATIOS_DISTORCH.items() if r > 0]
|
||||
nonzero_total_ratio = sum(DEVICE_RATIOS_DISTORCH[d] for d in nonzero_devices)
|
||||
device_assignments = {device: [] for device in DEVICE_RATIOS_DISTORCH.keys()}
|
||||
total_layers = len(layer_list)
|
||||
current_layer = 0
|
||||
|
||||
for idx, device in enumerate(nonzero_devices):
|
||||
ratio = DEVICE_RATIOS_DISTORCH[device]
|
||||
if idx == len(nonzero_devices) - 1:
|
||||
device_layer_count = total_layers - current_layer
|
||||
else:
|
||||
device_layer_count = int((ratio / nonzero_total_ratio) * total_layers)
|
||||
start_idx = current_layer
|
||||
end_idx = current_layer + device_layer_count
|
||||
device_assignments[device] = layer_list[start_idx:end_idx]
|
||||
current_layer += device_layer_count
|
||||
|
||||
logging.info(" DisTorch Final Device/Layer Assignments")
|
||||
logging.info(dash_line)
|
||||
fmt_assign = "{:<12}{:>10}{:>14}{:>10}"
|
||||
logging.info(fmt_assign.format("Device", "Layers", "Memory (MB)", "% Total"))
|
||||
logging.info(dash_line)
|
||||
total_assigned_memory = 0
|
||||
device_memories = {}
|
||||
for device, layers in device_assignments.items():
|
||||
device_memory = 0
|
||||
for layer_type in layer_summary:
|
||||
type_layers = sum(1 for _, _, lt in layers if lt == layer_type)
|
||||
if layer_summary[layer_type] > 0:
|
||||
mem_per_layer = memory_by_type[layer_type] / layer_summary[layer_type]
|
||||
device_memory += mem_per_layer * type_layers
|
||||
device_memories[device] = device_memory
|
||||
total_assigned_memory += device_memory
|
||||
|
||||
sorted_assignments = sorted(device_assignments.keys(), key=lambda d: (d == "cpu", d))
|
||||
|
||||
for dev in sorted_assignments:
|
||||
layers = device_assignments[dev]
|
||||
mem_mb = device_memories[dev] / (1024 * 1024)
|
||||
mem_percent = (device_memories[dev] / total_memory) * 100 if total_memory > 0 else 0
|
||||
logging.info(fmt_assign.format(dev,str(len(layers)),f"{mem_mb:.2f}",f"{mem_percent:.1f}%"))
|
||||
logging.info(dash_line)
|
||||
|
||||
return {"device_assignments": device_assignments}
|
||||
|
||||
|
||||
def override_class_with_distorch(cls):
|
||||
from . import register_patched_ggufmodelpatcher
|
||||
from . import get_device_list
|
||||
import copy
|
||||
import logging
|
||||
|
||||
class NodeOverrideDisTorch(cls):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.distorch_allocations = {}
|
||||
self.distorch_compute_device = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
inputs = copy.deepcopy(cls.INPUT_TYPES())
|
||||
devices = [d for d in get_device_list() if d != "cpu"]
|
||||
inputs["optional"] = inputs.get("optional", {})
|
||||
|
||||
inputs["required"]["compute_device"] = (devices, {"default": devices[0], "tooltip": "Device model will use for computation"})
|
||||
inputs["required"]["compute_device_alloc"] = ("FLOAT", {"default": 0.15, "step": 0.01, "tooltip": "Fraction of memory NOT allocated to active latent space computation, recommended <= 15%"})
|
||||
|
||||
for i in range(len(devices) - 1):
|
||||
inputs["optional"][f"distorch{i+1}_device"] = (devices, {"default": devices[i+1], "tooltip": f"Device for distorch{i+1} model layer VRAM allocation"})
|
||||
inputs["optional"][f"distorch{i+1}_alloc"] = ("FLOAT", {"default": 0.9, "step": 0.01, "tooltip": f"Fraction of memory allocated to distorch{i+1} model layer, recommended >= 90%"})
|
||||
|
||||
inputs["optional"]["distorch_cpu"] = (["cpu"], {"default": "cpu", "tooltip": "Device for distorch CPU memory allocation"})
|
||||
inputs["optional"]["distorch_cpu_alloc"] = ("FLOAT", {"default": 0.0, "step": 0.01, "tooltip": "Fraction of memory allocated to distorch CPU memory (potentially slower than cuda)"})
|
||||
|
||||
return inputs
|
||||
|
||||
CATEGORY = "multigpu"
|
||||
FUNCTION = "override"
|
||||
|
||||
|
||||
def override(self, *args, **kwargs):
|
||||
self.distorch_allocations = {}
|
||||
self.distorch_compute_device = kwargs.get("compute_device", None)
|
||||
if self.distorch_compute_device is not None:
|
||||
global current_device
|
||||
current_device = self.distorch_compute_device
|
||||
|
||||
register_patched_ggufmodelpatcher(self)
|
||||
|
||||
for key, value in list(kwargs.items()):
|
||||
if key not in {"unet_name", "clip_name1", "clip_name2", "clip_name2", "type"}:
|
||||
logging.info(f"MultiGPU: Removing {key} from kwargs")
|
||||
logging.info(f"MultiGPU: Value: {value}")
|
||||
self.distorch_allocations[key] = kwargs.pop(key)
|
||||
|
||||
fn = getattr(super(), cls.FUNCTION)
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return NodeOverrideDisTorch
|
||||
Reference in New Issue
Block a user