From f119787bee62b4e2bf73715334e61df7e48232b1 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sun, 19 Jan 2025 14:24:11 -0600 Subject: [PATCH] ggml learning --- __init__.py | 329 ++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 321 insertions(+), 8 deletions(-) diff --git a/__init__.py b/__init__.py index 30116ea..90cb6ee 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,7 @@ import time import copy import torch +import sys import comfy.model_management import os from pathlib import Path # Add this import @@ -11,31 +12,189 @@ import folder_paths logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logging.info("MultiGPU: Initialization started") +# Initialize the current device states and log them current_device = comfy.model_management.get_torch_device() current_offload_device = comfy.model_management.get_torch_device() +distorch_allocations = {} -logging.info(f"MultiGPU: Initial device {current_device}") -logging.info(f"MultiGPU: Initial offload device {current_offload_device}") +logging.info(f"MultiGPU: Initial device set to {current_device}") +logging.info(f"MultiGPU: Initial offload device set to {current_offload_device}") +# Define and patch the device logic def get_torch_device_patched(): + device = None if ( not torch.cuda.is_available() or comfy.model_management.cpu_state == comfy.model_management.CPUState.CPU or "cpu" in str(current_device).lower() ): - return torch.device("cpu") - return torch.device(current_device) + device = torch.device("cpu") + else: + device = torch.device(current_device) + + logging.info(f"MultiGPU: get_torch_device_patched invoked, returning {device}") + return device comfy.model_management.get_torch_device = get_torch_device_patched +logging.info(f"MultiGPU: Patched get_torch_device now returns {get_torch_device_patched()}") def unet_offload_device_patched(): - if (not torch.cuda.is_available() + device = None + if ( + not torch.cuda.is_available() or comfy.model_management.cpu_state == comfy.model_management.CPUState.CPU - or "cpu" in str(current_offload_device).lower()): - return torch.device("cpu") - return torch.device(current_offload_device) + or "cpu" in str(current_offload_device).lower() + ): + device = torch.device("cpu") + else: + device = torch.device(current_offload_device) + + logging.info(f"MultiGPU: unet_offload_device_patched invoked, returning {device}") + return device comfy.model_management.unet_offload_device = unet_offload_device_patched +logging.info(f"MultiGPU: Patched unet_offload_device now returns {unet_offload_device_patched()}") + +# Save the original patched logic for later restoration and log them +original_get_torch_device = comfy.model_management.get_torch_device +original_unet_offload_device = comfy.model_management.unet_offload_device +logging.info(f"MultiGPU: Saved original get_torch_device: {original_get_torch_device()}") +logging.info(f"MultiGPU: Saved original unet_offload_device: {original_unet_offload_device()}") + +logging.info("MultiGPU: Device management logic initialized") + + +def analyze_ggml_loading(model): + """ + Analyzes GGML model loading and determines device assignments. + Returns device assignments with accurate memory calculations. + """ + from collections import defaultdict + + # For testing - this would come from a global config in production + DEVICE_RATIOS = { + "cuda:0": 1, # 1/9 of layers + "cuda:1": 8 # 8/9 of layers + } + + # Step 1: Memory Analysis + device_properties = {} + for device in DEVICE_RATIOS.keys(): + if device.startswith("cuda"): + device_props = torch.cuda.get_device_properties(torch.device(device)) + device_properties[device] = { + "total_memory": device_props.total_memory, + "name": device_props.name + } + logging.info(f"ComfyUI-GGUF: Device {device} Memory: {device_props.total_memory / (1024 ** 3):.2f} GB") + + # Step 2: Layer Analysis + layer_summary = {} + layer_list = [] + memory_by_type = defaultdict(int) + total_memory = 0 + + # First pass: collect layers and calculate total memory + 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)) + + # Calculate memory for this layer + 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 + + # Step 3: Print Analysis Results as Tables + logging.info("\nGGML Layer Analysis") + logging.info("==================") + + # Layer Distribution Table + format_str = "{:<12} {:>8} {:>12} {:>8}" + logging.info("\nLayer Distribution:") + logging.info(format_str.format("Layer Type", "Layers", "Memory (MB)", "% Total")) + logging.info("-" * 42) + + for layer_type, count in layer_summary.items(): + mem = memory_by_type[layer_type] / (1024 * 1024) # MB + mem_percent = (memory_by_type[layer_type] / total_memory) * 100 if total_memory > 0 else 0 + logging.info(format_str.format( + layer_type, + str(count), + f"{mem:.2f}", + f"{mem_percent:.1f}%" + )) + + # Step 4: Calculate Device Assignments + total_ratio = sum(DEVICE_RATIOS.values()) + device_assignments = {device: [] for device in DEVICE_RATIOS.keys()} + + # Calculate layer counts for each device + total_layers = len(layer_list) + current_layer = 0 + + for device, ratio in DEVICE_RATIOS.items(): + if device == list(DEVICE_RATIOS.keys())[-1]: + # Last device gets all remaining layers + device_layer_count = total_layers - current_layer + else: + device_layer_count = int((ratio / 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 + + # Device Assignment Table with corrected memory calculations + format_str = "{:<10} {:>8} {:>16} {:>10}" + logging.info("\nDevice Assignments:") + logging.info(format_str.format("Device", "Layers", "Memory (MB)", "% Total")) + logging.info("-" * 46) + + total_assigned_memory = 0 + device_memories = {} + + # Calculate memory per device + for device, layers in device_assignments.items(): + device_memory = 0 + # Calculate memory per layer type for this device + for layer_type in layer_summary: + type_layers = sum(1 for _, _, lt in layers if lt == layer_type) + if layer_summary[layer_type] > 0: # Avoid div by zero + 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 + + # Print device assignments with memory percentages + for device, layers in device_assignments.items(): + mem_mb = device_memories[device] / (1024 * 1024) # Convert to MB + mem_percent = (device_memories[device] / total_memory) * 100 if total_memory > 0 else 0 + logging.info(format_str.format( + device, + str(len(layers)), + f"{mem_mb:.2f}", + f"{mem_percent:.1f}%" + )) + + # Verification log + total_mb = total_memory / (1024 * 1024) + assigned_mb = total_assigned_memory / (1024 * 1024) + logging.info(f"\nMemory Verification:") + logging.info(f"Total Model Memory: {total_mb:.2f} MB") + logging.info(f"Total Assigned Memory: {assigned_mb:.2f} MB") + if abs(total_mb - assigned_mb) > 0.01: # Allow for minor floating point differences + logging.warning(f"Memory assignment mismatch: {abs(total_mb - assigned_mb):.2f} MB difference") + + return { + "device_assignments": device_assignments + } def get_device_list(): import torch @@ -109,6 +268,47 @@ def override_class_with_offload(cls): return NodeOverrideDiffSynth +def override_class_with_distorch(cls): + class NodeOverrideDisTorch(cls): + @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, compute_device=None, **kwargs): + global current_device + global distorch_allocations + + current_device = compute_device + distorch_allocations = {} + + for key, value in list(kwargs.items()): + if key not in {"unet_name"}: + distorch_allocations[key] = kwargs.pop(key) + + logging.info(f"MultiGPU: DisTorch - distorch_allocations: {distorch_allocations}") + + fn = getattr(super(), cls.FUNCTION) + return fn(*args, **kwargs) + + return NodeOverrideDisTorch + NODE_CLASS_MAPPINGS = { "DeviceSelectorMultiGPU": DeviceSelectorMultiGPU } @@ -493,6 +693,118 @@ def register_UnetLoaderGGUFMultiGPU(): NODE_CLASS_MAPPINGS["UnetLoaderGGUFAdvancedMultiGPU"] = UnetLoaderGGUFAdvancedMultiGPU logging.info(f"MultiGPU: Registered UnetLoaderGGUFAdvancedMultiGPU") +def register_UnetLoaderGGUFDisTorchMultiGPU(): + global NODE_CLASS_MAPPINGS + + # First define the base UnetLoaderGGUFDisTorch class + class UnetLoaderGGUFDisTorch: + @classmethod + def INPUT_TYPES(s): + unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] + return { + "required": { + "unet_name": (unet_names,), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "load_unet" + CATEGORY = "bootleg" + TITLE = "Unet Loader (GGUFDisTorch)" + + def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None): + from nodes import NODE_CLASS_MAPPINGS + logging.info("MultiGPU: GGUFDisTorch - Starting GGUFDisTorch UNet load") + + # Get the correct module through the original loader + original_loader = NODE_CLASS_MAPPINGS["UnetLoaderGGUF"] + module = sys.modules[original_loader.__module__] + logging.info(f"MultiGPU: GGUFDisTorch - Got GGUF module: {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): + logging.info("MultiGPU: GGUFDisTorch - Entering patched GGUFDisTorch load function") + + # Save the current device states and logic + global current_device, current_offload_device + original_current_device = current_device + original_current_offload_device = current_offload_device + + try: + # Temporarily override the device logic for this load + current_device = torch.device("cuda:0") + current_offload_device = torch.device("cuda:1") + + logging.info(f"MultiGPU: GGUFDisTorch - Overriding current_device to {current_device}") + logging.info(f"MultiGPU: GGUFDisTorch - Overriding current_offload_device to {current_offload_device}") + + # Call the original load function with the temporary overrides + super(module.GGUFModelPatcher, self).load(*args, force_patch_weights=True, **kwargs) + + if not self.mmap_released: + logging.info("MultiGPU: GGUFDisTorch - Processing mmap release") + linked = [] + + # Debug the lowvram check + lowvram_value = kwargs.get("lowvram_model_memory", 0) + logging.info(f"MultiGPU: GGUFDisTorch - lowvram_model_memory value: {lowvram_value}") + + if lowvram_value > 0: + logging.info("MultiGPU: GGUFDisTorch - Entering module scanning") + 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 - Module {n} on device {device}, offload_device is {self.offload_device}") + if device == self.offload_device: + linked.append((n, m)) + logging.info(f"MultiGPU: GGUFDisTorch - Scanned {module_count} total modules") + else: + logging.info("MultiGPU: GGUFDisTorch - Skipped module scanning due to lowvram check") + + if linked: + logging.info(f"MultiGPU: GGUFDisTorch - Found {len(linked)} linked modules") + device_assignments = analyze_ggml_loading(self.model)['device_assignments'] + for device, layers in device_assignments.items(): + target_device = torch.device(device) + logging.info(f"MultiGPU: GGUFDisTorch - Moving {len(layers)} layers to {device}") + for n, m, _ in layers: + try: + m.to(self.load_device).to(target_device) + # logging.info(f"MultiGPU: GGUFDisTorch - Successfully moved layer {n} to {device}") + except Exception as e: + logging.error(f"MultiGPU: GGUFDisTorch - Error moving layer {n} to {device}: {str(e)}") + self.mmap_released = True + logging.info("MultiGPU: GGUFDisTorch - mmap release complete") + + finally: + # Restore the original device states + current_device = original_current_device + current_offload_device = original_current_offload_device + + logging.info(f"MultiGPU: GGUFDisTorch - Restored current_device to {current_device}") + logging.info(f"MultiGPU: GGUFDisTorch - Restored current_offload_device to {current_offload_device}") + + 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") + + logging.info("MultiGPU: GGUFDisTorch - Calling original GGUF loader") + loader_instance = original_loader() + return loader_instance.load_unet(unet_name, dequant_dtype, patch_dtype, patch_on_device) + + + # Create the MultiGPU version of the base class + UnetLoaderGGUFDisTorchMultiGPU = override_class_with_distorch(UnetLoaderGGUFDisTorch) + NODE_CLASS_MAPPINGS["UnetLoaderGGUFDisTorchMultiGPU"] = UnetLoaderGGUFDisTorchMultiGPU + logging.info(f"MultiGPU: Registered UnetLoaderGGUFDisTorchMultiGPU") + def register_CLIPLoaderGGUFMultiGPU(): global NODE_CLASS_MAPPINGS @@ -823,6 +1135,7 @@ 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()