From baa31a19617f794c01aedae34bd4477c1f48e16a Mon Sep 17 00:00:00 2001 From: John Pollock Date: Tue, 26 Aug 2025 14:08:30 -0500 Subject: [PATCH] refactor(distorch): Default unassigned tensor blocks to CPU Change the fallback assignment for tensor blocks that do not fit within any donor device's memory quota. These blocks are now assigned to the CPU instead of the primary compute device. This prevents potential VRAM Out-Of-Memory errors on the main GPU, improving the stability of the model loading process, especially under tight memory constraints. Additionally, the allocation log is now updated to include all available system devices, even those with zero allocation, to provide a more comprehensive report. --- distorch_2.py | 35 ++++++++++++----------------------- 1 file changed, 12 insertions(+), 23 deletions(-) diff --git a/distorch_2.py b/distorch_2.py index 3e74456..dd9f589 100644 --- a/distorch_2.py +++ b/distorch_2.py @@ -16,6 +16,7 @@ from collections import defaultdict import comfy.model_management as mm import comfy.model_patcher from . import current_device +from .nodes import get_device_list safetensor_allocation_store = {} safetensor_settings_store = {} @@ -148,6 +149,12 @@ def analyze_safetensor_loading(model_patcher, allocations_str): elif "%" in distorch_alloc: mode = "ratio" distorch_alloc = calculate_fraction_from_ratio_expert_string(model_patcher, distorch_alloc) + + all_devices = get_device_list() + present_devices = {item.split(',')[0] for item in distorch_alloc.split(';') if ',' in item} + for device in all_devices: + if device not in present_devices: + distorch_alloc += f";{device},0.0" logger.info(f"[MultiGPU_DisTorch2] Final Allocation String: {distorch_alloc}") @@ -169,7 +176,6 @@ def analyze_safetensor_loading(model_patcher, allocations_str): "alloc_gb": alloc_gb } - # Final Allocation Table logger.info(eq_line) logger.info(" DisTorch2 Model Device Allocations") logger.info(eq_line) @@ -180,7 +186,6 @@ def analyze_safetensor_loading(model_patcher, allocations_str): sorted_devices = sorted(device_table.keys(), key=lambda d: (d == "cpu", d)) - # Calculate total allocated model size for ratio calculation total_allocated_model_bytes = sum(d["alloc_gb"] * (1024**3) for d in device_table.values()) for dev in sorted_devices: @@ -188,7 +193,6 @@ def analyze_safetensor_loading(model_patcher, allocations_str): alloc_fraction = device_table[dev]["fraction"] alloc_gb = device_table[dev]["alloc_gb"] - # Calculate the distribution ratio percentage dist_ratio_percent = (alloc_gb * (1024**3) / total_allocated_model_bytes) * 100 if total_allocated_model_bytes > 0 else 0 logger.info(fmt_rosetta.format( @@ -208,15 +212,12 @@ def analyze_safetensor_loading(model_patcher, allocations_str): raw_block_list = model_patcher._load_list() - # Calculate total memory from ComfyUI's list (first pass replacement) total_memory = sum(module_size for module_size, _, _, _ in raw_block_list) - # Set the minimum block size threshold (0.01% of total model memory) MIN_BLOCK_THRESHOLD = total_memory * 0.0001 logger.debug(f"[MultiGPU_DisTorch2] Total model memory: {total_memory} bytes") logger.debug(f"[MultiGPU_DisTorch2] Tiny block threshold (0.01%): {MIN_BLOCK_THRESHOLD} bytes") - # Build all_blocks from ComfyUI's list (second pass replacement) all_blocks = [] for module_size, module_name, module_object, params in raw_block_list: block_type = type(module_object).__name__ @@ -225,7 +226,6 @@ def analyze_safetensor_loading(model_patcher, allocations_str): memory_by_type[block_type] += module_size all_blocks.append((module_name, module_object, block_type, module_size)) - # Filter out tiny blocks from the distribution list block_list = [b for b in all_blocks if b[3] >= MIN_BLOCK_THRESHOLD] tiny_block_list = [b for b in all_blocks if b[3] < MIN_BLOCK_THRESHOLD] @@ -247,15 +247,11 @@ def analyze_safetensor_loading(model_patcher, allocations_str): logger.info(dash_line) # Distribute blocks sequentially from the tail of the model + device_assignments = {device: [] for device in DEVICE_RATIOS_DISTORCH.keys()} block_assignments = {} - # Determine the primary compute device (first non-cpu device) - compute_device = "cuda:0" # Fallback - for dev in sorted_devices: - if dev != "cpu": - compute_device = dev - break + compute_device = str(current_device) # Create a memory quota for each donor device based on its calculated allocation. donor_devices = [d for d in sorted_devices if d != compute_device] @@ -267,7 +263,6 @@ def analyze_safetensor_loading(model_patcher, allocations_str): # Iterate from the TAIL of the model, assigning blocks to donors until their quotas are filled. for block_name, module, block_type, block_memory in reversed(block_list): assigned_to_donor = False - # Attempt to assign the block to a donor device that has quota remaining. for donor in donor_devices: if donor_quotas[donor] >= block_memory: block_assignments[block_name] = donor @@ -275,11 +270,9 @@ def analyze_safetensor_loading(model_patcher, allocations_str): assigned_to_donor = True break # Move to the next block - # If no donor had enough quota, assign it to the primary compute device. if not assigned_to_donor: - block_assignments[block_name] = compute_device + block_assignments[block_name] = "cpu" - # Explicitly assign tiny blocks to the compute device if tiny_block_list: for block_name, module, block_type, block_memory in tiny_block_list: block_assignments[block_name] = compute_device @@ -292,13 +285,11 @@ def analyze_safetensor_loading(model_patcher, allocations_str): device_assignments[device].append((b_name, b_module, b_type, b_mem)) break - # Log final assignments - IDENTICAL FORMAT TO GGML logger.info("DisTorch2 Model Final Device/Layer Assignments") logger.info(dash_line) logger.info(fmt_assign.format("Device", "Layers", "Memory (MB)", "% Total")) logger.info(dash_line) - # Calculate and log tiny blocks separately if tiny_block_list: tiny_block_memory = sum(b[3] for b in tiny_block_list) tiny_mem_mb = tiny_block_memory / (1024 * 1024) @@ -307,12 +298,10 @@ def analyze_safetensor_loading(model_patcher, allocations_str): logger.info(fmt_assign.format(device_label, str(len(tiny_block_list)), f"{tiny_mem_mb:.2f}", f"{tiny_mem_percent:.1f}%")) logger.debug(f"[MultiGPU_DisTorch2] Tiny block memory breakdown: {tiny_block_memory} bytes ({tiny_mem_mb:.2f} MB), which is {tiny_mem_percent:.4f}% of total model memory.") - # Log distributed blocks total_assigned_memory = 0 device_memories = {} for device, blocks in device_assignments.items(): - # Exclude tiny blocks from this calculation dist_blocks = [b for b in blocks if b[3] >= MIN_BLOCK_THRESHOLD] if not dist_blocks: continue @@ -420,7 +409,7 @@ def calculate_fraction_from_byte_expert_string(model_patcher, byte_str): allocation_parts.append(f"{dev},{fraction:.4f}") result_string = ";".join(allocation_parts) - logger.info(f"[MultiGPU_DisTorch2] Converted byte string '{byte_str}' to final fraction string: '{result_string}'") + return result_string def calculate_fraction_from_ratio_expert_string(model_patcher, ratio_str): @@ -470,7 +459,7 @@ def calculate_fraction_from_ratio_expert_string(model_patcher, ratio_str): logger.info(f"[MultiGPU_DisTorch2] Ratio(%) Mode - {ratio_str} -> {ratio_string} ratio, put {put_part}") result_string = ";".join(allocation_parts) - logger.info(f"[MultiGPU_DisTorch2] Converted ratio string to fraction string: {result_string}") + return result_string def calculate_safetensor_vvram_allocation(model_patcher, virtual_vram_str):