diff --git a/__init__.py b/__init__.py index d854750..8d1ef4a 100644 --- a/__init__.py +++ b/__init__.py @@ -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())}") \ No newline at end of file diff --git a/distorch.py b/distorch.py new file mode 100644 index 0000000..ddbc863 --- /dev/null +++ b/distorch.py @@ -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 \ No newline at end of file