401 lines
17 KiB
Python
401 lines
17 KiB
Python
"""
|
|
Compatibility module for SeedVR2
|
|
Contains FP8/FP16 compatibility layers and wrappers for different model architectures
|
|
|
|
Extracted from: seedvr2.py (lines 1045-1630)
|
|
"""
|
|
|
|
import torch
|
|
import platform
|
|
import types
|
|
from typing import List, Tuple, Union, Any, Optional
|
|
|
|
|
|
def call_rope_with_stability(method, *args, **kwargs):
|
|
"""
|
|
Call RoPE method with stability fixes:
|
|
1. Clear cache if available
|
|
2. Disable autocast to prevent numerical issues
|
|
This prevents artifacts in FP8/mixed precision models.
|
|
"""
|
|
if hasattr(method, 'cache_clear'):
|
|
method.cache_clear()
|
|
|
|
with torch.cuda.amp.autocast(enabled=False):
|
|
return method(*args, **kwargs)
|
|
|
|
class FP8CompatibleDiT(torch.nn.Module):
|
|
"""
|
|
Wrapper for DiT models with automatic compatibility management + advanced optimizations
|
|
- FP8: Keeps native FP8 parameters, converts inputs/outputs
|
|
- FP16: Uses native FP16
|
|
- Mixed Precision: Stabilizes RoPE for models with FP16 blocks
|
|
- RoPE: Converted from FP8 to BFloat16 only when detected as FP8
|
|
- Flash Attention: Automatic optimization of attention layers
|
|
"""
|
|
|
|
def __init__(self, dit_model, skip_conversion=False, debug=None):
|
|
super().__init__()
|
|
self.dit_model = dit_model
|
|
if debug is None:
|
|
raise ValueError("Debug instance must be provided to FP8CompatibleDiT")
|
|
self.debug = debug if debug is not None else None
|
|
self.model_dtype = self._detect_model_dtype()
|
|
self.is_fp8_model = self.model_dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
|
|
self.is_fp16_model = self.model_dtype == torch.float16
|
|
|
|
# Only convert if not already done (e.g., when reusing cached weights)
|
|
if not skip_conversion and self.is_fp8_model:
|
|
# Only FP8 models need RoPE frequency conversion
|
|
model_variant = "7B" if self._is_nadit_model() else "3B" if self._is_nadit_v2_model() else "Unknown"
|
|
self.debug.log(f"Detected NaDiT {model_variant} FP8 - Converting RoPE freqs for FP8 compatibility",
|
|
category="precision", force=True)
|
|
self._convert_rope_freqs()
|
|
|
|
# Apply RoPE stabilization for numerical stability
|
|
self._stabilize_rope_computations()
|
|
|
|
# 🚀 FLASH ATTENTION OPTIMIZATION (Phase 2)
|
|
self._apply_flash_attention_optimization()
|
|
|
|
def _detect_model_dtype(self) -> torch.dtype:
|
|
"""Detect main model dtype"""
|
|
try:
|
|
return next(self.dit_model.parameters()).dtype
|
|
except:
|
|
return torch.bfloat16
|
|
|
|
def _is_nadit_model(self) -> bool:
|
|
"""Detect if this is a NaDiT model (7B) with precise logic"""
|
|
# Check module path for dit (not dit_v2)
|
|
model_module = str(self.dit_model.__class__.__module__).lower()
|
|
return 'dit.nadit' in model_module and 'dit_v2' not in model_module
|
|
|
|
def _is_nadit_v2_model(self) -> bool:
|
|
"""Detect if this is a NaDiT v2 model (3B) with precise logic"""
|
|
# Check module path for dit_v2
|
|
model_module = str(self.dit_model.__class__.__module__).lower()
|
|
return 'dit_v2' in model_module
|
|
|
|
def _convert_rope_freqs(self) -> None:
|
|
"""Convert RoPE frequency buffers for FP8 compatibility"""
|
|
converted = 0
|
|
for module in self.dit_model.modules():
|
|
if 'RotaryEmbedding' in type(module).__name__:
|
|
if hasattr(module, 'rope') and hasattr(module.rope, 'freqs'):
|
|
if module.rope.freqs.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
|
|
module.rope.freqs.data = module.rope.freqs.to(torch.bfloat16)
|
|
converted += 1
|
|
self.debug.log(f"Converted {converted} RoPE frequency buffers", category="success")
|
|
|
|
def _stabilize_rope_computations(self):
|
|
"""
|
|
Add error handling to RoPE computations to prevent artifacts.
|
|
|
|
Wraps the get_axial_freqs method of RoPE modules with a try-except handler.
|
|
During normal operation, uses the original cached method for performance.
|
|
Only on exceptions (e.g., numerical instability, NaN propagation) does it
|
|
intervene by clearing the cache and retrying the computation through
|
|
call_rope_with_stability.
|
|
|
|
This prevents artifacts in FP8, mixed precision, and edge cases while
|
|
maintaining optimal performance for normal operations.
|
|
"""
|
|
if not hasattr(self.dit_model, 'blocks'):
|
|
return
|
|
|
|
self.debug.start_timer("stabilize_rope")
|
|
self.debug.log(f"Stabilizing RoPE computations for numerical stability", category="precision")
|
|
|
|
rope_count = 0
|
|
|
|
# Wrap RoPE modules to handle numerical instability
|
|
for name, module in self.dit_model.named_modules():
|
|
if "rope" in name.lower() and hasattr(module, "get_axial_freqs"):
|
|
# Check if already wrapped
|
|
if hasattr(module, '_rope_wrapped'):
|
|
continue
|
|
|
|
original_method = module.get_axial_freqs
|
|
|
|
# Mark as wrapped and store original
|
|
module._rope_wrapped = 'stability'
|
|
module._original_get_axial_freqs = original_method
|
|
|
|
# Error handler that prevents NaN propagation
|
|
def stable_rope_computation(self, *args, **kwargs):
|
|
try:
|
|
return original_method(*args, **kwargs)
|
|
except Exception:
|
|
return call_rope_with_stability(original_method, *args, **kwargs)
|
|
|
|
module.get_axial_freqs = types.MethodType(stable_rope_computation, module)
|
|
rope_count += 1
|
|
|
|
if rope_count > 0:
|
|
self.debug.log(f"Stabilized {rope_count} RoPE modules", category="success")
|
|
|
|
self.debug.end_timer("stabilize_rope", f"Stabilized {rope_count} RoPE modules")
|
|
|
|
def _apply_flash_attention_optimization(self) -> None:
|
|
"""🚀 FLASH ATTENTION OPTIMIZATION - 30-50% speedup of attention layers"""
|
|
attention_layers_optimized = 0
|
|
flash_attention_available = self._check_flash_attention_support()
|
|
|
|
for name, module in self.dit_model.named_modules():
|
|
# Identify all attention layers
|
|
if self._is_attention_layer(name, module):
|
|
# Apply optimization based on availability
|
|
if self._optimize_attention_layer(name, module, flash_attention_available):
|
|
attention_layers_optimized += 1
|
|
|
|
if not flash_attention_available:
|
|
self.debug.log("Flash Attention not available, using PyTorch SDPA as fallback", category="info", force=True)
|
|
|
|
def _check_flash_attention_support(self) -> bool:
|
|
"""Check if Flash Attention is available"""
|
|
try:
|
|
# Check PyTorch SDPA (includes Flash Attention on H100/A100)
|
|
if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
|
|
return True
|
|
|
|
# Check flash-attn package
|
|
import flash_attn
|
|
return True
|
|
except ImportError:
|
|
return False
|
|
|
|
def _is_attention_layer(self, name: str, module: torch.nn.Module) -> bool:
|
|
"""Identify if a module is an attention layer"""
|
|
attention_keywords = [
|
|
'attention', 'attn', 'self_attn', 'cross_attn', 'mhattn', 'multihead',
|
|
'transformer_block', 'dit_block'
|
|
]
|
|
|
|
# Check by name
|
|
if any(keyword in name.lower() for keyword in attention_keywords):
|
|
return True
|
|
|
|
# Check by module type
|
|
module_type = type(module).__name__.lower()
|
|
if any(keyword in module_type for keyword in attention_keywords):
|
|
return True
|
|
|
|
# Check by attributes (modules with q, k, v projections)
|
|
if hasattr(module, 'q_proj') or hasattr(module, 'qkv') or hasattr(module, 'to_q'):
|
|
return True
|
|
|
|
return False
|
|
|
|
def _optimize_attention_layer(self, name: str, module: torch.nn.Module, flash_attention_available: bool) -> bool:
|
|
"""Optimize a specific attention layer"""
|
|
try:
|
|
# Save original forward method
|
|
if not hasattr(module, '_original_forward'):
|
|
module._original_forward = module.forward
|
|
|
|
# Create new optimized forward method
|
|
if flash_attention_available:
|
|
optimized_forward = self._create_flash_attention_forward(module, name)
|
|
else:
|
|
optimized_forward = self._create_sdpa_forward(module, name)
|
|
|
|
# Replace forward method
|
|
module.forward = optimized_forward
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.debug.log(f"Failed to optimize attention layer '{name}': {e}", level="WARNING", category="model")
|
|
return False
|
|
|
|
def _create_flash_attention_forward(self, module: torch.nn.Module, layer_name: str):
|
|
"""Create optimized forward with Flash Attention"""
|
|
original_forward = module._original_forward
|
|
|
|
def flash_attention_forward(*args, **kwargs):
|
|
try:
|
|
# Try to use Flash Attention via SDPA
|
|
return self._sdpa_attention_forward(original_forward, module, *args, **kwargs)
|
|
except Exception as e:
|
|
# Fallback to original implementation
|
|
self.debug.log(f"Flash Attention failed for {layer_name}, using original: {e}", level="WARNING", category="model", force=True)
|
|
return original_forward(*args, **kwargs)
|
|
|
|
return flash_attention_forward
|
|
|
|
def _create_sdpa_forward(self, module: torch.nn.Module, layer_name: str):
|
|
"""Create optimized forward with PyTorch SDPA"""
|
|
original_forward = module._original_forward
|
|
|
|
def sdpa_forward(*args, **kwargs):
|
|
try:
|
|
return self._sdpa_attention_forward(original_forward, module, *args, **kwargs)
|
|
except Exception as e:
|
|
# Fallback to original implementation
|
|
return original_forward(*args, **kwargs)
|
|
|
|
return sdpa_forward
|
|
|
|
def _sdpa_attention_forward(self, original_forward, module: torch.nn.Module, *args, **kwargs):
|
|
"""Optimized forward pass using SDPA (Scaled Dot Product Attention)"""
|
|
# Detect if we can intercept and optimize this layer
|
|
if len(args) >= 1 and isinstance(args[0], torch.Tensor):
|
|
input_tensor = args[0]
|
|
|
|
# Check dimensions to ensure it's standard attention
|
|
if len(input_tensor.shape) >= 3: # [batch, seq_len, hidden_dim] or similar
|
|
try:
|
|
return self._optimized_attention_computation(module, input_tensor, *args[1:], **kwargs)
|
|
except:
|
|
pass
|
|
|
|
# Fallback to original implementation
|
|
return original_forward(*args, **kwargs)
|
|
|
|
def _optimized_attention_computation(self, module: torch.nn.Module, input_tensor: torch.Tensor, *args, **kwargs):
|
|
"""Optimized attention computation with SDPA"""
|
|
# Try to detect standard attention format
|
|
batch_size, seq_len = input_tensor.shape[:2]
|
|
|
|
# Check if module has standard Q, K, V projections
|
|
if hasattr(module, 'qkv') or (hasattr(module, 'q_proj') and hasattr(module, 'k_proj') and hasattr(module, 'v_proj')):
|
|
return self._compute_sdpa_attention(module, input_tensor, *args, **kwargs)
|
|
|
|
# If no standard format detected, use original
|
|
return module._original_forward(input_tensor, *args, **kwargs)
|
|
|
|
def _compute_sdpa_attention(self, module: torch.nn.Module, x: torch.Tensor, *args, **kwargs):
|
|
"""Optimized SDPA computation for standard attention modules"""
|
|
try:
|
|
# Case 1: Module with combined QKV projection
|
|
if hasattr(module, 'qkv'):
|
|
qkv = module.qkv(x)
|
|
# Reshape to separate Q, K, V
|
|
batch_size, seq_len, _ = qkv.shape
|
|
qkv = qkv.reshape(batch_size, seq_len, 3, -1)
|
|
q, k, v = qkv.unbind(dim=2)
|
|
|
|
# Case 2: Separate Q, K, V projections
|
|
elif hasattr(module, 'q_proj') and hasattr(module, 'k_proj') and hasattr(module, 'v_proj'):
|
|
q = module.q_proj(x)
|
|
k = module.k_proj(x)
|
|
v = module.v_proj(x)
|
|
else:
|
|
# Unsupported format, use original
|
|
return module._original_forward(x, *args, **kwargs)
|
|
|
|
# Detect number of heads
|
|
head_dim = getattr(module, 'head_dim', None)
|
|
num_heads = getattr(module, 'num_heads', None)
|
|
|
|
if head_dim is None or num_heads is None:
|
|
# Try to guess from dimensions
|
|
hidden_dim = q.shape[-1]
|
|
if hasattr(module, 'num_heads'):
|
|
num_heads = module.num_heads
|
|
head_dim = hidden_dim // num_heads
|
|
else:
|
|
# Reasonable defaults
|
|
head_dim = 64
|
|
num_heads = hidden_dim // head_dim
|
|
|
|
# Reshape for multi-head attention
|
|
batch_size, seq_len = q.shape[:2]
|
|
q = q.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
|
|
k = k.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
|
|
v = v.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
|
|
|
|
if platform.system() == "Darwin":
|
|
attn_output = torch.nn.functional.scaled_dot_product_attention(
|
|
q, k, v,
|
|
dropout_p=0.0,
|
|
is_causal=False
|
|
)
|
|
else:
|
|
# Use optimized SDPA
|
|
with torch.backends.cuda.sdp_kernel(
|
|
enable_flash=True,
|
|
enable_math=True,
|
|
enable_mem_efficient=True
|
|
):
|
|
attn_output = torch.nn.functional.scaled_dot_product_attention(
|
|
q, k, v,
|
|
dropout_p=0.0,
|
|
is_causal=False
|
|
)
|
|
|
|
# Reshape back
|
|
attn_output = attn_output.transpose(1, 2).contiguous().view(
|
|
batch_size, seq_len, num_heads * head_dim
|
|
)
|
|
|
|
# Output projection if it exists
|
|
if hasattr(module, 'out_proj') or hasattr(module, 'o_proj'):
|
|
proj = getattr(module, 'out_proj', None) or getattr(module, 'o_proj', None)
|
|
attn_output = proj(attn_output)
|
|
|
|
return attn_output
|
|
|
|
except Exception as e:
|
|
# In case of error, use original implementation
|
|
return module._original_forward(x, *args, **kwargs)
|
|
|
|
def forward(self, *args, **kwargs):
|
|
"""Forward pass with minimal dtype conversion overhead
|
|
|
|
Conversion strategy:
|
|
- FP16 models: Keep everything in FP16 (no conversion needed)
|
|
- FP8 models: Convert FP8 tensors to BFloat16 (required for arithmetic)
|
|
- BFloat16 models: No conversion needed
|
|
"""
|
|
|
|
# Only convert if we have an FP8 model for arithmetic operations
|
|
if self.is_fp8_model:
|
|
fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2)
|
|
|
|
# Convert args
|
|
converted_args = []
|
|
for arg in args:
|
|
if isinstance(arg, torch.Tensor) and arg.dtype in fp8_dtypes:
|
|
converted_args.append(arg.to(torch.bfloat16))
|
|
else:
|
|
converted_args.append(arg)
|
|
|
|
# Convert kwargs
|
|
converted_kwargs = {}
|
|
for key, value in kwargs.items():
|
|
if isinstance(value, torch.Tensor) and value.dtype in fp8_dtypes:
|
|
converted_kwargs[key] = value.to(torch.bfloat16)
|
|
else:
|
|
converted_kwargs[key] = value
|
|
|
|
args = tuple(converted_args)
|
|
kwargs = converted_kwargs
|
|
|
|
# Execute forward pass
|
|
try:
|
|
return self.dit_model(*args, **kwargs)
|
|
except Exception as e:
|
|
self.debug.log(f"Forward pass error: {e}", category="error", force=True)
|
|
if self.is_fp8_model:
|
|
self.debug.log(f"FP8 model - converted FP8 tensors to BFloat16", category="info", force=True)
|
|
else:
|
|
self.debug.log(f"{self.model_dtype} model - no conversion applied", category="info", force=True)
|
|
raise
|
|
|
|
def __getattr__(self, name):
|
|
"""Redirect all other attributes to original model"""
|
|
if name in ['dit_model', 'model_dtype', 'is_fp8_model', 'is_fp16_model']:
|
|
return super().__getattr__(name)
|
|
return getattr(self.dit_model, name)
|
|
|
|
def __setattr__(self, name, value):
|
|
"""Redirect assignments to original model except for our attributes"""
|
|
if name in ['dit_model', 'model_dtype', 'is_fp8_model', 'is_fp16_model']:
|
|
super().__setattr__(name, value)
|
|
else:
|
|
if hasattr(self, 'dit_model'):
|
|
setattr(self.dit_model, name, value)
|
|
else:
|
|
super().__setattr__(name, value)
|