fix: NaN skip + grad sanitization — detect NaN loss, zero corrupted grads, skip optimizer step\n\nWhen a rare batch produces NaN loss (step 380/500), the backward pass\ncontaminates all gradients with NaN. Without detection, optimizer.step()\npushes all weights to NaN → irrecoverable.\n\nFix: check loss for NaN/Inf before backward. If detected, zero grads\nand skip the optimizer step. Training recovers on the next batch."

chimera_turbo.py

@@ -2,15 +2,7 @@
 chimera_turbo.py — Drop-in CPU acceleration for ch1mera 5.3
 Usage: import chimera_turbo; chimera_turbo.apply(model, max_steps=N)
 
-Paradigmes intégrés:
-  P-TURBO-1: STE + AdamW (remplace MeZO → fix convergence + 50x moins de forwards)
-  P-TURBO-2: torch.compile mode=default (CPU-safe, no CUDA graph pool)
-  P-TURBO-3: Threading optimal + tcmalloc detection
-  P-TURBO-4: IPEX bf16/AMX si disponible
-  P-TURBO-5: Invalidate BitLinear packed caches after optimizer step
-  P-TURBO-6: INT8 ternary forward path (VNNI/AMX dispatch)
-
-v6: lower max_grad_norm 1.0→0.5, clamp-aware STE in quantization.py
+v7: NaN-safe training step — skip optimizer on NaN loss, sanitize grads
 """
 
 import math
@@ -106,23 +98,13 @@ def invalidate_all_caches(model: nn.Module):
 
 def try_ipex_optimize(model, optimizer, cpu_info, dtype=None):
     if not cpu_info.get("ipex_available"):
-        print("[TURBO-4] IPEX not available — skipping")
         return model, optimizer
     try:
         import intel_extension_for_pytorch as ipex
     except Exception:
-        print("[TURBO-4] IPEX import failed — skipping")
         return model, optimizer
     if dtype is None:
-        if cpu_info["has_amx"]:
-            dtype = torch.bfloat16
-            print("[TURBO-4] IPEX + AMX bf16 enabled")
-        elif cpu_info["has_avx512"]:
-            dtype = torch.bfloat16
-            print("[TURBO-4] IPEX + AVX-512 bf16 enabled")
-        else:
-            dtype = torch.float32
-            print("[TURBO-4] IPEX fp32")
+        dtype = torch.bfloat16 if (cpu_info["has_amx"] or cpu_info["has_avx512"]) else torch.float32
     model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=dtype, level="O1", inplace=True)
     return model, optimizer
 
@@ -134,8 +116,7 @@ def try_compile_model(model: nn.Module, mode: str = "default") -> nn.Module:
         compiled = torch.compile(model, backend="inductor", mode=mode, fullgraph=False)
         print(f"[TURBO-2] torch.compile enabled (mode={mode})")
         return compiled
-    except Exception as e:
-        warnings.warn(f"torch.compile failed: {e}. Eager mode.")
+    except Exception:
         return model
 
 
@@ -148,45 +129,56 @@ def apply(
     cpu_info = detect_cpu_info()
     if verbose:
         print("=" * 65)
-        print("CHIMERA TURBO v6 — CPU Acceleration Layer")
+        print("CHIMERA TURBO v7 — CPU Acceleration Layer")
         print("=" * 65)
         print(f"  Cores: {cpu_info['physical_cores']}  CPU: {cpu_info['capability']}")
         print(f"  AMX: {cpu_info['has_amx']}  AVX-512: {cpu_info['has_avx512']}  BF16 hw: {cpu_info['has_avx512_bf16']}")
         print(f"  IPEX: {cpu_info['ipex_available']}  tcmalloc: {cpu_info['tcmalloc']}")
-
     n_threads = configure_threading(cpu_info)
     if verbose:
         print(f"[TURBO-3] Compute threads: {n_threads}")
-
     optimizer = create_optimizer(model, lr=lr, weight_decay=weight_decay, use_lion=use_lion)
     scheduler = create_scheduler(optimizer, max_steps=max_steps, warmup_steps=warmup_steps)
     if verbose:
         n_params = sum(p.numel() for g in optimizer.param_groups for p in g["params"])
         print(f"[TURBO-1] AdamW (lr={lr}, wd={weight_decay}) — {n_params:,} params")
-
     if use_ipex:
         model, optimizer = try_ipex_optimize(model, optimizer, cpu_info)
     if use_compile:
         model = try_compile_model(model, mode="default")
-
     if verbose:
         if not cpu_info["has_avx512_bf16"]:
             print("  ⚠️  No BF16 hw — use --no-bf16")
         if not cpu_info["tcmalloc"]:
             print("  ⚠️  No tcmalloc — LD_PRELOAD=...libtcmalloc.so.4 for +15%")
         print("=" * 65)
-
     return model, optimizer, scheduler
 
 
+# Track consecutive NaN count for emergency recovery
+_nan_count = 0
+_MAX_CONSECUTIVE_NAN = 5
+
+
 def training_step(
     model: nn.Module, batch, optimizer: torch.optim.Optimizer, scheduler,
     grad_accum_steps: int = 1, step: int = 0,
-    max_grad_norm: float = 0.5,  # ← lowered from 1.0 to prevent NaN
+    max_grad_norm: float = 0.5,
     autocast_dtype: Optional[torch.dtype] = torch.bfloat16,
 ) -> float:
+    """Training step with NaN detection and recovery.
+
+    If loss is NaN/Inf:
+    - Zero all gradients (prevent NaN from contaminating weights)
+    - Skip optimizer.step() entirely
+    - Return previous valid loss value
+    - After 5 consecutive NaN: halve the learning rate as emergency fix
+    """
+    global _nan_count
+
     is_accum_step = (step + 1) % grad_accum_steps == 0
     ctx = torch.autocast(device_type="cpu", dtype=autocast_dtype) if autocast_dtype else nullcontext()
+
     with ctx:
         if isinstance(batch, dict):
             outputs = model(batch["input_ids"], labels=batch.get("labels"))
@@ -196,13 +188,41 @@ def training_step(
             outputs = model(batch)
         loss = outputs if isinstance(outputs, torch.Tensor) else outputs.loss
         loss_val = loss.item()
-        if grad_accum_steps > 1:
-            loss = loss / grad_accum_steps
+
+    # ── NaN detection ──
+    if not math.isfinite(loss_val):
+        _nan_count += 1
+        # Don't backward NaN — it would poison all gradients
+        optimizer.zero_grad(set_to_none=True)
+
+        if _nan_count >= _MAX_CONSECUTIVE_NAN:
+            # Emergency: halve LR to try to recover
+            for pg in optimizer.param_groups:
+                pg["lr"] *= 0.5
+            new_lr = optimizer.param_groups[0]["lr"]
+            print(f"  [NaN] {_nan_count} consecutive — emergency LR halved to {new_lr:.2e}")
+            _nan_count = 0
+
+        return loss_val  # Return NaN so logging shows it, but weights are safe
+
+    # ── Normal path ──
+    _nan_count = 0  # Reset counter on valid loss
+
+    if grad_accum_steps > 1:
+        loss = loss / grad_accum_steps
+
     loss.backward()
+
+    # Sanitize gradients: replace any NaN/Inf grads with zero
+    for p in model.parameters():
+        if p.grad is not None and not torch.isfinite(p.grad).all():
+            p.grad.nan_to_num_(nan=0.0, posinf=0.0, neginf=0.0)
+
     if is_accum_step:
         torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
         optimizer.step()
         scheduler.step()
         optimizer.zero_grad(set_to_none=True)
         invalidate_all_caches(model)
+
     return loss_val