1959 lines
82 KiB
Python
1959 lines
82 KiB
Python
PROFILE = False
|
||
|
||
import comfy.sd
|
||
import folder_paths
|
||
import datetime
|
||
from functools import wraps
|
||
from torch import nn
|
||
import time
|
||
import torch
|
||
import gc
|
||
import threading
|
||
import warnings
|
||
import psutil
|
||
import types
|
||
import torch._dynamo
|
||
from functools import partial
|
||
import torch.utils.checkpoint
|
||
from collections import defaultdict
|
||
import bitsandbytes as bnb
|
||
|
||
|
||
if PROFILE == True:
|
||
try:
|
||
import torch.profiler
|
||
import torch.cuda.nvtx as nvtx
|
||
except ImportError:
|
||
PROFILE = False
|
||
nvtx = None
|
||
"""
|
||
Dynamic GPU Layer Swapping for ComfyUI by obisin
|
||
======================================
|
||
Enhanced inference with automatic memory optimization for diffusion models.
|
||
Supports Cosmos, Flux, Wan2.1, Wan2.2, HunyuanVideo, and generic transformer blocks.
|
||
|
||
PERFORMANCE DESIGN NOTE:
|
||
This file uses module-level functions with local variable binding for optimal performance.
|
||
Hot path functions bind frequently-accessed attributes to locals at function start (LOAD_FAST)
|
||
to avoid Python's attribute lookup overhead (LOAD_ATTR). Call chains are minimized to
|
||
reduce function call overhead.
|
||
|
||
Architecture:
|
||
- Globals: Intentionally used for performance (documented below)
|
||
- Hot functions: Bind globals to locals once per call
|
||
- Minimal call depth: Direct function calls without abstraction layers
|
||
- Factory pattern: Clean call sites with fast execution
|
||
|
||
CRITICAL PATH: Functions marked "HOT PATH" - profile before changing!
|
||
|
||
Author: obisin
|
||
"""
|
||
|
||
class Args:
|
||
def __init__(self):
|
||
# Set defaults - these will be overridden by ComfyUI node
|
||
self.dynamic_swapping = True
|
||
self.initial_gpu_layers = 1 # Number of initial layers to keep permanently on GPU
|
||
self.final_gpu_layers = 1 # Number of final layers to keep permanently on GPU
|
||
self.gpu_layers = None # Comma-separated list of layer indices to keep permanently on GPU
|
||
self.prefetch = 0 # Number of layers to prefetch ahead (0=off, 1=single layer, 2+=multiple layers)
|
||
self.threading = False # Enable background threading for automatic layer management
|
||
self.cuda_streams = False # Enable CUDA streams for copy-compute overlap (needs more VRAM)
|
||
self.batch_move = False # Use batch layer moving (experimental, may cause device issues)
|
||
self.cast_target = None # Cast FROM dtype TO dtype at start-up (e.g., f32 bf16)
|
||
self.selective_packing = 0 # Size threshold in MB for using packed transfers (default: 64MB)
|
||
self.event_sync = False # Use CUDA events instead of torch.cuda.synchronize() for better performance
|
||
self.verbose = False # Enable verbose output with detailed timing and transfer information
|
||
self.device_sync_mode = 'off' # Not needed for inference
|
||
self.sync_only_on_resume = False
|
||
self.compile = False
|
||
self.autocast = None
|
||
self.mixed_precision = 'auto'
|
||
self.compute_casting = "disabled"
|
||
|
||
args = Args()
|
||
|
||
|
||
TOOLTIPS = {
|
||
"dynamic_swapping": "Smart dynamic layer swapping between GPU and CPU for optimal performance",
|
||
"initial_gpu_layers": "Number of initial layers to keep permanently on GPU. If not specified, uses reasonable defaults based on estimated VRAM",
|
||
"final_gpu_layers": "Number of final layers to keep permanently on GPU",
|
||
"gpu_layers": "Comma-separated list of layer indices to keep permanently on GPU (e.g., '0,1,2,14,18,19,20,21,22'). Overrides initial_gpu_layers and final_gpu_layers",
|
||
"prefetch": "Number of layers to prefetch ahead (0=off, 1=single layer, 2+=multiple layers), might not work with mixed layer type models",
|
||
"threading": "Enable background threading for automatic layer management",
|
||
"cuda_streams": "Enable CUDA streams for copy-compute overlap (needs more VRAM)",
|
||
"batch_move": "Use batch layer moving (experimental, may cause device issues)",
|
||
"cast_target": "Cast FROM dtype TO dtype at start-up (e.g., f32 bf16) choices=[f32, bf16, f16, f8_e4m3, f8_e5m2, nf4, fp4]", #DISABLED DUE TO COMFY AUTOCAST
|
||
"selective_packing": "Size threshold in MB for using packed transfers (default: 64MB)",
|
||
"event_sync": "Use CUDA events instead of torch.cuda.synchronize() for better performance",
|
||
"compile": "Enable torch.compile optimization for permanent GPU layers (experimental)",
|
||
"verbose": "Enable verbose output with detailed timing and transfer information"
|
||
}
|
||
|
||
# =============================================================================
|
||
# GLOBAL VARIABLES
|
||
# =============================================================================
|
||
|
||
swap_stats = {'to_gpu': 0, 'to_cpu': 0}
|
||
transfer_events = {}
|
||
transfer_stats = {
|
||
'to_gpu_times': [], 'to_cpu_times': [], 'to_gpu_speeds': [], 'to_cpu_speeds': [],
|
||
'current_step_gpu_times': [], 'current_step_cpu_times': [],
|
||
'current_step_gpu_speeds': [], 'current_step_cpu_speeds': []
|
||
}
|
||
device_cache = None
|
||
packed_layers = {}
|
||
layer_sizes_mb = {}
|
||
# VERBOSE = True
|
||
|
||
GPU_DEVICE = 'cuda'
|
||
CPU_DEVICE = 'cpu'
|
||
|
||
layers = None
|
||
|
||
cpu_swappable_layers = set()
|
||
gpu_resident_layers = set()
|
||
|
||
GPU_DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
CPU_DEVICE = 'cpu'
|
||
|
||
print(f"Device configuration: GPU={GPU_DEVICE}, CPU={CPU_DEVICE}")
|
||
casting_handler = None
|
||
|
||
# =============================================================================
|
||
# CORE FUNCTIONS by obisin
|
||
# =============================================================================
|
||
|
||
def fix_inference_tensor_parameters(layer):
|
||
"""Fix inference tensor parameters in a layer if needed"""
|
||
# Quick check if ANY parameters need fixing
|
||
needs_fixing = False
|
||
with torch.no_grad():
|
||
for name, param in layer.named_parameters():
|
||
if param is not None:
|
||
try:
|
||
_ = param._version
|
||
except (RuntimeError, AttributeError):
|
||
needs_fixing = True
|
||
break
|
||
|
||
# Only fix if needed
|
||
if needs_fixing:
|
||
with torch.no_grad():
|
||
for name, param in layer.named_parameters():
|
||
if param is not None:
|
||
try:
|
||
_ = param._version
|
||
except (RuntimeError, AttributeError):
|
||
param.data = param.data.detach().clone()
|
||
|
||
return needs_fixing
|
||
|
||
|
||
def get_cached_layer_size_mb(idx):
|
||
return layer_sizes_mb.get(idx, 0)
|
||
|
||
def event_based_sync(operation_name, idx=None):
|
||
"""Replace torch.cuda.synchronize() with specific event tracking"""
|
||
if args.event_sync:
|
||
event = torch.cuda.Event()
|
||
event.record()
|
||
|
||
key = f"{operation_name}_{idx}" if idx is not None else operation_name
|
||
transfer_events[key] = event
|
||
return event
|
||
else:
|
||
torch.cuda.synchronize()
|
||
return None
|
||
|
||
|
||
class LayerDeviceCache:
|
||
def __init__(self, model, layers): # Add layers parameter
|
||
self.cache = {}
|
||
self.dirty = set()
|
||
# Initialize cache using the layers list
|
||
for i, layer in enumerate(layers):
|
||
self.cache[i] = self.get_layer_device(layer)
|
||
|
||
def get_layer_device(self, layer):
|
||
"""Your existing function"""
|
||
try:
|
||
return next(layer.parameters()).device
|
||
except StopIteration:
|
||
return torch.device('cpu')
|
||
|
||
def get_device(self, layer_idx): # Simplified - no need for model parameter
|
||
"""Fast cached lookup"""
|
||
if layer_idx in self.dirty:
|
||
if layers is not None and layer_idx < len(layers):
|
||
self.cache[layer_idx] = self.get_layer_device(layers[layer_idx])
|
||
self.dirty.remove(layer_idx)
|
||
else:
|
||
return torch.device('cpu') # Safe fallback
|
||
return self.cache.get(layer_idx, torch.device('cpu'))
|
||
|
||
def mark_moved(self, layer_idx, new_device):
|
||
"""Update cache when we move a layer"""
|
||
self.cache[layer_idx] = new_device
|
||
|
||
|
||
|
||
def safe_move_to_cpu(layer, idx):
|
||
"""Move layer back to CPU and clean up GPU memory"""
|
||
if PROFILE:
|
||
nvtx.range_push(f"CPU_Transfer_L{idx}")
|
||
try:
|
||
try:
|
||
current_device = next(layer.parameters()).device
|
||
if current_device.type == 'cpu':
|
||
return # Already on CPU
|
||
except StopIteration:
|
||
# Layer has no parameters
|
||
pass
|
||
|
||
transfer_time = 0
|
||
|
||
if VERBOSE:
|
||
start_time = time.time()
|
||
layer_size_mb = get_cached_layer_size_mb(idx)
|
||
layer.to(CPU_DEVICE)
|
||
add_smart_swapping_to_layer.cleanup_stream = torch.cuda.Stream()
|
||
swap_stats['to_cpu'] += 1
|
||
if VERBOSE:
|
||
end_time = time.time()
|
||
transfer_time = end_time - start_time
|
||
|
||
if transfer_time > 0 and VERBOSE:
|
||
speed_mbps = layer_size_mb / transfer_time
|
||
|
||
# Track current step stats
|
||
transfer_stats['current_step_cpu_times'].append(transfer_time)
|
||
transfer_stats['current_step_cpu_speeds'].append(speed_mbps)
|
||
|
||
device_cache.mark_moved(idx, torch.device('cpu'))
|
||
|
||
# if idx % 10 == 0:
|
||
# print(f" Layer {idx} → CPU, memory: {torch.cuda.memory_allocated(0) / 1e9:.1f}GB")
|
||
|
||
return True
|
||
finally:
|
||
if PROFILE:
|
||
nvtx.range_pop()
|
||
|
||
|
||
def safe_move_to_gpu(layer, idx):
|
||
"""Move layer to GPU with dtype casting"""
|
||
|
||
if PROFILE:
|
||
nvtx.range_push(f"GPU_Transfer_L{idx}")
|
||
try:
|
||
try:
|
||
|
||
current_device = device_cache.get_device(idx)
|
||
if current_device.type == 'cuda':
|
||
return True
|
||
|
||
if VERBOSE:
|
||
start_time = time.time()
|
||
layer_size_mb = get_cached_layer_size_mb(idx)
|
||
|
||
layer.to(GPU_DEVICE, non_blocking=True)
|
||
|
||
if VERBOSE:
|
||
end_time = time.time()
|
||
transfer_time = end_time - start_time
|
||
|
||
if transfer_time > 0 and VERBOSE:
|
||
speed_mbps = layer_size_mb / transfer_time
|
||
# Track current step stats
|
||
transfer_stats['current_step_gpu_times'].append(transfer_time)
|
||
transfer_stats['current_step_gpu_speeds'].append(speed_mbps)
|
||
|
||
event_based_sync("gpu_transfer", idx)
|
||
device_cache.mark_moved(idx, torch.device('cuda'))
|
||
swap_stats['to_gpu'] += 1
|
||
return True
|
||
|
||
except RuntimeError as e:
|
||
if "out of memory" in str(e):
|
||
return False
|
||
raise e
|
||
finally:
|
||
if PROFILE:
|
||
nvtx.range_pop()
|
||
|
||
|
||
# THREADING: Thread-safe GPU operations
|
||
def safe_move_to_gpu_threaded(layer, idx):
|
||
"""Thread-safe move to GPU"""
|
||
with add_smart_swapping_to_layer.gpu_lock:
|
||
if device_cache.get_device(idx).type == 'cpu':
|
||
return safe_move_to_gpu(layer, idx)
|
||
return True
|
||
|
||
|
||
def safe_move_to_cpu_threaded(layer, idx):
|
||
"""Thread-safe move to CPU"""
|
||
with add_smart_swapping_to_layer.gpu_lock:
|
||
if layers is not None and device_cache is not None and device_cache.get_device(idx).type == 'cuda':
|
||
return safe_move_to_cpu(layer, idx)
|
||
return True
|
||
|
||
def stop_background_threading():
|
||
"""Stop the background threading system"""
|
||
if hasattr(add_smart_swapping_to_layer, 'training_active'):
|
||
add_smart_swapping_to_layer.training_active = False
|
||
if add_smart_swapping_to_layer.background_thread:
|
||
add_smart_swapping_to_layer.background_thread.join(timeout=1.0)
|
||
print(" Background threading stopped")
|
||
|
||
|
||
def calculate_needed_layers(layer_idx, prefetch):
|
||
needed = set()
|
||
needed.add(layer_idx)
|
||
|
||
if layers is None:
|
||
return needed
|
||
|
||
for i in range(1, prefetch + 1):
|
||
prefetch_idx = layer_idx + i
|
||
|
||
# CIRCULAR: If we go past the end, wrap to the beginning of swappable layers
|
||
if prefetch_idx >= len(layers):
|
||
prefetch_idx = min(cpu_swappable_layers) + (prefetch_idx - len(layers))
|
||
|
||
if prefetch_idx in cpu_swappable_layers:
|
||
needed.add(prefetch_idx)
|
||
|
||
return needed
|
||
|
||
|
||
def cleanup_excess_layers(keep_layers):
|
||
"""Remove layers from GPU that are not in keep_layers set"""
|
||
# current_step = getattr(add_smart_swapping_to_layer, 'current_sampling_step', 0)
|
||
|
||
if PROFILE:
|
||
nvtx.range_push(f"Cleanup_Excess_{len(cpu_swappable_layers - keep_layers)}layers")
|
||
try:
|
||
if args.batch_move:
|
||
# Batch approach (your current code)
|
||
layers_to_remove = []
|
||
for idx in cpu_swappable_layers:
|
||
if (idx < len(layers) and
|
||
idx not in keep_layers and
|
||
device_cache.get_device(idx).type == 'cuda'):
|
||
layers_to_remove.append(idx)
|
||
|
||
|
||
cleaned_count = batch_safe_move_to_cpu(layers_to_remove)
|
||
else:
|
||
# Individual approach (fallback)
|
||
cleaned_count = 0
|
||
layers_to_remove = []
|
||
for idx in cpu_swappable_layers:
|
||
if (idx < len(layers) and
|
||
idx not in keep_layers and
|
||
device_cache.get_device(idx).type == 'cuda'):
|
||
layers_to_remove.append(idx)
|
||
safe_move_to_cpu(layers[idx], idx)
|
||
cleaned_count += 1
|
||
|
||
return cleaned_count
|
||
finally:
|
||
if PROFILE:
|
||
nvtx.range_pop()
|
||
|
||
|
||
def fetch_missing_layers(needed_layers):
|
||
"""Ensure all needed layers are on GPU"""
|
||
current_step = getattr(add_smart_swapping_to_layer, 'current_sampling_step', 0)
|
||
|
||
layers_to_fetch = [idx for idx in needed_layers
|
||
if (idx < len(layers) and
|
||
idx in cpu_swappable_layers and
|
||
device_cache.get_device(idx).type == 'cpu')]
|
||
|
||
|
||
if PROFILE:
|
||
nvtx.range_push(f"Fetch_Missing_{len(needed_layers)}layers")
|
||
try:
|
||
# Your existing fetch logic...
|
||
fetched_count = 0
|
||
if args.batch_move:
|
||
fetched_count = batch_safe_move_to_gpu(layers_to_fetch)
|
||
else:
|
||
for idx in layers_to_fetch:
|
||
success = safe_move_to_gpu(layers[idx], idx)
|
||
if success:
|
||
fetched_count += 1
|
||
|
||
|
||
return fetched_count
|
||
finally:
|
||
if PROFILE:
|
||
nvtx.range_pop()
|
||
|
||
def batch_safe_move_to_gpu_packed(layer_indices, threshold_mb=64):
|
||
"""Batch move to GPU with selective packing"""
|
||
if not layer_indices:
|
||
return 0
|
||
|
||
moved_count = 0
|
||
large_layers = []
|
||
small_layers = []
|
||
|
||
# Separate layers by size
|
||
for idx in layer_indices:
|
||
if (idx < len(layers) and
|
||
device_cache.get_device(idx).type == 'cpu'):
|
||
|
||
layer_size_mb = get_cached_layer_size_mb(idx)
|
||
|
||
if layer_size_mb > threshold_mb and idx in packed_layers:
|
||
large_layers.append(idx)
|
||
else:
|
||
small_layers.append(idx)
|
||
|
||
# Handle large layers with packed transfer
|
||
if large_layers:
|
||
for idx in large_layers:
|
||
success = safe_move_to_gpu_packed(layers[idx], idx)
|
||
if success:
|
||
moved_count += 1
|
||
|
||
# Handle small layers with direct transfer
|
||
if small_layers:
|
||
for idx in small_layers:
|
||
success = safe_move_to_gpu(layers[idx], idx)
|
||
if success:
|
||
moved_count += 1
|
||
|
||
return moved_count
|
||
|
||
|
||
def batch_safe_move_to_cpu_packed(layer_indices, threshold_mb=64):
|
||
"""Batch move to CPU with selective packing"""
|
||
if not layer_indices:
|
||
return 0
|
||
|
||
moved_count = 0
|
||
large_layers = []
|
||
small_layers = []
|
||
|
||
# Separate layers by size
|
||
for idx in layer_indices:
|
||
if (idx < len(layers) and
|
||
device_cache.get_device(idx).type == 'cuda'):
|
||
|
||
layer_size_mb = get_cached_layer_size_mb(idx)
|
||
|
||
if layer_size_mb > threshold_mb:
|
||
large_layers.append(idx)
|
||
else:
|
||
small_layers.append(idx)
|
||
|
||
# Handle large layers with packed transfer
|
||
if large_layers:
|
||
for idx in large_layers:
|
||
success = safe_move_to_cpu_packed(layers[idx], idx)
|
||
if success:
|
||
moved_count += 1
|
||
|
||
# Handle small layers with direct transfer
|
||
if small_layers:
|
||
for idx in small_layers:
|
||
success = safe_move_to_cpu(layers[idx], idx)
|
||
if success:
|
||
moved_count += 1
|
||
|
||
return moved_count
|
||
|
||
|
||
def cleanup_excess_layers_packed(keep_layers, threshold_mb=64):
|
||
"""Remove layers from GPU with selective packing awareness"""
|
||
|
||
if args.batch_move:
|
||
# Collect all layers to remove, then use batch function
|
||
layers_to_remove = []
|
||
for idx in cpu_swappable_layers:
|
||
if (idx < len(layers) and
|
||
idx not in keep_layers and
|
||
device_cache.get_device(idx).type == 'cuda'):
|
||
layers_to_remove.append(idx)
|
||
return batch_safe_move_to_cpu_packed(layers_to_remove, threshold_mb)
|
||
|
||
else:
|
||
# Individual processing
|
||
moved_count = 0
|
||
for idx in cpu_swappable_layers:
|
||
if (idx < len(layers) and
|
||
idx not in keep_layers and
|
||
device_cache.get_device(idx).type == 'cuda'):
|
||
|
||
layer_size_mb = get_cached_layer_size_mb(idx)
|
||
|
||
if layer_size_mb > threshold_mb:
|
||
success = safe_move_to_cpu_packed(layers[idx], idx)
|
||
else:
|
||
success = safe_move_to_cpu(layers[idx], idx)
|
||
|
||
if success:
|
||
moved_count += 1
|
||
|
||
return moved_count
|
||
|
||
|
||
def fetch_missing_layers_packed(needed_layers, threshold_mb=64):
|
||
"""Use packed transfers only for large layers, direct for small ones"""
|
||
|
||
if args.batch_move:
|
||
# Use batch function
|
||
return batch_safe_move_to_gpu_packed(needed_layers, threshold_mb)
|
||
|
||
else:
|
||
# Individual processing
|
||
fetched_count = 0
|
||
for idx in needed_layers:
|
||
if (idx < len(layers) and
|
||
idx in cpu_swappable_layers and
|
||
device_cache.get_device(idx).type == 'cpu'):
|
||
|
||
layer_size_mb = get_cached_layer_size_mb(idx)
|
||
|
||
if layer_size_mb > threshold_mb and idx in packed_layers:
|
||
success = safe_move_to_gpu_packed(layers[idx], idx)
|
||
else:
|
||
success = safe_move_to_gpu(layers[idx], idx)
|
||
|
||
if success:
|
||
fetched_count += 1
|
||
|
||
return fetched_count
|
||
|
||
|
||
def background_layer_manager():
|
||
"""Background thread that maintains sliding window of layers"""
|
||
last_seen_step = -1
|
||
|
||
while add_smart_swapping_to_layer.training_active:
|
||
try:
|
||
# Detect new sampling step - reset and be extra aggressive
|
||
current_step = getattr(add_smart_swapping_to_layer, 'current_sampling_step', 0)
|
||
current_idx = add_smart_swapping_to_layer.current_layer_idx
|
||
|
||
if current_step != last_seen_step:
|
||
# print(f"🔄 Background thread detected step {current_step}, preloading early layers")
|
||
last_seen_step = current_step
|
||
|
||
# Preload for new step using actual prefetch setting
|
||
first_swappable = min(cpu_swappable_layers)
|
||
next_step_needed = calculate_needed_layers(first_swappable, args.prefetch)
|
||
|
||
for layer_idx in next_step_needed:
|
||
if device_cache.get_device(layer_idx).type == 'cpu':
|
||
safe_move_to_gpu_threaded(layers[layer_idx], layer_idx)
|
||
|
||
#OVERLY AGGRESSIVE DSEBUG VERSION
|
||
# first_swappable = min(cpu_swappable_layers)
|
||
# if device_cache.get_device(first_swappable).type == 'cpu':
|
||
# # print(f"🎯 Pre-warming first swappable layer {first_swappable}")
|
||
# safe_move_to_gpu_threaded(layers[first_swappable], first_swappable)
|
||
#
|
||
# # Aggressively preload first few layers for new step
|
||
# for early_idx in range(min(5, len(layers))): # First 5 layers
|
||
# if (early_idx in cpu_swappable_layers and
|
||
# device_cache.get_device(early_idx).type == 'cpu'):
|
||
# safe_move_to_gpu_threaded(layers[early_idx], early_idx)
|
||
|
||
if device_cache is None or layers is None:
|
||
time.sleep(0.1) # Wait for initialization
|
||
continue
|
||
|
||
if PROFILE:
|
||
nvtx.range_push("Background_Manager_Cycle")
|
||
|
||
current_idx = add_smart_swapping_to_layer.current_layer_idx
|
||
|
||
needed_layers = set()
|
||
for future_idx in range(current_idx, min(current_idx + args.prefetch + 2, len(layers))):
|
||
if future_idx in cpu_swappable_layers:
|
||
needed_layers.add(future_idx)
|
||
|
||
# Also include the standard prefetch calculation
|
||
standard_needed = calculate_needed_layers(current_idx, args.prefetch)
|
||
needed_layers.update(standard_needed)
|
||
|
||
# needed_layers = calculate_needed_layers(current_idx, args.prefetch)
|
||
|
||
# Collect layers to move instead of moving one by one
|
||
layers_to_cpu = []
|
||
layers_to_gpu = []
|
||
|
||
if args.batch_move:
|
||
# One in, one out: remove old layers, add new ones
|
||
for idx in cpu_swappable_layers:
|
||
if idx not in needed_layers and device_cache.get_device(idx).type == 'cuda':
|
||
layers_to_cpu.append(idx)
|
||
|
||
for idx in needed_layers:
|
||
if device_cache.get_device(idx).type == 'cpu':
|
||
layers_to_gpu.append(idx)
|
||
|
||
# Batch move with thread safety
|
||
if layers_to_cpu or layers_to_gpu:
|
||
with add_smart_swapping_to_layer.gpu_lock:
|
||
if layers_to_cpu:
|
||
batch_safe_move_to_cpu(layers_to_cpu)
|
||
if layers_to_gpu:
|
||
batch_safe_move_to_gpu(layers_to_gpu)
|
||
else:
|
||
# One in, one out: remove old layers, add new ones
|
||
for idx in cpu_swappable_layers:
|
||
if idx not in needed_layers and device_cache.get_device(idx).type == 'cuda':
|
||
safe_move_to_cpu_threaded(layers[idx], idx)
|
||
|
||
for idx in needed_layers:
|
||
if device_cache.get_device(idx).type == 'cpu':
|
||
safe_move_to_gpu_threaded(layers[idx], idx)
|
||
|
||
time.sleep(0.000001) #do not touch. very needed for sync
|
||
|
||
except Exception as e:
|
||
print(f" Background thread error: {e}")
|
||
time.sleep(0.1)
|
||
finally:
|
||
if PROFILE:
|
||
nvtx.range_pop()
|
||
|
||
|
||
def batch_safe_move_to_cpu(layer_indices):
|
||
"""Move multiple layers to CPU in batch"""
|
||
if not layer_indices:
|
||
return 0
|
||
if PROFILE:
|
||
nvtx.range_push(f"Batch_CPU_Transfer_{len(layer_indices)}layers")
|
||
try:
|
||
moved_count = 0
|
||
for idx in layer_indices:
|
||
if (idx < len(layers) and
|
||
device_cache.get_device(idx).type == 'cuda'):
|
||
|
||
# Measure transfer timing
|
||
transfer_time = 0
|
||
if VERBOSE:
|
||
start_time = time.time()
|
||
layer_size_mb = get_cached_layer_size_mb(idx)
|
||
layers[idx].to(CPU_DEVICE)
|
||
if VERBOSE:
|
||
end_time = time.time()
|
||
transfer_time = end_time - start_time
|
||
|
||
if transfer_time > 0 and VERBOSE:
|
||
speed_mbps = layer_size_mb / transfer_time
|
||
|
||
# Track current step stats
|
||
transfer_stats['current_step_cpu_times'].append(transfer_time)
|
||
transfer_stats['current_step_cpu_speeds'].append(speed_mbps)
|
||
device_cache.mark_moved(idx, torch.device('cpu'))
|
||
swap_stats['to_cpu'] += 1
|
||
moved_count += 1
|
||
|
||
return moved_count
|
||
finally:
|
||
if PROFILE:
|
||
nvtx.range_pop()
|
||
|
||
def batch_safe_move_to_gpu(layer_indices):
|
||
"""Batch move to GPU with dtype casting"""
|
||
if not layer_indices:
|
||
return 0
|
||
|
||
if PROFILE:
|
||
nvtx.range_push(f"Batch_GPU_Transfer_{len(layer_indices)}layers")
|
||
try:
|
||
|
||
moved_count = 0
|
||
for idx in layer_indices:
|
||
if (idx < len(layers) and
|
||
device_cache.get_device(idx).type == 'cpu'):
|
||
|
||
layer = layers[idx]
|
||
transfer_time = 0
|
||
|
||
if VERBOSE:
|
||
start_time = time.time()
|
||
layer_size_mb = get_cached_layer_size_mb(idx)
|
||
|
||
layer.to(GPU_DEVICE, non_blocking=True)
|
||
|
||
|
||
if VERBOSE:
|
||
end_time = time.time()
|
||
transfer_time = end_time - start_time
|
||
|
||
event_based_sync("gpu_transfer", idx)
|
||
if transfer_time > 0 and VERBOSE:
|
||
speed_mbps = layer_size_mb / transfer_time
|
||
# Track current step stats
|
||
transfer_stats['current_step_gpu_times'].append(transfer_time)
|
||
transfer_stats['current_step_gpu_speeds'].append(speed_mbps)
|
||
|
||
device_cache.mark_moved(idx, torch.device('cuda'))
|
||
swap_stats['to_gpu'] += 1
|
||
moved_count += 1
|
||
|
||
return moved_count
|
||
finally:
|
||
if PROFILE:
|
||
nvtx.range_pop()
|
||
|
||
# =============================================================================
|
||
# PACKED CODE by obisin
|
||
# =============================================================================
|
||
|
||
class PackedCPUBlock:
|
||
def __init__(self, layer):
|
||
self.packed_buffers = {} # One buffer per dtype
|
||
self.tensor_specs = []
|
||
self.total_elements = 0
|
||
self.gpu_blocks = {} # Keep GPU buffers alive when resident
|
||
self.gpu_events = {} # CUDA events for synchronization
|
||
self.pack_layer(layer)
|
||
|
||
def pack_layer(self, layer):
|
||
"""Pack all layer parameters/buffers into contiguous CPU blocks by dtype"""
|
||
dtype_groups = {}
|
||
|
||
# Build param/buffer maps once (avoid rebuilding in loops)
|
||
param_map = layer._parameters
|
||
buffer_map = layer._buffers
|
||
|
||
# Collect parameters by dtype
|
||
for name, param in layer.named_parameters(recurse=False):
|
||
if param is None:
|
||
continue
|
||
|
||
# Move to CPU first if needed (read-only copy during collection)
|
||
src_data = param.detach().to('cpu', copy=True) if param.device.type != 'cpu' else param.data
|
||
dtype = param.dtype
|
||
|
||
if dtype not in dtype_groups:
|
||
dtype_groups[dtype] = []
|
||
dtype_groups[dtype].append({
|
||
'name': name,
|
||
'data': src_data.flatten(),
|
||
'shape': param.shape,
|
||
'is_param': True,
|
||
'param_ref': param
|
||
})
|
||
|
||
# Collect buffers by dtype
|
||
for name, buffer in layer.named_buffers(recurse=True):
|
||
if buffer is None:
|
||
continue
|
||
|
||
# Move to CPU first if needed (read-only copy during collection)
|
||
src_data = buffer.detach().to('cpu', copy=True) if buffer.device.type != 'cpu' else buffer.data
|
||
dtype = buffer.dtype
|
||
|
||
if dtype not in dtype_groups:
|
||
dtype_groups[dtype] = []
|
||
dtype_groups[dtype].append({
|
||
'name': name,
|
||
'data': src_data.flatten(),
|
||
'shape': buffer.shape,
|
||
'is_param': False,
|
||
'buffer_ref': buffer
|
||
})
|
||
|
||
# Pack each dtype group into contiguous buffers
|
||
for dtype, tensors in dtype_groups.items():
|
||
if not tensors:
|
||
continue
|
||
|
||
total_size = sum(t['data'].numel() for t in tensors)
|
||
|
||
packed_buffer = torch.empty(total_size, dtype=dtype)
|
||
|
||
self.packed_buffers[dtype] = packed_buffer
|
||
|
||
# Pack data and create specs
|
||
offset = 0
|
||
for tensor_info in tensors:
|
||
data = tensor_info['data']
|
||
size = data.numel()
|
||
|
||
# Copy into packed buffer
|
||
packed_buffer[offset:offset + size] = data
|
||
|
||
# Store spec for later rebinding
|
||
spec = {
|
||
'name': tensor_info['name'],
|
||
'dtype': dtype,
|
||
'offset': offset,
|
||
'size': size,
|
||
'shape': tensor_info['shape'],
|
||
'is_param': tensor_info['is_param'],
|
||
'param_ref': tensor_info.get('param_ref'),
|
||
'buffer_ref': tensor_info.get('buffer_ref')
|
||
}
|
||
self.tensor_specs.append(spec)
|
||
offset += size
|
||
|
||
self.total_elements += total_size
|
||
|
||
# ATOMIC COMMIT: Rebind CPU params/buffers to views into packed_buffers
|
||
# This eliminates duplicate host RAM
|
||
for spec in self.tensor_specs:
|
||
dtype = spec['dtype']
|
||
start_idx = spec['offset']
|
||
end_idx = start_idx + spec['size']
|
||
cpu_view = self.packed_buffers[dtype][start_idx:end_idx].view(spec['shape'])
|
||
|
||
if spec['is_param']:
|
||
spec['param_ref'].data = cpu_view
|
||
else:
|
||
spec['buffer_ref'].data = cpu_view
|
||
|
||
def unpack_to_gpu(self, layer):
|
||
"""Move packed buffers to GPU and rebind layer parameters"""
|
||
# Group specs by dtype for efficient transfer
|
||
dtype_specs = {}
|
||
for spec in self.tensor_specs:
|
||
dtype = spec['dtype']
|
||
if dtype not in dtype_specs:
|
||
dtype_specs[dtype] = []
|
||
dtype_specs[dtype].append(spec)
|
||
|
||
# Transfer each dtype group
|
||
for dtype, specs in dtype_specs.items():
|
||
if dtype not in self.packed_buffers:
|
||
continue
|
||
|
||
packed_cpu = self.packed_buffers[dtype]
|
||
|
||
gpu_buffer = packed_cpu.to(GPU_DEVICE, non_blocking=True)
|
||
|
||
# Record event and keep buffer alive
|
||
event = torch.cuda.Event()
|
||
event.record()
|
||
self.gpu_blocks[dtype] = gpu_buffer
|
||
self.gpu_events[dtype] = event
|
||
|
||
# ATOMIC COMMIT: Rebind all params/buffers for this dtype
|
||
for spec in specs:
|
||
start_idx = spec['offset']
|
||
end_idx = start_idx + spec['size']
|
||
gpu_view = gpu_buffer[start_idx:end_idx].view(spec['shape'])
|
||
|
||
if spec['is_param']:
|
||
spec['param_ref'].data = gpu_view
|
||
else:
|
||
spec['buffer_ref'].data = gpu_view
|
||
|
||
def unpack_to_cpu(self, layer):
|
||
"""Move GPU data back to packed CPU buffers and rebind layer parameters"""
|
||
# Group specs by dtype
|
||
dtype_specs = {}
|
||
for spec in self.tensor_specs:
|
||
dtype = spec['dtype']
|
||
if dtype not in dtype_specs:
|
||
dtype_specs[dtype] = []
|
||
dtype_specs[dtype].append(spec)
|
||
|
||
# Copy GPU → CPU for each dtype
|
||
for dtype, specs in dtype_specs.items():
|
||
if dtype not in self.gpu_blocks:
|
||
continue
|
||
|
||
gpu_buffer = self.gpu_blocks[dtype]
|
||
packed_cpu = self.packed_buffers[dtype]
|
||
|
||
# Synchronous copy back to packed CPU buffer
|
||
packed_cpu.copy_(gpu_buffer.to('cpu'))
|
||
# packed_cpu.copy_(gpu_buffer)
|
||
|
||
# ATOMIC COMMIT: Rebind params/buffers back to CPU views
|
||
for spec in specs:
|
||
start_idx = spec['offset']
|
||
end_idx = start_idx + spec['size']
|
||
cpu_view = packed_cpu[start_idx:end_idx].view(spec['shape'])
|
||
|
||
if spec['is_param']:
|
||
spec['param_ref'].data = cpu_view
|
||
else:
|
||
spec['buffer_ref'].data = cpu_view
|
||
|
||
# Release GPU memory
|
||
self.gpu_blocks.clear()
|
||
self.gpu_events.clear()
|
||
|
||
def wait_for_gpu_transfer(self):
|
||
"""Wait for all async GPU transfers to complete"""
|
||
for event in self.gpu_events.values():
|
||
event.wait()
|
||
|
||
def get_memory_usage(self):
|
||
"""Return memory usage info"""
|
||
cpu_bytes = sum(buf.numel() * buf.element_size() for buf in self.packed_buffers.values())
|
||
gpu_bytes = sum(buf.numel() * buf.element_size() for buf in self.gpu_blocks.values())
|
||
return {
|
||
'cpu_bytes': cpu_bytes,
|
||
'gpu_bytes': gpu_bytes,
|
||
'total_elements': self.total_elements
|
||
}
|
||
|
||
|
||
def safe_move_to_gpu_packed(layer, idx):
|
||
"""Move layer to GPU using pre-packed buffer"""
|
||
if PROFILE:
|
||
nvtx.range_push(f"GPU_Transfer_Packed_L{idx}")
|
||
try:
|
||
current_device = device_cache.get_device(idx)
|
||
if current_device.type == 'cuda':
|
||
return True
|
||
|
||
if VERBOSE:
|
||
start_time = time.time()
|
||
layer_size_mb = get_cached_layer_size_mb(idx)
|
||
|
||
if idx in packed_layers:
|
||
# Use packed transfer (automatically uses pinned if enabled)
|
||
packed_layers[idx].unpack_to_gpu(layer, GPU_DEVICE)
|
||
else:
|
||
# Fallback to normal transfer
|
||
layer.to(GPU_DEVICE, non_blocking=True)
|
||
|
||
transfer_time = 0
|
||
|
||
event_based_sync("gpu_transfer", idx)
|
||
if VERBOSE:
|
||
end_time = time.time()
|
||
transfer_time = end_time - start_time
|
||
if transfer_time > 0 and VERBOSE:
|
||
speed_mbps = layer_size_mb / transfer_time
|
||
|
||
# Track current step stats
|
||
transfer_stats['current_step_gpu_times'].append(transfer_time)
|
||
transfer_stats['current_step_gpu_speeds'].append(speed_mbps)
|
||
device_cache.mark_moved(idx, torch.device('cuda'))
|
||
swap_stats['to_gpu'] += 1
|
||
return True
|
||
|
||
except RuntimeError as e:
|
||
if "out of memory" in str(e):
|
||
print(f" OOM moving packed layer {idx}")
|
||
return False
|
||
raise e
|
||
finally:
|
||
if PROFILE:
|
||
nvtx.range_pop()
|
||
|
||
|
||
def safe_move_to_cpu_packed(layer, idx):
|
||
"""Move layer to CPU - packed version already stored"""
|
||
if PROFILE:
|
||
nvtx.range_push(f"CPU_Transfer_Packed_L{idx}")
|
||
try:
|
||
try:
|
||
current_device = next(layer.parameters()).device
|
||
if current_device.type == 'cpu':
|
||
return True
|
||
except StopIteration:
|
||
pass
|
||
|
||
transfer_time = 0
|
||
|
||
if VERBOSE:
|
||
start_time = time.time()
|
||
layer_size_mb = get_cached_layer_size_mb(idx)
|
||
layer.to(CPU_DEVICE)
|
||
if VERBOSE:
|
||
end_time = time.time()
|
||
transfer_time = end_time - start_time
|
||
|
||
if transfer_time > 0 and VERBOSE:
|
||
speed_mbps = layer_size_mb / transfer_time
|
||
# Track current step stats
|
||
transfer_stats['current_step_cpu_times'].append(transfer_time)
|
||
transfer_stats['current_step_cpu_speeds'].append(speed_mbps)
|
||
device_cache.mark_moved(idx, torch.device('cpu'))
|
||
swap_stats['to_cpu'] += 1
|
||
return True
|
||
finally:
|
||
if PROFILE:
|
||
nvtx.range_pop()
|
||
|
||
|
||
# =============================================================================
|
||
# DYNAMIC LAYER SWAPPING INFERENCE for Comfy-UI - by obisin
|
||
# =============================================================================
|
||
|
||
def add_smart_swapping_to_layer(layer, layer_idx, layers_list, gpu_resident_layers, cpu_swappable_layers):
|
||
"""Add swapping capability with ComfyUI coordination"""
|
||
|
||
global layers, device_cache
|
||
layers = layers_list
|
||
|
||
original_forward = layer.forward
|
||
|
||
if not hasattr(add_smart_swapping_to_layer, 'prefetch_hits'):
|
||
add_smart_swapping_to_layer.prefetch_hits = 0
|
||
add_smart_swapping_to_layer.prefetch_misses = 0
|
||
|
||
if not hasattr(add_smart_swapping_to_layer, 'layer_compute_times'):
|
||
add_smart_swapping_to_layer.layer_compute_times = []
|
||
|
||
# Initialize coordination state (once) - this should match _patch_model_patcher
|
||
if not hasattr(add_smart_swapping_to_layer, 'current_sampling_step'):
|
||
add_smart_swapping_to_layer.current_sampling_step = 0
|
||
add_smart_swapping_to_layer.last_processed_step = -1
|
||
add_smart_swapping_to_layer.apply_model_call_count = 0
|
||
add_smart_swapping_to_layer.current_forward_logged = False
|
||
|
||
# Initialize coordination state (once)
|
||
if not hasattr(add_smart_swapping_to_layer, 'current_sampling_step'):
|
||
add_smart_swapping_to_layer.current_sampling_step = 0
|
||
add_smart_swapping_to_layer.last_processed_step = -1
|
||
|
||
# THREADING: Initialize threading components (once)
|
||
if not hasattr(add_smart_swapping_to_layer, 'threading_initialized'):
|
||
add_smart_swapping_to_layer.threading_initialized = True
|
||
add_smart_swapping_to_layer.gpu_lock = threading.Lock()
|
||
add_smart_swapping_to_layer.current_layer_idx = 0
|
||
add_smart_swapping_to_layer.training_active = True
|
||
add_smart_swapping_to_layer.background_thread = None
|
||
|
||
# CUDA STREAMS: Initialize streams (once) - only if enabled
|
||
if args.cuda_streams and not hasattr(add_smart_swapping_to_layer, 'streams_initialized'):
|
||
add_smart_swapping_to_layer.streams_initialized = True
|
||
add_smart_swapping_to_layer.copy_stream = torch.cuda.Stream()
|
||
add_smart_swapping_to_layer.compute_stream = torch.cuda.Stream()
|
||
print(" CUDA Streams enabled for copy-compute overlap")
|
||
|
||
# THREADING: Start background thread (once)
|
||
if args.threading and add_smart_swapping_to_layer.background_thread is None:
|
||
add_smart_swapping_to_layer.background_thread = threading.Thread(
|
||
target=background_layer_manager,
|
||
daemon=True
|
||
)
|
||
add_smart_swapping_to_layer.background_thread.start()
|
||
|
||
@wraps(original_forward)
|
||
def swapped_forward(*fwd_args, **kwargs):
|
||
|
||
# Detect Flux calling patterns
|
||
layer_type = type(layer).__name__
|
||
is_flux_call = 'DoubleStream' in layer_type or 'SingleStream' in layer_type
|
||
|
||
if is_flux_call:
|
||
# Flux patterns - preserve exact argument structure
|
||
if fwd_args:
|
||
# SingleStreamBlock: block(img, vec=vec, pe=pe)
|
||
x = fwd_args[0] # img
|
||
new_args = fwd_args[1:] # should be empty
|
||
new_kwargs = kwargs # vec=vec, pe=pe
|
||
else:
|
||
# DoubleStreamBlock: block(img=img, txt=txt, vec=vec, pe=pe)
|
||
x = kwargs.get('img')
|
||
new_args = ()
|
||
new_kwargs = kwargs
|
||
else:
|
||
# WAN/SDXL pattern
|
||
x = fwd_args[0] if fwd_args else None
|
||
new_args = fwd_args[1:]
|
||
new_kwargs = kwargs
|
||
|
||
if VERBOSE and not hasattr(add_smart_swapping_to_layer, 'step_timing_initialized'):
|
||
add_smart_swapping_to_layer.step_timing_initialized = True
|
||
add_smart_swapping_to_layer.step_start_time = time.time()
|
||
add_smart_swapping_to_layer.current_step = 0
|
||
|
||
if not hasattr(add_smart_swapping_to_layer, 'total_forward_calls'):
|
||
add_smart_swapping_to_layer.total_forward_calls = 0
|
||
add_smart_swapping_to_layer.current_step = 0
|
||
add_smart_swapping_to_layer.calls_this_step = 0
|
||
|
||
add_smart_swapping_to_layer.total_forward_calls += 1
|
||
add_smart_swapping_to_layer.calls_this_step += 1
|
||
|
||
# Detect new step when we hit layer 0 again (reset to early layers)
|
||
if layer_idx <= 2: # Early layer indicates new step
|
||
if add_smart_swapping_to_layer.calls_this_step > len(layers): # Processed full model
|
||
add_smart_swapping_to_layer.current_step += 1
|
||
add_smart_swapping_to_layer.calls_this_step = 0
|
||
|
||
if VERBOSE:
|
||
print(f"🔄 New sampling step {add_smart_swapping_to_layer.current_step} detected")
|
||
|
||
# At the start of swapped_forward
|
||
if VERBOSE and layer_idx % 10 == 0: # Every 10th layer
|
||
current_step = getattr(add_smart_swapping_to_layer, 'current_step', 0)
|
||
print(f"Layer {layer_idx} step {current_step}: device={device_cache.get_device(layer_idx)}")
|
||
|
||
# Step progress logging
|
||
if VERBOSE and layer_idx % 20 == 0:
|
||
current_step = getattr(add_smart_swapping_to_layer, 'current_step', 0)
|
||
if (add_smart_swapping_to_layer.prefetch_hits + add_smart_swapping_to_layer.prefetch_misses) > 0:
|
||
hit_rate = add_smart_swapping_to_layer.prefetch_hits / (
|
||
add_smart_swapping_to_layer.prefetch_hits + add_smart_swapping_to_layer.prefetch_misses) * 100
|
||
print(
|
||
f" Prefetch hit rate: {hit_rate:.1f}% ({add_smart_swapping_to_layer.prefetch_hits}/{add_smart_swapping_to_layer.prefetch_hits + add_smart_swapping_to_layer.prefetch_misses})")
|
||
|
||
|
||
# THREADING: Update current layer for background thread
|
||
if args.threading:
|
||
add_smart_swapping_to_layer.current_layer_idx = layer_idx
|
||
|
||
if PROFILE:
|
||
nvtx.range_push(f"Layer_{layer_idx}_{'Forward'}")
|
||
try:
|
||
|
||
# Initialize global flags if they don't exist
|
||
# if not hasattr(add_smart_swapping_to_layer, 'current_forward_logged'):
|
||
# add_smart_swapping_to_layer.current_forward_logged = False
|
||
|
||
# # Summary logging at the start - only once per pass type
|
||
if not add_smart_swapping_to_layer.current_forward_logged:
|
||
current_step = getattr(add_smart_swapping_to_layer, 'current_step', 0)
|
||
if VERBOSE:
|
||
print(
|
||
f" Forward pass step {current_step}: layers {min(cpu_swappable_layers)}-{max(cpu_swappable_layers)}")
|
||
add_smart_swapping_to_layer.current_forward_logged = True
|
||
|
||
if not hasattr(add_smart_swapping_to_layer, 'stats_initialized'):
|
||
add_smart_swapping_to_layer.stats_initialized = True
|
||
add_smart_swapping_to_layer.prefetch_hits = 0
|
||
add_smart_swapping_to_layer.prefetch_misses = 0
|
||
|
||
|
||
# 2. SMART LAYER MANAGEMENT
|
||
if layer_idx in cpu_swappable_layers:
|
||
if PROFILE:
|
||
nvtx.range_push(f"Smart_Management_L{layer_idx}")
|
||
try:
|
||
current_device = device_cache.get_device(layer_idx)
|
||
layer_already_on_gpu = (current_device.type == 'cuda')
|
||
|
||
if not layer_already_on_gpu:
|
||
# THREADING: If threading enabled, wait briefly for background thread
|
||
if args.threading:
|
||
print(f" Layer {layer_idx} not ready, waiting for background thread...")
|
||
for _ in range(50):
|
||
time.sleep(0.000001) #do not touch. very needed for sync
|
||
if device_cache.get_device(layer_idx).type == 'cuda':
|
||
layer_already_on_gpu = True
|
||
print(f" Background thread caught up!")
|
||
break
|
||
|
||
# If still not ready, run fallback
|
||
if not layer_already_on_gpu:
|
||
|
||
# THREADING: Use lock to prevent race condition
|
||
if args.threading:
|
||
with add_smart_swapping_to_layer.gpu_lock:
|
||
needed_layers = calculate_needed_layers(layer_idx, args.prefetch)
|
||
|
||
if args.selective_packing:
|
||
cleaned = cleanup_excess_layers_packed(needed_layers, args.selective_packing)
|
||
fetched = fetch_missing_layers_packed(needed_layers, args.selective_packing)
|
||
else:
|
||
cleaned = cleanup_excess_layers(needed_layers)
|
||
fetched = fetch_missing_layers(needed_layers)
|
||
else:
|
||
# No threading, run normally
|
||
needed_layers = calculate_needed_layers(layer_idx, args.prefetch)
|
||
if args.selective_packing:
|
||
cleaned = cleanup_excess_layers_packed(needed_layers, args.selective_packing)
|
||
fetched = fetch_missing_layers_packed(needed_layers, args.selective_packing)
|
||
else:
|
||
cleaned = cleanup_excess_layers(needed_layers)
|
||
fetched = fetch_missing_layers(needed_layers)
|
||
|
||
if layer_already_on_gpu:
|
||
add_smart_swapping_to_layer.prefetch_hits += 1
|
||
# if VERBOSE:
|
||
# print(f" ✅ Layer {layer_idx} prefetch hit")
|
||
else:
|
||
add_smart_swapping_to_layer.prefetch_misses += 1
|
||
# if VERBOSE:
|
||
# print(f" ⚠️ Layer {layer_idx} prefetch miss, fetching...")
|
||
|
||
device = device_cache.get_device(layer_idx)
|
||
gpu_success = device.type == 'cuda'
|
||
|
||
# if VERBOSE and layer_idx % 20 == 0:
|
||
# current_step = getattr(add_smart_swapping_to_layer, 'current_step', 0)
|
||
# if (
|
||
# add_smart_swapping_to_layer.prefetch_hits + add_smart_swapping_to_layer.prefetch_misses) > 0:
|
||
# hit_rate = add_smart_swapping_to_layer.prefetch_hits / (
|
||
# add_smart_swapping_to_layer.prefetch_hits + add_smart_swapping_to_layer.prefetch_misses) * 100
|
||
# print(
|
||
# f" Prefetch hit rate: {hit_rate:.1f}% ({add_smart_swapping_to_layer.prefetch_hits}/{add_smart_swapping_to_layer.prefetch_hits + add_smart_swapping_to_layer.prefetch_misses})")
|
||
|
||
finally:
|
||
if PROFILE:
|
||
nvtx.range_pop()
|
||
|
||
if layer_already_on_gpu:
|
||
add_smart_swapping_to_layer.prefetch_hits += 1
|
||
else:
|
||
add_smart_swapping_to_layer.prefetch_misses += 1
|
||
|
||
device = device_cache.get_device(layer_idx)
|
||
gpu_success = device.type == 'cuda'
|
||
else:
|
||
# Layer not in cpu_swappable_layers (permanent resident)
|
||
device = device_cache.get_device(layer_idx)
|
||
gpu_success = device.type == 'cuda'
|
||
|
||
# Handle GPU failure case
|
||
if not gpu_success:
|
||
print(f" Layer {layer_idx} failed to be on GPU, forcing aggressive cleanup...")
|
||
# Nuclear cleanup - evict ALL layers except this one
|
||
for cleanup_idx in cpu_swappable_layers:
|
||
if cleanup_idx != layer_idx and cleanup_idx < len(layers):
|
||
try:
|
||
layers[cleanup_idx].to('cpu')
|
||
except:
|
||
pass
|
||
|
||
gc.collect()
|
||
torch.cuda.empty_cache()
|
||
|
||
# Try again after cleanup
|
||
try:
|
||
layer.to(GPU_DEVICE)
|
||
device = torch.device(GPU_DEVICE)
|
||
gpu_success = True
|
||
print(f" Layer {layer_idx} moved to GPU after aggressive cleanup")
|
||
except RuntimeError:
|
||
print(f" CRITICAL: Layer {layer_idx} cannot fit on GPU, skipping computation!")
|
||
return x # Pass input through unchanged
|
||
|
||
# # Debug output every 10th layer
|
||
# if layer_idx % 10 == 0: # Every 10th layer
|
||
# current_step = getattr(add_smart_swapping_to_layer, 'current_step', 0)
|
||
# print(f"Layer {layer_idx} step {current_step}: device={device_cache.get_device(layer_idx)}")
|
||
|
||
|
||
def move_to_device(tensor, target_device):
|
||
# Only move device, don't touch dtype at all
|
||
if hasattr(tensor, 'device') and tensor.device != target_device:
|
||
tensor = tensor.to(target_device)
|
||
|
||
# Fix inference tensors (without dtype changes)
|
||
if isinstance(tensor, torch.Tensor):
|
||
try:
|
||
_ = tensor._version
|
||
return tensor
|
||
except (RuntimeError, AttributeError):
|
||
return tensor.detach().clone()
|
||
|
||
return tensor
|
||
|
||
|
||
move_to_device.current_layer = layer_idx
|
||
|
||
|
||
# DISABLE FLUX DETECTION FOR TESTING
|
||
# is_flux_call = False # Force to False
|
||
|
||
# Move tensors to device (respect the calling pattern you already set up)
|
||
if is_flux_call:
|
||
if fwd_args:
|
||
# SingleStreamBlock - move fwd_args and kwargs
|
||
new_args = tuple(move_to_device(arg, device) for arg in fwd_args)
|
||
new_kwargs = {k: move_to_device(v, device) for k, v in kwargs.items()}
|
||
else:
|
||
# DoubleStreamBlock - move only kwargs
|
||
new_args = ()
|
||
new_kwargs = {k: move_to_device(v, device) for k, v in kwargs.items()}
|
||
else:
|
||
x = move_to_device(x, device)
|
||
new_kwargs = {k: move_to_device(v, device) for k, v in kwargs.items()}
|
||
new_args = tuple(move_to_device(arg, device) for arg in new_args)
|
||
|
||
# Wait for any pending GPU transfers before executing the layer
|
||
if hasattr(layer, '_packed_block'):
|
||
for event in layer._packed_block.gpu_events.values():
|
||
torch.cuda.current_stream().wait_event(event)
|
||
|
||
# 4. COMPUTATION
|
||
if PROFILE:
|
||
nvtx.range_push(f"Compute_L{layer_idx}")
|
||
try:
|
||
|
||
if args.cuda_streams and hasattr(add_smart_swapping_to_layer, 'compute_stream'):
|
||
with torch.cuda.stream(add_smart_swapping_to_layer.compute_stream):
|
||
if VERBOSE:
|
||
layer_compute_start = time.time()
|
||
|
||
if args.compute_casting != "disabled":
|
||
# Compute casting: use specified dtype
|
||
dtype_map = {'fp32': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16}
|
||
cast_dtype = dtype_map[args.compute_casting]
|
||
|
||
with torch.autocast(device_type='cuda', dtype=cast_dtype, enabled=True):
|
||
fix_inference_tensor_parameters(layer)
|
||
# result = original_forward(x, *tuple(new_args), **new_kwargs)
|
||
# Call with preserved argument structure
|
||
if is_flux_call:
|
||
if fwd_args:
|
||
result = original_forward(*new_args, **new_kwargs)
|
||
return result
|
||
else:
|
||
return original_forward(**new_kwargs)
|
||
else:
|
||
# WAN/SDXL: original logic
|
||
return original_forward(x, *new_args, **new_kwargs)
|
||
|
||
elif args.autocast and args.autocast != 'fp32':
|
||
if args.autocast in ['f8_e4m3', 'f8_e5m2']:
|
||
# Use Transformer Engine for FP8
|
||
import transformer_engine.pytorch as te
|
||
with te.fp8_autocast():
|
||
fix_inference_tensor_parameters(layer)
|
||
result = original_forward(x, *tuple(new_args), **new_kwargs)
|
||
else:
|
||
# Regular PyTorch autocast for FP16/BF16
|
||
dtype_map = {'fp16': torch.float16, 'bf16': torch.bfloat16}
|
||
autocast_dtype = dtype_map[args.autocast]
|
||
with torch.autocast(device_type='cuda', dtype=autocast_dtype):
|
||
fix_inference_tensor_parameters(layer)
|
||
result = original_forward(x, *tuple(new_args), **new_kwargs)
|
||
else:
|
||
fix_inference_tensor_parameters(layer)
|
||
# result = original_forward(x, *tuple(new_args), **new_kwargs)
|
||
# Call with preserved argument structure
|
||
if is_flux_call:
|
||
if fwd_args:
|
||
result = original_forward(*new_args, **new_kwargs)
|
||
return result
|
||
else:
|
||
|
||
return original_forward(**new_kwargs)
|
||
else:
|
||
return original_forward(x, *new_args, **new_kwargs)
|
||
|
||
if VERBOSE:
|
||
layer_compute_end = time.time()
|
||
layer_compute_time = layer_compute_end - layer_compute_start
|
||
add_smart_swapping_to_layer.layer_compute_times.append(layer_compute_time)
|
||
return result
|
||
else:
|
||
if VERBOSE:
|
||
layer_compute_start = time.time()
|
||
|
||
if args.compute_casting != "disabled":
|
||
# Compute casting: use specified dtype
|
||
dtype_map = {'fp32': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16}
|
||
cast_dtype = dtype_map[args.compute_casting]
|
||
|
||
with torch.autocast(device_type='cuda', dtype=cast_dtype, enabled=True):
|
||
fix_inference_tensor_parameters(layer)
|
||
# result = original_forward(x, *tuple(new_args), **new_kwargs)
|
||
# Call with preserved argument structure
|
||
if is_flux_call:
|
||
if fwd_args:
|
||
result = original_forward(*new_args, **new_kwargs)
|
||
return result
|
||
else:
|
||
|
||
return original_forward(**new_kwargs)
|
||
else:
|
||
# WAN/SDXL: original logic
|
||
return original_forward(x, *new_args, **new_kwargs)
|
||
|
||
elif args.autocast and args.autocast != 'fp32':
|
||
if args.autocast in ['f8_e4m3', 'f8_e5m2']:
|
||
# Use Transformer Engine for FP8
|
||
import transformer_engine.pytorch as te
|
||
with te.fp8_autocast():
|
||
fix_inference_tensor_parameters(layer)
|
||
result = original_forward(x, *tuple(new_args), **new_kwargs)
|
||
else:
|
||
# Regular PyTorch autocast for FP16/BF16
|
||
dtype_map = {'fp16': torch.float16, 'bf16': torch.bfloat16}
|
||
autocast_dtype = dtype_map[args.autocast]
|
||
with torch.autocast(device_type='cuda', dtype=autocast_dtype):
|
||
fix_inference_tensor_parameters(layer)
|
||
result = original_forward(x, *tuple(new_args), **new_kwargs)
|
||
else:
|
||
fix_inference_tensor_parameters(layer)
|
||
# result = original_forward(x, *tuple(new_args), **new_kwargs)
|
||
# Call with preserved argument structure
|
||
if is_flux_call:
|
||
if fwd_args:
|
||
result = original_forward(*new_args, **new_kwargs)
|
||
return result
|
||
else:
|
||
return original_forward(**new_kwargs)
|
||
else:
|
||
# WAN/SDXL: original logic
|
||
return original_forward(x, *new_args, **new_kwargs)
|
||
|
||
if VERBOSE:
|
||
layer_compute_end = time.time()
|
||
layer_compute_time = layer_compute_end - layer_compute_start
|
||
add_smart_swapping_to_layer.layer_compute_times.append(layer_compute_time)
|
||
|
||
return result
|
||
finally:
|
||
if PROFILE:
|
||
nvtx.range_pop()
|
||
finally:
|
||
if PROFILE:
|
||
nvtx.range_pop()
|
||
|
||
# Replace forward method
|
||
layer.forward = swapped_forward
|
||
|
||
# =============================================================================
|
||
# COMFYUI NODE CLASS - by obisin
|
||
# =============================================================================
|
||
|
||
class DynamicSwappingLoader:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"model": ("MODEL",),
|
||
"layers": ("LAYERS",),
|
||
"initial_gpu_layers": ("INT",
|
||
{"default": 2, "min": 0, "max": 50, "tooltip": TOOLTIPS["initial_gpu_layers"]}),
|
||
"final_gpu_layers": ("INT",
|
||
{"default": 2, "min": 0, "max": 50, "tooltip": TOOLTIPS["final_gpu_layers"]}),
|
||
"prefetch": ("INT", {"default": 1, "min": 0, "max": 100, "tooltip": TOOLTIPS["prefetch"]}),
|
||
"threading": ("BOOLEAN", {"default": False, "tooltip": TOOLTIPS["threading"]}),
|
||
"event_sync": ("BOOLEAN", {"default": False, "tooltip": TOOLTIPS["event_sync"]}),
|
||
"cuda_streams": ("BOOLEAN", {"default": False, "tooltip": TOOLTIPS["cuda_streams"]}),
|
||
"batch_move": ("BOOLEAN", {"default": False, "tooltip": TOOLTIPS["batch_move"]}),
|
||
"selective_packing": ("INT",
|
||
{"default": 0, "min": 0, "max": 128, "tooltip": TOOLTIPS["selective_packing"]}),
|
||
"verbose": ("BOOLEAN", {"default": True, "tooltip": TOOLTIPS["verbose"]}),
|
||
"compile": ("BOOLEAN", {"default": False, "tooltip": TOOLTIPS["compile"]}),
|
||
},
|
||
"optional": {
|
||
"gpu_layer_indices": ("STRING", {"default": "", "tooltip": TOOLTIPS["gpu_layers"]}),
|
||
#Precision tools have been disabled due to confliction with comfy auto precision and casting
|
||
|
||
|
||
# "compute_casting": (["disabled", "fp32", "bf16", "fp16"],
|
||
# {"default": "disabled",
|
||
# "tooltip": "Cast to higher precision for computation only. Model stays in original dtype for storage/transfer, temporarily upcasts during forward pass for better accuracy"}),
|
||
# "cast_target": ("STRING", {"default": "", "tooltip": TOOLTIPS["cast_target"]}),
|
||
# "autocast": (["", 'auto', "fp16", "bf16", "f8_e4m3", "f8_e5m2"],
|
||
# {"default": "",}),
|
||
# "mixed_precision": (["auto", "f32", "bf16", "f16"],
|
||
# {"default": "auto", "tooltip": TOOLTIPS["mixed_precision"]}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("MODEL",)
|
||
RETURN_NAMES = ("model",)
|
||
|
||
FUNCTION = "load_model_with_swapping"
|
||
CATEGORY = "loaders"
|
||
TITLE = "Dynamic Swapping Loader"
|
||
|
||
|
||
def load_model_with_swapping(self, model, layers, initial_gpu_layers, final_gpu_layers,
|
||
prefetch, threading, event_sync, cuda_streams, batch_move,
|
||
selective_packing, verbose, compile, compute_casting=False, gpu_layer_indices="",
|
||
cast_target="", autocast="", mixed_precision="auto"):
|
||
|
||
# Force cleanup before starting
|
||
torch.cuda.empty_cache()
|
||
gc.collect()
|
||
|
||
global swap_stats, transfer_stats
|
||
swap_stats = {'to_gpu': 0, 'to_cpu': 0}
|
||
transfer_stats = {
|
||
'to_gpu_times': [], 'to_cpu_times': [], 'to_gpu_speeds': [], 'to_cpu_speeds': [],
|
||
'current_step_gpu_times': [], 'current_step_cpu_times': [],
|
||
'current_step_gpu_speeds': [], 'current_step_cpu_speeds': []
|
||
}
|
||
|
||
# Reset layer-specific stats
|
||
if hasattr(add_smart_swapping_to_layer, 'prefetch_hits'):
|
||
add_smart_swapping_to_layer.prefetch_hits = 0
|
||
add_smart_swapping_to_layer.prefetch_misses = 0
|
||
|
||
# Update ALL args with node parameters
|
||
args.dynamic_swapping = True
|
||
args.initial_gpu_layers = initial_gpu_layers
|
||
args.final_gpu_layers = final_gpu_layers
|
||
args.prefetch = prefetch
|
||
args.threading = threading
|
||
args.cuda_streams = cuda_streams
|
||
args.batch_move = batch_move
|
||
args.selective_packing = selective_packing
|
||
args.event_sync = event_sync
|
||
args.verbose = verbose
|
||
global VERBOSE
|
||
VERBOSE = args.verbose
|
||
args.compile = compile
|
||
|
||
args.device_sync_mode = 'off'
|
||
args.sync_only_on_resume = False
|
||
args.mixed_precision = mixed_precision
|
||
|
||
# Handle string parameters
|
||
args.gpu_layers = gpu_layer_indices.strip() if gpu_layer_indices.strip() else None
|
||
args.cast_target = False #cast_target.strip() if cast_target.strip() else None
|
||
args.autocast = False #autocast if autocast else None
|
||
|
||
args.compute_casting = "disabled"#compute_casting
|
||
|
||
if compute_casting:
|
||
args.autocast = False
|
||
|
||
|
||
if VERBOSE:
|
||
print(f"\n=== LAYER DTYPE INSPECTION ===")
|
||
print(f"Total layers: {len(layers)}")
|
||
|
||
for i, layer in enumerate(layers):
|
||
layer_type = type(layer).__name__
|
||
|
||
# Get parameter dtypes for this layer
|
||
param_dtypes = set()
|
||
param_count = 0
|
||
for name, param in layer.named_parameters():
|
||
param_dtypes.add(param.dtype)
|
||
param_count += 1
|
||
|
||
# Get buffer dtypes
|
||
buffer_dtypes = set()
|
||
buffer_count = 0
|
||
for name, buffer in layer.named_buffers():
|
||
buffer_dtypes.add(buffer.dtype)
|
||
buffer_count += 1
|
||
|
||
# Format dtype info
|
||
dtype_info = []
|
||
if param_dtypes:
|
||
dtype_info.append(f"params: {list(param_dtypes)}")
|
||
if buffer_dtypes:
|
||
dtype_info.append(f"buffers: {list(buffer_dtypes)}")
|
||
|
||
dtype_str = ", ".join(dtype_info) if dtype_info else "no tensors"
|
||
|
||
print(f" {i:2d}: {layer_type:<20} | {dtype_str}")
|
||
|
||
print("===============================\n")
|
||
|
||
if gpu_layer_indices and gpu_layer_indices.strip():
|
||
specified_gpu_layers = set(map(int, gpu_layer_indices.split(',')))
|
||
|
||
specified_gpu_layers.add(0) # First layer
|
||
specified_gpu_layers.add(len(layers) - 1) # Last layer
|
||
if VERBOSE:
|
||
print(f" Automatically added layers 0 and {len(layers) - 1} for stability")
|
||
else:
|
||
# Use initial/final logic, but ensure at least 1 initial and 1 final
|
||
initial_gpu_layers = max(1, args.initial_gpu_layers) # At least 1
|
||
final_gpu_layers = max(1, args.final_gpu_layers) # At least 1
|
||
|
||
if args.dynamic_swapping:
|
||
if VERBOSE:
|
||
print("Applying ENHANCED smart GPU allocation + dynamic swapping...")
|
||
print(f"Setting up enhanced swapping for {len(layers)} layers...")
|
||
|
||
if args.gpu_layers:
|
||
try:
|
||
specified_gpu_layers = set(map(int, args.gpu_layers.split(',')))
|
||
if VERBOSE:
|
||
print(f" Using specified GPU layers: {sorted(specified_gpu_layers)}")
|
||
|
||
# Validate layer indices
|
||
invalid_layers = [idx for idx in specified_gpu_layers if idx >= len(layers) or idx < 0]
|
||
if invalid_layers:
|
||
raise ValueError(f"Invalid layer indices: {invalid_layers}. Must be 0-{len(layers) - 1}")
|
||
|
||
except ValueError as e:
|
||
print(f" Error parsing --gpu_layers: {e}")
|
||
print("Example usage: --gpu_layers 0,1,2,14,18,19,20,21,22")
|
||
raise e
|
||
else:
|
||
specified_gpu_layers = None
|
||
if VERBOSE:
|
||
print(f"Using initial/final layer allocation: {initial_gpu_layers}/{final_gpu_layers}")
|
||
|
||
# PHASE 1: Determine which layers go where based on specified layers or initial/final counts
|
||
for i, layer in enumerate(layers):
|
||
if hasattr(layer, 'to'):
|
||
if specified_gpu_layers is not None:
|
||
# Use specified layer indices for GPU placement
|
||
if i in specified_gpu_layers:
|
||
try:
|
||
layer.to(GPU_DEVICE)
|
||
gpu_resident_layers.add(i)
|
||
if VERBOSE:
|
||
print(f"Layer {i} (specified) -> GPU permanent")
|
||
except RuntimeError as e:
|
||
print(f" CRITICAL: Cannot fit specified layer {i} on GPU!")
|
||
print(f"GPU memory may be insufficient. Consider removing layer {i} from --gpu_layers")
|
||
raise e
|
||
else:
|
||
# Not in specified list, make swappable
|
||
layer.to(CPU_DEVICE)
|
||
cpu_swappable_layers.add(i)
|
||
if VERBOSE:
|
||
print(f"Layer {i} (not specified) -> CPU swappable")
|
||
|
||
else:
|
||
# Use original initial/final logic as fallback
|
||
if i < initial_gpu_layers:
|
||
# Initial layers on GPU permanently
|
||
try:
|
||
layer.to(GPU_DEVICE)
|
||
gpu_resident_layers.add(i)
|
||
if VERBOSE:
|
||
print(f"Layer {i} (initial) -> GPU permanent")
|
||
except RuntimeError as e:
|
||
print(f"GPU exhausted at layer {i}, moving to CPU with swapping")
|
||
layer.to(CPU_DEVICE)
|
||
cpu_swappable_layers.add(i)
|
||
|
||
elif i >= (len(layers) - final_gpu_layers):
|
||
# Final layers on GPU permanently
|
||
try:
|
||
layer.to(GPU_DEVICE)
|
||
gpu_resident_layers.add(i)
|
||
if VERBOSE:
|
||
print(f"Layer {i} (final) -> GPU permanent")
|
||
except RuntimeError as e:
|
||
print(f"CRITICAL: Cannot fit final layer {i} on GPU!")
|
||
raise e
|
||
else:
|
||
# Middle layers on CPU with swapping capability
|
||
layer.to(CPU_DEVICE)
|
||
cpu_swappable_layers.add(i)
|
||
print(f"Layer {i} (middle) -> CPU swappable")
|
||
|
||
print(f"✓ {len(gpu_resident_layers)} layers permanently on GPU: {sorted(gpu_resident_layers)}")
|
||
print(f"✓ {len(cpu_swappable_layers)} layers on CPU with smart swapping: {sorted(cpu_swappable_layers)}")
|
||
|
||
|
||
if cast_target:
|
||
casting_handler = CastingHandler()
|
||
casting_handler.precast_all_layers(cpu_swappable_layers, layers)
|
||
|
||
global device_cache
|
||
device_cache = LayerDeviceCache(model, layers)
|
||
|
||
if args.selective_packing:
|
||
print(" Packing CPU-swappable layers for optimized transfers...")
|
||
for i in cpu_swappable_layers:
|
||
try:
|
||
# Pass pinned flag to PackedCPUBlock
|
||
packed_layers[i] = PackedCPUBlock(layers[i])
|
||
# print(f"✓ Layer {i} packed: {packed_layers[i].total_elements} elements")
|
||
except Exception as e:
|
||
print(f" Failed to pack layer {i}: {e}")
|
||
else:
|
||
print(" Selective packing disabled - using direct transfers only")
|
||
|
||
print(" Pre-calculating layer sizes...")
|
||
layer_sizes_mb = {}
|
||
for i, layer in enumerate(layers):
|
||
layer_sizes_mb[i] = sum(p.numel() * p.element_size() for p in layer.parameters()) / (1024 * 1024)
|
||
if VERBOSE:
|
||
print(f" Layer {i}: {layer_sizes_mb[i]:.1f}MB")
|
||
|
||
if args.cast_target:
|
||
print(f"\n=== LAYER CASTING DTYPES ===")
|
||
print(f"Total layers: {len(layers)}")
|
||
|
||
for i, layer in enumerate(layers):
|
||
layer_type = type(layer).__name__
|
||
|
||
# Get parameter dtypes for this layer
|
||
param_dtypes = set()
|
||
param_count = 0
|
||
for name, param in layer.named_parameters():
|
||
param_dtypes.add(param.dtype)
|
||
param_count += 1
|
||
|
||
# Get buffer dtypes
|
||
buffer_dtypes = set()
|
||
buffer_count = 0
|
||
for name, buffer in layer.named_buffers():
|
||
buffer_dtypes.add(buffer.dtype)
|
||
buffer_count += 1
|
||
|
||
# Format dtype info
|
||
dtype_info = []
|
||
if param_dtypes:
|
||
dtype_info.append(f"params: {list(param_dtypes)}")
|
||
if buffer_dtypes:
|
||
dtype_info.append(f"buffers: {list(buffer_dtypes)}")
|
||
|
||
dtype_str = ", ".join(dtype_info) if dtype_info else "no tensors"
|
||
|
||
print(f" {i:2d}: {layer_type:<20} | {dtype_str}")
|
||
|
||
print("===============================\n")
|
||
|
||
for layer_idx in cpu_swappable_layers:
|
||
add_smart_swapping_to_layer(
|
||
layers[layer_idx],
|
||
layer_idx,
|
||
layers,
|
||
gpu_resident_layers,
|
||
cpu_swappable_layers)
|
||
|
||
if VERBOSE:
|
||
print("Adding compute timing to permanent GPU layers...")
|
||
for layer_idx in gpu_resident_layers:
|
||
layer = layers[layer_idx]
|
||
original_forward = layer.forward
|
||
|
||
if args.compile:
|
||
# Try torch.compile version
|
||
@torch.compile
|
||
def compiled_forward( *layer_args, **layer_kwargs):
|
||
return original_forward( *layer_args, **layer_kwargs)
|
||
|
||
def create_resident_forward(layer_idx, original_forward, compiled_fn):
|
||
current_args = args
|
||
@wraps(original_forward)
|
||
def resident_forward( *args_tuple, **kwargs):
|
||
|
||
if VERBOSE:
|
||
layer_compute_start = time.time()
|
||
|
||
try:
|
||
|
||
if current_args.compute_casting != "disabled":
|
||
dtype_map = {'fp32': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16}
|
||
cast_dtype = dtype_map[args.compute_casting]
|
||
|
||
with torch.autocast(device_type='cuda', dtype=cast_dtype, enabled=True):
|
||
result = original_forward( *args_tuple, **kwargs)
|
||
else:
|
||
result = compiled_fn( *args_tuple, **kwargs)
|
||
except Exception as e:
|
||
print(f" Compile failed for layer {layer_idx}: {e}")
|
||
result = original_forward( *args_tuple, **kwargs)
|
||
|
||
if VERBOSE:
|
||
layer_compute_end = time.time()
|
||
layer_compute_time = layer_compute_end - layer_compute_start
|
||
add_smart_swapping_to_layer.layer_compute_times.append(layer_compute_time)
|
||
return result
|
||
|
||
return resident_forward
|
||
|
||
layer.forward = create_resident_forward(layer_idx, original_forward, compiled_forward)
|
||
if VERBOSE:
|
||
print(f" Added compiled timing to permanent GPU layer {layer_idx}")
|
||
|
||
else:
|
||
# Original working version (no compile)
|
||
def create_resident_forward(layer_idx, original_forward):
|
||
current_args = args
|
||
@wraps(original_forward)
|
||
def resident_forward( *args_tuple, **kwargs):
|
||
|
||
if VERBOSE:
|
||
layer_compute_start = time.time()
|
||
|
||
|
||
if current_args.compute_casting != "disabled":
|
||
dtype_map = {'fp32': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16}
|
||
cast_dtype = dtype_map[args.compute_casting]
|
||
|
||
with torch.autocast(device_type='cuda', dtype=cast_dtype, enabled=True):
|
||
result = original_forward( *args_tuple, **kwargs)
|
||
else:
|
||
result = original_forward( *args_tuple, **kwargs)
|
||
|
||
# result = original_forward( *args_tuple, **kwargs)
|
||
|
||
if VERBOSE:
|
||
layer_compute_end = time.time()
|
||
layer_compute_time = layer_compute_end - layer_compute_start
|
||
add_smart_swapping_to_layer.layer_compute_times.append(layer_compute_time)
|
||
return result
|
||
|
||
return resident_forward
|
||
|
||
layer.forward = create_resident_forward(layer_idx, original_forward)
|
||
if VERBOSE:
|
||
print(f"✓ Added timing to permanent GPU layer {layer_idx}")
|
||
|
||
# if step % 5 == 0:
|
||
# print(f"\n=== STEP {step} STATS ===")
|
||
# print(f"Expected total steps: {model_engine.total_steps}")
|
||
# if (add_smart_swapping_to_layer.prefetch_hits + add_smart_swapping_to_layer.prefetch_misses) > 0:
|
||
# hit_rate = add_smart_swapping_to_layer.prefetch_hits / (
|
||
# add_smart_swapping_to_layer.prefetch_hits + add_smart_swapping_to_layer.prefetch_misses) * 100
|
||
# print(f"Prefetch hit rate: {hit_rate:.1f}% ({add_smart_swapping_to_layer.prefetch_hits}/{add_smart_swapping_to_layer.prefetch_hits + add_smart_swapping_to_layer.prefetch_misses})")
|
||
# print("========================")
|
||
|
||
|
||
if VERBOSE:
|
||
print("✅ Dynamic swapping successfully integrated with ComfyUI inference")
|
||
|
||
return (model,)
|
||
|
||
|
||
|
||
# =============================================================================
|
||
# MIXED PRECISION FRAMEWORK CODE by obisin
|
||
# =============================================================================
|
||
|
||
"""
|
||
Universal Mixed Precision Handler by obisin
|
||
=================================
|
||
|
||
Drop-in utility to handle gradient operations with mixed precision (float8, bfloat16, float32)
|
||
across any model architecture. Patches PyTorch's gradient functions to handle unsupported
|
||
dtype operations gracefully.
|
||
|
||
Usage:
|
||
from utils.mixed_precision_handler import patch_mixed_precision
|
||
patch_mixed_precision()
|
||
|
||
# or with user control:
|
||
patch_mixed_precision(mixed_precision_target='f32')
|
||
"""
|
||
class CastingHandler:
|
||
def __init__(self):
|
||
self.original_clip_grad_norm = None
|
||
self.dtype_conversion_map = {
|
||
torch.float8_e4m3fn: torch.bfloat16,
|
||
torch.float8_e5m2: torch.bfloat16,
|
||
}
|
||
self.safe_dtypes = {torch.float32, torch.float16, torch.bfloat16}
|
||
self.is_patched = False
|
||
|
||
def detect_model_highest_precision(self, layers):
|
||
"""Find the single highest precision dtype in the model"""
|
||
all_dtypes = set()
|
||
for layer in layers:
|
||
for param in layer.parameters():
|
||
all_dtypes.add(param.dtype)
|
||
|
||
# Precision hierarchy (highest to lowest)
|
||
precision_order = [torch.float64, torch.float32, torch.bfloat16, torch.float16,
|
||
getattr(torch, 'float8_e4m3fn', None)]
|
||
|
||
# Find the highest precision actually used
|
||
for dtype in precision_order:
|
||
if dtype and dtype in all_dtypes:
|
||
return dtype
|
||
return torch.float32 # fallback
|
||
|
||
|
||
def cast_to_dtype(self, param, target_device, cast_target=None):
|
||
"""Cast specific source dtype to target dtype"""
|
||
source_dtype = param.dtype
|
||
|
||
if cast_target:
|
||
cast_from, cast_to = cast_target
|
||
|
||
# Map string to torch dtypes
|
||
dtype_map = {
|
||
'f32': torch.float32,
|
||
'bf16': torch.bfloat16,
|
||
'f16': torch.float16,
|
||
'f8_e4m3': torch.float8_e4m3fn,
|
||
'f8_e5m2': torch.float8_e5m2,
|
||
'nf4': 'nf4', # Special handling needed NOT TESTED
|
||
'fp4': 'fp4' # Special handling needed
|
||
}
|
||
|
||
# Only cast if source matches the FROM dtype
|
||
if source_dtype == dtype_map[cast_from]:
|
||
# Handle 4-bit quantization
|
||
if cast_to in ['nf4', 'fp4']:
|
||
try:
|
||
import bitsandbytes as bnb
|
||
quantized, quant_state = bnb.functional.quantize_4bit(
|
||
param.data, blocksize=64, quant_type=cast_to
|
||
)
|
||
# Return dequantized version on target device
|
||
return bnb.functional.dequantize_4bit(
|
||
quantized.to(target_device), quant_state.to(target_device)
|
||
)
|
||
except ImportError:
|
||
raise RuntimeError(
|
||
" bitsandbytes required for 4-bit quantization. Install with: pip install bitsandbytes")
|
||
except Exception as e:
|
||
raise RuntimeError(f" 4-bit quantization failed: {e}")
|
||
else:
|
||
# Regular dtype casting
|
||
return param.to(target_device, dtype=dtype_map[cast_to], non_blocking=True)
|
||
|
||
# No casting needed
|
||
return param.to(target_device, non_blocking=True)
|
||
|
||
def replace_with_4bit_layer(self, layer):
|
||
"""Replace torch.nn.Linear with bitsandbytes Linear4bit, skip LoRA layers"""
|
||
try:
|
||
# import bitsandbytes as bnb
|
||
|
||
for name, module in layer.named_modules():
|
||
if isinstance(module, torch.nn.Linear):
|
||
# Skip LoRA layers - they cause dropout issues with 4-bit
|
||
if any(lora_term in name.lower() for lora_term in ['lora', 'adapter']):
|
||
# print(f" Skipping LoRA layer: {name}")
|
||
continue
|
||
|
||
parent_names = name.split('.')
|
||
if any(any(lora_term in part.lower() for lora_term in ['lora', 'adapter'])
|
||
for part in parent_names):
|
||
continue
|
||
|
||
print(f" Converting layer: {name}")
|
||
linear_4bit = bnb.nn.Linear4bit(
|
||
module.in_features,
|
||
module.out_features,
|
||
bias=module.bias is not None,
|
||
compute_dtype=torch.float16,
|
||
quant_type='nf4'
|
||
)
|
||
|
||
linear_4bit.load_state_dict(module.state_dict())
|
||
linear_4bit = linear_4bit.to('cuda')
|
||
|
||
# Replace in parent module
|
||
parent_name = '.'.join(name.split('.')[:-1])
|
||
module_name = name.split('.')[-1]
|
||
parent = layer
|
||
for part in parent_name.split('.'):
|
||
if part:
|
||
parent = getattr(parent, part)
|
||
setattr(parent, module_name, linear_4bit)
|
||
|
||
return layer
|
||
except Exception as e:
|
||
print(f"️ 4-bit replacement failed: {e}")
|
||
return layer
|
||
|
||
|
||
def precast_all_layers(self, cpu_swappable_layers, layers):
|
||
"""Cast all layer dtypes once at startup using existing cast_to_dtype"""
|
||
print(" Pre-casting layer dtypes (one-time setup)...")
|
||
casted_layers = 0
|
||
|
||
for i, layer in enumerate(layers):
|
||
if i in cpu_swappable_layers: # Only cast swappable layers
|
||
layer_casted = False
|
||
|
||
for name, param in layer.named_parameters():
|
||
# Use your existing function but stay on current device
|
||
new_param = self.cast_to_dtype(param, param.device) # Keep on same device
|
||
if new_param.dtype != param.dtype:
|
||
param.data = new_param.data
|
||
layer_casted = True
|
||
|
||
if layer_casted:
|
||
casted_layers += 1
|
||
|
||
print(f"✓ Pre-casted {casted_layers} layers using existing cast_to_dtype function")
|