From d87fca3ee25e1e75e4d2c58f3b32fba2fb4c962d Mon Sep 17 00:00:00 2001 From: Adrien Toupet Date: Wed, 23 Jul 2025 17:10:59 -0400 Subject: [PATCH] Tentative fix for artifacts without blockswap --- src/optimization/compatibility.py | 41 +++++++++++++------------------ 1 file changed, 17 insertions(+), 24 deletions(-) diff --git a/src/optimization/compatibility.py b/src/optimization/compatibility.py index 5d3a246..68eba77 100644 --- a/src/optimization/compatibility.py +++ b/src/optimization/compatibility.py @@ -52,7 +52,7 @@ class FP8CompatibleDiT(torch.nn.Module): return torch.bfloat16 def _wrap_fp16_blocks_in_quantized_model(self): - """Apply minimal RoPE fixes for mixed precision models""" + """Clear RoPE caches for mixed precision models to prevent numerical instability""" if not hasattr(self.dit_model, 'blocks'): return @@ -70,33 +70,26 @@ class FP8CompatibleDiT(torch.nn.Module): if not has_fp16_blocks: return - print(f"🎯 Mixed Precision Model detected - applying stability fixes") + print(f"🎯 Mixed Precision Model detected - clearing RoPE caches") - # Apply minimal RoPE patching - only for mixed models - import types - rope_count = 0 + # Clear all RoPE LRU caches to prevent stale cached values + # This is the actual fix - cached RoPE freqs become invalid with mixed precision + cleared_count = 0 + for name, module in self.dit_model.named_modules(): + if "rope" in name.lower() and hasattr(module, "get_axial_freqs"): + if hasattr(module.get_axial_freqs, 'cache_clear'): + module.get_axial_freqs.cache_clear() + cleared_count += 1 + + # Also check for RotaryEmbedding modules (different naming convention) for module in self.dit_model.modules(): - if hasattr(module, 'get_axial_freqs') and 'rope' in module.__class__.__name__.lower(): - original = module.get_axial_freqs - - # Minimal wrapper - only catch the specific error that occurs - def stable_rope(self, *args, **kwargs): - try: - return original(*args, **kwargs) - except RuntimeError as e: - if "device" in str(e) or "memory" in str(e): - # Only fallback for memory/device errors - self.cpu() - result = original(*args, **kwargs) - self.cuda() - return result.cuda() if hasattr(result, 'cuda') else result - raise # Re-raise other errors - - module.get_axial_freqs = types.MethodType(stable_rope, module) - rope_count += 1 + if 'RotaryEmbedding' in type(module).__name__: + if hasattr(module, 'get_axial_freqs') and hasattr(module.get_axial_freqs, 'cache_clear'): + module.get_axial_freqs.cache_clear() + cleared_count += 1 - print(f" ✅ Stabilized {rope_count} RoPE modules for mixed precision") + print(f" ✅ Cleared {cleared_count} RoPE caches for mixed precision stability") def _is_nadit_model(self) -> bool: """Detect if this is a NaDiT model (7B) with precise logic"""