241 lines
10 KiB
Python
241 lines
10 KiB
Python
"""
|
|
Block Swap Module for SafeTensor Models
|
|
Contains all SafeTensor DisTorch code for block-swap memory optimization
|
|
"""
|
|
|
|
import torch
|
|
import logging
|
|
import copy
|
|
from collections import defaultdict
|
|
import comfy.model_management as mm
|
|
import torch.nn as nn
|
|
from .model_sig import get_model_type
|
|
|
|
|
|
def log_memory_usage(device, stage=""):
|
|
"""Logs the memory usage of a given device."""
|
|
if not isinstance(device, torch.device):
|
|
device = torch.device(device)
|
|
|
|
if device.type == 'cuda':
|
|
stats = torch.cuda.memory_stats(device)
|
|
total_mem = mm.get_total_memory(device)
|
|
allocated = stats['allocated_bytes.all.current']
|
|
reserved = stats['reserved_bytes.all.current']
|
|
|
|
logging.info(
|
|
f"[MemLog] {stage} - {device}: "
|
|
f"Allocated: {allocated / 1024**2:.2f}MB, "
|
|
f"Reserved: {reserved / 1024**2:.2f}MB, "
|
|
f"Total: {total_mem / 1024**3:.2f}GB"
|
|
)
|
|
elif device.type == 'cpu':
|
|
# Basic CPU memory logging (less detailed than CUDA)
|
|
# This requires psutil, which might not be a dependency.
|
|
# For now, we'll just log that it's a CPU.
|
|
logging.info(f"[MemLog] {stage} - {device}: CPU memory logging is not as detailed.")
|
|
|
|
|
|
class FluxBlockSwapManager:
|
|
"""
|
|
Manages block-swapping for FLUX models using the original, fast `forward` patching method.
|
|
"""
|
|
def __init__(self, model_patcher):
|
|
self.model_patcher = model_patcher
|
|
self.patched_blocks = {}
|
|
|
|
def apply_patch(self, compute_device, swap_device, blocks_to_swap):
|
|
logging.info(f"[FluxBlockSwapManager] Applying forward patch to {len(blocks_to_swap)} blocks.")
|
|
for i, block in enumerate(blocks_to_swap):
|
|
block.to(swap_device)
|
|
original_forward = block.forward
|
|
|
|
def create_patched_forward(original_f, b, block_index, cd, sd):
|
|
def patched_forward(*args, **kwargs):
|
|
b.to(cd, non_blocking=True)
|
|
result = original_f(*args, **kwargs)
|
|
b.to(sd, non_blocking=True)
|
|
return result
|
|
return patched_forward
|
|
|
|
block.forward = create_patched_forward(original_forward, block, i, torch.device(compute_device), torch.device(swap_device))
|
|
self.patched_blocks[block] = original_forward
|
|
|
|
def cleanup(self):
|
|
logging.info(f"[FluxBlockSwapManager] Cleaning up {len(self.patched_blocks)} patched blocks.")
|
|
for block, original_forward in self.patched_blocks.items():
|
|
block.forward = original_forward
|
|
self.patched_blocks = {}
|
|
|
|
|
|
def analyze_safetensor_distorch(model, compute_device, swap_device, virtual_vram_gb, all_blocks):
|
|
"""Provides a detailed analysis of the block swap configuration, mimicking the GGUF DisTorch style."""
|
|
|
|
eq_line = "=" * 60
|
|
dash_line = "-" * 60
|
|
|
|
logging.info(eq_line)
|
|
logging.info(" DisTorch SafeTensor Memory Analysis")
|
|
logging.info(eq_line)
|
|
|
|
# Device Allocation Table
|
|
fmt_assign = "{:<12}{:>15}{:>15}{:>15}"
|
|
logging.info(fmt_assign.format("Device", "Role", "Total Mem (GB)", "Config (GB)"))
|
|
logging.info(dash_line)
|
|
|
|
compute_total_gb = mm.get_total_memory(torch.device(compute_device)) / (1024**3)
|
|
swap_total_gb = mm.get_total_memory(torch.device(swap_device)) / (1024**3)
|
|
|
|
logging.info(fmt_assign.format(compute_device, "Compute", f"{compute_total_gb:.2f}", ""))
|
|
logging.info(fmt_assign.format(swap_device, "Swap", f"{swap_total_gb:.2f}", f"Offload: {virtual_vram_gb:.2f}"))
|
|
logging.info(dash_line)
|
|
|
|
# Block Analysis Table
|
|
block_summary = defaultdict(lambda: {'count': 0, 'memory': 0})
|
|
total_memory = 0
|
|
|
|
for block in all_blocks:
|
|
block_type = type(block).__name__
|
|
block_memory = sum(p.numel() * p.element_size() for p in block.parameters())
|
|
block_summary[block_type]['count'] += 1
|
|
block_summary[block_type]['memory'] += block_memory
|
|
total_memory += block_memory
|
|
|
|
logging.info(" DisTorch SafeTensor Block Analysis")
|
|
logging.info(dash_line)
|
|
fmt_layer = "{:<20}{:>10}{:>15}{:>12}"
|
|
logging.info(fmt_layer.format("Block Type", "Count", "Memory (MB)", "% Total"))
|
|
logging.info(dash_line)
|
|
|
|
sorted_blocks = sorted(block_summary.items(), key=lambda x: x[1]['memory'], reverse=True)
|
|
|
|
for block_type, data in sorted_blocks:
|
|
mem_mb = data['memory'] / (1024 * 1024)
|
|
mem_percent = (data['memory'] / total_memory) * 100 if total_memory > 0 else 0
|
|
logging.info(fmt_layer.format(block_type, str(data['count']), f"{mem_mb:.2f}", f"{mem_percent:.1f}%"))
|
|
logging.info(dash_line)
|
|
|
|
# Final Assignment Table
|
|
model_size_gb = total_memory / (1024**3)
|
|
block_size_gb = model_size_gb / len(all_blocks) if all_blocks else 0
|
|
blocks_to_offload = int(virtual_vram_gb / block_size_gb) if block_size_gb > 0 else 0
|
|
blocks_on_compute = len(all_blocks) - blocks_to_offload
|
|
|
|
logging.info(" DisTorch Final Block Assignments")
|
|
logging.info(dash_line)
|
|
fmt_final = "{:<20}{:>15}"
|
|
logging.info(fmt_final.format("Total Model Size (GB):", f"{model_size_gb:.2f}"))
|
|
logging.info(fmt_final.format("Average Block Size (MB):", f"{block_size_gb * 1024:.2f}" if all_blocks else "N/A"))
|
|
logging.info(dash_line)
|
|
logging.info(fmt_final.format("Blocks on Compute:", f"{blocks_on_compute}"))
|
|
logging.info(fmt_final.format("Blocks on Swap:", f"{blocks_to_offload}"))
|
|
logging.info(eq_line)
|
|
|
|
|
|
def apply_block_swap(model_patcher, compute_device="cuda:0", swap_device="cpu",
|
|
virtual_vram_gb=4.0, expert_mode_allocations=""):
|
|
"""
|
|
Identifies the model type and applies the appropriate block swapping strategy.
|
|
"""
|
|
model_type = get_model_type(model_patcher)
|
|
logging.info(f"[BlockSwap] Detected model type: {model_type}")
|
|
|
|
if model_type == "FLUX":
|
|
manager = FluxBlockSwapManager(model_patcher)
|
|
|
|
model_to_patch = model_patcher.model.diffusion_model
|
|
|
|
all_blocks = []
|
|
if hasattr(model_to_patch, 'double_blocks'):
|
|
all_blocks.extend(model_to_patch.double_blocks)
|
|
if hasattr(model_to_patch, 'single_blocks'):
|
|
all_blocks.extend(model_to_patch.single_blocks)
|
|
|
|
if not all_blocks:
|
|
logging.error("[BlockSwap] CRITICAL: No swappable blocks found for FLUX model.")
|
|
return
|
|
|
|
analyze_safetensor_distorch(model_to_patch, compute_device, swap_device, virtual_vram_gb, all_blocks)
|
|
|
|
log_memory_usage(compute_device, "Before Swap")
|
|
log_memory_usage(swap_device, "Before Swap")
|
|
|
|
model_size_gb = sum(p.numel() * p.element_size() for p in model_to_patch.parameters()) / (1024**3)
|
|
block_size_gb = model_size_gb / len(all_blocks) if all_blocks else 0
|
|
blocks_to_offload_count = int(virtual_vram_gb / block_size_gb) if block_size_gb > 0 else 0
|
|
|
|
blocks_to_swap = all_blocks[-blocks_to_offload_count:] if blocks_to_offload_count > 0 else []
|
|
|
|
if not blocks_to_swap:
|
|
logging.warning("[BlockSwap] No blocks designated for swapping.")
|
|
return
|
|
|
|
manager.apply_patch(compute_device, swap_device, blocks_to_swap)
|
|
|
|
if not hasattr(model_patcher, 'block_swap_managers'):
|
|
model_patcher.block_swap_managers = []
|
|
model_patcher.block_swap_managers.append(manager)
|
|
|
|
log_memory_usage(compute_device, "After Swap")
|
|
log_memory_usage(swap_device, "After Swap")
|
|
|
|
logging.info("[BlockSwap] FLUX block swap setup complete.")
|
|
else:
|
|
logging.warning(f"[BlockSwap] Model type '{model_type}' is not yet supported for block swapping.")
|
|
|
|
|
|
def override_class_with_distorch_safetensor(cls):
|
|
"""DisTorch 2.0 wrapper for SafeTensor models, providing block-swap memory optimization."""
|
|
from .nodes import get_device_list
|
|
|
|
class NodeOverrideDisTorchSafeTensorv2(cls):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
inputs = copy.deepcopy(cls.INPUT_TYPES())
|
|
devices = get_device_list()
|
|
compute_device = devices[1] if len(devices) > 1 else devices[0]
|
|
|
|
inputs["optional"] = inputs.get("optional", {})
|
|
|
|
inputs["optional"]["compute_device"] = (devices, {"default": compute_device})
|
|
inputs["optional"]["virtual_vram_gb"] = ("FLOAT", {"default": 4.0, "min": 0.0, "max": 128.0, "step": 0.1})
|
|
inputs["optional"]["donor_device"] = (devices, {"default": "cpu"})
|
|
inputs["optional"]["expert_mode_allocations"] = ("STRING", {"multiline": False, "default": ""})
|
|
return inputs
|
|
|
|
CATEGORY = "multigpu/distorch_2"
|
|
FUNCTION = "override"
|
|
|
|
def override(self, *args, compute_device=None, virtual_vram_gb=4.0,
|
|
donor_device="cpu", expert_mode_allocations="", **kwargs):
|
|
from . import set_current_device
|
|
|
|
logging.info(f"[DisTorch SafeTensor] Override called with: compute_device={compute_device}, donor_device={donor_device}, virtual_vram_gb={virtual_vram_gb}")
|
|
|
|
if compute_device is not None:
|
|
set_current_device(compute_device)
|
|
|
|
fn = getattr(super(), cls.FUNCTION)
|
|
out = fn(*args, **kwargs)
|
|
|
|
model = out[0]
|
|
if hasattr(model, 'model'):
|
|
logging.info("[DisTorch SafeTensor] Model has 'model' attribute, applying block swap.")
|
|
apply_block_swap(
|
|
model,
|
|
compute_device=compute_device,
|
|
swap_device=donor_device,
|
|
virtual_vram_gb=virtual_vram_gb,
|
|
expert_mode_allocations=expert_mode_allocations
|
|
)
|
|
else:
|
|
logging.warning("[DisTorch SafeTensor] Loaded object does not have a 'model' attribute, skipping block swap.")
|
|
|
|
return out
|
|
|
|
return NodeOverrideDisTorchSafeTensorv2
|
|
|
|
|
|
# For backwards compatibility, keep the old name pointing to the new safetensor wrapper
|
|
override_class_with_distorch_bs = override_class_with_distorch_safetensor
|