John Pollock
e288152dae
This commit introduces a major architectural refactoring, laying the groundwork for DisTorch V2. The changes focus on improving modularity, memory management, and diagnostics. Key changes include: - Renaming `distorch_safetensor.py` to `distorch_2.py` to house the new core logic. - Deleting the legacy `block_swap.py` module. - Adding `device_memory_audit.py` for more sophisticated analysis of GPU memory usage. - Implementing a centralized and configurable logging system in `__init__.py` to provide standardized and level-controlled (DEBUG/INFO) output for better debugging.
352 lines
12 KiB
Python
352 lines
12 KiB
Python
"""
|
|
Device Memory Audit Utility for ComfyUI MultiGPU
|
|
Provides tools to inspect and monitor GPU/CPU memory usage and model placement
|
|
"""
|
|
|
|
import torch
|
|
import gc
|
|
import psutil
|
|
import logging
|
|
from collections import defaultdict
|
|
from typing import Dict, List, Tuple, Any
|
|
import comfy.model_management as mm
|
|
|
|
def format_bytes(bytes_value: int) -> str:
|
|
"""Convert bytes to human-readable format"""
|
|
for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
|
|
if bytes_value < 1024.0:
|
|
return f"{bytes_value:.2f} {unit}"
|
|
bytes_value /= 1024.0
|
|
return f"{bytes_value:.2f} PB"
|
|
|
|
def get_device_memory_info() -> Dict[str, Dict[str, Any]]:
|
|
"""
|
|
Get current memory usage for all available devices
|
|
|
|
Returns:
|
|
Dict with device names as keys and memory info as values
|
|
"""
|
|
memory_info = {}
|
|
|
|
# Check CUDA devices
|
|
if torch.cuda.is_available():
|
|
for i in range(torch.cuda.device_count()):
|
|
device_name = f"cuda:{i}"
|
|
torch.cuda.synchronize(i) # Ensure accurate memory reading
|
|
|
|
allocated = torch.cuda.memory_allocated(i)
|
|
reserved = torch.cuda.memory_reserved(i)
|
|
total = torch.cuda.get_device_properties(i).total_memory
|
|
free = total - allocated
|
|
|
|
memory_info[device_name] = {
|
|
"allocated": allocated,
|
|
"allocated_str": format_bytes(allocated),
|
|
"reserved": reserved,
|
|
"reserved_str": format_bytes(reserved),
|
|
"total": total,
|
|
"total_str": format_bytes(total),
|
|
"free": free,
|
|
"free_str": format_bytes(free),
|
|
"usage_percent": (allocated / total * 100) if total > 0 else 0
|
|
}
|
|
|
|
# Check CPU/System memory
|
|
vm = psutil.virtual_memory()
|
|
memory_info["cpu"] = {
|
|
"allocated": vm.used,
|
|
"allocated_str": format_bytes(vm.used),
|
|
"total": vm.total,
|
|
"total_str": format_bytes(vm.total),
|
|
"free": vm.available,
|
|
"free_str": format_bytes(vm.available),
|
|
"usage_percent": vm.percent
|
|
}
|
|
|
|
return memory_info
|
|
|
|
def audit_torch_tensors() -> Dict[str, List[Dict[str, Any]]]:
|
|
"""
|
|
Find all torch tensors in memory and group by device
|
|
|
|
Returns:
|
|
Dict with device names as keys and list of tensor info as values
|
|
"""
|
|
tensors_by_device = defaultdict(list)
|
|
|
|
# Iterate through all objects in memory
|
|
for obj in gc.get_objects():
|
|
try:
|
|
if torch.is_tensor(obj):
|
|
device_str = str(obj.device)
|
|
size_bytes = obj.element_size() * obj.numel()
|
|
|
|
tensor_info = {
|
|
"id": id(obj),
|
|
"shape": list(obj.shape),
|
|
"dtype": str(obj.dtype),
|
|
"size_bytes": size_bytes,
|
|
"size_str": format_bytes(size_bytes),
|
|
"requires_grad": obj.requires_grad,
|
|
"is_leaf": obj.is_leaf
|
|
}
|
|
|
|
tensors_by_device[device_str].append(tensor_info)
|
|
except:
|
|
# Some objects might not be accessible
|
|
pass
|
|
|
|
# Sort tensors by size for each device
|
|
for device in tensors_by_device:
|
|
tensors_by_device[device].sort(key=lambda x: x["size_bytes"], reverse=True)
|
|
|
|
return dict(tensors_by_device)
|
|
|
|
def audit_model_placement(model) -> Dict[str, List[Tuple[str, str, int]]]:
|
|
"""
|
|
Audit where each layer of a model is placed
|
|
|
|
Args:
|
|
model: PyTorch model to audit
|
|
|
|
Returns:
|
|
Dict with device names as keys and list of (layer_name, layer_type, size) tuples
|
|
"""
|
|
device_layers = defaultdict(list)
|
|
|
|
for name, module in model.named_modules():
|
|
# Skip container modules
|
|
if len(list(module.children())) > 0:
|
|
continue
|
|
|
|
# Find device for this module
|
|
device = None
|
|
size_bytes = 0
|
|
|
|
# Check parameters
|
|
for param_name, param in module.named_parameters(recurse=False):
|
|
if param is not None:
|
|
device = str(param.device)
|
|
size_bytes += param.element_size() * param.numel()
|
|
|
|
if device:
|
|
layer_type = type(module).__name__
|
|
device_layers[device].append((name, layer_type, size_bytes))
|
|
|
|
return dict(device_layers)
|
|
|
|
def audit_comfy_models() -> Dict[str, Any]:
|
|
"""
|
|
Audit ComfyUI's loaded models and their placement
|
|
|
|
Returns:
|
|
Dict with model information
|
|
"""
|
|
model_info = {}
|
|
|
|
# Try to get loaded models from ComfyUI's model management
|
|
try:
|
|
# Access currently loaded models if available
|
|
if hasattr(mm, 'current_loaded_models'):
|
|
loaded_models = mm.current_loaded_models
|
|
|
|
for i, model_data in enumerate(loaded_models):
|
|
model_name = f"model_{i}"
|
|
if hasattr(model_data, 'model'):
|
|
model = model_data.model
|
|
placement = audit_model_placement(model)
|
|
|
|
total_size = 0
|
|
for device_layers in placement.values():
|
|
for _, _, size in device_layers:
|
|
total_size += size
|
|
|
|
model_info[model_name] = {
|
|
"type": type(model).__name__,
|
|
"device_placement": placement,
|
|
"total_size": total_size,
|
|
"total_size_str": format_bytes(total_size)
|
|
}
|
|
except Exception as e:
|
|
logging.warning(f"[MultiGPU] Could not audit ComfyUI models: {e}")
|
|
|
|
return model_info
|
|
|
|
def print_memory_audit(detailed=False):
|
|
"""
|
|
Print a formatted memory audit report
|
|
|
|
Args:
|
|
detailed: If True, include tensor-level details
|
|
"""
|
|
logging.info("\n" + "=" * 60)
|
|
logging.info(" DEVICE MEMORY AUDIT REPORT")
|
|
logging.info("=" * 60)
|
|
|
|
# Device memory overview
|
|
memory_info = get_device_memory_info()
|
|
logging.info("\n📊 MEMORY USAGE BY DEVICE:")
|
|
logging.info("-" * 60)
|
|
|
|
for device, info in memory_info.items():
|
|
logging.info(f"\n{device.upper()}:")
|
|
logging.info(f" Total: {info['total_str']:>12}")
|
|
logging.info(f" Allocated: {info['allocated_str']:>12} ({info['usage_percent']:.1f}%)")
|
|
logging.info(f" Free: {info['free_str']:>12}")
|
|
if 'reserved_str' in info:
|
|
logging.info(f" Reserved: {info['reserved_str']:>12}")
|
|
|
|
# Tensor audit
|
|
if detailed:
|
|
tensors = audit_torch_tensors()
|
|
logging.info("\n📦 TENSORS BY DEVICE:")
|
|
logging.info("-" * 60)
|
|
|
|
for device, tensor_list in tensors.items():
|
|
total_size = sum(t["size_bytes"] for t in tensor_list)
|
|
logging.info(f"\n{device}: {len(tensor_list)} tensors, {format_bytes(total_size)} total")
|
|
|
|
# Show top 5 largest tensors
|
|
for i, tensor in enumerate(tensor_list[:5]):
|
|
logging.info(f" #{i+1}: Shape {tensor['shape']}, {tensor['size_str']}, {tensor['dtype']}")
|
|
|
|
# ComfyUI model audit
|
|
models = audit_comfy_models()
|
|
if models:
|
|
logging.info("\n🤖 COMFYUI LOADED MODELS:")
|
|
logging.info("-" * 60)
|
|
|
|
for model_name, info in models.items():
|
|
logging.info(f"\n{model_name} ({info['type']}):")
|
|
logging.info(f" Total size: {info['total_size_str']}")
|
|
|
|
for device, layers in info['device_placement'].items():
|
|
device_size = sum(size for _, _, size in layers)
|
|
logging.info(f" {device}: {len(layers)} layers, {format_bytes(device_size)}")
|
|
|
|
logging.info("\n" + "=" * 60)
|
|
|
|
def get_memory_snapshot() -> Dict[str, Any]:
|
|
"""
|
|
Get a complete memory snapshot for benchmarking
|
|
|
|
Returns:
|
|
Dict containing all memory information
|
|
"""
|
|
return {
|
|
"device_memory": get_device_memory_info(),
|
|
"tensors": audit_torch_tensors(),
|
|
"models": audit_comfy_models()
|
|
}
|
|
|
|
def compare_snapshots(before: Dict[str, Any], after: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""
|
|
Compare two memory snapshots to see what changed
|
|
|
|
Args:
|
|
before: Snapshot taken before an operation
|
|
after: Snapshot taken after an operation
|
|
|
|
Returns:
|
|
Dict with differences between snapshots
|
|
"""
|
|
differences = {}
|
|
|
|
# Compare device memory
|
|
memory_diff = {}
|
|
for device in after["device_memory"]:
|
|
if device in before["device_memory"]:
|
|
before_mem = before["device_memory"][device]["allocated"]
|
|
after_mem = after["device_memory"][device]["allocated"]
|
|
diff = after_mem - before_mem
|
|
|
|
memory_diff[device] = {
|
|
"before": format_bytes(before_mem),
|
|
"after": format_bytes(after_mem),
|
|
"difference": format_bytes(abs(diff)),
|
|
"increased": diff > 0
|
|
}
|
|
|
|
differences["memory_changes"] = memory_diff
|
|
|
|
# Compare tensor counts
|
|
tensor_diff = {}
|
|
for device in after["tensors"]:
|
|
after_count = len(after["tensors"][device])
|
|
before_count = len(before["tensors"].get(device, []))
|
|
|
|
if after_count != before_count:
|
|
tensor_diff[device] = {
|
|
"before": before_count,
|
|
"after": after_count,
|
|
"difference": after_count - before_count
|
|
}
|
|
|
|
differences["tensor_count_changes"] = tensor_diff
|
|
|
|
return differences
|
|
|
|
def print_snapshot_comparison(before: Dict[str, Any], after: Dict[str, Any], label: str = "Operation"):
|
|
"""
|
|
Print a formatted comparison of two snapshots
|
|
|
|
Args:
|
|
before: Snapshot before operation
|
|
after: Snapshot after operation
|
|
label: Label for the operation being measured
|
|
"""
|
|
diff = compare_snapshots(before, after)
|
|
|
|
logging.info(f"\n📈 MEMORY CHANGES AFTER {label}:")
|
|
logging.info("-" * 60)
|
|
|
|
for device, changes in diff["memory_changes"].items():
|
|
symbol = "↑" if changes["increased"] else "↓"
|
|
logging.info(f"{device}: {changes['before']} → {changes['after']} ({symbol} {changes['difference']})")
|
|
|
|
if diff["tensor_count_changes"]:
|
|
logging.info("\nTensor count changes:")
|
|
for device, changes in diff["tensor_count_changes"].items():
|
|
logging.info(f"{device}: {changes['before']} → {changes['after']} ({changes['difference']:+d})")
|
|
|
|
# Example usage functions for benchmarking
|
|
def benchmark_memory_points():
|
|
"""
|
|
Example function showing how to capture the 4 key memory points
|
|
"""
|
|
logging.info("\n🔍 Starting Memory Benchmark...")
|
|
|
|
# 1. Pre-load memory
|
|
snapshot_pre = get_memory_snapshot()
|
|
print_memory_audit(detailed=False)
|
|
|
|
# User would load model here
|
|
logging.info("\n[Load your model now]")
|
|
input("Press Enter when model is loaded...")
|
|
|
|
# 2. Post-load memory
|
|
snapshot_post_load = get_memory_snapshot()
|
|
print_snapshot_comparison(snapshot_pre, snapshot_post_load, "MODEL LOAD")
|
|
|
|
# User would run inference here
|
|
logging.info("\n[Run inference now]")
|
|
input("Press Enter when inference is complete...")
|
|
|
|
# 3. Active inference memory (captured during)
|
|
snapshot_active = get_memory_snapshot()
|
|
print_snapshot_comparison(snapshot_post_load, snapshot_active, "INFERENCE")
|
|
|
|
# 4. Post-generation memory
|
|
logging.info("\n[Waiting for cleanup...]")
|
|
torch.cuda.empty_cache() # Force cleanup
|
|
gc.collect()
|
|
|
|
snapshot_post_gen = get_memory_snapshot()
|
|
print_snapshot_comparison(snapshot_active, snapshot_post_gen, "CLEANUP")
|
|
|
|
logging.info("\n✅ Benchmark complete!")
|
|
|
|
if __name__ == "__main__":
|
|
# Run a basic audit when script is executed directly
|
|
print_memory_audit(detailed=True)
|