Upload 6 files

nord_v4_700m-4.2/chat_v4.py

@@ -0,0 +1,625 @@
+"""
+╔══════════════════════════════════════════════════════════════════════════╗
+║         PROJECT NORD v4 — Interactive Chat v4.0                         ║
+║                                                                        ║
+║  Commands:                                                             ║
+║      /stdp on|off   — Toggle online learning                           ║
+║      /stats         — Show zone & MoE statistics                        ║
+║      /memory        — Show memory cortex state                          ║
+║      /reset         — Clear working memory                              ║
+║      /expert        — Show expert routing breakdown                     ║
+║      /tokens N      — Set max response tokens (default: 200)            ║
+║      /temp F        — Set temperature (default: 0.85)                   ║
+║      /rep F         — Set repetition penalty (default: 1.3)             ║
+║      /live on|off   — Toggle live spike visualization                   ║
+║      /quit          — Exit                                              ║
+╚══════════════════════════════════════════════════════════════════════════╝
+"""
+
+from __future__ import annotations
+
+import sys
+import time
+import torch
+import os
+from pathlib import Path
+
+from nord_core_700m import NordConfig, NordModel
+
+# ── ANSI Colors ──
+class C:
+    RESET   = "\033[0m"
+    BOLD    = "\033[1m"
+    DIM     = "\033[2m"
+    CYAN    = "\033[96m"
+    ORANGE  = "\033[38;5;208m"
+    PURPLE  = "\033[35m"
+    GREEN   = "\033[92m"
+    BLUE    = "\033[94m"
+    YELLOW  = "\033[93m"
+    RED     = "\033[91m"
+    WHITE   = "\033[97m"
+    GREY    = "\033[90m"
+
+SPARK_CHARS = " ░▒▓█"
+
+def spike_bar(rate, width=20, color=C.CYAN, max_rate=0.4):
+    """Colored bar with adjustable scale. max_rate=0.4 means 40% rate fills full bar"""
+    normalized = min(rate / max(max_rate, 0.001), 1.0)
+    filled = int(normalized * width)
+    bar = ""
+    for i in range(width):
+        if i < filled:
+            frac = normalized * width - i
+            intensity = min(4, int(frac * 4))
+            bar += color + SPARK_CHARS[min(intensity + 1, 4)]
+        else:
+            bar += C.DIM + "·"
+    return bar + C.RESET
+
+
+def render_live_spikes(stats, cfg):
+    spike_rates = stats.get("spike_rates", [])
+    if not spike_rates:
+        return
+
+    lines = []
+    lines.append(f"  {C.GREY}{'─' * 56}{C.RESET}")
+
+    ns = cfg.sensory_layers + 1
+    if len(spike_rates) > 0:
+        avg_s = sum(spike_rates[:ns]) / max(ns, 1)
+        bar = spike_bar(avg_s, 20, C.CYAN)
+        lines.append(f"  {C.CYAN}⚡ SEN{C.RESET} {bar} {C.CYAN}{avg_s*100:5.1f}%{C.RESET}")
+
+    na = cfg.association_layers
+    if len(spike_rates) > ns:
+        assoc_rates = spike_rates[ns:ns+na]
+        avg_a = sum(assoc_rates) / max(len(assoc_rates), 1) if assoc_rates else 0
+        bar = spike_bar(avg_a, 20, C.ORANGE)
+        lines.append(f"  {C.ORANGE}⚡ ASC{C.RESET} {bar} {C.ORANGE}{avg_a*100:5.1f}%{C.RESET}")
+
+    mem_rate = stats.get("memory_spike_rate", 0)
+    if isinstance(mem_rate, torch.Tensor):
+        mem_rate = mem_rate.item()
+    bar = spike_bar(mem_rate * 0.3, 20, C.PURPLE)
+    lines.append(f"  {C.PURPLE}⚡ MEM{C.RESET} {bar} {C.PURPLE}{mem_rate*100:5.1f}%{C.RESET}")
+
+    ne = cfg.executive_layers
+    offset = ns + na
+    if len(spike_rates) > offset:
+        exec_rates = spike_rates[offset:]
+        avg_e = sum(exec_rates) / max(len(exec_rates), 1) if exec_rates else 0
+        bar = spike_bar(avg_e, 20, C.GREEN)
+        lines.append(f"  {C.GREEN}⚡ EXE{C.RESET} {bar} {C.GREEN}{avg_e*100:5.1f}%{C.RESET}")
+
+    sp = stats.get("sparsity", 0)
+    if isinstance(sp, torch.Tensor):
+        sp = sp.item()
+    sp_color = C.GREEN if sp > 0.85 else C.YELLOW if sp > 0.7 else C.RED
+    lines.append(f"  {C.GREY}  SPR{C.RESET} {sp_color}{sp*100:.0f}%{C.RESET} {C.DIM}neurons silent{C.RESET}")
+    lines.append(f"  {C.GREY}{'─' * 56}{C.RESET}")
+
+    output = "\n".join(lines)
+    n_lines = len(lines)
+    sys.stdout.write(f"\033[{n_lines}A")
+    sys.stdout.write(output + "\n")
+    sys.stdout.flush()
+
+
+def init_live_display(cfg):
+    for _ in range(7):
+        print()
+
+
+def render_spike_panel(stats, cfg):
+    """Render a clean spike panel BELOW the generated text"""
+    spike_rates = stats.get("spike_rates", [])
+    if not spike_rates:
+        return
+
+    ns = cfg.sensory_layers + 1
+    na = cfg.association_layers
+
+    print(f"  {C.GREY}┌{'─' * 54}┐{C.RESET}")
+    print(f"  {C.GREY}│{C.RESET} {C.BOLD}Neural Activity{C.RESET}{' ' * 38}{C.GREY}│{C.RESET}")
+    print(f"  {C.GREY}├{'─' * 54}┤{C.RESET}")
+
+    # Sensory
+    if len(spike_rates) > 0:
+        avg_s = sum(spike_rates[:ns]) / max(ns, 1)
+        bar = spike_bar(avg_s, 25, C.CYAN)
+        print(f"  {C.GREY}│{C.RESET} {C.CYAN}⚡ Sensory    {C.RESET} {bar} {C.CYAN}{avg_s*100:5.1f}%{C.RESET} {C.GREY}│{C.RESET}")
+
+    # Association
+    if len(spike_rates) > ns:
+        assoc_rates = spike_rates[ns:ns+na]
+        avg_a = sum(assoc_rates) / max(len(assoc_rates), 1) if assoc_rates else 0
+        bar = spike_bar(avg_a, 25, C.ORANGE)
+        print(f"  {C.GREY}│{C.RESET} {C.ORANGE}⚡ Association{C.RESET} {bar} {C.ORANGE}{avg_a*100:5.1f}%{C.RESET} {C.GREY}│{C.RESET}")
+
+    # Memory
+    mem_rate = stats.get("memory_spike_rate", 0)
+    if isinstance(mem_rate, torch.Tensor):
+        mem_rate = mem_rate.item()
+    bar = spike_bar(min(mem_rate, 1.0), 25, C.PURPLE)
+    print(f"  {C.GREY}│{C.RESET} {C.PURPLE}⚡ Memory     {C.RESET} {bar} {C.PURPLE}{mem_rate*100:5.1f}%{C.RESET} {C.GREY}│{C.RESET}")
+
+    # Executive
+    offset = ns + na
+    if len(spike_rates) > offset:
+        exec_rates = spike_rates[offset:]
+        avg_e = sum(exec_rates) / max(len(exec_rates), 1) if exec_rates else 0
+        bar = spike_bar(avg_e, 25, C.GREEN)
+        print(f"  {C.GREY}│{C.RESET} {C.GREEN}⚡ Executive  {C.RESET} {bar} {C.GREEN}{avg_e*100:5.1f}%{C.RESET} {C.GREY}│{C.RESET}")
+
+    # Sparsity
+    sp = stats.get("sparsity", 0)
+    if isinstance(sp, torch.Tensor):
+        sp = sp.item()
+    sp_color = C.GREEN if sp > 0.85 else C.YELLOW if sp > 0.7 else C.RED
+    silent = int(sp * 100)
+    active = 100 - silent
+    print(f"  {C.GREY}├{'─' * 54}┤{C.RESET}")
+    print(f"  {C.GREY}│{C.RESET} {C.DIM}Sparsity:{C.RESET} {sp_color}{sp*100:.0f}%{C.RESET} silent  {C.DIM}({active}% neurons active per token){C.RESET}  {C.GREY}│{C.RESET}")
+    print(f"  {C.GREY}└{'─' * 54}┘{C.RESET}")
+
+
+def load_model(model_dir: str):
+    from transformers import AutoTokenizer
+
+    model_dir = Path(model_dir)
+
+    # ── Smart checkpoint search ──
+    # 1. If user gave a direct .pt file path
+    if model_dir.is_file() and model_dir.suffix == ".pt":
+        latest = model_dir
+    else:
+        latest = None
+
+        # 2. Search in the given directory
+        search_dirs = [model_dir]
+
+        # 3. Also search in current working directory (where the script is run from)
+        cwd = Path.cwd()
+        if cwd != model_dir:
+            search_dirs.append(cwd)
+
+        # 4. Also search in the script's own directory
+        script_dir = Path(__file__).resolve().parent
+        if script_dir != cwd and script_dir != model_dir:
+            search_dirs.append(script_dir)
+
+        # Search order: nord_v4_latest.pt, nord_v4_final.pt, step checkpoints, legacy names
+        checkpoint_names = [
+            "nord_v4_latest.pt",
+            "nord_v4_final.pt",
+            "nord_500m_latest.pt",
+            "nord_latest.pt",
+        ]
+
+        for search_dir in search_dirs:
+            if not search_dir.exists():
+                continue
+
+            # Try known names
+            for name in checkpoint_names:
+                p = search_dir / name
+                if p.exists():
+                    latest = p
+                    break
+
+            # Try step checkpoints
+            if latest is None:
+                ckpts = sorted(search_dir.glob("nord_v4_step_*.pt"))
+                if ckpts:
+                    latest = ckpts[-1]
+
+            # Try any .pt file
+            if latest is None:
+                all_pt = sorted(search_dir.glob("*.pt"))
+                if all_pt:
+                    latest = all_pt[-1]
+
+            if latest is not None:
+                break
+
+    if latest is None:
+        print(f"  {C.RED}[✗] No checkpoint found!{C.RESET}")
+        print(f"  {C.DIM}Searched in:{C.RESET}")
+        for d in search_dirs:
+            exists = "✓" if d.exists() else "✗"
+            print(f"    [{exists}] {d}")
+        print(f"\n  {C.DIM}Place your .pt file in the same folder as chat.py{C.RESET}")
+        print(f"  {C.DIM}Or give the full path: /path/to/nord_v4_latest.pt{C.RESET}")
+        sys.exit(1)
+
+    print(f"  [*] Loading: {latest.name}")
+    ckpt = torch.load(latest, map_location="cpu", weights_only=False)
+
+    saved_cfg = ckpt.get("config", {})
+    cfg = NordConfig(
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        dtype=torch.float16,
+    )
+    for k, v in saved_cfg.items():
+        if hasattr(cfg, k):
+            setattr(cfg, k, v)
+
+    tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer_id, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    if cfg.vocab_size < tokenizer.vocab_size:
+        cfg.vocab_size = tokenizer.vocab_size
+
+    model = NordModel(cfg)
+    state = ckpt["model_state_dict"]
+    filtered = {k: v for k, v in state.items()
+                if "_v_mem_state" not in k and "_i_syn_state" not in k}
+    model.load_state_dict(filtered, strict=False)
+    model = model.to(cfg.device)
+    model.eval()
+
+    total = sum(p.numel() for p in model.parameters())
+    print(f"  {C.GREEN}[✓]{C.RESET} Nord v4 loaded ({total/1e6:.1f}M params)")
+    print(f"  {C.GREEN}[✓]{C.RESET} {model.count_params()}")
+
+    return model, tokenizer, cfg
+
+
+@torch.no_grad()
+def generate_streaming(model, tokenizer, cfg, prompt: str,
+                       max_tokens: int = 200, temperature: float = 0.85,
+                       top_p: float = 0.9, repetition_penalty: float = 1.3,
+                       enable_stdp: bool = False, live_spikes: bool = False):
+
+    input_ids = tokenizer(
+        prompt, return_tensors="pt",
+        max_length=cfg.max_seq_len, truncation=True,
+    ).input_ids.to(cfg.device)
+
+    model.reset_state()
+
+    generated = input_ids.clone()
+    all_stats = {}
+    token_count = 0
+
+    t_start = time.time()
+
+    sys.stdout.write(f"  {C.BOLD}Nord:{C.RESET} ")
+    sys.stdout.flush()
+
+    for i in range(max_tokens):
+        context = generated[:, -cfg.max_seq_len:]
+
+        if torch.cuda.is_available():
+            with torch.amp.autocast(device_type="cuda", dtype=torch.float16,
+                                    enabled=(cfg.dtype == torch.float16)):
+                logits, stats = model(context, enable_stdp=enable_stdp)
+        else:
+            logits, stats = model(context, enable_stdp=enable_stdp)
+
+        next_logits = logits[:, -1, :].float()
+
+        if repetition_penalty != 1.0:
+            for token_id in generated[0].unique():
+                next_logits[0, token_id] /= repetition_penalty
+
+        next_logits = next_logits / max(temperature, 0.01)
+
+        probs = torch.softmax(next_logits, dim=-1)
+        sorted_probs, sorted_idx = torch.sort(probs, descending=True)
+        cumsum = sorted_probs.cumsum(dim=-1)
+        mask = cumsum - sorted_probs > top_p
+        sorted_probs[mask] = 0
+        sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)
+
+        token = sorted_idx[0, torch.multinomial(sorted_probs[0], 1)]
+        generated = torch.cat([generated, token.reshape(1, 1)], dim=1)
+        token_count += 1
+
+        if token.item() == tokenizer.eos_token_id:
+            break
+
+        # ── Stream token ──
+        decoded_token = tokenizer.decode([token.item()], skip_special_tokens=True)
+        sys.stdout.write(decoded_token)
+        sys.stdout.flush()
+
+        all_stats = stats
+
+    elapsed = time.time() - t_start
+    tps = token_count / elapsed if elapsed > 0 else 0
+
+    rep_score = 1.0
+    if token_count > 5:
+        out_ids = generated[0][input_ids.shape[1]:].tolist()
+        unique = len(set(out_ids))
+        rep_score = len(out_ids) / max(unique, 1)
+
+    sp = all_stats.get("sparsity", 0)
+    if isinstance(sp, torch.Tensor):
+        sp = sp.item()
+
+    print(f"\n  {C.GREY}[{token_count} tok, {elapsed:.1f}s, {tps:.1f} tok/s "
+          f"[REP {rep_score:.1f}] [SPR {sp:.0%}]]{C.RESET}")
+
+    if live_spikes and all_stats:
+        render_spike_panel(all_stats, cfg)
+
+    return all_stats
+
+
+def print_stats(stats: dict, cfg: NordConfig):
+    print(f"\n  {C.GREY}{'─' * 50}{C.RESET}")
+    print(f"  {C.BOLD}Zone Statistics:{C.RESET}")
+
+    spike_rates = stats.get("spike_rates", [])
+    if spike_rates:
+        print(f"    {C.DIM}Encoder:     {spike_rates[0]:.4f}{C.RESET}")
+        for i in range(min(cfg.sensory_layers, len(spike_rates)-1)):
+            rate = spike_rates[i+1]
+            bar = spike_bar(rate, 15, C.CYAN)
+            print(f"    {C.CYAN}Sensory[{i}]:{C.RESET}  {rate:.4f} {bar}")
+        offset = cfg.sensory_layers + 1
+        for i in range(cfg.association_layers):
+            if offset + i < len(spike_rates):
+                rate = spike_rates[offset+i]
+                bar = spike_bar(rate, 15, C.ORANGE)
+                print(f"    {C.ORANGE}Assoc[{i}]:{C.RESET}    {rate:.4f} {bar} {C.DIM}(MoE){C.RESET}")
+        offset += cfg.association_layers
+        for i in range(cfg.executive_layers):
+            if offset + i < len(spike_rates):
+                rate = spike_rates[offset+i]
+                bar = spike_bar(rate, 15, C.GREEN)
+                print(f"    {C.GREEN}Exec[{i}]:{C.RESET}     {rate:.4f} {bar}")
+
+    print(f"\n  {C.BOLD}MoE Routing:{C.RESET}")
+    expert_loads = stats.get("expert_loads", None)
+    moe_entropy = stats.get("moe_route_entropy", None)
+
+    # Also check for entropy with assoc_ prefix
+    if moe_entropy is None:
+        for key in stats:
+            if "route_entropy" in key:
+                moe_entropy = stats[key]
+                break
+
+    if expert_loads is not None:
+        if isinstance(expert_loads, torch.Tensor):
+            expert_loads = expert_loads.detach().cpu().tolist()
+            if isinstance(expert_loads, float):
+                expert_loads = [expert_loads]
+        for e, load in enumerate(expert_loads):
+            pct = load if isinstance(load, float) else float(load)
+            bar = spike_bar(pct, 30, C.YELLOW, max_rate=0.5)
+            print(f"    Expert {e}: {pct:.2%} {bar}")
+    else:
+        found = False
+        # Search with ALL possible key patterns including assoc_ prefix
+        for e in range(cfg.n_experts):
+            load = None
+            for key_pattern in [
+                f"expert_{e}_load",
+                f"expert_load_{e}",
+                f"moe_expert_{e}",
+            ]:
+                # Direct match
+                if key_pattern in stats:
+                    load = stats[key_pattern]
+                    break
+                # Prefixed match (assoc_0_expert_0_load, etc.)
+                for k, v in stats.items():
+                    if key_pattern in k:
+                        load = v
+                        break
+                if load is not None:
+                    break
+
+            if load is not None:
+                found = True
+                if isinstance(load, torch.Tensor): load = load.item()
+                bar = spike_bar(load, 30, C.YELLOW, max_rate=0.5)
+                print(f"    Expert {e}: {load:.2%} {bar}")
+
+        if not found:
+            # Last resort: scan all stats keys for anything with "expert" and "load"
+            expert_data = {k: v for k, v in stats.items() if "expert" in k and "load" in k}
+            if expert_data:
+                found = True
+                for k, v in sorted(expert_data.items()):
+                    if isinstance(v, torch.Tensor): v = v.item()
+                    bar = spike_bar(v, 30, C.YELLOW, max_rate=0.5)
+                    name = k.split("_expert_")[-1] if "_expert_" in k else k
+                    print(f"    {name}: {v:.2%} {bar}")
+
+            if not found:
+                moe_lb = stats.get("moe_lb_loss", None)
+                if moe_lb is None:
+                    for k, v in stats.items():
+                        if "load_balance" in k or "moe_lb" in k:
+                            moe_lb = v
+                            break
+                if moe_lb is not None:
+                    if isinstance(moe_lb, torch.Tensor): moe_lb = moe_lb.item()
+                    print(f"    {C.DIM}Load balance loss: {moe_lb:.4f}{C.RESET}")
+                print(f"    {C.DIM}Per-expert loads not in top-level stats.{C.RESET}")
+                print(f"    {C.DIM}They exist as assoc_N_expert_N_load — fixing...{C.RESET}")
+
+    if moe_entropy is not None:
+        if isinstance(moe_entropy, torch.Tensor): moe_entropy = moe_entropy.item()
+        print(f"    Entropy: {moe_entropy:.3f}")
+
+    mem_rate = stats.get("memory_spike_rate", None)
+    if mem_rate is not None:
+        if isinstance(mem_rate, torch.Tensor): mem_rate = mem_rate.item()
+        gate = stats.get("gate_activity", 0)
+        mix = stats.get("memory_mix", 0)
+        if isinstance(gate, torch.Tensor): gate = gate.item()
+        if isinstance(mix, torch.Tensor): mix = mix.item()
+        bar = spike_bar(mem_rate * 0.3, 15, C.PURPLE)
+        print(f"\n  {C.BOLD}Memory Cortex:{C.RESET}")
+        print(f"    {C.PURPLE}Spike rate:{C.RESET}  {mem_rate:.4f} {bar}")
+        print(f"    {C.PURPLE}Gate:{C.RESET}        {gate:.4f}")
+        print(f"    {C.PURPLE}Mix weight:{C.RESET}  {mix:.4f}")
+
+    sparsity = stats.get("sparsity", 0)
+    if isinstance(sparsity, torch.Tensor): sparsity = sparsity.item()
+    sp_color = C.GREEN if sparsity > 0.85 else C.YELLOW if sparsity > 0.7 else C.RED
+    print(f"\n  Overall Sparsity: {sp_color}{sparsity:.1%}{C.RESET}")
+    print(f"  {C.GREY}{'─' * 50}{C.RESET}")
+
+
+def main():
+    os.system('clear' if os.name != 'nt' else 'cls')
+
+    print(f"""
+  {C.CYAN}╔══════════════════════════════════════════════════════════╗{C.RESET}
+  {C.CYAN}║{C.RESET}  {C.BOLD}⚡ PROJECT NORD v4.2 — Brain-Inspired SNN Chat{C.RESET}         {C.CYAN}║{C.RESET}
+  {C.CYAN}║{C.RESET}  {C.DIM}618M params │ Spike-driven │ Zonal architecture{C.RESET}       {C.CYAN}║{C.RESET}
+  {C.CYAN}╚══════════════════════════════════════════════════════════╝{C.RESET}
+""")
+
+    default_dir = "nord_v4_700m"
+    print(f"  Model directory?")
+    print(f"  {C.DIM}(Enter = {default_dir}){C.RESET}")
+    model_input = input("  Path: ").strip()
+    model_dir = model_input if model_input else default_dir
+
+    model, tokenizer, cfg = load_model(model_dir)
+
+    stdp_enabled = False
+    live_spikes = False
+    max_tokens = 200
+    temperature = 0.85
+    top_p = 0.9
+    rep_penalty = 1.3
+    last_stats = {}
+
+    print(f"\n  {C.DIM}Type /help for commands{C.RESET}")
+    print(f"  {C.GREY}{'─' * 50}{C.RESET}\n")
+
+    while True:
+        try:
+            user = input(f"  {C.BOLD}You:{C.RESET} ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print(f"\n  {C.DIM}Goodbye!{C.RESET}")
+            break
+
+        if not user:
+            continue
+
+        cmd = user.lower().split()
+
+        if cmd[0] == "/quit":
+            break
+        elif cmd[0] == "/help":
+            print(f"""
+  {C.BOLD}Commands:{C.RESET}
+    {C.CYAN}/tokens N{C.RESET}     — Max response tokens (current: {max_tokens})
+    {C.CYAN}/temp F{C.RESET}       — Temperature (current: {temperature})
+    {C.CYAN}/rep F{C.RESET}        — Repetition penalty (current: {rep_penalty})
+    {C.CYAN}/stdp on|off{C.RESET}  — Toggle online learning ({C.GREEN if stdp_enabled else C.RED}{'ON' if stdp_enabled else 'OFF'}{C.RESET})
+    {C.CYAN}/live on|off{C.RESET}  — Live spike visualization ({C.GREEN if live_spikes else C.RED}{'ON' if live_spikes else 'OFF'}{C.RESET})
+    {C.CYAN}/stats{C.RESET}        — Zone & MoE statistics
+    {C.CYAN}/memory{C.RESET}       — Memory cortex state
+    {C.CYAN}/expert{C.RESET}       — Expert routing breakdown
+    {C.CYAN}/reset{C.RESET}        — Clear working memory
+    {C.CYAN}/quit{C.RESET}         — Exit""")
+            continue
+        elif cmd[0] == "/tokens":
+            if len(cmd) > 1:
+                try:
+                    max_tokens = int(cmd[1])
+                    print(f"  {C.GREEN}[✓]{C.RESET} Max tokens: {max_tokens}")
+                except ValueError:
+                    print(f"  {C.RED}[✗]{C.RESET} Usage: /tokens 300")
+            else:
+                print(f"  Max tokens: {max_tokens}")
+            continue
+        elif cmd[0] == "/temp":
+            if len(cmd) > 1:
+                try:
+                    temperature = float(cmd[1])
+                    print(f"  {C.GREEN}[✓]{C.RESET} Temperature: {temperature}")
+                except ValueError:
+                    print(f"  {C.RED}[✗]{C.RESET} Usage: /temp 0.7")
+            else:
+                print(f"  Temperature: {temperature}")
+            continue
+        elif cmd[0] == "/rep":
+            if len(cmd) > 1:
+                try:
+                    rep_penalty = float(cmd[1])
+                    print(f"  {C.GREEN}[✓]{C.RESET} Repetition penalty: {rep_penalty}")
+                except ValueError:
+                    print(f"  {C.RED}[✗]{C.RESET} Usage: /rep 1.3")
+            else:
+                print(f"  Repetition penalty: {rep_penalty}")
+            continue
+        elif cmd[0] == "/stdp":
+            if len(cmd) > 1 and cmd[1] == "on":
+                stdp_enabled = True
+                print(f"  {C.GREEN}[⚙] STDP enabled{C.RESET}")
+            elif len(cmd) > 1 and cmd[1] == "off":
+                stdp_enabled = False
+                print(f"  {C.YELLOW}[⚙] STDP disabled{C.RESET}")
+            else:
+                print(f"  STDP: {'ON' if stdp_enabled else 'OFF'}")
+            continue
+        elif cmd[0] == "/live":
+            if len(cmd) > 1 and cmd[1] == "on":
+                live_spikes = True
+                print(f"  {C.GREEN}[⚙] Live spike visualization ON{C.RESET}")
+            elif len(cmd) > 1 and cmd[1] == "off":
+                live_spikes = False
+                print(f"  {C.YELLOW}[⚙] Live spike visualization OFF{C.RESET}")
+            else:
+                print(f"  Live spikes: {'ON' if live_spikes else 'OFF'}")
+            continue
+        elif cmd[0] == "/stats":
+            print_stats(last_stats, cfg)
+            continue
+        elif cmd[0] == "/memory":
+            mem_rate = last_stats.get("memory_spike_rate", "N/A")
+            gate = last_stats.get("gate_activity", "N/A")
+            mix = last_stats.get("memory_mix", "N/A")
+            if isinstance(mem_rate, torch.Tensor): mem_rate = f"{mem_rate.item():.4f}"
+            if isinstance(gate, torch.Tensor): gate = f"{gate.item():.4f}"
+            if isinstance(mix, torch.Tensor): mix = f"{mix.item():.4f}"
+            print(f"  {C.PURPLE}Memory:{C.RESET} rate={mem_rate}, gate={gate}, mix={mix}")
+            continue
+        elif cmd[0] == "/expert":
+            # Search all stats keys for expert load data
+            expert_data = {k: v for k, v in last_stats.items() if "expert" in k and "load" in k}
+            if expert_data:
+                for k, v in sorted(expert_data.items()):
+                    if isinstance(v, torch.Tensor): v = v.item()
+                    bar = spike_bar(v, 30, C.YELLOW, max_rate=0.5)
+                    # Clean up key name for display
+                    display_name = k.replace("assoc_", "A").replace("_load", "")
+                    print(f"    {display_name}: {v:.2%} {bar}")
+            else:
+                print(f"    {C.DIM}No expert load data in stats{C.RESET}")
+                moe_keys = [k for k in last_stats.keys() if "moe" in k or "expert" in k]
+                if moe_keys:
+                    print(f"    {C.DIM}Related keys: {moe_keys}{C.RESET}")
+            continue
+        elif cmd[0] == "/reset":
+            model.reset_state()
+            print(f"  {C.GREEN}[⚙] Working memory cleared{C.RESET}")
+            continue
+
+        last_stats = generate_streaming(
+            model, tokenizer, cfg, user,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            repetition_penalty=rep_penalty,
+            enable_stdp=stdp_enabled,
+            live_spikes=live_spikes,
+        )
+
+
+if __name__ == "__main__":
+    main()

nord_v4_700m-4.2/download_data.py

@@ -0,0 +1,238 @@
+"""
+╔══════════════════════════════════════════════════════════════════════════╗
+║         PROJECT NORD — Завантаження датасетів                          ║
+║                                                                        ║
+║  Просто запусти:  python download_data.py                              ║
+║                                                                        ║
+║  Датасети для різних фаз навчання:                                     ║
+║    1. FineWeb-Edu     — загальні освітні тексти (база)                  ║
+║    2. OpenWebMath     — математика і reasoning                         ║
+║    3. The Stack v2    — код (Python, JS, C++ та інші)                  ║
+║    4. peS2o           — наукові статті                                  ║
+║    5. OpenHermes 2.5  — інструкції (chat/assistant формат)             ║
+║    6. SlimPajama      — різноманітний веб-текст                        ║
+║    7. Wikipedia       — енциклопедичні знання                          ║
+║    8. Cosmopedia      — синтетичні підручники                          ║
+║                                                                        ║
+║  Потрібно:  pip install datasets tqdm                                  ║
+╚══════════════════════════════════════════════════════════════════════════╝
+"""
+
+import json, os, sys, time
+
+DATASETS = {
+    "1": {
+        "name": "FineWeb-Edu (освітні тексти)",
+        "desc": "Високоякісні освітні тексти. Найкраще для базового навчання.",
+        "hf_id": "HuggingFaceFW/fineweb-edu", "hf_name": "sample-10BT",
+        "split": "train", "field": "text", "gb": 40,
+        "phase": "Фаза 1 — Базова мова",
+    },
+    "2": {
+        "name": "OpenWebMath (математика)",
+        "desc": "Математика: формули, докази, задачі. Для reasoning.",
+        "hf_id": "open-web-math/open-web-math", "hf_name": None,
+        "split": "train", "field": "text", "gb": 15,
+        "phase": "Фаза 2 — Reasoning і математика",
+    },
+    "3": {
+        "name": "StarCoder Data (код)",
+        "desc": "Код на Python, JS, C++, Java та інших мовах.",
+        "hf_id": "bigcode/starcoderdata", "hf_name": None,
+        "split": "train", "field": "content", "gb": 20,
+        "phase": "Фаза 3 — Програмування",
+    },
+    "4": {
+        "name": "peS2o (наукові статті)",
+        "desc": "Наукові papers від Semantic Scholar.",
+        "hf_id": "allenai/peS2o", "hf_name": "v2",
+        "split": "train", "field": "text", "gb": 20,
+        "phase": "Фаза 4 — Наукові тексти",
+    },
+    "5": {
+        "name": "OpenHermes 2.5 (інструкції)",
+        "desc": "Chat/Assistant формат. Перетворює base model в chat bot.",
+        "hf_id": "teknium/OpenHermes-2.5", "hf_name": None,
+        "split": "train", "field": "conversations", "gb": 2,
+        "phase": "Фаза 5 — Інструкції (chat)", "is_chat": True,
+    },
+    "6": {
+        "name": "SlimPajama (різноманітний текст)",
+        "desc": "Змішаний: веб, книги, Wikipedia, GitHub.",
+        "hf_id": "cerebras/SlimPajama-627B", "hf_name": None,
+        "split": "train", "field": "text", "gb": 30,
+        "phase": "Альтернатива — Різноманітний текст",
+    },
+    "7": {
+        "name": "Wikipedia (енциклопедія)",
+        "desc": "Вся англійська Wikipedia. Чисті факти.",
+        "hf_id": "wikimedia/wikipedia", "hf_name": "20231101.en",
+        "split": "train", "field": "text", "gb": 6,
+        "phase": "Додаток — Енциклопедичні знання",
+    },
+    "8": {
+        "name": "Cosmopedia (синтетичні підручники)",
+        "desc": "AI-згенеровані освітні тексти у форматі підручників.",
+        "hf_id": "HuggingFaceTB/cosmopedia", "hf_name": None,
+        "split": "train", "field": "text", "gb": 15,
+        "phase": "Додаток — Синтетичні освітні тексти",
+    },
+}
+
+def fmt(b):
+    for u in ["B","KB","MB","GB","TB"]:
+        if b < 1024: return f"{b:.1f} {u}"
+        b /= 1024
+    return f"{b:.1f} PB"
+
+def format_chat(convs):
+    if isinstance(convs, str): return convs
+    parts = []
+    for m in convs:
+        role = m.get("from", m.get("role", "user"))
+        text = m.get("value", m.get("content", ""))
+        if role in ("system","human","user"): parts.append(f"User: {text}")
+        elif role in ("gpt","assistant"): parts.append(f"Assistant: {text}")
+    return "\n".join(parts)
+
+def download_one(ds, save_dir, target_gb=None):
+    if target_gb is None: target_gb = ds["gb"]
+    target_bytes = int(target_gb * (1024**3))
+    safe = ds["hf_id"].split("/")[-1].replace("-","_").lower()
+    path = os.path.join(save_dir, f"{safe}.jsonl")
+
+    print(f"\n  {'═'*55}")
+    print(f"  📦 {ds['name']}")
+    print(f"  📁 {path}")
+    print(f"  🎯 {target_gb:.0f} GB")
+    print(f"  {'═'*55}")
+
+    os.makedirs(save_dir, exist_ok=True)
+    written = 0; count = 0; mode = "w"
+
+    if os.path.exists(path):
+        sz = os.path.getsize(path)
+        if sz >= target_bytes:
+            print(f"  [✓] Вже є! ({fmt(sz)})"); return path
+        if sz > 0:
+            written = sz
+            with open(path,"r",encoding="utf-8") as f: count = sum(1 for _ in f)
+            mode = "a"
+            print(f"  [*] Продовжуємо з {fmt(written)} ({count:,} зразків)")
+
+    print(f"  [*] Підключаємося до HuggingFace...")
+    try:
+        from datasets import load_dataset
+    except ImportError:
+        print("  [✗] pip install datasets"); return None
+
+    kw = {"path": ds["hf_id"], "split": ds["split"], "streaming": True}
+    if ds.get("hf_name"): kw["name"] = ds["hf_name"]
+
+    try:
+        data = load_dataset(**kw)
+    except Exception as e:
+        print(f"  [✗] Помилка: {e}"); return None
+
+    it = iter(data)
+    is_chat = ds.get("is_chat", False)
+    field = ds["field"]
+
+    if count > 0:
+        print(f"  [*] Пропускаємо {count:,} зразків...")
+        for _ in range(count):
+            try: next(it)
+            except StopIteration: break
+
+    print(f"  [*] Записуємо... (Ctrl+C = пауза)")
+    t0 = time.time(); lp = t0; start_b = written
+
+    try:
+        with open(path, mode, encoding="utf-8") as f:
+            for sample in it:
+                if is_chat:
+                    text = format_chat(sample.get(field, []))
+                else:
+                    text = sample.get(field, "")
+                if not text or len(text) < 50: continue
+
+                line = json.dumps({"text": text}, ensure_ascii=False) + "\n"
+                lb = len(line.encode("utf-8"))
+                f.write(line); written += lb; count += 1
+
+                now = time.time()
+                if now - lp >= 2.0:
+                    el = now - t0
+                    spd = (written - start_b) / el if el > 0 else 0
+                    pct = written / target_bytes * 100
+                    fl = int(30 * min(pct,100) / 100)
+                    bar = "█"*fl + "░"*(30-fl)
+                    eta = (target_bytes - written) / spd if spd > 0 else 0
+                    es = f"{eta/60:.0f}хв" if eta < 3600 else f"{eta/3600:.1f}год"
+                    print(f"\r  [{bar}] {pct:.1f}%  {fmt(written)}/{fmt(target_bytes)}  {count:,} зр.  {fmt(int(spd))}/s  ETA {es}    ", end="", flush=True)
+                    lp = now
+                    if count % 10000 == 0: f.flush()
+
+                if written >= target_bytes: break
+
+    except KeyboardInterrupt:
+        print(f"\n  [⏸] Пауза: {fmt(written)} ({count:,} зр.)")
+        return path
+
+    el = time.time() - t0
+    print(f"\n  [✓] {ds['name']}: {fmt(written)} | {count:,} зр. | {el/60:.0f}хв")
+    return path
+
+def main():
+    print("=" * 60)
+    print("  PROJECT NORD — Завантаження датасетів")
+    print("=" * 60)
+    print()
+    print("  ┌─────────────────────────────────────────────────────┐")
+    for k, ds in DATASETS.items():
+        print(f"  │ [{k}] {ds['name']:<45}│")
+        print(f"  │     {ds['phase']:<41} ~{ds['gb']:>2}GB │")
+    print(f"  │                                                     │")
+    print(f"  │ [A] Завантажити ВСЕ (Фази 1-5)                      │")
+    print(f"  │ [M] Кілька (через кому: 1,2,5)                      │")
+    print(f"  └─────────────────────────────────────────────────────┘")
+    print()
+
+    choice = input("  Вибери: ").strip().upper()
+
+    default_dir = os.path.join(os.sep, "nord_dataset")
+    print(f"\n  Папка? (Enter = {default_dir})")
+    di = input("  Папка: ").strip()
+    save_dir = di if di else default_dir
+
+    if choice == "A":
+        for k in ["1","2","4","5","7"]: download_one(DATASETS[k], save_dir)
+    elif choice == "M":
+        nums = input("  Номери (через кому): ").strip()
+        for k in [x.strip() for x in nums.split(",")]:
+            if k in DATASETS: download_one(DATASETS[k], save_dir)
+            else: print(f"  [!] Невідомий: {k}")
+    elif choice in DATASETS:
+        ds = DATASETS[choice]
+        print(f"\n  {ds['desc']}")
+        print(f"  Рекомендовано: {ds['gb']}GB")
+        print(f"  Скільки GB? (Enter = {ds['gb']})")
+        si = input("  GB: ").strip()
+        gb = float(si) if si else ds["gb"]
+        download_one(ds, save_dir, gb)
+    else:
+        print(f"  [!] Невідомий вибір: {choice}"); return
+
+    print(f"\n  {'═'*55}")
+    print(f"  [✓] ГОТОВО! Датасети в: {save_dir}")
+    print(f"  {'═'*55}")
+    print(f"\n  Як тренувати:")
+    print(f"    Базове:      python train_nord_700m.py --dataset {save_dir}/fineweb_edu.jsonl")
+    print(f"    Математика:  python train_nord_700m.py --dataset {save_dir}/open_web_math.jsonl --continued")
+    print(f"    Код:         python train_nord_700m.py --dataset {save_dir}/starcoderdata.jsonl --continued")
+    print(f"    Наука:       python train_nord_700m.py --dataset {save_dir}/pes2o.jsonl --continued")
+    print(f"    Chat:        python train_nord_700m.py --dataset {save_dir}/openhermes_2.5.jsonl --continued")
+    print()
+
+if __name__ == "__main__":
+    main()

nord_v4_700m-4.2/fast_tokenize.py

@@ -0,0 +1,201 @@
+"""
+╔══════════════════════════════════════════════════════════════════════════╗
+║         PROJECT NORD — Fast LMDB Tokenizer                             ║
+║                                                                        ║
+║  Usage:                                                                ║
+║      python build_lmdb.py                              (interactive)   ║
+║      python build_lmdb.py --src data.jsonl             (auto)          ║
+║      python build_lmdb.py --src data.jsonl --dst out_lmdb --seq 512   ║
+║                                                                        ║
+║  Batch tokenization with progress bar and resume support               ║
+╚══════════════════════════════════════════════════════════════════════════╝
+"""
+
+import argparse, json, struct, time, os, sys
+
+def build_lmdb(src, dst, seq_len=512, batch_size=1024):
+    import lmdb
+    import numpy as np
+    from transformers import AutoTokenizer
+
+    print("=" * 60, flush=True)
+    print("  PROJECT NORD — Fast LMDB Tokenizer", flush=True)
+    print("=" * 60, flush=True)
+    print(f"  Source:   {src}", flush=True)
+    print(f"  Output:   {dst}", flush=True)
+    print(f"  Seq len:  {seq_len}", flush=True)
+    print(f"  Batch:    {batch_size}", flush=True)
+    print(flush=True)
+
+    # ── Check if already exists ──
+    if os.path.exists(dst):
+        try:
+            env = lmdb.open(dst, readonly=True, lock=False)
+            with env.begin(write=False) as txn:
+                existing = struct.unpack("<Q", txn.get(b"__len__"))[0]
+                existing_tok = struct.unpack("<Q", txn.get(b"__total_tokens__"))[0]
+            env.close()
+            print(f"  [!] LMDB already exists: {existing:,} samples, {existing_tok/1e6:.0f}M tokens", flush=True)
+            print(f"  Overwrite? (y/n, Enter = n)")
+            choice = input("  > ").strip().lower()
+            if choice not in ("y", "yes"):
+                print("  [*] Skipped.", flush=True)
+                return dst
+            import shutil
+            shutil.rmtree(dst)
+            print("  [*] Deleted old LMDB.", flush=True)
+        except:
+            pass
+
+    # ── Init tokenizer ──
+    print("  [*] Loading tokenizer...", flush=True)
+    tok = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    PAD_ID = tok.pad_token_id
+    print(f"  [✓] Tokenizer ready (vocab={tok.vocab_size:,})", flush=True)
+
+    # ── Read all texts ──
+    print(f"\n  [1/3] Reading JSONL into memory...", flush=True)
+    t0 = time.time()
+    texts = []
+    with open(src, "r", encoding="utf-8") as f:
+        for i, line in enumerate(f):
+            if i % 1_000_000 == 0 and i > 0:
+                print(f"      read {i:,} lines...", flush=True)
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                obj = json.loads(line)
+            except:
+                continue
+            text = obj.get("text") or obj.get("content") or obj.get("passage", "")
+            if len(text) >= 30:
+                texts.append(text)
+    print(f"      {len(texts):,} valid texts in {time.time()-t0:.0f}s", flush=True)
+
+    if not texts:
+        print("  [✗] No valid texts found!", flush=True)
+        return None
+
+    # ── Batch tokenize ──
+    print(f"\n  [2/3] Batch tokenizing {len(texts):,} texts (batch={batch_size})...", flush=True)
+    t1 = time.time()
+
+    os.makedirs(os.path.dirname(dst) if os.path.dirname(dst) else ".", exist_ok=True)
+    env = lmdb.open(dst, map_size=80 * (1024**3))
+    txn = env.begin(write=True)
+
+    count = 0
+    total_tok = 0
+    total_batches = (len(texts) + batch_size - 1) // batch_size
+
+    for batch_idx in range(0, len(texts), batch_size):
+        batch = texts[batch_idx : batch_idx + batch_size]
+        batch_num = batch_idx // batch_size + 1
+
+        enc = tok(
+            batch,
+            max_length=seq_len,
+            truncation=True,
+            padding="max_length",
+            return_tensors="np",
+            return_attention_mask=False,
+        )
+        ids_np = enc.input_ids.astype(np.int32)
+
+        for j in range(ids_np.shape[0]):
+            row = ids_np[j]
+            non_pad = int(np.sum(row != PAD_ID))
+            if non_pad < 10:
+                continue
+            txn.put(f"sample_{count:010d}".encode(), row.tobytes())
+            count += 1
+            total_tok += non_pad
+
+        # Progress
+        if batch_num % 100 == 0 or batch_num == total_batches:
+            elapsed = time.time() - t1
+            pct = batch_num / total_batches * 100
+            eta = (elapsed / batch_num) * (total_batches - batch_num)
+            speed = count / elapsed if elapsed > 0 else 0
+            bar_len = 30
+            filled = int(bar_len * pct / 100)
+            bar = "█" * filled + "░" * (bar_len - filled)
+            print(
+                f"    [{bar}] {pct:5.1f}% | "
+                f"{count:,} samples | {total_tok/1e6:.0f}M tok | "
+                f"{speed:.0f} doc/s | ETA {eta:.0f}s",
+                flush=True,
+            )
+
+        # Commit every 500k
+        if count % 500_000 < batch_size and count >= 500_000:
+            txn.commit()
+            txn = env.begin(write=True)
+
+    # Save metadata
+    txn.put(b"__len__", struct.pack("<Q", count))
+    txn.put(b"__total_tokens__", struct.pack("<Q", total_tok))
+    txn.commit()
+    env.close()
+
+    elapsed = time.time() - t1
+    print(f"\n  [3/3] Done!", flush=True)
+    print(f"  {'═' * 50}", flush=True)
+    print(f"    Samples:  {count:,}", flush=True)
+    print(f"    Tokens:   {total_tok:,} ({total_tok/1e6:.0f}M)", flush=True)
+    print(f"    Time:     {elapsed:.0f}s ({elapsed/60:.1f} min)", flush=True)
+    print(f"    Speed:    {count/elapsed:.0f} doc/s", flush=True)
+    print(f"  {'═' * 50}", flush=True)
+    print(f"\n  Тепер тренуй:", flush=True)
+    print(f"    python train_nord_700m.py --dataset {src}", flush=True)
+    print(flush=True)
+    return dst
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Nord LMDB Tokenizer")
+    parser.add_argument("--src", type=str, default=None, help="Source JSONL file")
+    parser.add_argument("--dst", type=str, default=None, help="Output LMDB directory")
+    parser.add_argument("--seq", type=int, default=512, help="Max sequence length (default: 512)")
+    parser.add_argument("--batch", type=int, default=1024, help="Batch size (default: 1024)")
+    args = parser.parse_args()
+
+    # Interactive mode if no args
+    if args.src is None:
+        print("=" * 60)
+        print("  PROJECT NORD — Fast LMDB Tokenizer")
+        print("=" * 60)
+        print()
+        print("  Шлях до JSONL датасету?")
+        print("  (наприклад: /nord_dataset/train_data.jsonl)")
+        args.src = input("  Source: ").strip()
+        if not args.src:
+            print("  [✗] Потрібно вказати шлях!", flush=True)
+            sys.exit(1)
+
+    if not os.path.exists(args.src):
+        print(f"  [✗] Файл не знайдено: {args.src}", flush=True)
+        sys.exit(1)
+
+    if args.dst is None:
+        # Auto: same path but _lmdb suffix
+        args.dst = args.src.replace(".jsonl", "") + "_lmdb"
+        print(f"\n  Output LMDB? (Enter = {args.dst})")
+        user_dst = input("  Output: ").strip()
+        if user_dst:
+            args.dst = user_dst
+
+    print(f"\n  Sequence length? (Enter = {args.seq})")
+    seq_input = input("  Seq: ").strip()
+    if seq_input:
+        args.seq = int(seq_input)
+
+    print()
+    build_lmdb(args.src, args.dst, args.seq, args.batch)
+
+
+if __name__ == "__main__":
+    main()

nord_v4_700m-4.2/nord_core_700m.py

@@ -0,0 +1,634 @@
+"""
+╔══════════════════════════════════════════════════════════════════════════════╗
+║                PROJECT NORD — Core Engine  v4.2  (700M)                 ║
+║          Spiking Neural Network LLM with Brain-Inspired Architecture       ║
+║                                                                            ║
+║  v4.1 CRITICAL FIXES (from code review):                                   ║
+║    FIX A: Vectorized MoE dispatch — no Python loops over experts           ║
+║    FIX B: Temporal attention memory — multi-head read over ALL timesteps   ║
+║    FIX C: Differentiable spike loss — proper gradient flow                 ║
+║    FIX D: LIF stability — clamped tau/threshold, warmup freeze             ║
+║    FIX E: Temporal mixing in attention (no naive T*Dh flattening)         ║
+║    FIX F: STDP isolation — only executive zone, bounded magnitude          ║
+║    FIX G: MoE load balancing loss — prevents expert collapse               ║
+║    FIX H: Gradient checkpointing support — VRAM control                    ║
+║    FIX I: Fused LIF operations — reduced kernel launch overhead            ║
+║    FIX J: Realistic training estimates in docs                             ║
+║                                                                            ║
+║  v4.2 FIXES (from 13K step training analysis):                             ║
+║    FIX K: Block outputs spike-only — clamp negative before spike_ts        ║
+║    FIX L: Stronger spike regulator — adaptive weight, per-layer targeting  ║
+║    FIX M: Executive clamp floor=0 — prevent negative spike propagation     ║
+╚══════════════════════════════════════════════════════════════════════════════╝
+"""
+
+from __future__ import annotations
+import math, torch, torch.nn as nn, torch.nn.functional as F
+from torch import Tensor
+from torch.utils.checkpoint import checkpoint as grad_checkpoint
+from dataclasses import dataclass
+from typing import Dict, Tuple, Optional, List
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# §0  CONFIG
+# ═══════════════════════════════════════════════════════════════════════════════
+@dataclass
+class NordConfig:
+    tokenizer_id:str="meta-llama/Llama-3.2-1B"
+    # ── 700M Architecture ──
+    vocab_size:int=128_256; d_model:int=1536; n_heads:int=24; n_layers:int=10
+    d_ff:int=4096; max_seq_len:int=512
+    T:int=8; T_slow:int=2; persistent_mem:bool=True
+    # LIF — FIX D: constrained ranges
+    tau_mem:float=0.9; tau_mem_min:float=0.8; tau_mem_max:float=0.98
+    tau_syn:float=0.50; v_threshold:float=0.12
+    v_thresh_min:float=0.05; v_thresh_max:float=0.5
+    v_reset:float=-0.1; refractory_t:int=2; threshold_lr:float=0.01
+    lif_freeze_steps:int=1000
+    n_clusters:int=128; cascade_radius:int=3; cascade_gain:float=0.8
+    # STDP — FIX F: bounded
+    stdp_a_plus:float=0.005; stdp_a_minus:float=0.005
+    stdp_tau_plus:float=20.0; stdp_tau_minus:float=20.0
+    stdp_w_max:float=0.5; stdp_w_min:float=-0.15
+    stdp_reward_scale:float=1.0; stdp_layers:Optional[List[str]]=None
+    resonance_top_k:int=64; clamp_floor:float=-0.1; surrogate_alpha:float=4.0
+    rope_theta:float=10000.0
+    # MoE — FIX A+G
+    n_experts:int=4; top_k_experts:int=2; moe_capacity_factor:float=1.25
+    moe_load_balance_weight:float=0.01; moe_route_temperature:float=1.0
+    # Spike loss — FIX L
+    target_spike_rate:float=0.03; spike_loss_weight:float=0.5
+    # Zones: 3 sensory + 3 association(MoE) + 4 executive = 10
+    sensory_layers:int=3; association_layers:int=3; executive_layers:int=4
+    # Memory — FIX B
+    memory_tau_mem:float=0.99; memory_size:int=256
+    memory_gate_threshold:float=0.3; memory_n_read_heads:int=8
+    # FIX H
+    gradient_checkpointing:bool=False
+    # Training
+    batch_size:int=1; grad_accum:int=32; lr:float=2e-4; min_lr:float=1e-5
+    weight_decay:float=0.01; warmup_steps:int=1000; max_steps:int=50_000
+    save_every:int=1000; log_every:int=10; max_grad_norm:float=1.0
+    dtype:torch.dtype=torch.float16; device:str="cuda"
+    @property
+    def T_total(self)->int: return self.T+self.T_slow
+    @property
+    def n_layers_total(self)->int: return self.sensory_layers+self.association_layers+self.executive_layers
+    def __post_init__(self):
+        if self.stdp_layers is None:
+            self.stdp_layers=[f"executive_{i}" for i in range(self.executive_layers)]
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# §1  SURROGATE GRADIENT
+# ═══════════════════════════════════════════════════════════════════════════════
+class ATanSurrogate(torch.autograd.Function):
+    alpha=2.0
+    @staticmethod
+    def forward(ctx,membrane:Tensor,threshold:Tensor)->Tensor:
+        ctx.save_for_backward(membrane,threshold)
+        return(membrane>=threshold).to(membrane.dtype)
+    @staticmethod
+    def backward(ctx,grad_output:Tensor)->Tuple[Tensor,Tensor]:
+        membrane,threshold=ctx.saved_tensors
+        x=(membrane.float()-threshold.float())
+        grad=ATanSurrogate.alpha/(2.0*math.pi*(1.0+(ATanSurrogate.alpha*x)**2))
+        grad_v=(grad_output.float()*grad).to(membrane.dtype)
+        return grad_v,-grad_v
+
+def spike_fn(v:Tensor,th:Tensor,alpha:float=2.0)->Tensor:
+    ATanSurrogate.alpha=alpha; return ATanSurrogate.apply(v,th)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# §2  ASSOCIATIVE LIF — FIX D: Stability + FIX I: Fused ops
+# ═══════════════════════════════════════════════════════════════════════════════
+class AssociativeLIF(nn.Module):
+    def __init__(self,d:int,cfg:NordConfig,persistent:bool=False,
+                 tau_mem_override:Optional[float]=None):
+        super().__init__()
+        self.cfg=cfg; self.d=d; self.persistent=persistent
+        self.threshold_raw=nn.Parameter(torch.full((d,),cfg.v_threshold))
+        tau_mem=tau_mem_override if tau_mem_override is not None else cfg.tau_mem
+        self.beta_mem_raw=nn.Parameter(torch.tensor(math.log(tau_mem/(1-tau_mem+1e-6))))
+        self.beta_syn_raw=nn.Parameter(torch.tensor(math.log(cfg.tau_syn/(1-cfg.tau_syn+1e-6))))
+        nc=cfg.n_clusters
+        self.register_buffer("cluster_ids",torch.arange(d)%nc)
+        r=cfg.cascade_radius; idx=torch.arange(nc)
+        iw=torch.zeros(nc,nc)
+        for offset in range(-r,r+1):
+            if offset!=0: iw[idx,(idx+offset)%nc]=1.0-abs(offset)/(r+1)
+        self.neighbor_weights=nn.Parameter(iw)
+        self.cluster_gain=nn.Parameter(torch.full((nc,),cfg.cascade_gain))
+        if persistent:
+            self.register_buffer("_v_mem_state",torch.zeros(1,d))
+            self.register_buffer("_i_syn_state",torch.zeros(1,d))
+        self.register_buffer("_firing_rate_ema",torch.full((d,),cfg.target_spike_rate))
+        self.register_buffer("_step_counter",torch.tensor(0,dtype=torch.long))
+
+    @property
+    def threshold(self)->Tensor:
+        return self.threshold_raw.clamp(self.cfg.v_thresh_min,self.cfg.v_thresh_max)
+    @property
+    def beta_mem(self)->Tensor:
+        return torch.sigmoid(self.beta_mem_raw).clamp(self.cfg.tau_mem_min,self.cfg.tau_mem_max)
+    @property
+    def beta_syn(self)->Tensor: return torch.sigmoid(self.beta_syn_raw)
+
+    def _cascade_amplify(self,spikes:Tensor)->Tensor:
+        B,D=spikes.shape; nc=self.cfg.n_clusters
+        cid=self.cluster_ids.unsqueeze(0).expand(B,-1)
+        cf=torch.zeros(B,nc,device=spikes.device,dtype=spikes.dtype)
+        cf.scatter_add_(1,cid,spikes); cf=cf/max(D//nc,1)
+        W=torch.sigmoid(self.neighbor_weights)
+        ns=(W.to(cf.dtype)@cf.T).T*self.cluster_gain.to(cf.dtype).unsqueeze(0)
+        return ns.gather(1,cid)
+
+    def reset_state(self):
+        if self.persistent: self._v_mem_state.zero_(); self._i_syn_state.zero_()
+
+    def forward(self,current_in:Tensor)->Tuple[Tensor,Tensor]:
+        T,B,D=current_in.shape; device=current_in.device; dtype=current_in.dtype
+        bm=self.beta_mem; bs=self.beta_syn; thresh=self.threshold
+        if self.persistent and self._v_mem_state.shape[0]==B:
+            v_mem=self._v_mem_state.clone(); i_syn=self._i_syn_state.clone()
+        else:
+            v_mem=torch.zeros(B,D,device=device,dtype=dtype)
+            i_syn=torch.zeros(B,D,device=device,dtype=dtype)
+            if self.persistent:
+                self._v_mem_state=torch.zeros(B,D,device=device,dtype=dtype)
+                self._i_syn_state=torch.zeros(B,D,device=device,dtype=dtype)
+        refrac=torch.zeros(B,D,device=device,dtype=torch.int32)
+        spikes_out=[]; v_trace=[]
+        refractory_val=torch.full_like(v_mem,self.cfg.v_reset)
+        ref_t=self.cfg.refractory_t; alpha=self.cfg.surrogate_alpha
+        for t in range(T):
+            i_syn=bs*i_syn+current_in[t]
+            rmask=(refrac>0)
+            new_v=bm*v_mem+(1.0-bm)*i_syn
+            v_mem=torch.where(rmask,refractory_val,new_v)
+            s=spike_fn(v_mem,thresh,alpha)
+            if s.sum()>0: i_syn=i_syn+self._cascade_amplify(s)
+            v_mem=v_mem-s*thresh.detach()
+            refrac=torch.where(s.bool(),torch.full_like(refrac,ref_t),(refrac-1).clamp(min=0))
+            spikes_out.append(s); v_trace.append(v_mem)
+        if self.persistent:
+            self._v_mem_state=v_mem.detach(); self._i_syn_state=i_syn.detach()
+        ss=torch.stack(spikes_out)
+        with torch.no_grad():
+            self._firing_rate_ema.lerp_(ss.mean(dim=(0,1)),0.01)
+            self._step_counter+=1
+        return ss,torch.stack(v_trace)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# §3  TEMPORAL ENCODER
+# ═══════════════════════════════════════════════════════════════════════════════
+class TemporalSpikeEncoder(nn.Module):
+    def __init__(self,cfg:NordConfig):
+        super().__init__(); self.cfg=cfg; D=cfg.d_model
+        self.embed=nn.Embedding(cfg.vocab_size,D)
+        nn.init.kaiming_uniform_(self.embed.weight,a=math.sqrt(5))
+        self.temporal_proj=nn.Linear(D,D,bias=False)
+        self.drive_scale=nn.Parameter(torch.tensor(25.0))
+        self.fast_basis=nn.Parameter(torch.randn(cfg.T,D)*0.02)
+        self.slow_basis=nn.Parameter(torch.randn(cfg.T_slow,D)*0.02)
+        self.slow_scale=nn.Parameter(torch.tensor(8.0))
+    def forward(self,token_ids:Tensor)->Tensor:
+        B,S=token_ids.shape; D=self.cfg.d_model
+        x=self.temporal_proj(self.embed(token_ids)).reshape(B*S,D)
+        fast=torch.sigmoid(self.fast_basis).unsqueeze(1)*x.unsqueeze(0)*self.drive_scale
+        slow=torch.sigmoid(self.slow_basis).unsqueeze(1)*x.unsqueeze(0)*self.slow_scale
+        return torch.cat([fast,slow],dim=0)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# §4  RoPE
+# ═══════════════════════════════════════════════════════════════════════════════
+class RotaryPositionEmbedding(nn.Module):
+    def __init__(self,dim:int,max_seq_len:int=2048,theta:float=10000.0):
+        super().__init__()
+        inv_freq=1.0/(theta**(torch.arange(0,dim,2).float()/dim))
+        self.register_buffer("inv_freq",inv_freq)
+        t=torch.arange(max_seq_len).float(); freqs=torch.outer(t,inv_freq)
+        self.register_buffer("cos_cached",freqs.cos())
+        self.register_buffer("sin_cached",freqs.sin())
+    def forward(self,x:Tensor,seq_len:int)->Tuple[Tensor,Tensor]:
+        return self.cos_cached[:seq_len].to(x.dtype),self.sin_cached[:seq_len].to(x.dtype)
+
+def apply_rope(x:Tensor,cos:Tensor,sin:Tensor)->Tensor:
+    d=cos.shape[-1]; x1=x[...,:d]; x2=x[...,d:2*d]
+    c=cos.unsqueeze(0).unsqueeze(0); s=sin.unsqueeze(0).unsqueeze(0)
+    rot=torch.cat([x1*c-x2*s,x1*s+x2*c],dim=-1)
+    return torch.cat([rot,x[...,2*d:]],dim=-1) if x.shape[-1]>2*d else rot
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# §5  SYNAPTIC RESONANCE — FIX E: Temporal mixing (not flattening)
+# ═══════════════════════════════════════════════════════════════════════════════
+class SpikingSynapticResonance(nn.Module):
+    def __init__(self,cfg:NordConfig):
+        super().__init__(); self.cfg=cfg
+        self.n_heads=cfg.n_heads; self.d_head=cfg.d_model//cfg.n_heads
+        self.top_k=cfg.resonance_top_k; D=cfg.d_model; T_t=cfg.T_total
+        self.W_q=nn.Linear(D,D,bias=False); self.W_k=nn.Linear(D,D,bias=False)
+        self.W_v=nn.Linear(D,D,bias=False); self.W_o=nn.Linear(D,D,bias=False)
+        self.lif_q=AssociativeLIF(D,cfg); self.lif_k=AssociativeLIF(D,cfg)
+        self.resonance_temp=nn.Parameter(torch.tensor(1.0/math.sqrt(self.d_head)))
+        # FIX E: Learned temporal mixing weights (not concatenation)
+        self.temporal_mix_q=nn.Parameter(torch.ones(T_t)/T_t)
+        self.temporal_mix_k=nn.Parameter(torch.ones(T_t)/T_t)
+        self.rope=RotaryPositionEmbedding(self.d_head,cfg.max_seq_len,cfg.rope_theta)
+
+    def forward(self,x_spikes:Tensor)->Tensor:
+        T_t,B,S,D=x_spikes.shape; H=self.n_heads; Dh=self.d_head
+        xf=x_spikes.reshape(T_t*B*S,D)
+        qc=self.W_q(xf).reshape(T_t,B*S,D)
+        kc=self.W_k(xf).reshape(T_t,B*S,D)
+        vr=self.W_v(xf).reshape(T_t,B,S,D)
+        qs,_=self.lif_q(qc); ks,_=self.lif_k(kc)
+        qs=qs.reshape(T_t,B,S,H,Dh); ks=ks.reshape(T_t,B,S,H,Dh)
+        # FIX E: Weighted sum over time, preserves spike timing semantics
+        twq=F.softmax(self.temporal_mix_q,dim=0).reshape(T_t,1,1,1,1)
+        twk=F.softmax(self.temporal_mix_k,dim=0).reshape(T_t,1,1,1,1)
+        qm=(qs*twq).sum(0).permute(0,2,1,3) # (B,H,S,Dh)
+        km=(ks*twk).sum(0).permute(0,2,1,3)
+        cos,sin=self.rope(qm,S)
+        qm=apply_rope(qm,cos,sin); km=apply_rope(km,cos,sin)
+        res=torch.matmul(qm,km.transpose(-2,-1))*self.resonance_temp
+        cmask=torch.triu(torch.ones(S,S,device=x_spikes.device,dtype=torch.bool),diagonal=1)
+        res.masked_fill_(cmask.unsqueeze(0).unsqueeze(0),float("-inf"))
+        K=min(self.top_k,S)
+        if K<S:
+            tv,ti=torch.topk(res,K,dim=-1)
+            sr=torch.full_like(res,float("-inf")); sr.scatter_(-1,ti,tv); res=sr
+        attn=F.softmax(res.float(),dim=-1).to(res.dtype)
+        vm=vr.mean(dim=0).reshape(B,S,H,Dh).permute(0,2,1,3)
+        ctx=torch.matmul(attn,vm).permute(0,2,1,3).reshape(B,S,D)
+        return self.W_o(ctx).unsqueeze(0).expand(T_t,-1,-1,-1)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# §6  SPIKE-DRIVEN MoE — FIX A: Vectorized + FIX G: Load Balance
+# ═══════════════════════════════════════════════════════════════════════════════
+class SpikingExpertGroup(nn.Module):
+    """FIX A: Memory-efficient expert dispatch using per-expert Linear + masking.
+    Instead of bmm with (N,ef,D) tensors, we loop over experts (not tokens).
+    With 4 experts this is 4 iterations — much better than 2048-token bmm."""
+    def __init__(self,cfg:NordConfig):
+        super().__init__()
+        self.n_experts=cfg.n_experts; self.expert_ff=cfg.d_ff//cfg.n_experts
+        D=cfg.d_model; ef=self.expert_ff
+        # Standard Linear layers per expert — memory efficient
+        self.up=nn.ModuleList([nn.Linear(D,ef,bias=False) for _ in range(cfg.n_experts)])
+        self.down=nn.ModuleList([nn.Linear(ef,D,bias=False) for _ in range(cfg.n_experts)])
+        self.lif1=AssociativeLIF(ef,cfg); self.lif2=AssociativeLIF(D,cfg)
+
+    def forward(self,x:Tensor,expert_indices:Tensor,expert_weights:Tensor)->Tensor:
+        """x:(T,N,D), expert_indices:(N,top_k), expert_weights:(N,top_k)"""
+        T,N,D=x.shape; top_k=expert_indices.shape[1]
+        output=torch.zeros_like(x)
+        # Loop over experts (4 iterations), not tokens (2048)
+        for e in range(self.n_experts):
+            # Find which tokens use this expert and with what weight
+            mask=torch.zeros(N,device=x.device,dtype=x.dtype)
+            for k in range(top_k):
+                is_e=(expert_indices[:,k]==e).to(x.dtype)
+                mask=mask+is_e*expert_weights[:,k]
+            if mask.sum()==0: continue
+            # Which tokens actually route here
+            active=(mask>0)
+            if not active.any(): continue
+            # Extract active tokens across all timesteps
+            active_x=x[:,active,:] # (T, n_active, D)
+            Ta,Na,Da=active_x.shape
+            # Up projection + LIF
+            h=self.up[e](active_x.reshape(Ta*Na,Da)).reshape(Ta,Na,-1)
+            h,_=self.lif1(h)
+            # Down projection + LIF
+            o=self.down[e](h.reshape(Ta*Na,-1)).reshape(Ta,Na,Da)
+            o,_=self.lif2(o)
+            # Weighted scatter back
+            w=mask[active].unsqueeze(0).unsqueeze(-1) # (1,n_active,1)
+            output[:,active,:]+=o*w
+        return output
+
+class SpikeDrivenMoE(nn.Module):
+    def __init__(self,cfg:NordConfig):
+        super().__init__(); self.cfg=cfg
+        self.n_experts=cfg.n_experts; self.top_k=cfg.top_k_experts
+        self.clusters_per_expert=cfg.n_clusters//cfg.n_experts
+        self.expert_group=SpikingExpertGroup(cfg)
+        self.route_lif=AssociativeLIF(cfg.d_model,cfg)
+        self.expert_bias=nn.Parameter(torch.zeros(cfg.n_experts))
+        self.register_buffer("expert_counts_ema",torch.ones(cfg.n_experts)/cfg.n_experts)
+
+    def _compute_expert_scores(self,spikes:Tensor)->Tensor:
+        fr=spikes.mean(dim=0); N,D=fr.shape; nc=self.cfg.n_clusters
+        cid=torch.arange(D,device=fr.device)%nc
+        cr=torch.zeros(N,nc,device=fr.device,dtype=fr.dtype)
+        cr.scatter_add_(1,cid.unsqueeze(0).expand(N,-1),fr)
+        cr=cr/max(D//nc,1)
+        es=cr.reshape(N,self.n_experts,self.clusters_per_expert).mean(dim=-1)
+        es=es/max(self.cfg.moe_route_temperature,0.01)
+        return es+self.expert_bias.to(es.dtype)
+
+    def _load_balance_loss(self,scores:Tensor,top_idx:Tensor)->Tensor:
+        N=scores.shape[0]
+        ef=torch.zeros(self.n_experts,device=scores.device)
+        for e in range(self.n_experts):
+            ef[e]=(top_idx==e).float().sum()/(N*self.top_k)
+        rp=F.softmax(scores,dim=-1).mean(dim=0)
+        loss=self.n_experts*(ef*rp).sum()
+        with torch.no_grad(): self.expert_counts_ema.lerp_(ef,0.01)
+        return loss
+
+    def forward(self,x:Tensor)->Tuple[Tensor,Dict]:
+        T,B,S,D=x.shape; N=B*S
+        xf=x.reshape(T,N,D); rs,_=self.route_lif(xf)
+        es=self._compute_expert_scores(rs)
+        ts,ti=torch.topk(es,self.top_k,dim=-1)
+        tw=F.softmax(ts.float(),dim=-1).to(x.dtype)
+        output=self.expert_group(xf,ti,tw).reshape(T,B,S,D)
+        lb=self._load_balance_loss(es,ti)
+        stats={"moe_route_entropy":-(F.softmax(es,dim=-1)*F.log_softmax(es+1e-8,dim=-1)).sum(-1).mean().item(),
+               "moe_load_balance_loss":lb}
+        with torch.no_grad():
+            for e in range(self.n_experts): stats[f"expert_{e}_load"]=self.expert_counts_ema[e].item()
+        return output,stats
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# §7  MEMORY CORTEX — FIX B: Temporal attention readout
+# ═══════════════════════════════════════════════════════════════════════════════
+class MemoryCortex(nn.Module):
+    def __init__(self,cfg:NordConfig):
+        super().__init__(); self.cfg=cfg; D=cfg.d_model; M=cfg.memory_size
+        self.to_memory=nn.Linear(D,M,bias=False)
+        self.from_memory=nn.Linear(M,D,bias=False)
+        self.memory_lif=AssociativeLIF(M,cfg,persistent=True,tau_mem_override=cfg.memory_tau_mem)
+        self.gate_lif=AssociativeLIF(M,cfg)
+        self.gate_proj=nn.Linear(D,M,bias=False)
+        self.gate_threshold=nn.Parameter(torch.tensor(cfg.memory_gate_threshold))
+        # FIX B: Multi-head temporal attention for memory readout
+        H=cfg.memory_n_read_heads; hd=M//H
+        self.n_read_heads=H
+        self.read_query=nn.Parameter(torch.randn(H,hd)*0.02)
+        self.read_key_proj=nn.Linear(M,M,bias=False)
+        self.read_scale=1.0/math.sqrt(hd)
+        self.mem_norm=nn.LayerNorm(D)
+        self.memory_mix=nn.Parameter(torch.tensor(0.1))
+
+    def reset_state(self): self.memory_lif.reset_state()
+
+    def forward(self,x:Tensor)->Tuple[Tensor,Dict[str,float]]:
+        T,B,S,D=x.shape; M=self.cfg.memory_size; N=B*S; H=self.n_read_heads; hd=M//H
+        xf=x.reshape(T,N,D)
+        mi=self.to_memory(xf.reshape(T*N,D)).reshape(T,N,M)
+        ms,mv=self.memory_lif(mi)
+        gi=self.gate_proj(xf.reshape(T*N,D)).reshape(T,N,M)
+        gs,_=self.gate_lif(gi)
+        gate_sig=gs.mean(dim=0)
+        gate_mask=torch.sigmoid((gate_sig-self.gate_threshold)*10.0)
+        # FIX B: Temporal attention over ALL timesteps
+        mvh=mv.reshape(T,N,H,hd)
+        mk=self.read_key_proj(mv.reshape(T*N,M)).reshape(T,N,H,hd)
+        q=self.read_query.unsqueeze(0).unsqueeze(0) # (1,1,H,hd)
+        attn_s=(q*mk).sum(-1)*self.read_scale # (T,N,H)
+        attn_w=F.softmax(attn_s.float(),dim=0).to(mv.dtype) # (T,N,H)
+        mem_read=(mvh*attn_w.unsqueeze(-1)).sum(0).reshape(N,M) # (N,M)
+        mem_read=mem_read*gate_mask
+        mem_out=self.mem_norm(self.from_memory(mem_read).float()).to(x.dtype)
+        mix=torch.sigmoid(self.memory_mix)
+        x_e=x+mix*mem_out.reshape(1,B,S,D).expand_as(x)
+        stats={"memory_spike_rate":ms.mean().item(),"gate_activity":gate_sig.mean().item(),
+               "memory_mix":mix.item(),
+               "memory_attn_entropy":-(attn_w.float()*(attn_w.float()+1e-8).log()).sum(0).mean().item()}
+        return x_e,stats
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# §8  BLOCKS — FIX H: Gradient checkpointing
+# ═══════════════════════════════════════════════════════════════════════════════
+class SpikingFeedForward(nn.Module):
+    def __init__(self,cfg:NordConfig):
+        super().__init__()
+        self.up=nn.Linear(cfg.d_model,cfg.d_ff,bias=False)
+        self.down=nn.Linear(cfg.d_ff,cfg.d_model,bias=False)
+        self.lif1=AssociativeLIF(cfg.d_ff,cfg); self.lif2=AssociativeLIF(cfg.d_model,cfg)
+    def forward(self,x:Tensor)->Tensor:
+        T,B,S,D=x.shape
+        h=self.up(x.reshape(T*B*S,D)).reshape(T,B*S,-1); h,_=self.lif1(h)
+        h=self.down(h.reshape(T*B*S,-1)).reshape(T,B*S,D); h,_=self.lif2(h)
+        return h.reshape(T,B,S,D)
+
+class LeakyClamp(nn.Module):
+    def __init__(self,d:int,floor_init:float=-0.1,leak_init:float=0.1,force_nonneg:bool=False):
+        super().__init__()
+        # FIX M: force_nonneg=True for executive blocks — no negative spikes
+        self.force_nonneg=force_nonneg
+        if force_nonneg:
+            floor_init=0.0
+        self.floor=nn.Parameter(torch.full((d,),floor_init))
+        self.leak_raw=nn.Parameter(torch.full((d,),math.log(leak_init/(1-leak_init+1e-6))))
+    @property
+    def leak(self)->Tensor: return torch.sigmoid(self.leak_raw)
+    def forward(self,x:Tensor)->Tensor:
+        if self.force_nonneg:
+            # Executive: no negative values allowed
+            return F.relu(x)
+        return torch.where(x>=0,x,(self.leak*x).clamp(min=self.floor))
+
+class NordBlock(nn.Module):
+    def __init__(self,cfg:NordConfig,layer_idx:int=0,use_moe:bool=False,zone:str="sensory"):
+        super().__init__(); D=cfg.d_model; self.use_moe=use_moe; self.zone=zone
+        self.layer_idx=layer_idx; self.use_checkpoint=cfg.gradient_checkpointing
+        self.norm1=nn.LayerNorm(D); self.norm2=nn.LayerNorm(D)
+        self.resonance=SpikingSynapticResonance(cfg)
+        if use_moe: self.moe=SpikeDrivenMoE(cfg)
+        else: self.ffn=SpikingFeedForward(cfg)
+        sc=0.1/max(cfg.n_layers_total,1)
+        self.gamma_attn=nn.Parameter(torch.full((D,),sc))
+        self.gamma_ffn=nn.Parameter(torch.full((D,),sc))
+        # FIX M: Executive blocks force non-negative output
+        self.clamp=LeakyClamp(D,floor_init=cfg.clamp_floor,
+                              force_nonneg=(zone=="executive"))
+    @staticmethod
+    def _sn(nl:nn.LayerNorm,x:Tensor)->Tensor:
+        od=x.dtype
+        return F.layer_norm(x.float(),nl.normalized_shape,
+            nl.weight.float() if nl.weight is not None else None,
+            nl.bias.float() if nl.bias is not None else None,nl.eps).to(od)
+    def _forward_inner(self,x:Tensor)->Tuple[Tensor,Dict]:
+        stats={}
+        x=x+self.gamma_attn*self.resonance(self._sn(self.norm1,x))
+        xn=self._sn(self.norm2,x)
+        if self.use_moe: fo,ms=self.moe(xn); stats.update(ms)
+        else: fo=self.ffn(xn)
+        return self.clamp(x+self.gamma_ffn*fo),stats
+    def forward(self,x:Tensor)->Tuple[Tensor,Dict]:
+        if self.use_checkpoint and self.training:
+            x=grad_checkpoint(lambda inp:self._forward_inner(inp)[0],x,use_reentrant=False)
+            return x,{}
+        return self._forward_inner(x)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# §9  SPIKE REGULATOR — FIX C: Differentiable
+# ═══════════════════════════════════════════════════════════════════════════════
+class AuxiliarySpikeRegulator(nn.Module):
+    """FIX L: Adaptive spike regulator.
+    - Stronger weight (0.5 default)
+    - Extra penalty when any layer drops below min_rate (anti-death)
+    - Asymmetric: penalizes too-low firing 3x more than too-high"""
+    def __init__(self,cfg:NordConfig):
+        super().__init__(); self.target=cfg.target_spike_rate
+        self.weight=cfg.spike_loss_weight
+        self.min_rate=0.01  # absolute minimum — below this = dead layer
+    def forward(self,spike_tensors:List[Tensor])->Tensor:
+        if not spike_tensors: return torch.tensor(0.0)
+        loss=torch.tensor(0.0,device=spike_tensors[0].device,dtype=torch.float32)
+        for s in spike_tensors:
+            # FIX K: Only count non-negative values as spikes
+            rate=s.float().clamp(min=0).mean()
+            diff=self.target-rate
+            # Asymmetric: penalize too-low firing 3x more
+            if diff>0:
+                loss=loss+3.0*diff**2
+            else:
+                loss=loss+diff**2
+            # Anti-death penalty: heavy penalty if rate < min_rate
+            if rate<self.min_rate:
+                loss=loss+10.0*(self.min_rate-rate)**2
+        return self.weight*loss/len(spike_tensors)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# §10  STDP — FIX F: Bounded + Isolated
+# ═══════════════════════════════════════════════════════════════════════════════
+class STDPEngine:
+    def __init__(self,cfg:NordConfig):
+        self.cfg=cfg; self.a_plus=cfg.stdp_a_plus; self.a_minus=cfg.stdp_a_minus
+        self.tau_plus=cfg.stdp_tau_plus; self.tau_minus=cfg.stdp_tau_minus
+        self.w_max=cfg.stdp_w_max; self.w_min=cfg.stdp_w_min
+        self.reward_scale=cfg.stdp_reward_scale
+        self.allowed=set(cfg.stdp_layers or [])
+        self._loss_ema=10.0; self._ema_decay=0.99; self.max_update_norm=0.01
+
+    def update_reward(self,cl:float): self._loss_ema=self._ema_decay*self._loss_ema+(1-self._ema_decay)*cl
+    def _compute_reward(self,cl:float)->float:
+        return float(torch.sigmoid(torch.tensor((self._loss_ema-cl)*self.reward_scale)).item())
+    def is_allowed(self,name:str)->bool: return name in self.allowed
+
+    @torch.no_grad()
+    def compute_stdp_update(self,pre:Tensor,post:Tensor)->Tensor:
+        T=pre.shape[0]; d=pre.device
+        tp=torch.zeros_like(pre[0]); tpo=torch.zeros_like(post[0])
+        dp=math.exp(-1.0/self.tau_plus); dm=math.exp(-1.0/self.tau_minus)
+        dW=torch.zeros(post.shape[1],pre.shape[1],device=d,dtype=pre.dtype)
+        for t in range(T):
+            tp=tp*dp+pre[t]; tpo=tpo*dm+post[t]
+            if post[t].any(): dW+=self.a_plus*torch.outer(post[t],tp)
+            if pre[t].any(): dW-=self.a_minus*torch.outer(tpo,pre[t])
+        n=dW.norm()
+        if n>self.max_update_norm: dW=dW*(self.max_update_norm/n)
+        return dW
+
+    @torch.no_grad()
+    def apply_to_layer(self,layer:nn.Linear,pre:Tensor,post:Tensor,
+                       cl:Optional[float]=None,name:str=""):
+        if name and not self.is_allowed(name): return
+        if pre.dim()==3: pre=pre.mean(dim=1)
+        if post.dim()==3: post=post.mean(dim=1)
+        dW=self.compute_stdp_update(pre,post)
+        if cl is not None:
+            r=self._compute_reward(cl); dW=dW*(2.0*r-1.0); self.update_reward(cl)
+        o,i=layer.weight.shape; dW=dW[:o,:i]
+        layer.weight.data=(layer.weight.data+dW).clamp(self.w_min,self.w_max)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# §11  NORD MODEL v4.1
+# ═══════════════════════════════════════════════════════════════════════════════
+class NordModel(nn.Module):
+    def __init__(self,cfg:NordConfig):
+        super().__init__(); self.cfg=cfg
+        self.encoder=TemporalSpikeEncoder(cfg)
+        self.input_lif=AssociativeLIF(cfg.d_model,cfg,persistent=cfg.persistent_mem)
+        self.sensory_blocks=nn.ModuleList([NordBlock(cfg,i,False,zone="sensory") for i in range(cfg.sensory_layers)])
+        self.association_blocks=nn.ModuleList([NordBlock(cfg,cfg.sensory_layers+i,True,zone="association") for i in range(cfg.association_layers)])
+        self.memory_cortex=MemoryCortex(cfg)
+        self.executive_blocks=nn.ModuleList([NordBlock(cfg,cfg.sensory_layers+cfg.association_layers+i,False,zone="executive") for i in range(cfg.executive_layers)])
+        self.readout_lif=AssociativeLIF(cfg.d_model,cfg,persistent=cfg.persistent_mem)
+        self.readout_ema_raw=nn.Parameter(torch.tensor(1.4))
+        self.readout_norm=nn.LayerNorm(cfg.d_model)
+        self.lm_head=nn.Linear(cfg.d_model,cfg.vocab_size,bias=False)
+        self.stdp=STDPEngine(cfg); self._last_loss=None
+        self.spike_regulator=AuxiliarySpikeRegulator(cfg)
+
+    @property
+    def readout_ema_decay(self)->Tensor: return torch.sigmoid(self.readout_ema_raw)
+    def reset_state(self):
+        self.input_lif.reset_state(); self.readout_lif.reset_state()
+        self.memory_cortex.reset_state()
+
+    def forward(self,token_ids:Tensor,enable_stdp:bool=False)->Tuple[Tensor,Dict]:
+        B,S=token_ids.shape; T_t=self.cfg.T_total; D=self.cfg.d_model
+        cur=self.encoder(token_ids); isp,_=self.input_lif(cur)
+        isp=isp.reshape(T_t,B,S,D)
+        spike_ts=[isp]; stats={}; moe_lb=torch.tensor(0.0,device=token_ids.device)
+
+        x=isp
+        for i,bl in enumerate(self.sensory_blocks):
+            x,bs=bl(x); spike_ts.append(x)
+            for k,v in bs.items(): stats[f"sensory_{i}_{k}"]=v
+
+        for i,bl in enumerate(self.association_blocks):
+            x,bs=bl(x); spike_ts.append(x)
+            lb=bs.pop("moe_load_balance_loss",None)
+            if lb is not None: moe_lb=moe_lb+lb
+            for k,v in bs.items(): stats[f"assoc_{i}_{k}"]=v
+
+        x,ms=self.memory_cortex(x); stats.update(ms)
+
+        for i,bl in enumerate(self.executive_blocks):
+            x,bs=bl(x); spike_ts.append(x)
+            for k,v in bs.items(): stats[f"exec_{i}_{k}"]=v
+
+        xf=x.reshape(T_t,B*S,D); rsp,vm=self.readout_lif(xf)
+        a=self.readout_ema_decay
+        ema=torch.zeros(B*S,D,device=x.device,dtype=vm.dtype)
+        for t in range(T_t): ema=a*ema+(1-a)*vm[t]
+        vs=ema.reshape(B,S,D)
+        sm=rsp.mean(dim=0).reshape(B,S,D)
+        ro=vs+sm
+        xn=F.layer_norm(ro.float(),self.readout_norm.normalized_shape,
+            self.readout_norm.weight.float() if self.readout_norm.weight is not None else None,
+            self.readout_norm.bias.float() if self.readout_norm.bias is not None else None,
+            self.readout_norm.eps).to(ro.dtype)
+        logits=self.lm_head(xn)
+
+        out_rate=rsp.detach().mean().item()
+        # FIX K: clamp negatives — spike rates cannot be negative
+        sr=[s.detach().clamp(min=0).mean().item() for s in spike_ts]
+
+        # Convert ALL stats to tensors for DataParallel gather compatibility
+        dev = token_ids.device
+        tensor_stats = {}
+        tensor_stats["sparsity"] = torch.tensor(1.0 - out_rate, device=dev)
+        tensor_stats["avg_spike_rate"] = torch.tensor(sum(sr)/len(sr), device=dev)
+        tensor_stats["spike_loss"] = self.spike_regulator(spike_ts)
+        tensor_stats["moe_lb_loss"] = moe_lb
+        # Pack spike_rates as a single tensor
+        tensor_stats["spike_rates_tensor"] = torch.tensor(sr, device=dev)
+        # Convert any float stats from blocks/memory to tensors
+        for k, v in stats.items():
+            if isinstance(v, (int, float)):
+                tensor_stats[k] = torch.tensor(v, device=dev)
+            elif isinstance(v, torch.Tensor):
+                tensor_stats[k] = v.to(dev) if v.device != dev else v
+            # skip lists and other non-tensor types
+        return logits, tensor_stats
+
+    def set_last_loss(self,l:float): self._last_loss=l
+    def count_params(self)->str:
+        total=sum(p.numel() for p in self.parameters())
+        train=sum(p.numel() for p in self.parameters() if p.requires_grad)
+        se=sum(p.numel() for n,p in self.named_parameters() if 'sensory' in n)
+        a=sum(p.numel() for n,p in self.named_parameters() if 'association' in n)
+        m=sum(p.numel() for n,p in self.named_parameters() if 'memory' in n)
+        e=sum(p.numel() for n,p in self.named_parameters() if 'executive' in n)
+        return(f"Total: {total/1e6:.1f}M | Trainable: {train/1e6:.1f}M\n"
+               f"  Sensory:     {se/1e6:.1f}M ({self.cfg.sensory_layers} blocks)\n"
+               f"  Association: {a/1e6:.1f}M ({self.cfg.association_layers} blocks, MoE)\n"
+               f"  Memory:      {m/1e6:.1f}M\n"
+               f"  Executive:   {e/1e6:.1f}M ({self.cfg.executive_layers} blocks)")

nord_v4_700m-4.2/train_nord_700m.py

@@ -0,0 +1,644 @@
+"""
+╔══════════════════════════════════════════════════════════════════════════╗
+║         PROJECT NORD v4.2 — Training Script (700M)                     ║
+║                                                                        ║
+║  Usage:                                                                ║
+║      python train_nord_700m.py                                         ║
+║                                                                        ║
+║  v4.2 (700M) — Scaling test on L40/A100                                ║
+║      - Auxiliary spike loss (homeostatic regulation)                    ║
+║      - MoE routing stats (expert load, entropy)                        ║
+║      - Memory cortex monitoring                                        ║
+║      - Zone-aware logging (sensory/association/executive)               ║
+║      - Combined loss: L_total = L_CE + λ_spike * L_spike + λ_lb * L_lb║
+║                                                                        ║
+║  Hardware:                                                              ║
+║      - RTX 5070 (8GB) — batch=2, ~3GB VRAM                             ║
+║      - RTX 3090/4090 (24GB) — batch=4                                  ║
+║      - A100/L40 (48-80GB) — batch=8-16                                 ║
+╚══════════════════════════════════════════════════════════════════════════╝
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import shutil
+import struct
+import sys
+import time
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+from torch.amp import autocast
+from torch.utils.data import Dataset, DataLoader
+from torch.nn.parallel import DataParallel
+
+# Use local nord_core_700m (fix for gradient checkpoint + MoE metadata mismatch)
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from nord_core_700m import NordConfig, NordModel
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# TOKENIZER
+# ─────────────────────────────────────────────────────────────────────────────
+
+class NordTokenizer:
+    def __init__(self, cfg: NordConfig):
+        from transformers import AutoTokenizer
+
+        print(f"  [*] Loading Llama-3.2 tokenizer...", flush=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            cfg.tokenizer_id, trust_remote_code=True,
+        )
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
+        self.max_len = cfg.max_seq_len
+        self.vocab_size = self.tokenizer.vocab_size
+        if cfg.vocab_size < self.vocab_size:
+            cfg.vocab_size = self.vocab_size
+
+        print(f"  [✓] Tokenizer ready (vocab={self.vocab_size:,})", flush=True)
+
+    def encode(self, text: str) -> torch.Tensor:
+        enc = self.tokenizer(
+            text, return_tensors="pt",
+            max_length=self.max_len, truncation=True, padding="max_length",
+        )
+        return enc.input_ids
+
+    def decode(self, ids) -> str:
+        return self.tokenizer.decode(ids, skip_special_tokens=True)
+
+    @property
+    def pad_id(self) -> int:
+        return self.tokenizer.pad_token_id
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# LMDB DATASET
+# ─────────────────────────────────────────────────────────────────────────────
+
+class LMDBDataset(Dataset):
+    def __init__(self, db_path: str, max_seq_len: int):
+        import lmdb
+        self.db_path = db_path
+        self.max_seq_len = max_seq_len
+        self._env = None
+
+        env = lmdb.open(db_path, readonly=True, lock=False, readahead=False, meminit=False)
+        with env.begin(write=False) as txn:
+            raw = txn.get(b"__len__")
+            self.length = struct.unpack("<Q", raw)[0]
+        env.close()
+        print(f"  [✓] LMDB: {self.length:,} samples", flush=True)
+
+    def _get_env(self):
+        if self._env is None:
+            import lmdb
+            self._env = lmdb.open(
+                self.db_path, readonly=True, lock=False,
+                readahead=True, meminit=False, max_readers=64,
+            )
+        return self._env
+
+    def __len__(self): return self.length
+
+    def __getitem__(self, idx):
+        env = self._get_env()
+        with env.begin(write=False) as txn:
+            raw = txn.get(f"sample_{idx:010d}".encode())
+        ids = torch.frombuffer(bytearray(raw), dtype=torch.int32).long()
+        S = self.max_seq_len
+        return ids[:S] if ids.shape[0] >= S else F.pad(ids, (0, S - ids.shape[0]))
+
+
+def build_lmdb(jsonl_path: str, db_path: str, tokenizer: NordTokenizer,
+               max_seq_len: int, map_size_gb: float = 80.0):
+    import lmdb
+    import numpy as np
+
+    print(f"\n  [*] Building LMDB database (fast batch mode)...", flush=True)
+    print(f"      Source:  {jsonl_path}", flush=True)
+    print(f"      Target:  {db_path}", flush=True)
+
+    # Read all texts into memory
+    print(f"  [*] Reading JSONL into memory...", flush=True)
+    t0 = time.time()
+    texts = []
+    with open(jsonl_path, "r", encoding="utf-8") as f:
+        for i, line in enumerate(f):
+            if i % 1_000_000 == 0 and i > 0:
+                print(f"      read {i:,} lines...", flush=True)
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                obj = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            text = obj.get("text") or obj.get("content") or obj.get("passage", "")
+            if len(text) >= 30:
+                texts.append(text)
+
+    print(f"      {len(texts):,} valid texts in {time.time()-t0:.0f}s", flush=True)
+
+    # Batch tokenize
+    print(f"  [*] Batch tokenizing...", flush=True)
+    t1 = time.time()
+    BATCH = 1024
+    PAD_ID = tokenizer.pad_id
+
+    env = lmdb.open(db_path, map_size=int(map_size_gb * (1024**3)))
+    txn = env.begin(write=True)
+    count = 0
+    total_tokens = 0
+    total_batches = (len(texts) + BATCH - 1) // BATCH
+
+    for batch_idx in range(0, len(texts), BATCH):
+        batch = texts[batch_idx : batch_idx + BATCH]
+        batch_num = batch_idx // BATCH + 1
+
+        enc = tokenizer.tokenizer(
+            batch, max_length=max_seq_len, truncation=True,
+            padding="max_length", return_tensors="np",
+            return_attention_mask=False,
+        )
+        ids_np = enc.input_ids.astype(np.int32)
+
+        for j in range(ids_np.shape[0]):
+            row = ids_np[j]
+            non_pad = int(np.sum(row != PAD_ID))
+            if non_pad < 10:
+                continue
+            txn.put(f"sample_{count:010d}".encode(), row.tobytes())
+            count += 1
+            total_tokens += non_pad
+
+        if batch_num % 100 == 0 or batch_num == total_batches:
+            elapsed = time.time() - t1
+            pct = batch_num / total_batches * 100
+            eta = (elapsed / batch_num) * (total_batches - batch_num)
+            print(
+                f"      [{pct:5.1f}%] {count:,} samples | "
+                f"{total_tokens/1e6:.0f}M tok | ETA {eta:.0f}s",
+                flush=True,
+            )
+
+        if count % 500_000 < BATCH and count >= 500_000:
+            txn.commit()
+            txn = env.begin(write=True)
+
+    txn.put(b"__len__", struct.pack("<Q", count))
+    txn.put(b"__total_tokens__", struct.pack("<Q", total_tokens))
+    txn.commit()
+    env.close()
+
+    elapsed = time.time() - t1
+    print(f"\n  [✓] LMDB ready!", flush=True)
+    print(f"      Samples:  {count:,}", flush=True)
+    print(f"      Tokens:   {total_tokens:,} ({total_tokens/1e6:.1f}M)", flush=True)
+    print(f"      Time:     {elapsed:.0f}s ({elapsed/60:.1f} min)", flush=True)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# LR SCHEDULE
+# ─────────────────────────────────────────────────────────────────────────────
+
+def get_lr(step: int, cfg: NordConfig) -> float:
+    """Warmup → cosine decay to min_lr.
+    
+    Phase 1 (0 → warmup_steps): linear warmup from 0 → lr
+    Phase 2 (warmup_steps → max_steps): cosine decay from lr → min_lr
+    """
+    if step < cfg.warmup_steps:
+        return cfg.lr * (step + 1) / cfg.warmup_steps
+    
+    # Cosine decay phase
+    decay_steps = cfg.max_steps - cfg.warmup_steps
+    progress = (step - cfg.warmup_steps) / max(decay_steps, 1)
+    progress = min(progress, 1.0)  # clamp at 1.0
+    
+    # Cosine annealing: lr → min_lr
+    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+    return cfg.min_lr + (cfg.lr - cfg.min_lr) * cosine
+
+
+# ─────────────────────��───────────────────────────────────────────────────────
+# CHECKPOINT MANAGER
+# ─────────────────────────────────────────────────────────────────────────────
+
+class CheckpointManager:
+    def __init__(self, save_dir: str, keep_last: int = 5):
+        self.save_dir = Path(save_dir)
+        self.save_dir.mkdir(parents=True, exist_ok=True)
+        self.keep_last = keep_last
+
+    def save(self, model, optimizer, scaler, step, loss, cfg):
+        path = self.save_dir / f"nord_v4_step_{step:07d}.pt"
+        # Handle DataParallel: save inner model
+        model_to_save = model.module if hasattr(model, 'module') else model
+        torch.save({
+            "step": step, "loss": loss,
+            "version": "v4.1",
+            "model_state_dict": model_to_save.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "scaler_state_dict": scaler.state_dict(),
+            "config": {k: v for k, v in cfg.__dict__.items()
+                       if not k.startswith("_") and k != "dtype"},
+        }, path)
+
+        latest = self.save_dir / "nord_v4_latest.pt"
+        if latest.exists():
+            latest.unlink()
+        shutil.copy2(path, latest)
+
+        ckpts = sorted(self.save_dir.glob("nord_v4_step_*.pt"),
+                        key=lambda p: p.stat().st_mtime)
+        for old in ckpts[:max(0, len(ckpts) - self.keep_last)]:
+            old.unlink()
+
+        print(f"  [💾] Saved: {path.name} (loss={loss:.4f})", flush=True)
+
+    def load(self, model, optimizer, scaler, device) -> int:
+        latest = self.save_dir / "nord_v4_latest.pt"
+        if not latest.exists():
+            ckpts = sorted(self.save_dir.glob("nord_v4_step_*.pt"))
+            latest = ckpts[-1] if ckpts else None
+        if latest is None:
+            return 0
+
+        print(f"  [*] Resuming from: {latest.name}", flush=True)
+        ckpt = torch.load(latest, map_location=device, weights_only=False)
+        # Handle DataParallel: load into inner model
+        model_to_load = model.module if hasattr(model, 'module') else model
+        # Filter out persistent LIF state buffers — they resize with batch
+        state = ckpt["model_state_dict"]
+        filtered = {k: v for k, v in state.items()
+                    if "_v_mem_state" not in k and "_i_syn_state" not in k}
+        model_to_load.load_state_dict(filtered, strict=False)
+        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+        scaler.load_state_dict(ckpt["scaler_state_dict"])
+        step = ckpt["step"]
+        print(f"  [✓] Resumed at step {step:,} (loss={ckpt.get('loss', '?')})", flush=True)
+        return step
+
+    def save_final(self, model, cfg):
+        path = self.save_dir / "nord_v4_final.pt"
+        model_to_save = model.module if hasattr(model, 'module') else model
+        torch.save({
+            "version": "v4.1",
+            "model_state_dict": model_to_save.state_dict(),
+            "config": {k: v for k, v in cfg.__dict__.items()
+                       if not k.startswith("_") and k != "dtype"},
+        }, path)
+        print(f"  [⭐] Final model: {path}", flush=True)
+        return path
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# TRAINING
+# ─────────────────────────────────────────────────────────────────────────────
+
+# ─────────────────────────────────────────────────────────────────────────────
+# TRAINING
+# ─────────────────────────────────────────────────────────────────────────────
+
+def train(dataset_path: str, model_dir: str):
+    # ── Config — v4.2 (700M) ──
+    cfg = NordConfig(
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        dtype=torch.float16,
+
+        d_model=1536,
+        n_heads=24,
+        d_ff=4096,
+        n_clusters=128,
+        max_seq_len=192,  # 23.5GB без checkpoint: 384 OOM → 192 вміщається
+
+        sensory_layers=3,
+        association_layers=3,
+        executive_layers=4,
+
+        T=8,
+        T_slow=2,
+        persistent_mem=False,  # FIX для градієнтів
+
+        n_experts=4,
+        top_k_experts=2,
+
+        memory_size=256,
+        memory_tau_mem=0.99,
+        memory_n_read_heads=8,
+
+        target_spike_rate=0.03,
+        spike_loss_weight=0.5,
+
+        v_threshold=0.12,
+        tau_mem=0.9,
+        lif_freeze_steps=1000,
+
+        gradient_checkpointing=False,  # OFF: MoE+SNN несумісні з checkpoint (219 vs 220 metadata)
+
+        batch_size=1,
+        grad_accum=64,
+        lr=2e-4,
+        warmup_steps=1000,
+        max_steps=50_000,
+        save_every=1000,
+        log_every=10,
+    )
+    
+
+    print(flush=True)
+    print("═" * 60, flush=True)
+    print("  PROJECT NORD v4.2 — 700M SNN Training", flush=True)
+    print("═" * 60, flush=True)
+
+    if torch.cuda.is_available():
+        n_gpus = torch.cuda.device_count()
+        print(f"  GPU:            {torch.cuda.get_device_name()}", flush=True)
+        vram = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+        print(f"  VRAM:           {vram:.1f} GB" + (f" × {n_gpus} GPUs" if n_gpus > 1 else ""), flush=True)
+
+        # Auto-adjust batch size for 700M SNN
+        if vram >= 80:
+            cfg.batch_size = 8
+            cfg.grad_accum = 4
+            print(f"  [Auto] batch=8, accum=4 (A100 80GB)", flush=True)
+        elif vram >= 40:
+            cfg.batch_size = 4
+            cfg.grad_accum = 8
+            print(f"  [Auto] batch=4, accum=8 (L40 48GB)", flush=True)
+        elif vram >= 20:
+            cfg.batch_size = 1
+            cfg.grad_accum = 32
+            print(f"  [Auto] batch=1, accum=32 (24GB VRAM)", flush=True)
+        else:
+            print(f"  [ERROR] 700M model needs minimum 24GB VRAM!", flush=True)
+            sys.exit(1)
+    else:
+        print("  CPU mode (not recommended!)", flush=True)
+
+    print(f"  Architecture:   d={cfg.d_model}, heads={cfg.n_heads}, clusters={cfg.n_clusters}", flush=True)
+    print(f"  Zones:          Sensory({cfg.sensory_layers}) → Association({cfg.association_layers},MoE) → Memory → Executive({cfg.executive_layers})", flush=True)
+    print(f"  MoE:            {cfg.n_experts} experts, top-{cfg.top_k_experts}", flush=True)
+    print(f"  Memory:         {cfg.memory_size} neurons (τ={cfg.memory_tau_mem})", flush=True)
+    print(f"  Spike target:   {cfg.target_spike_rate:.0%} firing rate (λ={cfg.spike_loss_weight})", flush=True)
+    print(f"  Effective batch: {cfg.batch_size} × {cfg.grad_accum} = {cfg.batch_size * cfg.grad_accum}", flush=True)
+    print(f"  LR:             {cfg.lr} → {cfg.min_lr} (cosine decay, {cfg.warmup_steps} warmup)", flush=True)
+    print(f"  Max steps:      {cfg.max_steps:,}", flush=True)
+    print(f"  Dataset:        {dataset_path}", flush=True)
+    print(f"  Model dir:      {model_dir}", flush=True)
+    print(flush=True)
+
+    # ── Tokenizer ──
+    tokenizer = NordTokenizer(cfg)
+
+    # ── LMDB ──
+    db_path = str(Path(dataset_path).with_suffix("")) + "_lmdb"
+    if not Path(db_path).exists():
+        build_lmdb(dataset_path, db_path, tokenizer, cfg.max_seq_len)
+
+    dataset = LMDBDataset(db_path, cfg.max_seq_len)
+    dataloader = DataLoader(
+        dataset, batch_size=cfg.batch_size, shuffle=True,
+        num_workers=2, pin_memory=True, drop_last=True, persistent_workers=True,
+    )
+
+    # ── Model ──
+    print(f"\n  [*] Building Nord v4 model...", flush=True)
+    model = NordModel(cfg).to(cfg.device)
+    print(f"  [✓] {model.count_params()}", flush=True)
+
+    # ── Multi-GPU support ──
+    n_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0
+    if n_gpus > 1:
+        print(f"  [⚡] {n_gpus} GPUs detected! Using DataParallel", flush=True)
+        for i in range(n_gpus):
+            name = torch.cuda.get_device_name(i)
+            vram_i = torch.cuda.get_device_properties(i).total_memory / (1024**3)
+            print(f"       GPU {i}: {name} ({vram_i:.1f} GB)", flush=True)
+        model = DataParallel(model)
+        # Scale batch size by number of GPUs
+        cfg.batch_size = cfg.batch_size * n_gpus
+        cfg.grad_accum = max(1, cfg.grad_accum // n_gpus)
+        print(f"  [Auto] Scaled: batch={cfg.batch_size}, accum={cfg.grad_accum} "
+              f"(effective={cfg.batch_size * cfg.grad_accum})", flush=True)
+
+    # ── Gradient checkpointing for OOM prevention ──
+    if cfg.gradient_checkpointing:
+        print(f"  [*] Gradient checkpointing: ON (saves VRAM)", flush=True)
+
+    if torch.cuda.is_available():
+        allocated = torch.cuda.memory_allocated() / (1024**3)
+        print(f"  [*] Model VRAM: {allocated:.2f} GB", flush=True)
+
+    # ── Optimizer ──
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=cfg.lr,
+        weight_decay=cfg.weight_decay, betas=(0.9, 0.95),
+    )
+    scaler = torch.amp.GradScaler("cuda", enabled=(cfg.dtype == torch.float16))
+
+    # ── Checkpoints ──
+    ckpt_mgr = CheckpointManager(model_dir)
+    start_step = ckpt_mgr.load(model, optimizer, scaler, cfg.device)
+
+    # ── Training loop ──
+    model.train()
+    data_iter = iter(dataloader)
+    running_loss = 0.0
+    running_spike_loss = 0.0
+    tokens_seen = 0
+    t_start = time.time()
+
+    print(f"\n  {'─' * 55}", flush=True)
+    print(f"  Starting from step {start_step:,}  |  {len(dataset):,} samples", flush=True)
+    print(f"  Ctrl+C = stop (model will be saved!)", flush=True)
+    print(f"  {'─' * 55}\n", flush=True)
+
+    try:
+        for step in range(start_step, cfg.max_steps):
+            accum_loss = 0.0
+            accum_spike_loss = 0.0
+            stats = {}
+
+            for _ in range(cfg.grad_accum):
+                try:
+                    input_ids = next(data_iter)
+                except StopIteration:
+                    data_iter = iter(dataloader)
+                    input_ids = next(data_iter)
+
+                input_ids = input_ids.to(cfg.device, non_blocking=True)
+
+                with autocast(device_type="cuda", dtype=torch.float16,
+                              enabled=(cfg.dtype == torch.float16)):
+                    logits, stats = model(input_ids)
+
+                    shift_logits = logits[:, :-1, :].contiguous()
+                    shift_labels = input_ids[:, 1:].contiguous()
+
+                    # Main loss: cross entropy
+                    ce_loss = F.cross_entropy(
+                        shift_logits.reshape(-1, cfg.vocab_size),
+                        shift_labels.reshape(-1),
+                        ignore_index=tokenizer.pad_id,
+                    )
+
+                    # v4.1: Auxiliary spike loss
+                    spike_loss = stats.get("spike_loss", torch.tensor(0.0))
+                    if isinstance(spike_loss, torch.Tensor):
+                        spike_loss = spike_loss.to(ce_loss.device)
+                    else:
+                        spike_loss = torch.tensor(0.0, device=ce_loss.device)
+
+                    # v4.1: MoE load balance loss
+                    moe_lb_loss = stats.get("moe_lb_loss", torch.tensor(0.0))
+                    if isinstance(moe_lb_loss, torch.Tensor):
+                        moe_lb_loss = moe_lb_loss.to(ce_loss.device)
+                    else:
+                        moe_lb_loss = torch.tensor(0.0, device=ce_loss.device)
+
+                    # Combined loss: CE + spike homeostasis + MoE load balance
+                    loss = (ce_loss + spike_loss + 0.01 * moe_lb_loss) / cfg.grad_accum
+
+                scaler.scale(loss).backward()
+                accum_loss += ce_loss.item() / cfg.grad_accum
+                accum_spike_loss += spike_loss.item() / cfg.grad_accum
+                tokens_seen += input_ids.numel()
+
+            scaler.unscale_(optimizer)
+            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm)
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad(set_to_none=True)
+
+            # LR schedule
+            lr = get_lr(step, cfg)
+            for pg in optimizer.param_groups:
+                pg["lr"] = lr
+
+            running_loss += accum_loss
+            running_spike_loss += accum_spike_loss
+
+            if step % cfg.log_every == 0 and step > start_step:
+                avg = running_loss / cfg.log_every
+                avg_spike = running_spike_loss / cfg.log_every
+                elapsed = time.time() - t_start
+                tps = tokens_seen / elapsed / 1000 if elapsed > 0 else 0
+                sp = stats.get("sparsity", 0)
+
+                # VRAM monitoring
+                vram_used = ""
+                if torch.cuda.is_available():
+                    vram_gb = torch.cuda.memory_allocated() / (1024**3)
+                    vram_used = f" | VRAM {vram_gb:.1f}G"
+
+                # MoE routing info
+                moe_info = ""
+                entropy = stats.get("moe_route_entropy", None)
+                if entropy is not None:
+                    moe_info = f" | MoE H={entropy:.2f}"
+
+                # Memory info
+                mem_info = ""
+                mem_rate = stats.get("memory_spike_rate", None)
+                if mem_rate is not None:
+                    mem_info = f" | mem={mem_rate:.3f}"
+
+                print(
+                    f"  step {step:>7,} │ "
+                    f"loss {avg:.4f} │ "
+                    f"spike_L {avg_spike:.4f} │ "
+                    f"lr {lr:.1e} │ "
+                    f"grad {grad_norm:.1f} │ "
+                    f"sparsity {sp:.0%} │ "
+                    f"{tps:.1f}k tok/s"
+                    f"{moe_info}{mem_info}{vram_used}",
+                    flush=True,
+                )
+                running_loss = 0.0
+                running_spike_loss = 0.0
+
+            # Detailed stats every 100 steps
+            if step % 100 == 0 and step > start_step:
+                print(f"  {'·' * 50}", flush=True)
+                # Spike rates per zone
+                spike_rates = stats.get("spike_rates", [])
+                if spike_rates:
+                    s_rates = spike_rates[:cfg.sensory_layers + 1]
+                    a_rates = spike_rates[cfg.sensory_layers + 1:
+                                          cfg.sensory_layers + 1 + cfg.association_layers]
+                    e_rates = spike_rates[cfg.sensory_layers + 1 + cfg.association_layers:]
+
+                    print(f"    Sensory spike rates:     {[f'{r:.4f}' for r in s_rates]}", flush=True)
+                    print(f"    Association spike rates:  {[f'{r:.4f}' for r in a_rates]}", flush=True)
+                    print(f"    Executive spike rates:    {[f'{r:.4f}' for r in e_rates]}", flush=True)
+
+                # Expert load balance
+                loads = [stats.get(f"expert_{e}_load", 0) for e in range(cfg.n_experts)]
+                if any(l > 0 for l in loads):
+                    print(f"    Expert loads: {[f'{l:.2f}' for l in loads]}", flush=True)
+
+                # Memory stats
+                gate = stats.get("gate_activity", None)
+                mix = stats.get("memory_mix", None)
+                if gate is not None:
+                    print(f"    Memory gate={gate:.4f} mix={mix:.4f}", flush=True)
+
+                print(f"  {'·' * 50}", flush=True)
+
+            if step > 0 and step % cfg.save_every == 0:
+                ckpt_mgr.save(model, optimizer, scaler, step, accum_loss, cfg)
+
+    except KeyboardInterrupt:
+        print(f"\n\n  [⏸] Stopped at step {step:,}", flush=True)
+        ckpt_mgr.save(model, optimizer, scaler, step, accum_loss, cfg)
+        print(f"  To resume — just run the script again.", flush=True)
+
+    ckpt_mgr.save_final(model, cfg)
+
+    print(f"\n  {'═' * 55}", flush=True)
+    print(f"  Training complete!", flush=True)
+    print(f"  Model saved in: {model_dir}", flush=True)
+    print(f"  {'═' * 55}", flush=True)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# ENTRY POINT
+# ─────────────────────────────────────────────────────────────────────────────
+
+def main():
+    print("=" * 60, flush=True)
+    print("  PROJECT NORD v4 — Brain-Inspired SNN Training", flush=True)
+    print("=" * 60, flush=True)
+
+    default_data = "train_data.jsonl"
+    print(f"\n  Dataset path? (JSONL file)", flush=True)
+    print(f"  (Enter = {default_data})", flush=True)
+    data_input = input("  Dataset: ").strip()
+    dataset_path = data_input if data_input else default_data
+
+    if not Path(dataset_path).exists():
+        print(f"\n  [✗] File not found: {dataset_path}", flush=True)
+        sys.exit(1)
+
+    default_model = "nord_v4_700m"
+    print(f"\n  Model save directory?", flush=True)
+    print(f"  (Enter = {default_model})", flush=True)
+    model_input = input("  Model dir: ").strip()
+    model_dir = model_input if model_input else default_model
+
+    train(dataset_path, model_dir)
+
+
+if __name__ == "__main__":
+    main()

nord_v4_700m-4.2/train_nord_tpu_700m.py

@@ -0,0 +1,487 @@
+"""
+╔══════════════════════════════════════════════════════════════════════════╗
+║         PROJECT NORD v4.2 — Training Script (700M)                     ║
+║                                                                        ║
+║  Usage:                                                                ║
+║      CUDA:  python train_nord_700m.py                                  ║
+║      TPU:   python train_nord_700m.py --tpu                            ║
+║                                                                        ║
+║  v4.2 (700M) — Supports CUDA GPU and Google Cloud TPU                  ║
+╚══════════════════════════════════════════════════════════════════════════╝
+"""
+
+from __future__ import annotations
+import argparse, json, math, os, shutil, struct, sys, time
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+
+# ── Backend globals ──
+USE_TPU = False
+xm = None
+
+def detect_backend(force_tpu=False):
+    global USE_TPU, xm
+    if force_tpu:
+        try:
+            import torch_xla.core.xla_model as _xm
+            xm = _xm; USE_TPU = True
+            print("  [✓] TPU backend: torch_xla loaded", flush=True); return
+        except ImportError:
+            print("  [!] --tpu but torch_xla not found, fallback CUDA", flush=True)
+    if torch.cuda.is_available():
+        print(f"  [✓] CUDA backend: {torch.cuda.get_device_name()}", flush=True)
+    else:
+        try:
+            import torch_xla.core.xla_model as _xm
+            xm = _xm; USE_TPU = True
+            print("  [✓] TPU backend (auto-detected)", flush=True)
+        except ImportError:
+            print("  [!] CPU mode (very slow!)", flush=True)
+
+def get_device():
+    if USE_TPU: return xm.xla_device()
+    if torch.cuda.is_available(): return torch.device("cuda")
+    return torch.device("cpu")
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from nord_core_700m import NordConfig, NordModel
+
+# ── Tokenizer ──
+class NordTokenizer:
+    def __init__(self, cfg):
+        from transformers import AutoTokenizer
+        print(f"  [*] Loading Llama-3.2 tokenizer...", flush=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer_id, trust_remote_code=True)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        self.max_len = cfg.max_seq_len
+        self.vocab_size = self.tokenizer.vocab_size
+        if cfg.vocab_size < self.vocab_size: cfg.vocab_size = self.vocab_size
+        print(f"  [✓] Tokenizer ready (vocab={self.vocab_size:,})", flush=True)
+    def encode(self, text):
+        return self.tokenizer(text, return_tensors="pt", max_length=self.max_len, truncation=True, padding="max_length").input_ids
+    def decode(self, ids): return self.tokenizer.decode(ids, skip_special_tokens=True)
+    @property
+    def pad_id(self): return self.tokenizer.pad_token_id
+
+# ── LMDB Dataset ──
+class LMDBDataset(Dataset):
+    def __init__(self, db_path, max_seq_len):
+        import lmdb
+        self.db_path = db_path; self.max_seq_len = max_seq_len; self._env = None
+        env = lmdb.open(db_path, readonly=True, lock=False, readahead=False, meminit=False)
+        with env.begin(write=False) as txn: self.length = struct.unpack("<Q", txn.get(b"__len__"))[0]
+        env.close()
+        print(f"  [✓] LMDB: {self.length:,} samples", flush=True)
+    def _get_env(self):
+        if self._env is None:
+            import lmdb
+            self._env = lmdb.open(self.db_path, readonly=True, lock=False, readahead=True, meminit=False, max_readers=64)
+        return self._env
+    def __len__(self): return self.length
+    def __getitem__(self, idx):
+        env = self._get_env()
+        with env.begin(write=False) as txn: raw = txn.get(f"sample_{idx:010d}".encode())
+        ids = torch.frombuffer(bytearray(raw), dtype=torch.int32).long()
+        S = self.max_seq_len
+        return ids[:S] if ids.shape[0] >= S else F.pad(ids, (0, S - ids.shape[0]))
+
+def build_lmdb(jsonl_path, db_path, tokenizer, max_seq_len, map_size_gb=80.0):
+    import lmdb, numpy as np
+    print(f"\n  [*] Building LMDB...", flush=True)
+    t0 = time.time(); texts = []
+    with open(jsonl_path, "r", encoding="utf-8") as f:
+        for i, line in enumerate(f):
+            if i % 1_000_000 == 0 and i > 0: print(f"      read {i:,} lines...", flush=True)
+            line = line.strip()
+            if not line: continue
+            try: obj = json.loads(line)
+            except: continue
+            text = obj.get("text") or obj.get("content") or obj.get("passage", "")
+            if len(text) >= 30: texts.append(text)
+    print(f"      {len(texts):,} texts in {time.time()-t0:.0f}s", flush=True)
+    t1 = time.time(); BATCH = 1024; PAD_ID = tokenizer.pad_id
+    env = lmdb.open(db_path, map_size=int(map_size_gb * (1024**3))); txn = env.begin(write=True)
+    count = 0; total_tokens = 0; total_batches = (len(texts) + BATCH - 1) // BATCH
+    for batch_idx in range(0, len(texts), BATCH):
+        batch = texts[batch_idx:batch_idx+BATCH]; batch_num = batch_idx // BATCH + 1
+        enc = tokenizer.tokenizer(batch, max_length=max_seq_len, truncation=True, padding="max_length", return_tensors="np", return_attention_mask=False)
+        ids_np = enc.input_ids.astype(np.int32)
+        for j in range(ids_np.shape[0]):
+            row = ids_np[j]; non_pad = int(np.sum(row != PAD_ID))
+            if non_pad < 10: continue
+            txn.put(f"sample_{count:010d}".encode(), row.tobytes()); count += 1; total_tokens += non_pad
+        if batch_num % 100 == 0 or batch_num == total_batches:
+            pct = batch_num / total_batches * 100
+            print(f"      [{pct:5.1f}%] {count:,} samples | {total_tokens/1e6:.0f}M tok", flush=True)
+        if count % 500_000 < BATCH and count >= 500_000: txn.commit(); txn = env.begin(write=True)
+    txn.put(b"__len__", struct.pack("<Q", count)); txn.put(b"__total_tokens__", struct.pack("<Q", total_tokens))
+    txn.commit(); env.close()
+    print(f"  [✓] LMDB: {count:,} samples, {total_tokens/1e6:.1f}M tokens in {time.time()-t1:.0f}s", flush=True)
+
+# ── LR Schedule ──
+def get_lr(step, cfg):
+    if step < cfg.warmup_steps: return cfg.lr * (step + 1) / cfg.warmup_steps
+    progress = min((step - cfg.warmup_steps) / max(cfg.max_steps - cfg.warmup_steps, 1), 1.0)
+    return cfg.min_lr + (cfg.lr - cfg.min_lr) * 0.5 * (1.0 + math.cos(math.pi * progress))
+
+# ── Checkpoint Manager ──
+class CheckpointManager:
+    def __init__(self, save_dir, keep_last=5):
+        self.save_dir = Path(save_dir); self.save_dir.mkdir(parents=True, exist_ok=True); self.keep_last = keep_last
+
+    def save(self, model, optimizer, step, loss, cfg, scaler=None):
+        path = self.save_dir / f"nord_v4_step_{step:07d}.pt"
+        m = model.module if hasattr(model, 'module') else model
+        d = {"step": step, "loss": loss, "version": "v4.2", "model_state_dict": m.state_dict(),
+             "optimizer_state_dict": optimizer.state_dict(),
+             "config": {k: v for k, v in cfg.__dict__.items() if not k.startswith("_") and k != "dtype"}}
+        if scaler: d["scaler_state_dict"] = scaler.state_dict()
+        if USE_TPU: xm.save(d, str(path))
+        else: torch.save(d, path)
+        latest = self.save_dir / "nord_v4_latest.pt"
+        if latest.exists(): latest.unlink()
+        shutil.copy2(path, latest)
+        ckpts = sorted(self.save_dir.glob("nord_v4_step_*.pt"), key=lambda p: p.stat().st_mtime)
+        for old in ckpts[:max(0, len(ckpts) - self.keep_last)]: old.unlink()
+        print(f"  [💾] Saved: {path.name} (loss={loss:.4f})", flush=True)
+
+    def load(self, model, optimizer, device, scaler=None):
+        latest = self.save_dir / "nord_v4_latest.pt"
+        if not latest.exists():
+            ckpts = sorted(self.save_dir.glob("nord_v4_step_*.pt"))
+            latest = ckpts[-1] if ckpts else None
+        if latest is None: return 0
+        print(f"  [*] Resuming from: {latest.name}", flush=True)
+        ckpt = torch.load(latest, map_location="cpu", weights_only=False)
+        m = model.module if hasattr(model, 'module') else model
+        filtered = {k: v for k, v in ckpt["model_state_dict"].items() if "_v_mem_state" not in k and "_i_syn_state" not in k}
+        m.load_state_dict(filtered, strict=False)
+        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+        if scaler and "scaler_state_dict" in ckpt: scaler.load_state_dict(ckpt["scaler_state_dict"])
+        print(f"  [✓] Resumed at step {ckpt['step']:,} (loss={ckpt.get('loss', '?')})", flush=True)
+        return ckpt["step"]
+
+    def save_final(self, model, cfg):
+        path = self.save_dir / "nord_v4_final.pt"
+        m = model.module if hasattr(model, 'module') else model
+        d = {"version": "v4.2", "model_state_dict": m.state_dict(),
+             "config": {k: v for k, v in cfg.__dict__.items() if not k.startswith("_") and k != "dtype"}}
+        if USE_TPU: xm.save(d, str(path))
+        else: torch.save(d, path)
+        print(f"  [⭐] Final model: {path}", flush=True)
+
+# ── Training ──
+def train(dataset_path, model_dir, lr_override=None, continued=False):
+    device = get_device()
+
+    # Determine LR: continued pretraining uses lower LR
+    base_lr = 2e-4
+    if continued:
+        base_lr = 5e-5
+        print("  [*] Continued pretraining mode: LR=5e-5, warmup=200", flush=True)
+    if lr_override is not None:
+        base_lr = lr_override
+        print(f"  [*] LR override: {base_lr}", flush=True)
+
+    warmup = 200 if continued else 1000
+
+    cfg = NordConfig(
+        device=str(device), dtype=torch.bfloat16 if USE_TPU else torch.float16,
+        d_model=1536, n_heads=24, d_ff=4096, n_clusters=128, max_seq_len=192,
+        sensory_layers=3, association_layers=3, executive_layers=4,
+        T=8, T_slow=2, persistent_mem=False,
+        n_experts=4, top_k_experts=2,
+        memory_size=256, memory_tau_mem=0.99, memory_n_read_heads=8,
+        target_spike_rate=0.03, spike_loss_weight=0.5,
+        v_threshold=0.12, tau_mem=0.9, lif_freeze_steps=1000,
+        gradient_checkpointing=False,
+        batch_size=1, grad_accum=64, lr=base_lr, min_lr=1e-5,
+        warmup_steps=warmup, max_steps=50_000,
+        save_every=1000, log_every=10,
+    )
+
+    print(flush=True); print("═" * 60, flush=True)
+    print("  PROJECT NORD v4.2 — 700M SNN Training", flush=True); print("═" * 60, flush=True)
+
+    # ── Auto-adjust batch size ──
+    if USE_TPU:
+        print(f"  Device:    TPU ({device})", flush=True)
+        print(f"  Precision: bfloat16 (native)", flush=True)
+        cfg.batch_size = 8; cfg.grad_accum = 4
+        print(f"  [Auto] batch=8, accum=4 (TPU)", flush=True)
+    elif torch.cuda.is_available():
+        vram = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+        n_gpus = torch.cuda.device_count()
+        total_vram = vram * n_gpus
+        print(f"  GPU:  {torch.cuda.get_device_name()} ({vram:.1f}GB)" + (f" × {n_gpus} = {total_vram:.1f}GB total" if n_gpus > 1 else ""), flush=True)
+        if vram < 16:
+            print("  [ERROR] Need ≥16GB VRAM per GPU!", flush=True); sys.exit(1)
+
+    print(f"  Arch:      d={cfg.d_model}, h={cfg.n_heads}, ff={cfg.d_ff}", flush=True)
+    print(f"  Zones:     S({cfg.sensory_layers})→A({cfg.association_layers},MoE)→M→E({cfg.executive_layers})", flush=True)
+
+    tokenizer = NordTokenizer(cfg)
+    db_path = str(Path(dataset_path).with_suffix("")) + "_lmdb"
+    if not Path(db_path).exists(): build_lmdb(dataset_path, db_path, tokenizer, cfg.max_seq_len)
+    dataset = LMDBDataset(db_path, cfg.max_seq_len)
+
+    print(f"\n  [*] Building Nord v4 model...", flush=True)
+    model = NordModel(cfg).to(device)
+    print(f"  [✓] {model.count_params()}", flush=True)
+
+    # Multi-GPU (CUDA only)
+    n_gpus = 1
+    if not USE_TPU and torch.cuda.is_available() and torch.cuda.device_count() > 1:
+        from torch.nn.parallel import DataParallel
+        n_gpus = torch.cuda.device_count()
+        print(f"  [⚡] {n_gpus} GPUs → DataParallel", flush=True)
+        model = DataParallel(model)
+
+    # ── Smart VRAM auto-tuning: probe batch sizes to fill 85% VRAM ──
+    if not USE_TPU and torch.cuda.is_available():
+        TARGET_VRAM_PCT = 0.85  # Fill 85% of VRAM
+        EFF_BATCH_TARGET = 32   # Target effective batch size
+
+        vram_total = torch.cuda.get_device_properties(0).total_memory
+        vram_after_model = torch.cuda.memory_allocated()
+        vram_free = vram_total - vram_after_model
+        print(f"\n  [*] Smart VRAM auto-tuning...", flush=True)
+        print(f"      Total VRAM (per GPU): {vram_total/(1024**3):.1f}GB", flush=True)
+        print(f"      Model + optimizer:    {vram_after_model/(1024**3):.1f}GB", flush=True)
+        print(f"      Available:            {vram_free/(1024**3):.1f}GB", flush=True)
+        print(f"      Target fill:          {TARGET_VRAM_PCT:.0%}", flush=True)
+
+        # Probe increasing batch sizes with a dummy forward+backward
+        best_batch = 1
+        test_seq_len = cfg.max_seq_len
+        model.train()
+
+        # Create temporary optimizer for probing
+        temp_optim = torch.optim.AdamW(model.parameters(), lr=1e-4)
+        temp_scaler = torch.amp.GradScaler("cuda", enabled=(cfg.dtype == torch.float16))
+
+        for test_batch in [1, 2, 3, 4, 6, 8, 10, 12, 16]:
+            # For DataParallel, total batch = test_batch, split across GPUs
+            # Each GPU gets test_batch // n_gpus, need at least 1 per GPU
+            per_gpu = test_batch // n_gpus if n_gpus > 1 else test_batch
+            if per_gpu < 1: continue
+
+            torch.cuda.empty_cache()
+            torch.cuda.reset_peak_memory_stats()
+
+            try:
+                dummy_ids = torch.randint(0, 1000, (test_batch, test_seq_len), device=device)
+                with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=(cfg.dtype == torch.float16)):
+                    logits, stats = model(dummy_ids)
+                    loss = logits[:, :-1, :].contiguous().reshape(-1, cfg.vocab_size).mean()
+                temp_scaler.scale(loss).backward()
+                temp_scaler.unscale_(temp_optim)
+                temp_scaler.step(temp_optim)
+                temp_scaler.update()
+                temp_optim.zero_grad(set_to_none=True)
+
+                peak = torch.cuda.max_memory_allocated()
+                pct = peak / vram_total
+                print(f"      batch={test_batch:>2} → peak {peak/(1024**3):.1f}GB ({pct:.0%})", flush=True)
+
+                if pct <= TARGET_VRAM_PCT:
+                    best_batch = test_batch
+                else:
+                    # Exceeded target, stop probing
+                    break
+            except RuntimeError as e:
+                if "out of memory" in str(e).lower():
+                    torch.cuda.empty_cache()
+                    print(f"      batch={test_batch:>2} → OOM!", flush=True)
+                    break
+                else:
+                    raise
+
+        # Clean up probe state
+        del temp_optim, temp_scaler
+        torch.cuda.empty_cache()
+
+        # Re-initialize model weights since probe corrupted them
+        model_to_reinit = model.module if hasattr(model, 'module') else model
+        model_to_reinit.__init__(cfg)
+        model_to_reinit.to(device)
+        if hasattr(model, 'module'):
+            # Re-wrap in DataParallel
+            model = DataParallel(model_to_reinit)
+
+        cfg.batch_size = best_batch
+        cfg.grad_accum = max(1, EFF_BATCH_TARGET // best_batch)
+        eff = cfg.batch_size * cfg.grad_accum
+
+        print(f"\n  [✓] Auto-tuned: batch={cfg.batch_size}, accum={cfg.grad_accum}, effective={eff}", flush=True)
+        print(f"      VRAM utilization: ~{TARGET_VRAM_PCT:.0%} target", flush=True)
+
+    # Rebuild dataloader with tuned batch size
+    if USE_TPU:
+        dataloader = DataLoader(dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=4, drop_last=True)
+    else:
+        dataloader = DataLoader(dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=2, pin_memory=True, drop_last=True, persistent_workers=True)
+
+    print(f"  Eff batch: {cfg.batch_size}×{cfg.grad_accum}={cfg.batch_size*cfg.grad_accum}", flush=True)
+    print(f"  LR:        {cfg.lr}→{cfg.min_lr} (cosine, {cfg.warmup_steps} warmup)", flush=True)
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay, betas=(0.9, 0.95))
+    scaler = None
+    if not USE_TPU and cfg.dtype == torch.float16:
+        scaler = torch.amp.GradScaler("cuda", enabled=True)
+
+    ckpt_mgr = CheckpointManager(model_dir)
+    start_step = ckpt_mgr.load(model, optimizer, device, scaler)
+
+    model.train(); data_iter = iter(dataloader)
+    running_loss = 0.0; running_spike_loss = 0.0; tokens_seen = 0; t_start = time.time()
+
+    print(f"\n  {'─'*55}", flush=True)
+    print(f"  Start step {start_step:,} | {len(dataset):,} samples | {'TPU' if USE_TPU else 'CUDA'}", flush=True)
+    print(f"  Ctrl+C = save & stop", flush=True)
+    print(f"  {'─'*55}\n", flush=True)
+
+    try:
+        for step in range(start_step, cfg.max_steps):
+            accum_loss = 0.0; accum_spike_loss = 0.0; stats = {}
+            for _ in range(cfg.grad_accum):
+                try: input_ids = next(data_iter)
+                except StopIteration: data_iter = iter(dataloader); input_ids = next(data_iter)
+                input_ids = input_ids.to(device)
+
+                if USE_TPU:
+                    with torch.autocast(device_type="xla", dtype=torch.bfloat16):
+                        logits, stats = model(input_ids)
+                        ce_loss = F.cross_entropy(logits[:, :-1, :].contiguous().reshape(-1, cfg.vocab_size),
+                                                  input_ids[:, 1:].contiguous().reshape(-1), ignore_index=tokenizer.pad_id)
+                        spike_loss = stats.get("spike_loss", torch.tensor(0.0, device=device))
+                        if isinstance(spike_loss, torch.Tensor):
+                            if spike_loss.dim() > 0: spike_loss = spike_loss.mean()
+                        else:
+                            spike_loss = torch.tensor(float(spike_loss), device=device)
+                        moe_lb = stats.get("moe_lb_loss", torch.tensor(0.0, device=device))
+                        if isinstance(moe_lb, torch.Tensor):
+                            if moe_lb.dim() > 0: moe_lb = moe_lb.mean()
+                        else:
+                            moe_lb = torch.tensor(float(moe_lb), device=device)
+                        loss = (ce_loss + spike_loss + 0.01 * moe_lb) / cfg.grad_accum
+                    loss.backward()
+                else:
+                    with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=(cfg.dtype == torch.float16)):
+                        logits, stats = model(input_ids)
+                        ce_loss = F.cross_entropy(logits[:, :-1, :].contiguous().reshape(-1, cfg.vocab_size),
+                                                  input_ids[:, 1:].contiguous().reshape(-1), ignore_index=tokenizer.pad_id)
+                        spike_loss = stats.get("spike_loss", torch.tensor(0.0, device=device))
+                        if isinstance(spike_loss, torch.Tensor):
+                            if spike_loss.dim() > 0: spike_loss = spike_loss.mean()
+                        else:
+                            spike_loss = torch.tensor(float(spike_loss), device=device)
+                        moe_lb = stats.get("moe_lb_loss", torch.tensor(0.0, device=device))
+                        if isinstance(moe_lb, torch.Tensor):
+                            if moe_lb.dim() > 0: moe_lb = moe_lb.mean()
+                        else:
+                            moe_lb = torch.tensor(float(moe_lb), device=device)
+                        loss = (ce_loss + spike_loss + 0.01 * moe_lb) / cfg.grad_accum
+                    scaler.scale(loss).backward()
+
+                accum_loss += ce_loss.item() / cfg.grad_accum
+                sp_item = spike_loss.item() if isinstance(spike_loss, torch.Tensor) else float(spike_loss)
+                accum_spike_loss += sp_item / cfg.grad_accum
+                tokens_seen += input_ids.numel()
+
+            # Optimizer step
+            if USE_TPU:
+                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm)
+                xm.optimizer_step(optimizer)
+                optimizer.zero_grad(set_to_none=True)
+            else:
+                scaler.unscale_(optimizer)
+                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm)
+                scaler.step(optimizer); scaler.update()
+                optimizer.zero_grad(set_to_none=True)
+
+            lr = get_lr(step, cfg)
+            for pg in optimizer.param_groups: pg["lr"] = lr
+            running_loss += accum_loss; running_spike_loss += accum_spike_loss
+
+            # Logging
+            if step % cfg.log_every == 0 and step > start_step:
+                avg = running_loss / cfg.log_every; avg_sp = running_spike_loss / cfg.log_every
+                tps = tokens_seen / (time.time() - t_start) / 1000
+                # Handle both tensor and float stats (DataParallel returns averaged tensors)
+                sp = stats.get("sparsity", 0)
+                if isinstance(sp, torch.Tensor): sp = sp.mean().item()
+                mem_r = stats.get("memory_spike_rate", None)
+                if isinstance(mem_r, torch.Tensor): mem_r = mem_r.mean().item()
+                mem_s = f" | mem={mem_r:.3f}" if mem_r is not None else ""
+                gn = grad_norm.item() if isinstance(grad_norm, torch.Tensor) else grad_norm
+                dev = " | TPU" if USE_TPU else (f" | VRAM {torch.cuda.memory_allocated()/(1024**3):.1f}G" if torch.cuda.is_available() else "")
+                print(f"  step {step:>7,} │ loss {avg:.4f} │ spike_L {avg_sp:.4f} │ lr {lr:.1e} │ grad {gn:.1f} │ sparsity {sp:.0%} │ {tps:.1f}k tok/s{mem_s}{dev}", flush=True)
+                running_loss = 0.0; running_spike_loss = 0.0
+
+            if step % 100 == 0 and step > start_step:
+                print(f"  {'·'*50}", flush=True)
+                # Handle spike_rates as tensor (DataParallel) or list
+                sr = stats.get("spike_rates_tensor", stats.get("spike_rates", []))
+                if isinstance(sr, torch.Tensor):
+                    sr = sr.float()
+                    if sr.dim() > 1: sr = sr.mean(dim=0)  # average across DataParallel replicas
+                    sr = sr.tolist()
+                if sr:
+                    ns = cfg.sensory_layers + 1; na = cfg.association_layers
+                    print(f"    Sensory spike rates:     {[f'{r:.4f}' for r in sr[:ns]]}", flush=True)
+                    print(f"    Association spike rates:  {[f'{r:.4f}' for r in sr[ns:ns+na]]}", flush=True)
+                    print(f"    Executive spike rates:    {[f'{r:.4f}' for r in sr[ns+na:]]}", flush=True)
+                gate = stats.get("gate_activity"); mix = stats.get("memory_mix")
+                if isinstance(gate, torch.Tensor): gate = gate.mean().item()
+                if isinstance(mix, torch.Tensor): mix = mix.mean().item()
+                if gate is not None: print(f"    Memory gate={gate:.4f} mix={mix:.4f}", flush=True)
+                print(f"  {'·'*50}", flush=True)
+
+            if step > 0 and step % cfg.save_every == 0:
+                ckpt_mgr.save(model, optimizer, step, accum_loss, cfg, scaler)
+
+    except KeyboardInterrupt:
+        print(f"\n\n  [⏸] Stopped at step {step:,}", flush=True)
+        ckpt_mgr.save(model, optimizer, step, accum_loss, cfg, scaler)
+
+    ckpt_mgr.save_final(model, cfg)
+    print(f"\n  {'═'*55}\n  Training complete! Model: {model_dir}\n  {'═'*55}", flush=True)
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--tpu", action="store_true", help="Force TPU backend")
+    parser.add_argument("--dataset", type=str, default=None)
+    parser.add_argument("--model_dir", type=str, default=None)
+    parser.add_argument("--lr", type=float, default=None, help="Override learning rate (e.g. 5e-5 for continued pretraining)")
+    parser.add_argument("--continued", action="store_true", help="Continued pretraining mode: auto LR=5e-5, shorter warmup")
+    args = parser.parse_args()
+
+    print("=" * 60, flush=True)
+    print("  PROJECT NORD v4.2 — Brain-Inspired SNN Training", flush=True)
+    print("=" * 60, flush=True)
+    detect_backend(force_tpu=args.tpu)
+
+    if args.dataset: dataset_path = args.dataset
+    else:
+        d = "train_data.jsonl"
+        print(f"\n  Dataset? (Enter = {d})", flush=True)
+        inp = input("  Dataset: ").strip(); dataset_path = inp if inp else d
+    if not Path(dataset_path).exists(): print(f"  [✗] Not found: {dataset_path}", flush=True); sys.exit(1)
+
+    if args.model_dir: model_dir = args.model_dir
+    else:
+        d = "nord_v4_700m"
+        print(f"\n  Model dir? (Enter = {d})", flush=True)
+        inp = input("  Model dir: ").strip(); model_dir = inp if inp else d
+
+    train(dataset_path, model_dir, lr_override=args.lr, continued=args.continued)
+
+if __name__ == "__main__":
+    main()