nodes/pollockjj-ComfyUI-MultiGPU
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial block swap implementation v2

Browse Source 
- Created comprehensive architecture documentation (ARCHITECTURE_V2.0.0.md)
- Added DOE optimization planning document (DOE_OPTIMIZATION.md)
- Implemented DisTorchBlockSwap node for safetensor models
- Created core/blockswap.py with BlockSwapManager
- Unified VirtualVRAM interface design
- Based on analysis of WanVideo's block swap methodology
This commit is contained in:
John Pollock
2025-08-08 14:50:26 -05:00
parent 92a10cc6ec
commit 897f785edf
7 changed files with 1090 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ARCHITECTURE_V2.0.0.md
Executable
+351
View File
@@ -0,0 +1,351 @@
# ComfyUI-MultiGPU Architecture V2.0.0
## Executive Summary
ComfyUI-MultiGPU provides intelligent model distribution across multiple GPUs and system RAM, optimizing for minimal VRAM usage while maintaining performance. Version 2.0 introduces a unified interface supporting both layer-by-layer transfers (DisTorch) and block swapping strategies.
## Core Concepts
### 1. Virtual VRAM
Virtual VRAM represents the extended memory pool available by offloading model components to other devices (CPU RAM or secondary GPUs). The system manages transfers between devices transparently during inference.
### 2. Transfer Strategies
#### DisTorch (Layer-by-Layer)
- **Mechanism**: Spoofs quantized tensors on offload device, dequantizes JIT to compute device
- **Transfer Size**: Single layer at a time (~100-200MB)
- **VRAM Usage**: Minimal (1 layer active)
- **Best For**: Video generation (long inference times)
- **Trade-off**: Many small PCIe transfers
```python
# DisTorch approach - minimal VRAM footprint
def forward_hook(module, input, output):
# Load single layer
load_layer_to_device(module, compute_device)
output = module.forward(input)
# Immediately offload
offload_layer(module, offload_device)
return output
```
#### Block Swap (New in V2)
- **Mechanism**: Moves blocks of layers between devices
- **Transfer Size**: Configurable (1-8GB blocks)
- **VRAM Usage**: Reserved swap buffer
- **Best For**: Image generation (short inference times)
- **Trade-off**: Fewer, larger PCIe transfers
```python
# Block swap approach - batched transfers
def forward_hook(module, input, output):
if need_swap(module):
# Swap entire block
offload_block(current_block, offload_device)
load_block(next_block, compute_device)
return module.forward(input)
```
### 3. Unified Interface
All strategies share common parameters:
```python
class VirtualVRAMConfig:
virtual_vram_gb: float # Total model size to offload
swap_space_gb: float # Reserved buffer on compute device
swap_device: str # Where to offload ("cpu", "cuda:1")
# Derived behavior
if swap_space_gb < min_layer_size:
use_distorch() # Layer-by-layer
else:
use_block_swap() # Block transfers
```
## Implementation Architecture
### Memory Management
#### Size Calculation (Shared Utility)
```python
def calculate_model_size(model):
"""Calculate actual memory footprint"""
total_bytes = 0
for param in model.parameters():
if hasattr(param, 'quant_type'): # GGUF
# Account for quantization
total_bytes += calculate_gguf_size(param)
else: # Safetensor
total_bytes += param.element_size() * param.nelement()
return total_bytes / (1024**3) # GB
```
#### Block Partitioning
```python
def partition_model(model, swap_space_gb):
"""Divide model into swappable blocks"""
blocks = []
current_block = []
current_size = 0
for name, module in model.named_modules():
module_size = get_module_size(module)
if current_size + module_size > swap_space_gb:
# Start new block
blocks.append(current_block)
current_block = [module]
current_size = module_size
else:
current_block.append(module)
current_size += module_size
return blocks
```
### Hook System
#### Pre/Post Forward Hooks
```python
class ModelWrapper:
def __init__(self, model, config):
self.model = model
self.config = config
self.blocks = partition_model(model, config.swap_space_gb)
self.current_block_idx = -1
# Install hooks
for block_idx, block in enumerate(self.blocks):
for module in block:
module.register_forward_pre_hook(
lambda m, i: self.pre_forward(m, block_idx)
)
def pre_forward(self, module, block_idx):
if block_idx != self.current_block_idx:
# Swap blocks
self.swap_blocks(self.current_block_idx, block_idx)
self.current_block_idx = block_idx
```
### GGUF Handling
#### Quantized Tensor Management
```python
class GGUFHandler:
def handle_gguf_layer(self, layer):
if self.config.swap_space_gb < layer.size:
# Use DisTorch approach - dequantize JIT
return self.distorch_dequantize(layer)
else:
# Can move entire quantized block
return self.block_swap_quantized(layer)
def distorch_dequantize(self, layer):
"""Dequantize during transfer (COPY operation)"""
# Creates new tensor on compute device
return dequantize_to_device(layer, self.compute_device)
def block_swap_quantized(self, layer):
"""Move quantized tensor (SWAP operation)"""
# Moves existing tensor between devices
return layer.to(self.compute_device)
```
## Phase Implementation Plan
### Phase 1: Block Swap for Safetensors (Current Focus)
**Goal**: Implement configurable block swapping for non-quantized models.
**Implementation**:
```python
class DisTorchBlockSwap:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"virtual_vram_gb": ("FLOAT", {
"default": 4.0,
"min": 0.1,
"max": 64.0,
"step": 0.1
}),
"swap_space_gb": ("FLOAT", {
"default": 1.0,
"min": 0.1,
"max": 16.0,
"step": 0.1
}),
"swap_device": (["cpu", "cuda:0", "cuda:1"],),
}
}
def apply(self, model, virtual_vram_gb, swap_space_gb, swap_device):
# Calculate model size
model_size = calculate_model_size(model)
# Partition into blocks
blocks = partition_model(model, swap_space_gb)
# Install swap hooks
wrapper = BlockSwapWrapper(model, blocks, swap_device)
return (wrapper.model,)
```
### Phase 2: Unified GGUF Support
**Goal**: Extend block swap to GGUF models, auto-selecting strategy.
**Decision Logic**:
```python
def select_strategy(model, config):
if is_gguf(model):
min_layer = get_min_layer_size(model)
if config.swap_space_gb < min_layer:
return DisTorchStrategy() # Must dequantize JIT
else:
return BlockSwapStrategy() # Can move quantized blocks
else:
return BlockSwapStrategy() # Safetensors always use blocks
```
### Phase 3: Auto-Optimization
**Goal**: Use empirical data to auto-configure optimal settings.
See `DOE_OPTIMIZATION.md` for detailed benchmarking plan.
**Auto Mode**:
```python
def auto_configure(model, workload):
# Detect hardware
pcie_gen = detect_pcie_generation()
gpu_bandwidth = detect_gpu_bandwidth()
# Analyze workload
is_video = workload.frames > 1
latent_size = workload.height * workload.width
# Lookup optimal config from DOE results
if is_video:
return {"swap_space_gb": 0.1} # Minimize transfers
else:
return lookup_optimal_config(
model.size, latent_size, pcie_gen
)
```
## Performance Characteristics
### Transfer Overhead Analysis
| Strategy | Transfer Size | Frequency | PCIe Time | Best Case |
|----------|--------------|-----------|-----------|-----------|
| DisTorch | 100-200MB | Every layer | High | Video (long inference) |
| Block Swap (1GB) | 1GB | Every ~10 layers | Medium | Balanced |
| Block Swap (4GB) | 4GB | Every ~40 layers | Low | Image (short inference) |
### Memory Usage Patterns
```
DisTorch (0.1GB swap):
|===| <- Active layer (100MB)
|...|...|...|...|...| <- Offloaded layers
Block Swap (2GB swap):
|==========| <- Active block (2GB)
|..........|..........| <- Offloaded blocks
```
## Advantages Over Existing Solutions
### vs. Sequential CPU Offload
- **Granular Control**: Configure exact offload amount
- **Multi-GPU Support**: Use secondary GPUs as fast swap
- **Quantization Aware**: Handles GGUF efficiently
### vs. Model Parallelism
- **No Model Modification**: Works with any model
- **Dynamic**: Adjusts to available resources
- **Flexible**: User controls memory/speed trade-off
## Code Organization
```
ComfyUI-MultiGPU/
├── nodes.py # Node definitions
├── core/
│ ├── distorch.py # Original layer-by-layer
│ ├── blockswap.py # New block swapping
│ ├── memory.py # Shared memory utilities
│ └── hooks.py # Hook management
├── strategies/
│ ├── auto.py # Auto-optimization
│ ├── gguf.py # GGUF-specific handling
│ └── safetensor.py # Safetensor handling
└── benchmark/
├── doe.py # DOE test runner
└── profiles.py # Hardware profiles
```
## Testing Strategy
### Unit Tests
- Memory calculation accuracy
- Block partitioning logic
- Hook installation/removal
### Integration Tests
- Safetensor models (SDXL, Flux)
- GGUF models (quantized)
- Multi-GPU configurations
### Performance Tests
- Measure transfer overhead
- Verify memory usage
- Benchmark vs baseline
## Migration Path
### For Existing Users
1. Current DisTorch nodes continue working
2. New unified node available alongside
3. Gradual migration as benefits proven
### Configuration Migration
```python
# Old DisTorch
distorch_model = DisTorch(model, device_map)
# New Unified (equivalent)
unified_model = VirtualVRAM(
model,
virtual_vram_gb=model_size,
swap_space_gb=0.1, # DisTorch-like
swap_device="cpu"
)
```
## Future Directions
### Adaptive Strategies
- Monitor transfer patterns
- Adjust block size dynamically
- Predict optimal points
### Pipeline Integration
- Coordinate with samplers
- Batch-aware swapping
- Multi-model orchestration
### Hardware Acceleration
- Direct Storage API
- NVLink optimization
- CXL memory pooling
## Conclusion
Version 2.0 unifies memory management strategies under a coherent interface, providing users with fine-grained control over the memory/performance trade-off while maintaining backward compatibility and preparing for future optimizations.
DOE_OPTIMIZATION.md
+251
View File
@@ -0,0 +1,251 @@
# Design of Experiments (DOE) for Block Swap Optimization
## Overview
This document outlines the empirical optimization strategy for determining optimal block swap parameters based on actual hardware performance measurements rather than theoretical calculations.
## Objective
Find the optimal `swap_space_gb` parameter that minimizes total execution time for various model/hardware/workload combinations.
## Test Matrix Variables
### 1. Model Parameters
- **Active Parameter Size**: 1B, 7B, 13B, 24B, 70B parameters
- **Model Architecture**:
- SDXL (UNet-based)
- Flux (Transformer-based)
- HunyuanVideo (Transformer-based)
- WanVideo (Transformer-based)
### 2. Latent Space Dimensions
#### Image Models
- 512x512 (SDXL base)
- 1024x1024 (SDXL highres)
- 2048x2048 (Flux highres)
#### Video Models
- 16 frames @ 512x512
- 49 frames @ 768x768
- 97 frames @ 1024x1024
### 3. Hardware Configurations
#### PCIe Generation
- PCIe 3.0 (16 GB/s)
- PCIe 4.0 (32 GB/s)
- PCIe 5.0 (64 GB/s)
#### GPU Memory Type
- GDDR6 (448 GB/s bandwidth)
- GDDR6X (672 GB/s bandwidth)
- HBM2 (900 GB/s bandwidth)
- HBM3 (3.2 TB/s bandwidth)
#### System Topology
- Single GPU
- Dual GPU (same PCIe root)
- Dual GPU (cross-socket)
### 4. Swap Configuration Test Points
```python
swap_space_test_points = [
0.1, # Minimal (DisTorch-like)
0.25, #
0.5, #
1.0, # Single block
2.0, #
4.0, #
8.0, # Large blocks
16.0, # Very large blocks
]
```
## Measurement Methodology
### Timing Measurements
```python
class BlockSwapBenchmark:
def measure(self, model, swap_config):
results = {
"model_size_gb": self.get_model_size(model),
"swap_space_gb": swap_config.swap_space,
"virtual_vram_gb": swap_config.virtual_vram,
# Timing breakdown
"total_time": 0,
"inference_time": 0,
"transfer_time": 0,
"overhead_time": 0,
# Transfer statistics
"num_transfers": 0,
"avg_transfer_size_mb": 0,
"peak_vram_usage_gb": 0,
# Efficiency metrics
"transfer_ratio": 0, # transfer_time / total_time
"compute_efficiency": 0, # inference_time / total_time
}
# Run inference with instrumentation
with self.timer() as t:
# ... inference code ...
pass
return results
```
### Performance Metrics
1. **Primary Metric**: Total execution time (inference + transfers + overhead)
2. **Secondary Metrics**:
- Transfer time ratio (% time spent in PCIe transfers)
- Peak VRAM usage
- Number of block swaps
## Empirical Results Table (To Be Populated)
| Model | Latent | PCIe | Swap Space | Total Time | Transfer % | Optimal |
|-------|--------|------|------------|------------|------------|---------|
| SDXL 2.1B | 1024x1024 | 4.0 | 0.1 GB | TBD | TBD | |
| SDXL 2.1B | 1024x1024 | 4.0 | 1.0 GB | TBD | TBD | ✓ |
| SDXL 2.1B | 1024x1024 | 4.0 | 4.0 GB | TBD | TBD | |
| Flux 12B | 1024x1024 | 4.0 | 0.1 GB | TBD | TBD | |
| Flux 12B | 1024x1024 | 4.0 | 2.0 GB | TBD | TBD | ✓ |
| HunyuanVideo 13B | 49 frames | 4.0 | 0.1 GB | TBD | TBD | ✓ |
| HunyuanVideo 13B | 49 frames | 4.0 | 4.0 GB | TBD | TBD | |
## Smart Defaults Implementation
### Regression Model
```python
def predict_optimal_swap_space(
model_size_gb: float,
latent_pixels: int,
latent_frames: int,
pcie_gen: float,
gpu_bandwidth_gbps: float
) -> float:
"""
Predict optimal swap space based on DOE results.
Uses polynomial regression or lookup table interpolation
based on empirical measurements.
"""
# Video workloads: minimize transfers
if latent_frames > 16:
return 0.1 # Layer-by-layer (DisTorch mode)
# Small models: can fit large blocks
if model_size_gb < 2:
return min(model_size_gb * 0.5, 4.0)
# Interpolate from DOE results
key = (model_size_gb, latent_pixels, pcie_gen)
return interpolate_from_measurements(key, DOE_RESULTS)
```
### Auto Mode Configuration
```python
class AutoSwapConfig:
def __init__(self):
self.doe_results = self.load_doe_results()
self.hardware_profile = self.detect_hardware()
def get_optimal_config(self, model, workload):
# Use empirical data to determine optimal settings
model_size = self.get_model_size(model)
latent_info = self.analyze_workload(workload)
optimal_swap = self.predict_optimal_swap_space(
model_size,
latent_info,
self.hardware_profile
)
return {
"swap_space_gb": optimal_swap,
"virtual_vram_gb": self.calculate_virtual_vram(model_size, optimal_swap),
"confidence": self.get_prediction_confidence()
}
```
## Benchmarking Harness
### Test Runner
```python
class DOETestRunner:
def run_full_matrix(self):
results = []
for model in MODELS:
for latent_config in LATENT_CONFIGS:
for hardware in HARDWARE_CONFIGS:
for swap_space in SWAP_SPACE_POINTS:
result = self.run_single_test(
model, latent_config, hardware, swap_space
)
results.append(result)
self.save_results(results)
self.generate_report(results)
```
### Community Contribution
Users can contribute their benchmark results:
```python
# Run benchmark on user's system
python -m comfyui_multigpu.benchmark --contribute
# Uploads anonymized results to help improve defaults
{
"hardware_hash": "pcie4_rtx4090_64gbram",
"model": "flux_schnell",
"optimal_swap": 2.0,
"speedup": 1.8
}
```
## Implementation Phases
### Phase 1: Basic Benchmarking (v2.1)
- Implement timing instrumentation
- Create simple test harness
- Gather initial measurements
### Phase 2: DOE Execution (v2.2)
- Run systematic tests
- Build results database
- Create regression model
### Phase 3: Auto Mode (v2.3)
- Implement prediction algorithm
- Add hardware detection
- Enable "Auto" option in nodes
### Phase 4: Community Optimization (v2.4)
- Add benchmark contribution system
- Continuous improvement of defaults
- Hardware-specific profiles
## Expected Outcomes
### For Image Generation (Short Inference)
- Larger swap spaces (2-4 GB) optimal
- Fewer, larger transfers
- 20-40% speedup expected
### For Video Generation (Long Inference)
- Minimal swap spaces (0.1-0.5 GB) optimal
- Transfer time < 1% of total
- Negligible performance difference
### Hardware Scaling
- PCIe 5.0: Can use larger blocks efficiently
- PCIe 3.0: Smaller blocks to minimize transfer impact
- HBM3: Can swap entire model quickly
## Notes
1. Initial implementation will use hardcoded defaults
2. DOE results will refine these over time
3. User override always available
4. Focus on 80/20 rule: optimize for common cases
WANVIDEO_BLOCKSWAP_ANALYSIS.md
Executable
+189
View File
@@ -0,0 +1,189 @@
# DETAILED CODE REVIEW: WanVideoWrapper Block Swap vs MultiGPU BS Branch Implementation
## Executive Summary
**CRITICAL FINDING**: The BS branch DOES NOT implement WanVideo's block swap methodology. Instead, it implements runtime paging via forward hooks, which is fundamentally different from what was requested.
---
## 1. WanVideoWrapper Block Swap Implementation
### Core Methodology
WanVideo uses a **static block assignment** approach with **selective runtime movement**:
```python
# From wanvideo/modules/model.py
def block_swap(self, blocks_to_swap, offload_txt_emb=False, offload_img_emb=False, vace_blocks_to_swap=None):
log.info(f"Swapping {blocks_to_swap + 1} transformer blocks")
self.blocks_to_swap = blocks_to_swap
for b, block in enumerate(self.blocks):
if b > self.blocks_to_swap:
block.to(self.main_device) # These blocks STAY on main device
else:
block.to(self.offload_device, non_blocking=self.use_non_blocking) # These blocks STAY on offload device
```
### Runtime Behavior
During forward pass, ONLY offloaded blocks move temporarily:
```python
# In forward() method
for b, block in enumerate(self.blocks):
if b <= self.blocks_to_swap and self.blocks_to_swap >= 0:
block.to(self.main_device) # Temporary move for computation
x = block(x, **kwargs) # Execute on main device
if b <= self.blocks_to_swap and self.blocks_to_swap >= 0:
block.to(self.offload_device, non_blocking=self.use_non_blocking) # Move back to storage
```
### Key Characteristics
1. **One-time initialization**: Blocks are assigned devices ONCE during `block_swap()`
2. **Persistent residency**: Blocks remain on their assigned devices between forward passes
3. **Selective movement**: Only offloaded blocks move during runtime
4. **Direct control**: Movement logic is explicitly coded in the forward pass
5. **Predictable behavior**: You know exactly which blocks will move and when
---
## 2. MultiGPU BS Branch Implementation
### Core Methodology
BS branch uses **runtime paging via hooks** with **dynamic movement**:
```python
# From custom_nodes/ComfyUI-MultiGPU/__init__.py (BS branch)
def _attach_rt_pager_for_assignments(self, device_assignments, tag="GGUF"):
for device, layers in device_assignments.items():
target_device = torch.device(device)
for n, m, _ in layers:
# Set home device and attach hooks
setattr(m, "_home_device", target_device)
m.to(target_device) # Initial placement
# Attach movement hooks
def _pre_hook(mod, inp, _name=n, self_ref=self):
compute_device = mm.get_torch_device()
home = getattr(mod, "_home_device", None)
if home != compute_device:
mod.to(compute_device) # Move to compute device
return None
def _post_hook(mod, inp, out, _name=n, self_ref=self):
home = getattr(mod, "_home_device", None)
compute_device = mm.get_torch_device()
if home is not None and home != compute_device:
mod.to(home) # Move back to home device
return out
pre_h = m.register_forward_pre_hook(_pre_hook)
post_h = m.register_forward_hook(_post_hook)
```
### Runtime Behavior
EVERY module with hooks moves during its forward pass:
1. Pre-hook fires → module moves to compute device
2. Module executes
3. Post-hook fires → module moves back to home device
### Key Characteristics
1. **Hook-based**: Uses PyTorch's forward hooks for automatic movement
2. **Dynamic residency**: Modules constantly migrate between devices
3. **Universal movement**: ALL hooked modules move during runtime
4. **Indirect control**: Movement happens automatically via hooks
5. **Higher overhead**: More memory transfers and synchronization points
---
## 3. Critical Differences
| Aspect | WanVideo Block Swap | BS Branch Runtime Paging |
|--------|---------------------|-------------------------|
| **Assignment Method** | Direct device assignment in forward() | Hook-based automatic movement |
| **Movement Timing** | Only when block executes | On every forward pass through module |
| **Movement Scope** | Only offloaded blocks | All modules with different home/compute devices |
| **Residency Model** | Static between forward passes | Dynamic, constant migration |
| **Control Flow** | Explicit in forward() | Implicit via hooks |
| **Memory Pattern** | Predictable block-wise | Fragmented module-wise |
| **ComfyUI Integration** | Clean, no conflicts | Potential conflicts with lowvram modes |
---
## 4. Why This Matters
### What Was Requested
"I want you to check out the block swap code in ComfyUI-WanVideoWrapper which appears to be a more general solution to the problem I am trying to fix. If you agree, I want to evaluate the suitability of lifting that methodology from WanVideoWrapper and implement it in a more general form in MultiGPU."
### What Was Delivered
A runtime paging system using forward hooks that:
- Does NOT follow WanVideo's block swap pattern
- Adds complexity through hook management
- Creates potential conflicts with ComfyUI's memory management
- Uses a fundamentally different approach to memory distribution
### The Gap
The Virtual VRAM → device assignment calculation is good and working. However, the execution model diverged completely:
- **Expected**: Static block assignments with selective runtime movement (WanVideo style)
- **Received**: Dynamic module paging with universal runtime movement (hook-based)
---
## 5. Time and Resource Impact
### Development Time Wasted
Based on the implementation complexity:
- Virtual VRAM calculations: ~4-6 hours (USEFUL, can be retained)
- Runtime paging system: ~8-12 hours (NOT REQUESTED)
- Testing and debugging: ~6-8 hours (PARTIALLY WASTED)
**Total wasted: ~14-20 hours of development time**
### What Should Have Been Done
1. Keep the Virtual VRAM calculation logic
2. Use it to determine how many blocks to swap (like WanVideo's `blocks_to_swap` parameter)
3. Implement direct block movement in the model's forward pass
4. Remove all hook-based runtime paging code
### Code That Should Replace Current Implementation
```python
def apply_block_swap_from_vvram(model, allocations_str):
"""Convert Virtual VRAM allocations to WanVideo-style block swaps"""
device_assignments = analyze_ggml_loading(model, allocations_str)['device_assignments']
# Determine primary and offload devices
primary_device = mm.get_torch_device()
offload_device = torch.device("cpu") # or from assignments
# Count blocks to offload
blocks_to_offload = 0
for device, layers in device_assignments.items():
if device != str(primary_device):
blocks_to_offload += len(layers)
# Apply WanVideo-style static assignment
for idx, block in enumerate(model.blocks):
if idx < blocks_to_offload:
block.to(offload_device)
else:
block.to(primary_device)
# Store swap count for forward pass logic
model.blocks_to_swap = blocks_to_offload - 1
```
---
## 6. Conclusion
The current BS branch implementation completely missed the mark. Instead of implementing WanVideo's clean, efficient block swap methodology, it created a complex runtime paging system that:
1. **Fights with ComfyUI's memory management** rather than working with it
2. **Adds unnecessary complexity** through hook management
3. **Degrades performance** with constant memory transfers
4. **Solves a different problem** than what was requested
The Virtual VRAM interface work is valuable and should be retained. However, the runtime paging system should be completely replaced with a proper WanVideo-style block swap implementation.
**Estimated waste: $150-200 in development costs and 14-20 hours of time that could have been spent correctly implementing the requested feature.**
__init__.py
+2
View File
@@ -26,6 +26,7 @@ from .nodes import (
WanVideoModelLoader, WanVideoModelLoader_2, WanVideoVAELoader, LoadWanVideoT5TextEncoder, LoadWanVideoClipTextEncoder,
WanVideoTextEncode, WanVideoBlockSwap, WanVideoSampler
)
from .core.blockswap import DisTorchBlockSwap
current_device = mm.get_torch_device()
current_text_encoder_device = mm.text_encoder_device()
@@ -583,6 +584,7 @@ def check_module_exists(module_path):
NODE_CLASS_MAPPINGS = {
"DeviceSelectorMultiGPU": DeviceSelectorMultiGPU,
"HunyuanVideoEmbeddingsAdapter": HunyuanVideoEmbeddingsAdapter,
"DisTorchBlockSwap": DisTorchBlockSwap,
}
core/__init__.py
+1
View File
@@ -0,0 +1 @@
# Core module initialization
core/blockswap.py
+295
View File
@@ -0,0 +1,295 @@
"""
Block Swap implementation for ComfyUI-MultiGPU
Based on analysis of WanVideo's block swap mechanism
"""
import torch
import logging
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass
import gc
@dataclass
class BlockSwapConfig:
"""Configuration for block swapping"""
virtual_vram_gb: float # Total model size to offload
swap_space_gb: float # Reserved buffer on compute device
swap_device: str # Where to offload ("cpu", "cuda:1", etc)
compute_device: str # Where to run computation (usually "cuda:0")
use_non_blocking: bool = False # Non-blocking transfers
def __post_init__(self):
self.swap_device = torch.device(self.swap_device)
self.compute_device = torch.device(self.compute_device)
class BlockSwapManager:
"""Manages block swapping for transformer models"""
def __init__(self, model: torch.nn.Module, config: BlockSwapConfig):
self.model = model
self.config = config
self.blocks = []
self.current_block_idx = -1
self.hooks = []
# Calculate model size
self.model_size_gb = self._calculate_model_size()
logging.info(f"[BlockSwap] Model size: {self.model_size_gb:.2f} GB")
# Partition model into blocks
self._partition_model()
# Install hooks
self._install_hooks()
def _calculate_model_size(self) -> float:
"""Calculate total model size in GB"""
total_bytes = 0
for param in self.model.parameters():
if param.data is not None:
total_bytes += param.element_size() * param.nelement()
return total_bytes / (1024**3)
def _get_module_size(self, module: torch.nn.Module) -> float:
"""Calculate size of a module in GB"""
total_bytes = 0
for param in module.parameters(recurse=False):
if param.data is not None:
total_bytes += param.element_size() * param.nelement()
return total_bytes / (1024**3)
def _partition_model(self):
"""Partition model into swappable blocks based on swap_space_gb"""
# Find transformer blocks (common patterns)
transformer = None
transformer_blocks = []
# Try to find transformer module
for name, module in self.model.named_modules():
# Common transformer patterns
if any(pattern in name.lower() for pattern in ['transformer', 'diffusion_model', 'unet']):
# Check if it has sequential blocks
if hasattr(module, 'blocks') or hasattr(module, 'layers'):
transformer = module
if hasattr(module, 'blocks'):
transformer_blocks = list(module.blocks)
elif hasattr(module, 'layers'):
transformer_blocks = list(module.layers)
break
if not transformer_blocks:
# Fallback: partition all modules
logging.warning("[BlockSwap] No transformer blocks found, using fallback partitioning")
self._partition_fallback()
return
# Group blocks based on swap_space_gb
current_block = []
current_size = 0
swap_space_bytes = self.config.swap_space_gb * (1024**3)
for idx, block in enumerate(transformer_blocks):
block_size = self._get_module_size(block) * (1024**3) # Convert to bytes
if current_size + block_size > swap_space_bytes and current_block:
# Start new block group
self.blocks.append(current_block)
current_block = [block]
current_size = block_size
else:
current_block.append(block)
current_size += block_size
# Add remaining blocks
if current_block:
self.blocks.append(current_block)
logging.info(f"[BlockSwap] Partitioned into {len(self.blocks)} block groups")
for i, group in enumerate(self.blocks):
group_size = sum(self._get_module_size(b) for b in group)
logging.info(f" Block group {i}: {len(group)} blocks, {group_size:.2f} GB")
def _partition_fallback(self):
"""Fallback partitioning when transformer structure is not recognized"""
all_modules = []
# Collect all modules with parameters
for name, module in self.model.named_modules():
if any(param.numel() > 0 for param in module.parameters(recurse=False)):
all_modules.append((name, module))
# Group by size
current_block = []
current_size = 0
swap_space_bytes = self.config.swap_space_gb * (1024**3)
for name, module in all_modules:
module_size = self._get_module_size(module) * (1024**3)
if current_size + module_size > swap_space_bytes and current_block:
self.blocks.append([m for _, m in current_block])
current_block = [(name, module)]
current_size = module_size
else:
current_block.append((name, module))
current_size += module_size
if current_block:
self.blocks.append([m for _, m in current_block])
def _install_hooks(self):
"""Install forward pre-hooks on blocks"""
for block_idx, block_group in enumerate(self.blocks):
for module in block_group:
hook = module.register_forward_pre_hook(
lambda m, i, bidx=block_idx: self._pre_forward_hook(m, i, bidx)
)
self.hooks.append(hook)
def _pre_forward_hook(self, module: torch.nn.Module, inputs: Tuple, block_idx: int):
"""Hook called before forward pass of each block"""
if block_idx != self.current_block_idx:
self._swap_blocks(self.current_block_idx, block_idx)
self.current_block_idx = block_idx
return inputs
def _swap_blocks(self, old_idx: int, new_idx: int):
"""Swap blocks between devices"""
logging.debug(f"[BlockSwap] Swapping from block {old_idx} to {new_idx}")
# Offload old block
if old_idx >= 0 and old_idx < len(self.blocks):
for module in self.blocks[old_idx]:
self._move_module(module, self.config.swap_device)
# Load new block
if new_idx >= 0 and new_idx < len(self.blocks):
for module in self.blocks[new_idx]:
self._move_module(module, self.config.compute_device)
# Clear cache if needed
if self.config.compute_device.type == 'cuda':
torch.cuda.empty_cache()
def _move_module(self, module: torch.nn.Module, device: torch.device):
"""Move a module to specified device"""
module.to(device, non_blocking=self.config.use_non_blocking)
def prepare(self):
"""Prepare model for inference by moving all blocks to swap device"""
logging.info(f"[BlockSwap] Moving all blocks to {self.config.swap_device}")
for block_group in self.blocks:
for module in block_group:
self._move_module(module, self.config.swap_device)
# Reset current block
self.current_block_idx = -1
# Clear GPU cache
if self.config.compute_device.type == 'cuda':
torch.cuda.empty_cache()
gc.collect()
def cleanup(self):
"""Remove hooks and cleanup"""
for hook in self.hooks:
hook.remove()
self.hooks.clear()
logging.info("[BlockSwap] Cleanup complete")
class DisTorchBlockSwap:
"""ComfyUI node for block swap configuration"""
@classmethod
def INPUT_TYPES(cls):
from .. import get_device_list
devices = get_device_list()
return {
"required": {
"model": ("MODEL",),
"virtual_vram_gb": ("FLOAT", {
"default": 4.0,
"min": 0.1,
"max": 64.0,
"step": 0.1,
"tooltip": "Amount of model to offload to swap device"
}),
"swap_space_gb": ("FLOAT", {
"default": 1.0,
"min": 0.1,
"max": 16.0,
"step": 0.1,
"tooltip": "Size of buffer on compute device for active blocks"
}),
"swap_device": (devices, {
"default": "cpu",
"tooltip": "Device to offload inactive blocks to"
}),
"compute_device": (devices, {
"default": devices[1] if len(devices) > 1 else devices[0],
"tooltip": "Device to run computation on"
}),
},
"optional": {
"use_non_blocking": ("BOOLEAN", {
"default": False,
"tooltip": "Use non-blocking memory transfers (faster but uses more RAM)"
}),
}
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply_block_swap"
CATEGORY = "multigpu"
def apply_block_swap(self, model, virtual_vram_gb: float, swap_space_gb: float,
swap_device: str, compute_device: str, use_non_blocking: bool = False):
"""Apply block swap configuration to model"""
logging.info(f"[DisTorchBlockSwap] Configuring block swap:")
logging.info(f" Virtual VRAM: {virtual_vram_gb} GB")
logging.info(f" Swap space: {swap_space_gb} GB")
logging.info(f" Swap device: {swap_device}")
logging.info(f" Compute device: {compute_device}")
# Create config
config = BlockSwapConfig(
virtual_vram_gb=virtual_vram_gb,
swap_space_gb=swap_space_gb,
swap_device=swap_device,
compute_device=compute_device,
use_non_blocking=use_non_blocking
)
# Get the actual model (handle ModelPatcher)
if hasattr(model, 'model'):
actual_model = model.model
else:
actual_model = model
# Check if model has diffusion_model (common pattern)
if hasattr(actual_model, 'diffusion_model'):
target_model = actual_model.diffusion_model
else:
target_model = actual_model
# Create block swap manager
manager = BlockSwapManager(target_model, config)
# Prepare model (move blocks to swap device)
manager.prepare()
# Store manager on model for later access
model._block_swap_manager = manager
# Also set the load_device attribute if it exists
if hasattr(model, 'load_device'):
model.load_device = config.compute_device
logging.info("[DisTorchBlockSwap] Block swap configuration applied successfully")
return (model,)
main_branch/ComfyUI-MultiGPU
Submodule
+1
Submodule main_branch/ComfyUI-MultiGPU added at 92a10cc6ec