486 lines
21 KiB
Python
486 lines
21 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 time
|
||
import torch
|
||
from typing import List, Tuple, Union, Any, Optional
|
||
|
||
|
||
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
|
||
- RoPE: ALWAYS forced to BFloat16 for maximum compatibility
|
||
- Flash Attention: Automatic optimization of attention layers
|
||
"""
|
||
|
||
def __init__(self, dit_model, skip_conversion=False):
|
||
super().__init__()
|
||
self.dit_model = dit_model
|
||
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:
|
||
# Detect model type
|
||
is_nadit_7b = self._is_nadit_model() # NaDiT 7B (dit/nadit)
|
||
is_nadit_v2_3b = self._is_nadit_v2_model() # NaDiT v2 3B (dit_v2/nadit)
|
||
|
||
if is_nadit_7b:
|
||
# 🎯 CRITICAL FIX: ALL NaDiT 7B models (FP8 AND FP16) require BFloat16 conversion
|
||
# 7B architecture has dtype compatibility issues regardless of storage format
|
||
if self.is_fp8_model:
|
||
print("🎯 Detected NaDiT 7B FP8 - Converting all parameters to BFloat16")
|
||
self._force_nadit_bfloat16()
|
||
else:
|
||
print("🎯 Detected NaDiT 7B FP16")
|
||
|
||
|
||
elif self.is_fp8_model and is_nadit_v2_3b:
|
||
# For NaDiT v2 3B FP8: Convert ALL model to BFloat16
|
||
print("🎯 Detected NaDiT v2 3B FP8 - Converting all parameters to BFloat16")
|
||
self._force_nadit_bfloat16()
|
||
|
||
# 🚀 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"""
|
||
# 🎯 PRIMARY METHOD: Check emb_scale attribute (specific to 7B)
|
||
# This is the most reliable criterion to distinguish 7B vs 3B
|
||
if hasattr(self.dit_model, 'emb_scale'):
|
||
return True
|
||
|
||
# 🎯 SECONDARY METHOD: Check module path for NaDiT 7B (dit/nadit, not dit_v2)
|
||
model_module = str(self.dit_model.__class__.__module__).lower()
|
||
if 'dit.nadit' in model_module and 'dit_v2' not in model_module:
|
||
return True
|
||
|
||
return False
|
||
|
||
def _is_nadit_v2_model(self) -> bool:
|
||
"""Detect if this is a NaDiT v2 model (3B) with precise logic"""
|
||
# 🎯 PRIMARY METHOD: Check module path for NaDiT v2 (dit_v2/nadit)
|
||
model_module = str(self.dit_model.__class__.__module__).lower()
|
||
if 'dit_v2' in model_module:
|
||
return True
|
||
|
||
# 🎯 SECONDARY METHOD: Check specific 3B structure
|
||
# NaDiT v2 3B has vid_in, txt_in, emb_in but NO emb_scale
|
||
if (hasattr(self.dit_model, 'vid_in') and
|
||
hasattr(self.dit_model, 'txt_in') and
|
||
hasattr(self.dit_model, 'emb_in') and
|
||
not hasattr(self.dit_model, 'emb_scale')): # Absence of emb_scale = 3B
|
||
return True
|
||
|
||
return False
|
||
|
||
def _force_rope_bfloat16(self) -> None:
|
||
"""🎯 Force ALL RoPE modules to BFloat16 for maximum compatibility"""
|
||
rope_count = 0
|
||
for name, module in self.dit_model.named_modules():
|
||
# Identify RoPE modules by name or type
|
||
if any(keyword in name.lower() for keyword in ['rope', 'rotary', 'embedding']):
|
||
# Convert all parameters of this module to BFloat16
|
||
for param_name, param in module.named_parameters():
|
||
if param.dtype != torch.bfloat16:
|
||
param.data = param.data.to(torch.bfloat16)
|
||
rope_count += 1
|
||
|
||
# Also convert buffers (non-trainable parameters)
|
||
for buffer_name, buffer in module.named_buffers():
|
||
if buffer.dtype != torch.bfloat16:
|
||
buffer.data = buffer.data.to(torch.bfloat16)
|
||
rope_count += 1
|
||
|
||
def _force_nadit_bfloat16(self) -> None:
|
||
"""🎯 Force ALL NaDiT parameters to BFloat16 to avoid promotion errors"""
|
||
print("🔧 Converting ALL NaDiT parameters to BFloat16 for type compatibility...")
|
||
t = time.time()
|
||
converted_count = 0
|
||
original_dtype = None
|
||
|
||
# Convert ALL parameters to BFloat16 (FP8, FP16, etc.)
|
||
for name, param in self.dit_model.named_parameters():
|
||
if original_dtype is None:
|
||
original_dtype = param.dtype
|
||
if param.dtype != torch.bfloat16:
|
||
param.data = param.data.to(torch.bfloat16)
|
||
converted_count += 1
|
||
|
||
# Also convert buffers
|
||
for name, buffer in self.dit_model.named_buffers():
|
||
if buffer.dtype != torch.bfloat16:
|
||
buffer.data = buffer.data.to(torch.bfloat16)
|
||
converted_count += 1
|
||
|
||
print(f" ✅ Converted {converted_count} parameters/buffers from {original_dtype} to BFloat16")
|
||
|
||
# Update detected dtype
|
||
self.model_dtype = torch.bfloat16
|
||
self.is_fp8_model = False # Model is no longer FP8 after conversion
|
||
|
||
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:
|
||
print(" ℹ️ Flash Attention not available, using PyTorch SDPA as fallback")
|
||
|
||
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:
|
||
print(f" ⚠️ Failed to optimize attention layer '{name}': {e}")
|
||
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
|
||
print(f" ⚠️ Flash Attention failed for {layer_name}, using original: {e}")
|
||
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)
|
||
|
||
# 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 intelligent type management according to architecture"""
|
||
is_nadit_7b = self._is_nadit_model()
|
||
is_nadit_v2_3b = self._is_nadit_v2_model()
|
||
|
||
# Input conversion according to architecture
|
||
if is_nadit_7b or is_nadit_v2_3b:
|
||
# For NaDiT models (7B and v2 3B): Everything to BFloat16
|
||
converted_args = []
|
||
for arg in args:
|
||
if isinstance(arg, torch.Tensor):
|
||
if arg.dtype in (torch.float32, torch.float8_e4m3fn, torch.float8_e5m2):
|
||
converted_args.append(arg.to(torch.bfloat16))
|
||
else:
|
||
converted_args.append(arg)
|
||
else:
|
||
converted_args.append(arg)
|
||
|
||
converted_kwargs = {}
|
||
for key, value in kwargs.items():
|
||
if isinstance(value, torch.Tensor):
|
||
if value.dtype in (torch.float32, torch.float8_e4m3fn, torch.float8_e5m2):
|
||
converted_kwargs[key] = value.to(torch.bfloat16)
|
||
else:
|
||
converted_kwargs[key] = value
|
||
else:
|
||
converted_kwargs[key] = value
|
||
|
||
args = tuple(converted_args)
|
||
kwargs = converted_kwargs
|
||
else:
|
||
# For standard models: Conversion according to model dtype
|
||
if self.is_fp8_model:
|
||
# Convert FP8 → BFloat16 for calculations
|
||
converted_args = []
|
||
for arg in args:
|
||
if isinstance(arg, torch.Tensor) and arg.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
|
||
converted_args.append(arg.to(torch.bfloat16))
|
||
else:
|
||
converted_args.append(arg)
|
||
|
||
converted_kwargs = {}
|
||
for key, value in kwargs.items():
|
||
if isinstance(value, torch.Tensor) and value.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
|
||
converted_kwargs[key] = value.to(torch.bfloat16)
|
||
else:
|
||
converted_kwargs[key] = value
|
||
|
||
args = tuple(converted_args)
|
||
kwargs = converted_kwargs
|
||
elif self.is_fp16_model:
|
||
# Convert Float32 → FP16 for FP16 models
|
||
converted_args = []
|
||
for arg in args:
|
||
if isinstance(arg, torch.Tensor) and arg.dtype == torch.float32:
|
||
converted_args.append(arg.to(torch.float16))
|
||
else:
|
||
converted_args.append(arg)
|
||
|
||
converted_kwargs = {}
|
||
for key, value in kwargs.items():
|
||
if isinstance(value, torch.Tensor) and value.dtype == torch.float32:
|
||
converted_kwargs[key] = value.to(torch.float16)
|
||
else:
|
||
converted_kwargs[key] = value
|
||
|
||
args = tuple(converted_args)
|
||
kwargs = converted_kwargs
|
||
|
||
try:
|
||
return self.dit_model(*args, **kwargs)
|
||
except Exception as e:
|
||
print(f"❌ Error in forward pass: {e}")
|
||
print(f" Model type: NaDiT 7B={is_nadit_7b}, NaDiT v2 3B={is_nadit_v2_3b}")
|
||
print(f" Args dtypes: {[arg.dtype if isinstance(arg, torch.Tensor) else type(arg) for arg in args]}")
|
||
print(f" Kwargs dtypes: {[(k, v.dtype if isinstance(v, torch.Tensor) else type(v)) for k, v in kwargs.items()]}")
|
||
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', '_forward_count']:
|
||
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', '_forward_count']:
|
||
super().__setattr__(name, value)
|
||
else:
|
||
if hasattr(self, 'dit_model'):
|
||
setattr(self.dit_model, name, value)
|
||
else:
|
||
super().__setattr__(name, value)
|
||
|
||
|
||
def apply_fp8_compatibility_hooks(model: torch.nn.Module) -> List[Tuple[str, Any]]:
|
||
"""
|
||
Hook system to intercept problematic FP8 modules
|
||
Alternative if the wrapper is not sufficient.
|
||
|
||
Args:
|
||
model: Model to apply hooks to
|
||
|
||
Returns:
|
||
List of (module_name, hook) tuples for cleanup
|
||
"""
|
||
def create_fp8_safe_hook(original_dtype: torch.dtype):
|
||
def hook_fn(module, input, output):
|
||
# Convert FP8 output → BFloat16 if necessary for compatibility
|
||
if isinstance(output, torch.Tensor) and output.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
|
||
# Temporarily keep in BFloat16 to avoid downstream errors
|
||
return output.to(torch.bfloat16)
|
||
elif isinstance(output, (tuple, list)):
|
||
# Handle multiple outputs
|
||
converted_output = []
|
||
for item in output:
|
||
if isinstance(item, torch.Tensor) and item.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
|
||
converted_output.append(item.to(torch.bfloat16))
|
||
else:
|
||
converted_output.append(item)
|
||
return type(output)(converted_output)
|
||
return output
|
||
return hook_fn
|
||
|
||
# Apply hooks to critical modules
|
||
problematic_modules = []
|
||
for name, module in model.named_modules():
|
||
# Identify RoPE and attention modules that cause FP8 problems
|
||
if any(keyword in name.lower() for keyword in ['rope', 'rotary', 'attention', 'mmattn']):
|
||
if hasattr(module, 'register_forward_hook'):
|
||
hook = module.register_forward_hook(create_fp8_safe_hook(torch.float8_e4m3fn))
|
||
problematic_modules.append((name, hook))
|
||
|
||
print(f"🔧 Applied FP8 compatibility hooks to {len(problematic_modules)} modules")
|
||
return problematic_modules
|
||
|
||
|
||
def remove_compatibility_hooks(hooks: List[Tuple[str, Any]]) -> None:
|
||
"""
|
||
Remove previously applied compatibility hooks
|
||
|
||
Args:
|
||
hooks: List of (module_name, hook) tuples from apply_fp8_compatibility_hooks
|
||
"""
|
||
removed_count = 0
|
||
for name, hook in hooks:
|
||
try:
|
||
hook.remove()
|
||
removed_count += 1
|
||
except Exception as e:
|
||
print(f"⚠️ Failed to remove hook from {name}: {e}")
|
||
|
||
print(f"🧹 Removed {removed_count}/{len(hooks)} compatibility hooks")
|
||
|