perf: re-enable torch.compile now that STE uses detach() trick (zero graph breaks)"

chimera_turbo.py

@@ -4,30 +4,27 @@ 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 regional — DISABLED (84 graph breaks from _RoundTernarySTE)
+  P-TURBO-2: torch.compile (now possible — STE uses detach() trick, zero graph breaks)
   P-TURBO-3: Threading optimal + tcmalloc detection
   P-TURBO-4: IPEX bf16/AMX si disponible
-  P-TURBO-5: Cache poids quantifiés inter micro-batch (via BitLinear existing cache)
+  P-TURBO-5: Invalidate BitLinear packed caches after optimizer step
   P-TURBO-6: INT8 ternary forward path (VNNI/AMX dispatch)
-  P-TURBO-7: Arrow mmap dataset
-
-v2 changes:
-  - torch.compile DISABLED by default: _RoundTernarySTE (autograd.Function) causes
-    84+ graph breaks (28 layers × 3 BitLinear each). Net effect is SLOWER than eager.
-    Re-enable only after migrating STE to functional torch (torch.round + custom_vjp).
-  - Fix grad_accum_steps logic: DataLoader already provides eff_batch, don't double-accumulate.
-  - Add profile_bottleneck() for quick diagnosis.
-  - Better bf16 autocast handling: skip autocast if CPU has no AMX/AVX512-BF16.
+
+v3 changes (after quantization.py STE migration):
+  - torch.compile RE-ENABLED: _RoundTernarySTE replaced with detach() trick
+    in quantization.py — zero graph breaks, Inductor can fuse quantize+linear.
+  - Compile uses fullgraph=False as safety net for any remaining breaks
+    in non-BitLinear modules (evolution engine, loop controller, etc.)
+  - grad_accum_steps fix from v2 preserved.
 """
 
 import math
 import os
-import sys
 import warnings
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from typing import Optional, Dict, Any, Tuple, List
+from typing import Optional, Dict, Any, Tuple
 from contextlib import nullcontext
 
 # ═══════════════════════════════════════════════════════════
@@ -38,7 +35,6 @@ def detect_cpu_info() -> Dict[str, Any]:
     """Detect CPU capabilities for optimal configuration."""
     info = {}
 
-    # Physical cores (not hyperthreads)
     try:
         physical = len(os.sched_getaffinity(0))
         import multiprocessing
@@ -50,7 +46,6 @@ def detect_cpu_info() -> Dict[str, Any]:
         info["logical_cores"] = multiprocessing.cpu_count()
         info["physical_cores"] = info["logical_cores"] // 2
 
-    # CPU capability
     try:
         info["capability"] = torch.backends.cpu.get_cpu_capability()
     except Exception:
@@ -62,7 +57,6 @@ def detect_cpu_info() -> Dict[str, Any]:
     info["has_avx512_bf16"] = "avx512_bf16" in cap or info["has_amx"]
     info["has_vnni"] = info["has_avx512"]
 
-    # IPEX available?
     try:
         import intel_extension_for_pytorch
         info["ipex_available"] = True
@@ -70,7 +64,6 @@ def detect_cpu_info() -> Dict[str, Any]:
     except ImportError:
         info["ipex_available"] = False
 
-    # tcmalloc loaded?
     info["tcmalloc"] = "tcmalloc" in os.environ.get("LD_PRELOAD", "")
 
     return info
@@ -79,14 +72,9 @@ def detect_cpu_info() -> Dict[str, Any]:
 def configure_threading(cpu_info: Dict[str, Any], reserve_for_io: int = 1):
     """Set optimal threading for CPU training."""
     n_compute = max(1, cpu_info["physical_cores"] - reserve_for_io)
-
-    # Only set num_threads — interop threads can only be set once before
-    # any tensor ops, and train_hyper.py already sets them at import time.
     torch.set_num_threads(n_compute)
-
     os.environ["OMP_NUM_THREADS"] = str(n_compute)
     os.environ["MKL_NUM_THREADS"] = str(n_compute)
-
     return n_compute
 
 
@@ -105,7 +93,7 @@ def create_optimizer(
     Create optimizer for STE-based ternary training (replaces MeZO).
 
     Based on BitNet b1.58 Reloaded (2407.09527):
-    - lr=1e-3 for <300M params (NOT 1e-2, that's for 3B+)
+    - lr=1e-3 for <300M params
     - weight_decay=0.05
     - AdamW with β=(0.9, 0.95)
     """
@@ -149,16 +137,11 @@ def create_scheduler(optimizer, max_steps: int, warmup_steps: int = 500):
 
 
 # ═══════════════════════════════════════════════════════════
-# P-TURBO-5 : Invalidate BitLinear packed caches after optimizer step
+# P-TURBO-5 : Invalidate BitLinear packed caches
 # ═══════════════════════════════════════════════════════════
 
 def invalidate_all_caches(model: nn.Module):
-    """Call after optimizer.step() to force BitLinear re-quantization.
-
-    In training mode, BitLinear._forward_train() recomputes quantized
-    weights every call via STE, so the packed cache is not used.
-    This is still good practice for eval steps between training.
-    """
+    """Call after optimizer.step() to force BitLinear re-quantization."""
     from chimera.quantization import BitLinear
     for m in model.modules():
         if isinstance(m, BitLinear):
@@ -194,41 +177,43 @@ def try_ipex_optimize(
             print("[TURBO-4] IPEX fp32 (no bf16 hardware support detected)")
 
     model, optimizer = ipex.optimize(
-        model,
-        optimizer=optimizer,
-        dtype=dtype,
-        level="O1",
-        inplace=True,
+        model, optimizer=optimizer, dtype=dtype, level="O1", inplace=True,
     )
-
     return model, optimizer
 
 
 # ═══════════════════════════════════════════════════════════
-# P-TURBO-2 : torch.compile — DISABLED by default
+# P-TURBO-2 : torch.compile
 # ═══════════════════════════════════════════════════════════
 
 def try_compile_model(model: nn.Module, mode: str = "reduce-overhead") -> nn.Module:
     """
-    Attempt torch.compile with graceful fallback.
-
-    CURRENTLY DISABLED: _RoundTernarySTE (torch.autograd.Function) causes
-    84+ graph breaks across 28 layers × 3 BitLinear. This makes torch.compile
-    slower than eager mode due to recompilation overhead.
+    Compile model with torch.compile for kernel fusion.
 
-    To re-enable: migrate STE to use torch library custom ops:
-        @torch.library.custom_op("chimera::ste_ternary", mutates_args=())
-        def ste_ternary(w: torch.Tensor) -> torch.Tensor:
-            return torch.round(torch.clamp(w, -1.0, 1.0))
+    Now possible because STE uses the detach() trick (zero graph breaks
+    in BitLinear). Uses fullgraph=False as safety net for any remaining
+    breaks in non-BitLinear modules (evolution engine, grammar FST, etc.)
 
-        @ste_ternary.register_fake
-        def _(w): return torch.empty_like(w)
-
-        @torch.library.register_autograd("chimera::ste_ternary", ...)
+    Expected speedup: 1.5-3x from fusing quantize + linear + activation.
+    First call is slow (compilation); subsequent calls are fast.
     """
-    print("[TURBO-2] torch.compile SKIPPED (84 graph breaks from STE autograd.Function)")
-    print("          To enable: migrate _RoundTernarySTE to torch.library.custom_op")
-    return model
+    if not hasattr(torch, "compile"):
+        warnings.warn("torch.compile not available (PyTorch < 2.0)")
+        return model
+
+    try:
+        compiled = torch.compile(
+            model,
+            backend="inductor",
+            mode=mode,
+            fullgraph=False,  # safety net for non-BitLinear graph breaks
+        )
+        print(f"[TURBO-2] torch.compile enabled (backend=inductor, mode={mode})")
+        print(f"          First few steps will be slow (compilation). Then 1.5-3x speedup.")
+        return compiled
+    except Exception as e:
+        warnings.warn(f"torch.compile failed: {e}. Running in eager mode.")
+        return model
 
 
 # ═══════════════════════════════════════════════════════════
@@ -259,7 +244,7 @@ def ternary_matmul_int8(
 
 
 # ═══════════════════════════════════════════════════════════
-# MAIN: apply() — Point d'entrée unique
+# MAIN: apply()
 # ═══════════════════════════════════════════════════════════
 
 def apply(
@@ -268,7 +253,7 @@ def apply(
     lr: float = 1e-3,
     weight_decay: float = 0.05,
     warmup_steps: int = 500,
-    use_compile: bool = False,  # ← DISABLED by default (was True)
+    use_compile: bool = True,   # ← RE-ENABLED (STE detach trick = zero graph breaks)
     use_ipex: bool = True,
     use_lion: bool = False,
     verbose: bool = True,
@@ -282,7 +267,7 @@ def apply(
 
     if verbose:
         print("=" * 65)
-        print("CHIMERA TURBO v2 — CPU Acceleration Layer")
+        print("CHIMERA TURBO v3 — CPU Acceleration Layer")
         print("=" * 65)
         print(f"  Physical cores: {cpu_info['physical_cores']}")
         print(f"  CPU capability: {cpu_info['capability']}")
@@ -312,14 +297,13 @@ def apply(
     if use_compile:
         model = try_compile_model(model)
 
-    # ── Autocast recommendation ──
+    # ── Warnings ──
     if verbose:
         if not cpu_info["has_avx512_bf16"]:
             print()
             print("  ⚠️  No hardware BF16 support detected (need AVX512-BF16 or AMX).")
             print("     BF16 autocast may be SLOWER than fp32 on this CPU.")
             print("     Consider --no-bf16 flag if training is slow.")
-
         if not cpu_info["tcmalloc"]:
             print()
             print("  ⚠️  tcmalloc not detected. For +10-25% speedup:")
@@ -347,11 +331,8 @@ def training_step(
     """
     Single training step with all turbo optimizations active.
 
-    Handles: autocast, gradient accumulation, clipping, cache invalidation.
-
     IMPORTANT: grad_accum_steps should be 1 if the DataLoader already provides
-    the full effective batch. Set >1 only if you want to split a large batch
-    across multiple forward passes.
+    the full effective batch. Set >1 only for memory-constrained scenarios.
     """
     is_accum_step = (step + 1) % grad_accum_steps == 0
 
@@ -412,7 +393,11 @@ def profile_model(model: nn.Module, dummy_input: torch.Tensor, steps: int = 5):
 
 
 def count_compile_graph_breaks(model: nn.Module, dummy_input: torch.Tensor):
-    """Count how many graph breaks torch.compile would produce."""
+    """Count how many graph breaks torch.compile would produce.
+
+    After the STE detach() migration, BitLinear should produce ZERO breaks.
+    Remaining breaks come from non-BitLinear modules (evolution, grammar, etc.)
+    """
     try:
         import torch._dynamo as dynamo
         explanation = dynamo.explain(model)(dummy_input)
@@ -422,6 +407,8 @@ def count_compile_graph_breaks(model: nn.Module, dummy_input: torch.Tensor):
             print(f"  [{i+1}] {reason}")
         if n_breaks > 10:
             print(f"  ... and {n_breaks - 10} more")
+        if n_breaks == 0:
+            print("  ✅ Zero graph breaks — full model is compilable!")
         return n_breaks
     except Exception as e:
         print(f"[TURBO-DIAG] dynamo.explain failed: {e}")