CEDL.py

#!/usr/bin/env python3
"""
CEDL: A Hippocampal-Inspired Architecture for Language Modeling
===============================================================

Reference implementation for the CEDL architecture at ~100M parameter scale,
trained on WikiText-103 and evaluated via perplexity and LAMBADA accuracy.

Architecture:
    C-Stage (EC)  — Grid-cell periodic attention with per-layer AdaLN
    E-Stage (DG)  — Expansion + top-k with AHSD and CSR pattern separation
    D-Stage (CA3) — Dual memory: attractor refinement + 256-slot memory bank
    L-Stage (CA1) — Two-channel comparator with learned fusion gate
    Feedback      — Loop 1 (L→C via AdaLN), Loop 2 (D→E additive)
    NeuromodGate  — Output-level gain modulation

Baselines (parameter-matched):
    Transformer, Transformer-XL, RetNet, Mamba, LSTM

Usage:
    python cedl_release.py --model CEDL --batch-size 4 --max-steps 30000
    python cedl_release.py --model all --batch-size 8 --max-steps 30000
    python cedl_release.py --model all --eval-only

Paper: "CEDL: A Hippocampal-Inspired Architecture for Advancing LLMs"
Author: Dian Jiao (University of Pennsylvania)
License: MIT
"""

from __future__ import annotations

import argparse
import math
import os
import random
import sys
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset


LOGIT_SCALAR_MODES = frozenset({
    "v3", "v3m", "v3m2", "v3m2_nll", "v4c", "v4c_gate", "v4d",
})


@dataclass
class Config:
    """Hyperparameters for 100M-scale training."""
    dataset: str = "wikitext103"
    vocab_size: int = 50257
    max_seq: int = 1024
    batch_size: int = 4
    grad_accum: int = 4
    lr: float = 3e-4
    min_lr: float = 1e-5
    warmup_steps: int = 1000
    max_steps: int = 30_000
    weight_decay: float = 0.1
    grad_clip: float = 1.0
    bfloat16: bool = True
    eval_interval: int = 1000
    eval_steps: int = 200
    save_interval: int = 5000
    save_dir: str = "checkpoints_100m"
    tpu: bool = False
    seed: int = 42
    resume_checkpoint: Optional[str] = None
    resume_start_step: int = 0
    save_trainstate: bool = False
    run_until_step: Optional[int] = None
    v5_grad_isolation: bool = False
    v5_trunk_aux_frac: float = 0.0
    v5_aux_scale_early: float = 0.50
    v5_aux_scale_mid:   float = 0.25
    v5_aux_scale_late:  float = 0.10
    v5_cadence_early:   int   = 8
    v5_cadence_mid:     int   = 16
    v5_cadence_late:    int   = 32
    v5_family_reweight: bool = False
    v6_mixture: bool = False
    v6_lambda_init: float = -4.0
    v6_lambda_a_init: float = 1.0
    v6_aux_weight: float = 0.0
    v6_margin_target: float = 1.0
    v6_mix_weight: float = 0.5
    v6_gate_weight: float = 0.01
    v6_lambda_floor: float = 0.05
    v6_lambda_head: bool = False
    v6_lambda_head_hidden: int = 160
    v6_lambda_head_bias_init: float = -7.0
    v6_bg_weight: float = 1.0
    v6_bg_target: float = 0.01
    v6_bce_objective: bool = False
    v6_sel_weight: float = 1.0
    v6_lambda_head_w_init_std: float = 1e-3
    v6_wt_sparsity_weight: float = 0.0
    v6_wt_sparsity_target: float = 0.0
    v6_mem_head_bank: bool = False
    v6_bank_ce_weight: float = 0.0
    v6_bank_pair_ce_weight: float = 0.0
    v6_bank_head_lr: float = 0.0
    v6_bank_query_source: str = "h_d"
    v6_bank_readout_mode: str = "bank"
    v6_source_adapter: bool = False
    v6_context_adapter: bool = False
    v6_specialist_from_b0: bool = False
    v6_specialist_freeze: str = "none"
    v6_specialist_noinject: bool = False

@dataclass
class CEDLConfig:
    d_model: int = 640
    n_heads: int = 10
    c_layers: int = 8
    ffn_dim: int = 2816
    e_expand: int = 4
    e_sparsity: float = 0.10
    d_refine: int = 3
    d_slots: int = 256
    dropout: float = 0.1
    n_feedback_iters: int = 1
    feedback_decay: float = 0.5
    feedback_warmup_start: int = 2000
    feedback_warmup_end: int = 5000
    sparsity_final: float = 0.05
    sparsity_anneal_frac: float = 0.80
    use_salience: bool = True
    salience_mode: str = "v0"
    lex_anchor_weight: float = 0.0
    v4c_variant: str = "base"
    v4c_z_dim: int = 64
    v4c_temperature: float = 0.05
    v4c_proj_hidden: int = 256
    v4c_nce_weight: float = 0.05
    v4c_bce_weight: float = 0.02
    v4c_norm_cap: float = 0.3
    v4c_every: int = 4
    v4c_batch_size: int = 32
    v4c_max_seq: int = 128
    v4c_warm_start_sigma: float = 0.05
    v4c_margin_weight: float = 0.0
    v4c_margin_target: float = 1.0
    v4d_w_sal_sigma: float = 0.02
    v4d_logit_cap:   float = 4.0
    v4d_causal_weight: float = 0.0
    v4d_causal_gap:    float = 0.25
    v4d_causal_z_weight: float = 0.0
    v4d_causal_z_gap:    float = 0.15
    v4d_w_sal_rank: int = 0
    v4d_swap_consistency_weight: float = 0.0
    v4d_role_sep_weight: float = 0.0
    v4d_noinject:    bool  = False
    v6_mixture: bool = False
    v6_lambda_init: float = -4.0
    v6_lambda_a_init: float = 1.0
    v6_aux_weight: float = 0.0
    v6_margin_target: float = 1.0
    v6_mix_weight: float = 0.5
    v6_gate_weight: float = 0.01
    v6_lambda_floor: float = 0.05
    v6_lambda_head: bool = False
    v6_lambda_head_hidden: int = 160
    v6_lambda_head_bias_init: float = -7.0
    v6_bg_weight: float = 1.0
    v6_bg_target: float = 0.01
    v6_bce_objective: bool = False
    v6_sel_weight: float = 1.0
    v6_lambda_head_w_init_std: float = 1e-3
    v6_wt_sparsity_weight: float = 0.0
    v6_wt_sparsity_target: float = 0.0
    v6_mem_head_bank: bool = False
    v6_bank_ce_weight: float = 0.0
    v6_bank_pair_ce_weight: float = 0.0
    v6_bank_query_source: str = "h_d"
    v6_bank_readout_mode: str = "bank"
    v6_source_adapter: bool = False
    v6_context_adapter: bool = False
    v4c_neg_weights: Dict[str, float] = field(default_factory=lambda: {
        "stale": 2.0, "distractor": 1.0, "neutral": 1.0})


@dataclass
class TransformerConfig:
    d_model: int = 768
    n_heads: int = 12
    n_layers: int = 9
    ffn_dim: int = 3072
    dropout: float = 0.1


@dataclass
class TransformerXLConfig:
    d_model: int = 768
    n_heads: int = 12
    n_layers: int = 9
    ffn_dim: int = 3072
    mem_len: int = 256
    dropout: float = 0.1


@dataclass
class RetNetConfig:
    d_model: int = 640
    n_heads: int = 10
    n_layers: int = 12
    ffn_dim: int = 2880
    dropout: float = 0.1


@dataclass
class MambaConfig:
    d_model: int = 768
    n_layers: int = 18
    d_state: int = 16
    d_conv: int = 4
    expand: int = 2
    dropout: float = 0.0


@dataclass
class LSTMConfig:
    d_model: int = 640
    hidden_size: int = 1760
    n_layers: int = 3
    dropout: float = 0.1



class MultiScaleRetention(nn.Module):
    """Parallel multi-scale retention (RetNet-style)."""

    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        assert d_model % n_heads == 0

        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_out = nn.Linear(d_model, d_model)
        self.out_gate_proj = nn.Linear(d_model, d_model)
        self.gn = nn.GroupNorm(n_heads, d_model)

        gammas = torch.linspace(0.80, 0.99, n_heads)
        self.gamma_log = nn.Parameter(
            torch.log(gammas / (1.0 - gammas))
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        K = self.n_heads
        dh = self.d_head

        Q = self.w_q(x).view(B, T, K, dh).transpose(1, 2)
        Kmat = self.w_k(x).view(B, T, K, dh).transpose(1, 2)
        V = self.w_v(x).view(B, T, K, dh).transpose(1, 2)

        gamma = torch.sigmoid(self.gamma_log)
        log_gamma = torch.log(gamma)
        positions = torch.arange(T, device=x.device, dtype=x.dtype)
        diff = (positions.unsqueeze(1) - positions.unsqueeze(0)).clamp(min=0)
        D = torch.exp(log_gamma.view(K, 1, 1) * diff.unsqueeze(0))
        causal = torch.tril(torch.ones(T, T, device=x.device, dtype=x.dtype))
        D = D * causal.unsqueeze(0)

        retention = (Q @ Kmat.transpose(-1, -2)) * D.unsqueeze(0)
        row_sum = retention.sum(dim=-1, keepdim=True).abs().clamp(min=1.0)
        retention = retention / row_sum
        out = retention @ V

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        out = self.gn(out.transpose(1, 2)).transpose(1, 2)
        gate = F.silu(self.out_gate_proj(x))
        out = self.w_out(self.dropout(out * gate))
        return out


class MultiScalePeriodicRetention(nn.Module):
    """Grid-cell-inspired periodic retention for CEDL C-stage.

    Replaces RetNet's exponential decay D[i,j] = γ^(i-j) with a damped
    periodic kernel D[i,j] = γ^(i-j) · cos(2π(i-j)/λ_k + φ_k), inspired
    by entorhinal cortex grid cells that fire at regular spatial intervals
    at multiple scales. No existing sequence model uses periodic attention
    weights — this is the key novelty of the C-stage.
    """

    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        assert d_model % n_heads == 0

        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_out = nn.Linear(d_model, d_model)
        self.out_gate_proj = nn.Linear(d_model, d_model)
        self.gn = nn.GroupNorm(n_heads, d_model)

        gammas = torch.linspace(0.85, 0.995, n_heads)
        self.gamma_log = nn.Parameter(
            torch.log(gammas / (1.0 - gammas)))

        self.log_lambda = nn.Parameter(
            torch.linspace(math.log(4.0), math.log(256.0), n_heads))
        self.phi = nn.Parameter(torch.zeros(n_heads))

        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        K = self.n_heads
        dh = self.d_head

        Q = self.w_q(x).view(B, T, K, dh).transpose(1, 2)
        Kmat = self.w_k(x).view(B, T, K, dh).transpose(1, 2)
        V = self.w_v(x).view(B, T, K, dh).transpose(1, 2)

        gamma = torch.sigmoid(self.gamma_log)
        log_gamma = torch.log(gamma)

        lam = torch.exp(self.log_lambda)
        positions = torch.arange(T, device=x.device, dtype=x.dtype)
        diff = positions.unsqueeze(1) - positions.unsqueeze(0)
        diff_pos = diff.clamp(min=0)

        decay = torch.exp(log_gamma.view(K, 1, 1) * diff_pos.unsqueeze(0))
        periodic = torch.cos(
            2 * math.pi * diff.unsqueeze(0) / lam.view(K, 1, 1)
            + self.phi.view(K, 1, 1))
        D = decay * periodic
        causal = torch.tril(torch.ones(T, T, device=x.device, dtype=x.dtype))
        D = D * causal.unsqueeze(0)

        retention = (Q @ Kmat.transpose(-1, -2)) * D.unsqueeze(0)
        row_sum = retention.abs().sum(dim=-1, keepdim=True).clamp(min=1.0)
        retention = retention / row_sum
        out = retention @ V

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        out = self.gn(out.transpose(1, 2)).transpose(1, 2)
        gate = F.silu(self.out_gate_proj(x))
        out = self.w_out(self.dropout(out * gate))
        return out


class PeriodicRetentionLayer(nn.Module):
    """Pre-norm periodic retention layer with hierarchical neuromodulatory gating.

    Each layer's LayerNorm is modulated by feedback from the L-stage
    mismatch detector via learned scale+shift (Adaptive LayerNorm).
    This distributes neuromodulation across ALL C-stage layers, matching
    how cholinergic projections from the medial septum innervate every
    level of the hippocampal circuit — not just a single output stage.

    Zero-initialized so modulation starts as identity (no warmup needed).
    The model learns to use feedback gradually through training.
    """

    def __init__(self, d_model: int, n_heads: int, ffn_dim: int,
                 dropout: float = 0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model, elementwise_affine=False)
        self.retention = MultiScalePeriodicRetention(d_model, n_heads, dropout)
        self.ln2 = nn.LayerNorm(d_model, elementwise_affine=False)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, ffn_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_dim, d_model),
            nn.Dropout(dropout),
        )
        bottleneck = d_model // 4
        self.neuro_proj = nn.Sequential(
            nn.Linear(d_model, bottleneck),
            nn.GELU(),
            nn.Linear(bottleneck, 4 * d_model),
        )
        nn.init.zeros_(self.neuro_proj[2].weight)
        nn.init.zeros_(self.neuro_proj[2].bias)

    def forward(self, x: torch.Tensor,
                feedback: Optional[torch.Tensor] = None) -> torch.Tensor:
        if feedback is not None:
            params = self.neuro_proj(feedback)
            s1, b1, s2, b2 = params.chunk(4, dim=-1)
            s1 = s1.clamp(-0.5, 0.5)
            b1 = b1.clamp(-0.5, 0.5)
            s2 = s2.clamp(-0.5, 0.5)
            b2 = b2.clamp(-0.5, 0.5)
            h = (1 + s1) * self.ln1(x) + b1
            x = x + self.retention(h)
            h = (1 + s2) * self.ln2(x) + b2
            x = x + self.ffn(h)
        else:
            x = x + self.retention(self.ln1(x))
            x = x + self.ffn(self.ln2(x))
        return x


class RetentionLayer(nn.Module):
    """Pre-norm retention layer with FFN."""

    def __init__(self, d_model: int, n_heads: int, ffn_dim: int,
                 dropout: float = 0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.retention = MultiScaleRetention(d_model, n_heads, dropout)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, ffn_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_dim, d_model),
            nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.retention(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x



class CStageRetention(nn.Module):
    """C-Stage: grid-cell-inspired periodic retention encoder."""

    def __init__(self, vocab: int, max_seq: int, d_model: int, n_heads: int,
                 n_layers: int, ffn_dim: int, dropout: float = 0.1):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab, d_model)
        self.pos_emb = nn.Embedding(max_seq, d_model)
        self.layers = nn.ModuleList([
            PeriodicRetentionLayer(d_model, n_heads, ffn_dim, dropout)
            for _ in range(n_layers)
        ])
        self.ln = nn.LayerNorm(d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, ids: torch.Tensor,
                feedback: Optional[torch.Tensor] = None) -> torch.Tensor:
        B, T = ids.shape
        pos = torch.arange(T, device=ids.device).unsqueeze(0)
        x = self.drop(self.tok_emb(ids) + self.pos_emb(pos))
        for layer in self.layers:
            x = layer(x, feedback=feedback)
        return self.ln(x)


class EStage(nn.Module):
    """DG: expansion + pattern separation via AHSD + CSR.

    Anti-Hebbian Support Drift (AHSD): Tracks per-neuron activation
    frequencies and suppresses overused neurons before top-k selection.
    Prevents mode collapse and ensures full use of the sparse code space.
    Inspired by homeostatic plasticity and intrinsic excitability regulation
    that maintain the DG's characteristically low population activity
    (~2-5% of granule cells active; Chawla et al. 2005).

    Contrastive Support Repulsion (CSR): Auxiliary loss that penalizes
    support overlap between similar inputs. The E-stage doesn't just
    sparsify — it actively maximizes code dissimilarity for similar
    inputs, which is the computational definition of pattern separation.
    """

    def __init__(self, d_model: int, expansion: int = 4,
                 sparsity: float = 0.10, ema_decay: float = 0.99,
                 inhibition_strength: float = 0.5):
        super().__init__()
        self.d_expand = d_model * expansion
        self.k = max(1, int(self.d_expand * sparsity))
        self.expand = nn.Linear(d_model, self.d_expand)
        self.contract = nn.Linear(self.d_expand, d_model)
        self.ln = nn.LayerNorm(d_model)

        self.register_buffer('neuron_freq', torch.zeros(self.d_expand))
        self.ema_decay = ema_decay
        self.inhibition_strength = inhibition_strength
        self.inhibition_temp = nn.Parameter(torch.tensor(1.0))

    def set_sparsity(self, sparsity: float):
        """Update sparsity level (for annealing during training)."""
        self.k = max(1, int(self.d_expand * sparsity))

    def forward(self, h_C: torch.Tensor,
                feedback: Optional[torch.Tensor] = None
                ) -> Tuple[torch.Tensor, torch.Tensor]:
        if feedback is not None:
            h_C = h_C + feedback
        expanded = F.gelu(self.expand(h_C))

        if self.training and self.neuron_freq.sum() > 0:
            penalty = (self.neuron_freq * self.inhibition_strength
                       * F.softplus(self.inhibition_temp))
            inhibited = expanded - penalty
        else:
            inhibited = expanded

        topk_vals, topk_idx = inhibited.topk(self.k, dim=-1)
        original_vals = expanded.gather(-1, topk_idx)
        sparse = torch.zeros_like(expanded)
        sparse.scatter_(-1, topk_idx, original_vals)

        if self.training:
            with torch.no_grad():
                batch_freq = (sparse != 0).float().mean(dim=(0, 1))
                self.neuron_freq.mul_(self.ema_decay).add_(
                    batch_freq, alpha=1 - self.ema_decay)

        contracted = self.contract(sparse)
        return self.ln(contracted + h_C), sparse

    @staticmethod
    def csr_loss(h_C: torch.Tensor, sparse: torch.Tensor,
                 n_sample: int = 128) -> torch.Tensor:
        """Contrastive Support Repulsion: penalize support overlap between
        similar inputs. This IS pattern separation as an explicit objective.

        Args:
            h_C: [B, T, d] pre-expansion input
            sparse: [B, T, d_expand] sparse codes
            n_sample: subsample tokens for O(n^2) efficiency
        """
        B, T, _ = h_C.shape
        if T > n_sample:
            idx = torch.randperm(T, device=h_C.device)[:n_sample]
            h_sub = h_C[:, idx]
            s_sub = sparse[:, idx]
        else:
            h_sub, s_sub = h_C, sparse
            n_sample = T

        h_norm = F.normalize(h_sub, dim=-1)
        input_sim = F.relu(torch.bmm(h_norm, h_norm.transpose(1, 2)))

        mask = (s_sub != 0).float()
        k = mask.sum(dim=-1, keepdim=True).clamp(min=1)
        mask_norm = mask / k.sqrt()
        overlap = torch.bmm(mask_norm, mask_norm.transpose(1, 2))

        eye = torch.eye(n_sample, device=h_C.device).unsqueeze(0)
        loss = ((input_sim * overlap) * (1 - eye)).sum()
        return loss / (B * n_sample * max(n_sample - 1, 1))


class LogitInjectingMHA(nn.Module):
    """Multi-head cross-attention with optional pre-softmax logit injection.

    Used by V4d to bypass the lossy query-perturbation geometry of V0/V3m/
    V3m2/V4c/V4c-M (where v_vec was added to the query). With this module,
    v_vec is projected to a per-slot bias `sal_bias [B, T, S]` and added
    DIRECTLY to the pre-softmax attention logits. Softmax cannot flatten
    logit-space perturbations — it exponentiates them — so small biases
    produce measurable retrieval-distribution shifts.

    Parameter names + shapes MATCH `nn.MultiheadAttention`'s standard
    parameterization (`in_proj_weight`, `in_proj_bias`, `out_proj.weight`,
    `out_proj.bias`) so legacy strict-load of an nn.MultiheadAttention
    state_dict would round-trip into this module (although V4d's gated
    conditional replacement means we don't rely on this in practice).

    Forward output is `(out, None)` to mirror `nn.MultiheadAttention`'s
    `(output, attn_weights)` return shape; callers `_, _ = self.cross_attn(
    q, k, v)` work unchanged.

    Side-channel storage (`_last_attn`, `_last_sal_bias`, scalar logit-stat
    fields) is gated on `store=True` per-call to avoid wasting GPU memory
    on the [B, H, T, S] attn tensor every forward. Probes and the V4d
    periodic logger toggle this flag externally before calling.
    """

    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % n_heads == 0, \
            f"d_model {d_model} not divisible by n_heads {n_heads}"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head  = d_model // n_heads
        self.scale   = self.d_head ** -0.5
        self.in_proj_weight = nn.Parameter(torch.empty(3 * d_model, d_model))
        self.in_proj_bias   = nn.Parameter(torch.empty(3 * d_model))
        self.out_proj       = nn.Linear(d_model, d_model, bias=True)
        nn.init.xavier_uniform_(self.in_proj_weight)
        nn.init.constant_(self.in_proj_bias, 0.0)
        self.dropout = nn.Dropout(dropout)
        self._last_attn              = None
        self._last_sal_bias          = None
        self._last_base_logits_std   = None
        self._last_attn_logits_std   = None
        self._last_sal_logits_std    = None

    def forward(self, q, k, v, sal_bias=None, store: bool = False):
        B, T, _ = q.shape
        S = k.shape[1]
        Wq, Wk, Wv = self.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = self.in_proj_bias.chunk(3, dim=0)
        Q = F.linear(q, Wq, bq).view(
            B, T, self.n_heads, self.d_head).transpose(1, 2)
        K = F.linear(k, Wk, bk).view(
            B, S, self.n_heads, self.d_head).transpose(1, 2)
        V = F.linear(v, Wv, bv).view(
            B, S, self.n_heads, self.d_head).transpose(1, 2)
        attn_logits = (Q @ K.transpose(-1, -2)) * self.scale
        if sal_bias is not None:
            attn_logits = attn_logits + sal_bias.unsqueeze(1)
        attn = F.softmax(attn_logits, dim=-1)
        if store:
            self._last_attn     = attn.detach()
            self._last_sal_bias = (sal_bias.detach()
                                    if sal_bias is not None else None)
            with torch.no_grad():
                base_logits = (Q @ K.transpose(-1, -2)) * self.scale
                self._last_base_logits_std = float(base_logits.std().item())
                self._last_attn_logits_std = float(attn_logits.std().item())
                self._last_sal_logits_std  = (float(sal_bias.std().item())
                                                if sal_bias is not None else 0.0)
        attn = self.dropout(attn)
        out = (attn @ V).transpose(1, 2).contiguous().view(B, T, self.d_model)
        out = self.out_proj(out)
        return out, None


class DStage(nn.Module):
    """CA3: dual-memory — attractor refinement + explicit memory bank + Loop 2.

    Combines two complementary memory mechanisms, matching real CA3:
    1. Attractor refinement (implicit): patterns stored in shared QKV weights,
       retrieved via iterative self-attention settling to fixed points.
       Analog: CA3 recurrent collateral autoassociation.
    2. Memory bank (explicit): learned key-value slots accessed via
       cross-attention after attractor settling.
       Analog: CA3 specific learned synaptic associations.

    The attractor completes patterns from partial cues (implicit recall).
    The memory bank stores specific associations (explicit lookup).
    A learned gate integrates both sources.
    """

    def __init__(self, d_model: int, d_expand: int, n_heads: int,
                 refine_steps: int = 3, num_slots: int = 256,
                 dropout: float = 0.1, use_salience: bool = True,
                 salience_mode: str = "v0",
                 salience_bias_init_sigma: float = 0.0,
                 v4c_norm_cap: float = 0.3,
                 v4d_w_sal_sigma: float = 0.02,
                 v4d_logit_cap: float = 4.0,
                 v4d_noinject: bool = False,
                 v4d_w_sal_rank: int = 0):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.refine_steps = refine_steps
        self.scale = self.d_head ** -0.5
        self.use_salience = use_salience
        assert salience_mode in ("v0", "v2a", "v3", "v3m", "v3m2", "v3m2_nll",
                                  "v4c", "v4c_gate", "v4d"), \
            f"salience_mode must be one of 'v0'/'v2a'/'v3'/'v3m'/'v3m2'/" \
            f"'v3m2_nll'/'v4c'/'v4c_gate'/'v4d', got {salience_mode!r}"
        self.salience_mode = salience_mode
        self._salience_bias_init_sigma = salience_bias_init_sigma
        self.v4c_norm_cap = v4c_norm_cap
        self.num_slots = num_slots
        self.v4d_logit_cap  = v4d_logit_cap
        self.v4d_noinject   = v4d_noinject
        self.v4d_w_sal_rank = v4d_w_sal_rank
        self.store_v4d_attn = False
        self.store_v4d_code = False
        self.store_v6_bank_sources = False
        self._last_q_attractor_v6 = None
        self._last_q_mem_v6 = None

        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_out = nn.Linear(d_model, d_model)

        self.step_sizes = nn.ParameterList([
            nn.Parameter(torch.tensor(0.1)) for _ in range(refine_steps)])
        self.step_lns = nn.ModuleList([
            nn.LayerNorm(d_model) for _ in range(refine_steps)])

        self.dropout = nn.Dropout(dropout)

        self.mem_keys = nn.Parameter(torch.randn(num_slots, d_model) * 0.02)
        self.mem_vals = nn.Parameter(torch.randn(num_slots, d_model) * 0.02)
        if salience_mode == "v4d":
            self.cross_attn = LogitInjectingMHA(
                d_model, n_heads, dropout=dropout)
        else:
            self.cross_attn = nn.MultiheadAttention(
                d_model, n_heads, dropout=dropout, batch_first=True)
        self.ln_mem = nn.LayerNorm(d_model)

        self.mem_gate_linear = nn.Linear(2 * d_model, d_model)

        self.loop2_q_proj = nn.Linear(d_model, d_model)
        self.loop2_k_proj = nn.Linear(d_expand, d_model)
        self.loop2_v_proj = nn.Linear(d_expand, d_model)
        self.loop2_attn = nn.MultiheadAttention(
            d_model, n_heads, dropout=dropout, batch_first=True)
        self.loop2_gate = nn.Sequential(
            nn.Linear(2 * d_model, d_model),
            nn.Sigmoid(),
        )
        self.ln_loop2 = nn.LayerNorm(d_model)

        if use_salience:
            with torch.random.fork_rng(devices=[]):
                self.salience_bias_down = nn.Linear(d_model, 16, bias=False)
                self.salience_bias_up = nn.Linear(16, d_model, bias=False)
                if self._salience_bias_init_sigma > 0.0:
                    nn.init.normal_(self.salience_bias_up.weight,
                                    std=self._salience_bias_init_sigma)
                else:
                    nn.init.zeros_(self.salience_bias_up.weight)
                if salience_mode == "v4d":
                    if v4d_w_sal_rank and v4d_w_sal_rank > 0:
                        k = v4d_w_sal_rank
                        self.w_sal_A = nn.Linear(d_model, k, bias=False)
                        self.w_sal_B = nn.Linear(k, num_slots, bias=False)
                        nn.init.normal_(self.w_sal_A.weight,
                                        std=(1.0 / d_model) ** 0.5)
                        nn.init.normal_(self.w_sal_B.weight,
                                        std=v4d_w_sal_sigma * (d_model / k) ** 0.5)
                    else:
                        self.w_sal = nn.Linear(d_model, num_slots, bias=False)
                        nn.init.normal_(self.w_sal.weight, std=v4d_w_sal_sigma)

    def attractor_settle(self, h: torch.Tensor,
                         causal: torch.Tensor) -> torch.Tensor:
        """Iterative attractor settling via self-attention with learned step sizes."""
        B, T, D = h.shape
        q_state = h
        for step in range(self.refine_steps):
            Q = self.w_q(q_state).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
            K = self.w_k(q_state).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
            V = self.w_v(q_state).view(B, T, self.n_heads, self.d_head).transpose(1, 2)

            attn = (Q @ K.transpose(-1, -2)) * self.scale
            attn = attn.masked_fill(causal.unsqueeze(0).unsqueeze(0), -1e9)
            attn = self.dropout(F.softmax(attn, dim=-1))
            attn_out = (attn @ V).transpose(1, 2).contiguous().view(B, T, D)
            attn_out = self.w_out(attn_out)

            eta = torch.sigmoid(self.step_sizes[step])
            q_state = self.step_lns[step](q_state + eta * attn_out)
        return q_state

    def forward(self, h_E: torch.Tensor,
                h_E_sparse: torch.Tensor,
                v_vec: Optional[torch.Tensor] = None,
                ) -> Tuple[torch.Tensor, torch.Tensor]:
        B, T = h_E.size(0), h_E.size(1)
        causal = torch.triu(
            torch.ones(T, T, device=h_E.device, dtype=torch.bool), 1)

        q_attractor = self.attractor_settle(h_E, causal)
        if self.store_v6_bank_sources:
            self._last_q_attractor_v6 = q_attractor

        mk = self.mem_keys.unsqueeze(0).expand(B, -1, -1)
        mv = self.mem_vals.unsqueeze(0).expand(B, -1, -1)
        if (v_vec is not None and self.use_salience
                and self.salience_mode in ("v2a", "v4c")):
            bias_raw = self.salience_bias_up(self.salience_bias_down(v_vec))
            if self.salience_mode == "v4c":
                qa_norm = q_attractor.detach().norm(dim=-1, keepdim=True)
                bias_norm = bias_raw.norm(dim=-1, keepdim=True).clamp(min=1e-6)
                scale_factor = (qa_norm * self.v4c_norm_cap / bias_norm
                                ).clamp(max=1.0)
                bias_raw = bias_raw * scale_factor
            q_for_retrieval = q_attractor + bias_raw
            retrieved, _ = self.cross_attn(q_for_retrieval, mk, mv)
        elif (v_vec is not None and self.use_salience
                and self.salience_mode == "v4d"):
            if not self.v4d_noinject:
                if getattr(self, "v4d_w_sal_rank", 0) and hasattr(self, "w_sal_A"):
                    g = self.w_sal_A(v_vec)
                    if self.store_v4d_code:
                        self._last_sal_code = g
                    sal_bias = self.w_sal_B(g)
                else:
                    sal_bias = self.w_sal(v_vec)
                    if self.store_v4d_code:
                        self._last_sal_code = v_vec
                sal_bias = sal_bias - sal_bias.mean(dim=-1, keepdim=True)
                sal_bias = sal_bias.clamp(-self.v4d_logit_cap,
                                            +self.v4d_logit_cap)
            else:
                sal_bias = None
            retrieved, _ = self.cross_attn(q_attractor, mk, mv,
                                            sal_bias=sal_bias,
                                            store=self.store_v4d_attn)
            if self.store_v4d_attn:
                self._last_cross_attn = self.cross_attn._last_attn
                self._last_sal_bias   = self.cross_attn._last_sal_bias
        else:
            retrieved, _ = self.cross_attn(q_attractor, mk, mv)
        q_mem = self.ln_mem(q_attractor + retrieved)
        if self.salience_mode == "v4d" and self.store_v4d_attn:
            self._last_q_mem = q_mem.detach()
        if self.store_v6_bank_sources:
            self._last_q_mem_v6 = q_mem

        gate_logits = self.mem_gate_linear(torch.cat([q_attractor, q_mem], dim=-1))
        if (v_vec is not None and self.use_salience
                and self.salience_mode in ("v0", "v3", "v3m", "v3m2",
                                            "v3m2_nll", "v4c_gate")):
            gate_bias = self.salience_bias_up(self.salience_bias_down(v_vec))
            if self.salience_mode == "v4c_gate":
                gl_norm = gate_logits.detach().norm(dim=-1, keepdim=True)
                gb_norm = gate_bias.norm(dim=-1, keepdim=True).clamp(min=1e-6)
                gate_bias = gate_bias * (gl_norm * self.v4c_norm_cap
                                          / gb_norm).clamp(max=1.0)
            gate_logits = gate_logits + gate_bias
        gate = torch.sigmoid(gate_logits)
        q = gate * q_mem + (1 - gate) * q_attractor
        if self.salience_mode == "v4d" and self.store_v4d_attn:
            self._last_gate_logits = gate_logits.detach()

        loop2_q = self.loop2_q_proj(q)
        loop2_k = self.loop2_k_proj(h_E_sparse)
        loop2_v = self.loop2_v_proj(h_E_sparse)
        loop2_out, _ = self.loop2_attn(
            loop2_q, loop2_k, loop2_v, attn_mask=causal)
        gate = self.loop2_gate(torch.cat([q, loop2_out], dim=-1))
        h_D = self.ln_loop2(q + gate * loop2_out)

        return h_D, loop2_out


class LStage(nn.Module):
    """CA1: two-channel comparator."""

    def __init__(self, d_model: int, vocab: int, n_heads: int,
                 dropout: float = 0.1):
        super().__init__()
        self.mem_head = nn.Linear(d_model, vocab)
        self.per_head = nn.Linear(d_model, vocab, bias=False)

        self.gate_net = nn.Sequential(
            nn.Linear(2 * d_model, 128),
            nn.GELU(),
            nn.Linear(128, 1),
            nn.Sigmoid(),
        )

        self.fuse_gate = nn.Sequential(
            nn.Linear(2 * d_model, d_model),
            nn.Sigmoid(),
        )

        self.loop1_attn = nn.MultiheadAttention(
            d_model, n_heads, dropout=dropout, batch_first=True)
        self.loop1_gate = nn.Sequential(
            nn.Linear(2 * d_model, d_model),
            nn.Sigmoid(),
        )
        self.ln_loop1 = nn.LayerNorm(d_model)

    def forward(self, h_D: torch.Tensor, h_C: torch.Tensor):
        logits_mem = self.mem_head(h_D)

        g = self.fuse_gate(torch.cat([h_D, h_C], dim=-1))
        fused = g * h_D + (1 - g) * h_C
        T = h_C.size(1)
        causal = torch.triu(torch.ones(T, T, device=h_C.device, dtype=torch.bool), 1)
        attn_out, _ = self.loop1_attn(fused, h_C, h_C, attn_mask=causal)
        gate = self.loop1_gate(torch.cat([fused, attn_out], dim=-1))
        loop1_fb = self.ln_loop1(fused + gate * attn_out)

        return logits_mem, loop1_fb


class NeuromodulatorGate(nn.Module):
    """Output-level gain modulation of C-stage by L and D signals.

    Distinct from the per-layer cholinergic AdaLN in PeriodicRetentionLayer:
    this gate models dopaminergic/noradrenergic modulation that gates the
    final hippocampal output (Lisman & Grace 2005), while AdaLN models
    cholinergic modulation during encoding (Hasselmo 1999).
    """

    def __init__(self, d_model: int):
        super().__init__()
        self.gain_from_L = nn.Sequential(nn.Linear(d_model, d_model), nn.Sigmoid())
        self.gain_from_D = nn.Sequential(nn.Linear(d_model, d_model), nn.Sigmoid())
        self.bias_from_L = nn.Linear(d_model, d_model)
        self.ln = nn.LayerNorm(d_model)

    def forward(self, h_C, loop1_fb, h_D):
        g_L = self.gain_from_L(loop1_fb)
        g_D = self.gain_from_D(h_D)
        bias = self.bias_from_L(loop1_fb)
        return self.ln(g_L * g_D * h_C + bias)



DISCOURSE_MARKERS = [
    "however", "but", "although", "nevertheless", "despite", "instead",
    "whereas", "conversely",
    "previously", "currently", "originally", "initially", "formerly",
    "recently", "now",
]


def build_marker_token_ids() -> torch.Tensor:
    """Tokenize the discourse-marker lexicon with GPT-2 BPE.
    Returns a sorted 1-D LongTensor of unique FIRST-token IDs (sub-word
    continuations are excluded; the lead BPE chunk carries the marker signal).
    Returns an empty tensor if transformers is unavailable so the marker BCE
    loss degrades gracefully."""
    try:
        from transformers import GPT2TokenizerFast
        tok = GPT2TokenizerFast.from_pretrained("gpt2")
    except Exception:
        return torch.empty(0, dtype=torch.long)

    ids = set()
    for w in DISCOURSE_MARKERS:
        for form in (w, " " + w, w.capitalize(), " " + w.capitalize()):
            toks = tok.encode(form, add_special_tokens=False)
            if toks:
                ids.add(toks[0])
    return torch.tensor(sorted(ids), dtype=torch.long)


def build_marker_mask(ids: torch.Tensor, marker_ids: torch.Tensor) -> torch.Tensor:
    """[B, T] -> [B, T] float mask, 1 at any position within a +/-1-token
    window of a marker. Smooths supervision and matches phrase-boundary
    semantics of "however,", " but ", etc."""
    if marker_ids.numel() == 0:
        return torch.zeros_like(ids, dtype=torch.float)
    in_set = torch.isin(ids, marker_ids).float().unsqueeze(1)
    kernel = torch.ones(1, 1, 3, device=ids.device, dtype=in_set.dtype)
    dilated = F.conv1d(in_set, kernel, padding=1).squeeze(1)
    return (dilated > 0).float()


class SalienceLoop(nn.Module):
    """Causal, low-rank salience scorer + contrast detector + d->r->d gate.

    Computed from h_E (E-stage output) so it can bias D-stage memory fusion
    before final D output is committed. (SL injects v_vec into the mem_gate
    logits — the blend of attractor output and memory-bank retrieval — which
    runs AFTER attractor_settle(). It does not modify the attractor settling
    iterations themselves.) Modality-agnostic by construction.

    Outputs:
        v_vec    [B, T, d]    vector signal for D-stage mem-fusion bias
        v_scalar [B, T, 1]    scalar signal for aux losses (NLL-surprise + marker BCE)
    """

    def __init__(self, d: int = 640, H: int = 8, kernel: int = 5,
                 gate_rank: int = 16, scalar_from_contrast: bool = False,
                 contrastive_head: bool = False, z_dim: int = 64,
                 proj_hidden: int = 256):
        super().__init__()
        self._d = d
        self.U = nn.Parameter(torch.randn(H, d) * (d ** -0.5))
        self.b = nn.Parameter(torch.zeros(H))
        self.contrast = nn.Conv1d(H, H, kernel_size=kernel,
                                  groups=H, padding=0, bias=False)
        nn.init.normal_(self.contrast.weight, std=0.1)
        self.pad = kernel - 1
        self.proj = nn.Linear(2 * H, d, bias=True)
        self.ln = nn.LayerNorm(d)
        self.gate_down = nn.Linear(d, gate_rank, bias=False)
        self.gate_up = nn.Linear(gate_rank, d, bias=False)
        self.scale = nn.Parameter(torch.tensor(0.1))
        self.scalar_from_contrast = scalar_from_contrast
        if scalar_from_contrast:
            self.scalar_head = nn.Linear(2 * H, 1, bias=True)
            nn.init.constant_(self.scalar_head.bias, -2.0)
        self.contrastive_head = contrastive_head
        if contrastive_head:
            self.sl_embed_head = nn.Sequential(
                nn.Linear(d, proj_hidden),
                nn.GELU(),
                nn.Linear(proj_hidden, z_dim),
            )
            self._z_log_inv_temperature = nn.Parameter(
                torch.tensor(math.log(1.0 / 0.05)))

    def forward(self, h_E: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        s = F.softplus(h_E @ self.U.t() + self.b)
        s_c = F.pad(s.transpose(1, 2), (self.pad, 0))
        c = self.contrast(s_c).transpose(1, 2)
        v_feat = torch.cat([s, c.abs()], dim=-1)
        v_raw = self.ln(self.proj(v_feat))
        v_gate = self.gate_up(self.gate_down(v_raw))
        v_vec = torch.tanh(self.scale) * v_gate
        if self.scalar_from_contrast:
            v_scalar = self.scalar_head(v_feat)
        else:
            v_scalar = s.mean(-1, keepdim=True)
        return v_vec, v_scalar

    def embed_z(self, h_D: torch.Tensor) -> torch.Tensor:
        """V4c-only: compute the L2-normalized contrastive embedding from
        POST-DStage state h_D. Used by `_v4c_contrastive_loss` to supervise
        a contrastive manifold whose gradient backprops through h_D -> v_vec
        -> salience_bias_up -> everything upstream.

        Args:  h_D [B, T, d]
        Returns: z [B, T, z_dim], with ||z[..., :]||_2 = 1 per row.
        """
        assert self.contrastive_head, "embed_z called on a SalienceLoop " \
            "built without contrastive_head=True (V4c-only API)."
        return F.normalize(self.sl_embed_head(h_D), dim=-1)


def anchor_ortho_reg(U: torch.Tensor) -> torch.Tensor:
    """Anchor orthogonality regularizer: keeps the H salience anchors
    pointing in distinct directions instead of collapsing to redundancy.
    """
    G = U @ U.t()
    I = torch.eye(U.size(0), device=U.device, dtype=U.dtype)
    return ((G - I) ** 2).mean()


def _ema_zscore(x: torch.Tensor,
                ema_mean_buf: torch.Tensor,
                ema_var_buf: torch.Tensor,
                init_buf: torch.Tensor,
                decay: float,
                training: bool) -> torch.Tensor:
    """Initialize-from-first-batch EMA z-score for V3/V3m.

    EMA updates ONLY when training=True. The first observed batch initializes
    the EMA buffers from its own mean/var (sentinel `init_buf` < 0.5) so we
    avoid the start-at-0 vs NLL-around-6 explosion when aux loss first
    activates. Subsequent batches update with `decay` (≈0.99 → ~100-batch
    effective window). Returns the z-scored tensor (not clamped — callers
    clamp to ±sl_z_clip before combining channels).
    """
    if training:
        with torch.no_grad():
            if init_buf.item() < 0.5:
                ema_mean_buf.copy_(x.mean())
                ema_var_buf.copy_(x.var().clamp(min=1e-6))
                init_buf.fill_(1.0)
            else:
                ema_mean_buf.mul_(decay).add_((1 - decay) * x.mean())
                ema_var_buf.mul_(decay).add_((1 - decay) * x.var())
    return (x - ema_mean_buf) / (ema_var_buf.sqrt() + 1e-6)


def sigreg_loss(z: torch.Tensor) -> torch.Tensor:
    """Soft variance regularizer for D-stage outputs. Adapted from
    LeWorldModel's SIGReg (arXiv:2603.19312), but softened: only penalizes
    variance COLLAPSE (< 0.1), not high variance. This prevents
    representational collapse without fighting attractor basin structure
    (which is inherently multimodal/high-variance).
    """
    if z.dim() == 3:
        z = z.reshape(-1, z.size(-1))
    var = z.var(dim=0)
    return F.relu(0.1 - var).mean()


def predictive_latent_loss(h_D: torch.Tensor, h_C: torch.Tensor) -> torch.Tensor:
    """JEPA-inspired predictive loss: D-stage output at position t should
    predict C-stage encoding at position t+1 in representation space.

    This makes the D-stage a forward model (not just a pattern completer),
    directly implementing the hippocampal predictive coding hypothesis:
    CA3 generates a prediction via Schaffer collaterals, CA1 compares it
    against actual EC input (Hasselmo 2005, Lisman & Redish 2009).

    Stop-gradient on target prevents representational collapse (JEPA principle).
    """
    pred = h_D[:, :-1, :]
    target = h_C[:, 1:, :].detach()
    pred_n = F.normalize(pred, dim=-1)
    tgt_n = F.normalize(target, dim=-1)
    return F.smooth_l1_loss(pred_n, tgt_n)



V3M2_ANCHOR_NEGATIVE_WORDS = [
    "<|endoftext|>",
]

V3M2_LABEL_NEGATIVE_PUNCT = [
    ".", ",", ";", ":", "!", "?", "—", "-", '"', "'",
    "(", ")", "[", "]", "{", "}", "/", "\\",
    "`", "*", "_", "~", "|", "<", ">", "=", "+", "&", "%", "$", "#", "@",
]
V3M2_LABEL_NEGATIVE_WORDS = [
    "a", "an", "the",
    "is", "was", "were", "be", "been", "being", "am", "are",
    "have", "has", "had",
    "do", "does", "did",
    "will", "would", "shall", "should",
    "may", "might", "can", "could", "must", "ought",
    "he", "she", "it", "they", "we", "you", "I",
    "me", "him", "her", "us", "them",
    "his", "hers", "its", "their", "our", "your", "my",
    "of", "in", "on", "at", "to", "for", "with", "by", "from",
    "into", "onto", "upon", "about", "over", "under", "between", "through",
]


def _bpe_token_ids_for_forms(words_or_punct, tokenizer):
    """For each word/punctuation entry, return the FIRST GPT-2 BPE token id
    for both bare and leading-space forms (and capitalization variants for
    words). Returns sorted unique list of int ids."""
    ids = set()
    for entry in words_or_punct:
        forms = [entry, " " + entry]
        if entry and entry[0].isalpha():
            forms.append(entry.capitalize())
            forms.append(" " + entry.capitalize())
        for f in forms:
            try:
                toks = tokenizer.encode(f, add_special_tokens=False)
            except Exception:
                continue
            if toks:
                ids.add(int(toks[0]))
    return sorted(ids)


def build_v3m2_lexicons():
    """Return (anchor_negative_ids_tensor, label_negative_ids_tensor) as
    sorted 1-D LongTensors. Built once at model __init__ from a hand-curated
    list (committed to the repo above); empty tensors if transformers is
    unavailable so the masks degrade to "allow everything"."""
    try:
        from transformers import GPT2TokenizerFast
        tok = GPT2TokenizerFast.from_pretrained("gpt2")
    except Exception:
        return (torch.empty(0, dtype=torch.long),
                torch.empty(0, dtype=torch.long))

    anchor_ids = set()
    eot = tok.encode("<|endoftext|>", add_special_tokens=False)
    for t in eot:
        anchor_ids.add(int(t))
    try:
        space_id = tok.encode(" ", add_special_tokens=False)
        for t in space_id:
            anchor_ids.add(int(t))
    except Exception:
        pass

    label_ids = set()
    for tok_id in _bpe_token_ids_for_forms(V3M2_LABEL_NEGATIVE_PUNCT, tok):
        label_ids.add(tok_id)
    for tok_id in _bpe_token_ids_for_forms(V3M2_LABEL_NEGATIVE_WORDS, tok):
        label_ids.add(tok_id)

    anchor_t = torch.tensor(sorted(anchor_ids), dtype=torch.long)
    label_t  = torch.tensor(sorted(label_ids),  dtype=torch.long)
    return anchor_t, label_t



def _curvature_torch(h):
    """[B, T, D] -> [B, T-1] per-token curvature 1 - cos(d1, d2),
    padded with one leading 0 to align with downstream [:, :-1] slicing.
    h is detached + float internally; caller passes Pass-0 tensors only."""
    h_n = F.normalize(h.detach().float(), dim=-1)
    d1 = F.normalize(h_n[:, 1:-1] - h_n[:, :-2], dim=-1)
    d2 = F.normalize(h_n[:, 2:]   - h_n[:, 1:-1], dim=-1)
    c  = 1.0 - (d1 * d2).sum(-1)
    return F.pad(c, (1, 0))


def _topk_nms_with_threshold(scores, max_anchors, z_threshold, min_gap):
    """Per-row top-K with NMS min-gap and absolute threshold. Routine rows
    (no scores above threshold) get zero anchors. Returns bool [B, T].

    Sync-conscious: pulls sorted scores + indices to CPU ONCE per call so
    the inner double loop runs as pure Python over numpy arrays rather than
    one GPU sync per `.item()`. Single CPU->GPU sync at the end to set the
    bool mask. NMS itself is inherently sequential (per-row top-K with a
    min-gap exclusion) so loops are unavoidable, but they are no longer
    per-iter GPU syncs."""
    B, T = scores.shape
    sorted_scores, sorted_idx = scores.sort(dim=-1, descending=True)
    sorted_scores_cpu = sorted_scores.detach().cpu().numpy()
    sorted_idx_cpu    = sorted_idx.detach().cpu().numpy()

    rows, cols = [], []
    NEG_INF = float('-inf')
    for b in range(B):
        accepted = []
        for j in range(T):
            s = float(sorted_scores_cpu[b, j])
            if s == NEG_INF or s < z_threshold:
                break
            pos = int(sorted_idx_cpu[b, j])
            ok = True
            for ap in accepted:
                if abs(pos - ap) < min_gap:
                    ok = False
                    break
            if ok:
                accepted.append(pos)
                rows.append(b); cols.append(pos)
                if len(accepted) >= max_anchors:
                    break

    out = torch.zeros_like(scores, dtype=torch.bool)
    if rows:
        out[torch.as_tensor(rows, device=scores.device, dtype=torch.long),
            torch.as_tensor(cols, device=scores.device, dtype=torch.long)] = True
    return out


def _build_decay_target_tensorized(anchors, strength, is_label_eligible,
                                    window=12, tau=2.0, k_max=6):
    """Build per-position graded labels via exponential decay over the
    content-window AFTER each anchor.

    Args:
      anchors:           [B, T] bool, True at anchor positions
      strength:          [B, T] float in [0, 1], anchor strength
      is_label_eligible: [B, T] bool, True for content tokens
      window:            int, max position offset to consider after anchor
      tau:               float, decay constant in content-rank units
      k_max:             int, cap on content rank assigned a positive label

    Returns: target [B, T] float; target[t] = max over anchors of
             strength[t_a] * exp(-(content_rank - 1) / tau), for label-
             eligible positions within `window` of t_a and within k_max
             content-eligible positions of t_a (zero otherwise).
    """
    B, T = anchors.shape
    device = anchors.device

    cum_le = is_label_eligible.long().cumsum(dim=-1)

    anchor_counts = anchors.sum(dim=-1)
    A = int(anchor_counts.max().item()) if anchor_counts.max() > 0 else 0
    if A == 0:
        return torch.zeros(B, T, device=device,
                            dtype=strength.dtype)

    anchor_pos = torch.full((B, A), -1, dtype=torch.long, device=device)
    for b in range(B):
        idx = anchors[b].nonzero(as_tuple=True)[0]
        anchor_pos[b, :idx.size(0)] = idx

    offsets = torch.arange(1, window + 1, device=device).view(1, 1, window)
    cand_pos = anchor_pos.unsqueeze(-1) + offsets
    cand_valid = (anchor_pos.unsqueeze(-1) >= 0) & (cand_pos < T)
    cand_pos_clamped = cand_pos.clamp(min=0, max=T - 1)

    le_at_cand = is_label_eligible.gather(
        1, cand_pos_clamped.view(B, -1)
    ).view(B, A, window)
    cum_at_cand = cum_le.gather(
        1, cand_pos_clamped.view(B, -1)
    ).view(B, A, window)

    anchor_pos_clamped = anchor_pos.clamp(min=0)
    cum_at_anchor = cum_le.gather(1, anchor_pos_clamped)
    content_rank = cum_at_cand - cum_at_anchor.unsqueeze(-1)

    valid_label = cand_valid & le_at_cand & (content_rank >= 1) & (content_rank <= k_max)

    rank_for_decay = (content_rank - 1).clamp(min=0).float()
    decay = torch.exp(-rank_for_decay / tau)
    label_vals = strength.gather(1, anchor_pos_clamped).unsqueeze(-1) * decay
    label_vals = label_vals * valid_label.float()

    target = torch.zeros(B, T, device=device, dtype=strength.dtype)
    flat_pos = cand_pos_clamped.view(B, -1)
    flat_val = label_vals.view(B, -1)
    target = target.scatter_reduce(
        dim=1, index=flat_pos, src=flat_val, reduce="amax",
        include_self=True,
    )
    return target


def _build_v3m2_targets(
    *, ids, h_C_p0, h_E_p0, logits_mem_p0, model, salience_mode,
    update_ema, with_nll,
):
    """Build the V3m2 supervision target. Callable from training (update_ema
    =True) and from --preview-only (update_ema=False).

    Returns (target, w, debug_info) where target is [B, T-1] in [0, 1],
    w is per-token weight [B, T-1] normalized to mean 1, and debug_info is
    a dict with intermediate tensors useful for inspection.
    """
    decay = model.nll_ema_decay

    curve_C = _curvature_torch(h_C_p0)
    curve_E = _curvature_torch(h_E_p0)
    h_C_n = F.normalize(h_C_p0.detach().float(), dim=-1)
    h_E_n = F.normalize(h_E_p0.detach().float(), dim=-1)
    cosmis = (1.0 - (h_C_n * h_E_n).sum(-1))[:, :-1]

    def _z(x, m_buf, v_buf, i_buf):
        if update_ema:
            return _ema_zscore(x, m_buf, v_buf, i_buf, decay, training=True
                               ).clamp(-4.0, 4.0)
        return ((x - m_buf) / (v_buf.sqrt() + 1e-6)).clamp(-4.0, 4.0)

    curve_C_z = _z(curve_C,
                    model.curve_C_ema_mean, model.curve_C_ema_var,
                    model.curve_C_ema_initialized)
    curve_E_z = _z(curve_E,
                    model.curve_E_ema_mean, model.curve_E_ema_var,
                    model.curve_E_ema_initialized)
    cosmis_z  = _z(cosmis,
                    model.cosmis_ema_mean,  model.cosmis_ema_var,
                    model.cosmis_ema_initialized)

    event_score = 0.5 * curve_C_z + 0.5 * curve_E_z + 0.7 * cosmis_z
    if with_nll and logits_mem_p0 is not None:
        flat_logits = logits_mem_p0[:, :-1].reshape(
            -1, logits_mem_p0.size(-1)).float()
        nll = F.cross_entropy(
            flat_logits, ids[:, 1:].reshape(-1), reduction='none'
        ).view(ids.size(0), -1)
        nll_z = _z(nll, model.nll_ema_mean, model.nll_ema_var,
                    model.nll_ema_initialized)
        event_score = event_score + 0.7 * nll_z

    ANCHOR_Z = 1.5
    TEMP = 0.75
    strength = torch.sigmoid((event_score - ANCHOR_Z) / TEMP)

    ids_for_mask = ids[:, :-1]
    is_anchor_eligible = ~torch.isin(ids_for_mask, model.anchor_negative_ids)
    is_label_eligible  = ~torch.isin(ids_for_mask, model.label_negative_ids)
    es_for_anchor = event_score.masked_fill(
        ~is_anchor_eligible, float('-inf'))

    max_anchors = min(8, max(1, ids.size(1) // 128))
    anchors = _topk_nms_with_threshold(
        es_for_anchor, max_anchors=max_anchors,
        z_threshold=ANCHOR_Z, min_gap=6,
    )

    target = _build_decay_target_tensorized(
        anchors=anchors, strength=strength,
        is_label_eligible=is_label_eligible,
        window=12, tau=2.0, k_max=6,
    )

    w = 1.0 + 4.0 * target.detach()
    w = w / w.mean().clamp(min=1e-6)

    debug = {
        "event_score": event_score,
        "strength": strength,
        "is_anchor_eligible": is_anchor_eligible,
        "is_label_eligible":  is_label_eligible,
        "anchors": anchors,
        "target": target,
    }
    return target, w, debug



def _pool_span(z: torch.Tensor, span: Tuple[int, int]) -> torch.Tensor:
    """Mean-pool then L2-normalize z over (start, end) — used to extract
    a single span vector for InfoNCE. z is [T, z_dim]; returns [z_dim]."""
    s, e = span
    if e <= s:
        return None
    pooled = z[s:e].mean(dim=0)
    return F.normalize(pooled, dim=-1)


def _v4c_contrastive_loss(*, z, v_scalar, spans, cfg, salience_module,
                           bce_alpha, nce_alpha,
                           randomize_positive_mapping: bool = False):
    """Compute the V4c contrastive aux loss.

    InfoNCE on `z` pooled over (query, current/positive, stale, distractor,
    neutral) spans with weighted negatives + in-batch false-negative masking.
    Span-level BCE on `v_scalar` (positive mask = current span; negative mask
    = stale ∪ distractor ∪ neutral; other tokens ignored).

    Args:
      z:                 [B, T, z_dim] L2-normalized contrastive embedding.
      v_scalar:          [B, T, 1] LOGIT from SalienceLoop.scalar_head.
      spans:             list of len-B V4cItem-like dicts (or objects).
      cfg:               CEDLConfig.
      salience_module:   the SalienceLoop instance.
      bce_alpha, nce_alpha: ramp buffers in [0, 1].
      randomize_positive_mapping:
                          V4c-randlabel control. When True, each anchor's
                          positive is REPLACED by a different anchorable
                          item's pooled `current` z (cross-item rotation).
                          BCE is also DISABLED in this mode — the per-token
                          BCE target is meaningless when the positive comes
                          from another sequence. This is the proper random-
                          label control: same prompt distribution, but the
                          (anchor → positive) mapping is shuffled, destroying
                          the contrastive signal-to-supervision link.

    Returns: scalar tensor.
    """
    device = z.device
    B = z.size(0)
    inv_temp = salience_module._z_log_inv_temperature.exp().clamp(min=10.0, max=100.0)
    neg_weights = cfg.v4c_neg_weights

    anchors = []
    positives = []
    item_neg_pool = []
    global_neg_pool = []
    item_in_batch_currents = []

    for b in range(B):
        item = spans[b]
        family = item.family if hasattr(item, "family") else item.get("family")
        ids_b = item.ids if hasattr(item, "ids") else item.get("ids")
        if family == "neutral_control":
            for sp in (item.neutral if hasattr(item, "neutral")
                       else item.get("neutral", [])):
                v = _pool_span(z[b], sp)
                if v is not None:
                    global_neg_pool.append((b, v, neg_weights.get("neutral", 1.0), "neutral"))
            continue
        q_spans = item.query if hasattr(item, "query") else item.get("query", [])
        c_spans = item.current if hasattr(item, "current") else item.get("current", [])
        s_spans = item.stale if hasattr(item, "stale") else item.get("stale", [])
        d_spans = item.distractor if hasattr(item, "distractor") else item.get("distractor", [])
        n_spans = item.neutral if hasattr(item, "neutral") else item.get("neutral", [])
        pp_spans = (item.paraphrase_positives if hasattr(item, "paraphrase_positives")
                     else item.get("paraphrase_positives", []))
        if not q_spans or not c_spans:
            continue
        a = _pool_span(z[b], q_spans[0])
        if a is None:
            continue
        anchors.append((b, a))
        pos_for_b = []
        cv = _pool_span(z[b], c_spans[0])
        if cv is not None:
            pos_for_b.append(cv)
        for pp in pp_spans:
            ppv = _pool_span(z[b], pp)
            if ppv is not None:
                pos_for_b.append(ppv)
        positives.append((b, pos_for_b))
        for sp in s_spans:
            v = _pool_span(z[b], sp)
            if v is not None:
                item_neg_pool.append((b, v, neg_weights.get("stale", 2.0), "stale"))
        for sp in d_spans:
            v = _pool_span(z[b], sp)
            if v is not None:
                item_neg_pool.append((b, v, neg_weights.get("distractor", 1.0), "distractor"))
        for sp in n_spans:
            v = _pool_span(z[b], sp)
            if v is not None:
                item_neg_pool.append((b, v, neg_weights.get("neutral", 1.0), "neutral"))
        current_tok_ids = frozenset(ids_b[c_spans[0][0]:c_spans[0][1]])
        query_tok_ids = frozenset(ids_b[q_spans[0][0]:q_spans[0][1]])
        if pos_for_b:
            item_in_batch_currents.append(
                (b, pos_for_b[0], current_tok_ids, query_tok_ids))

    if randomize_positive_mapping and len(positives) > 1:
        rotated = positives[1:] + positives[:1]
        positives_for_loss = [(b, plist) for (_, _), (b, plist)
                               in zip(anchors, rotated)]
    else:
        positives_for_loss = positives

    L_nce = z.new_zeros(())
    n_anchored = 0
    for (b, a), (b_pos, pos_list) in zip(anchors, positives_for_loss):
        if not pos_list:
            continue
        candidates = []
        for p in pos_list:
            candidates.append((p, 1.0, True))
        for (b_n, v, w, kind) in item_neg_pool:
            if b_n == b:
                candidates.append((v, w, False))
        for (b_origin, v, w, kind) in global_neg_pool:
            candidates.append((v, w, False))
        own_current_ids = None
        own_query_ids = None
        for (bb, _v, cids, qids) in item_in_batch_currents:
            if bb == b:
                own_current_ids = cids
                own_query_ids = qids
                break
        for (bb, v_other, c_ids_other, q_ids_other) in item_in_batch_currents:
            if bb == b:
                continue
            if randomize_positive_mapping and bb == b_pos:
                continue
            if own_current_ids is not None:
                if c_ids_other & own_current_ids:
                    continue
                if q_ids_other == own_query_ids:
                    continue
            candidates.append((v_other, 1.0, False))
        if len(candidates) < 2:
            continue
        cand_vecs = torch.stack([c[0] for c in candidates], dim=0)
        cand_weights = torch.tensor([c[1] for c in candidates],
                                     device=device, dtype=z.dtype)
        is_pos = torch.tensor([c[2] for c in candidates],
                               device=device, dtype=torch.bool)
        logits = (cand_vecs @ a) * inv_temp
        max_l = logits.max().detach()
        log_denom = (cand_weights * (logits - max_l).exp()).sum().clamp(min=1e-12).log() + max_l
        pos_mask = is_pos.float()
        pos_count = pos_mask.sum().clamp(min=1.0)
        L_i = -((logits - log_denom) * pos_mask).sum() / pos_count
        L_nce = L_nce + L_i
        n_anchored += 1
    if n_anchored > 0:
        L_nce = L_nce / n_anchored
    else:
        L_nce = z.sum() * 0.0

    if randomize_positive_mapping:
        L_bce = z.sum() * 0.0
    else:
        v_logit = v_scalar[:, :, 0]
        target_mask = torch.zeros_like(v_logit)
        target_value = torch.zeros_like(v_logit)
        for b in range(B):
            item = spans[b]
            family = item.family if hasattr(item, "family") else item.get("family")
            if family == "neutral_control":
                for sp in (item.neutral if hasattr(item, "neutral")
                           else item.get("neutral", [])):
                    s, e = sp
                    target_mask[b, s:e] = 1.0
                    target_value[b, s:e] = 0.0
                continue
            c_spans = item.current if hasattr(item, "current") else item.get("current", [])
            s_spans = item.stale if hasattr(item, "stale") else item.get("stale", [])
            d_spans = item.distractor if hasattr(item, "distractor") else item.get("distractor", [])
            pp_spans = (item.paraphrase_positives if hasattr(item, "paraphrase_positives")
                         else item.get("paraphrase_positives", []))
            for s, e in c_spans + pp_spans:
                target_mask[b, s:e] = 1.0
                target_value[b, s:e] = 1.0
            for s, e in s_spans + d_spans:
                target_mask[b, s:e] = 1.0
                target_value[b, s:e] = 0.0
        bce = F.binary_cross_entropy_with_logits(
            v_logit, target_value, reduction='none')
        L_bce = (bce * target_mask).sum() / target_mask.sum().clamp(min=1.0)

    nce_w = float(getattr(cfg, "v4c_nce_weight", 0.05))
    bce_w = float(getattr(cfg, "v4c_bce_weight", 0.02))
    return nce_alpha * nce_w * L_nce + bce_alpha * bce_w * L_bce


def _v4d_compute_per_item_margins(logits_mem, ids, spans):
    """Compute per-item answer margins (log_p_cur − log_p_stale) at the last
    prompt position. Returns:
      margins:    [N_anchorable] tensor — one margin per item that passed all
                  validity checks (non-neutral, has spans, valid length,
                  current_tok != stale_tok). May be empty if all skipped.
      item_idxs:  list[int] — the batch index `b` of each kept item, so the
                  caller can ALIGN margins across two forwards (normal +
                  clamped) by checking that item_idxs match.

    Used by both _v4c_answer_margin_loss (hinge over each margin) and the
    V4d-Causal loss (per-item margin gap between normal/clamped forwards).
    """
    B = logits_mem.size(0)
    margins, item_idxs = [], []
    for b in range(B):
        item = spans[b]
        family = item.family if hasattr(item, "family") else item.get("family")
        if family == "neutral_control":
            continue
        c_spans = item.current if hasattr(item, "current") else item.get("current", [])
        s_spans = item.stale if hasattr(item, "stale") else item.get("stale", [])
        ids_b = item.ids if hasattr(item, "ids") else item.get("ids")
        if not c_spans or not s_spans:
            continue
        orig_len = (item.original_length if hasattr(item, "original_length")
                    else item.get("original_length", 0))
        if orig_len <= 0 or orig_len > logits_mem.size(1):
            continue
        last_pos = orig_len - 1
        cur_tok = ids_b[c_spans[0][0]]
        stale_tok = ids_b[s_spans[0][0]]
        if cur_tok == stale_tok:
            continue
        logp = F.log_softmax(logits_mem[b, last_pos].float(), dim=-1)
        margins.append(logp[cur_tok] - logp[stale_tok])
        item_idxs.append(b)
    if margins:
        return torch.stack(margins), item_idxs
    return logits_mem.new_zeros((0,)), []


def _v4d_compute_per_item_z_margins(z, spans):
    """Compute per-item z-cosine margins matching the held-out probe's formula:

        m_z = cos(z_q, z_current) − max(cos(z_q, z_s), cos(z_q, z_d), cos(z_q, z_n))

    where the neutral span is the first 8 tokens of the prompt (probe convention).
    z is [B, T, z_dim], already L2-normalized per row (SalienceLoop.embed_z
    output). For each anchorable item, pool z over spans, re-normalize, take
    cosines, return the margin tensor + item indices.

    Returns:
      margins:    [N_anchorable] tensor of m_z per item
      item_idxs:  list[int] of batch indices aligned with margins (so two
                  calls under normal/clamped forwards can be aligned).
    """
    B = z.size(0)
    margins, item_idxs = [], []

    def _pool(z_seq, span):
        s, e = span
        if e <= s:
            return None
        return F.normalize(z_seq[s:e].mean(dim=0), dim=-1)

    for b in range(B):
        item = spans[b]
        family = item.family if hasattr(item, "family") else item.get("family")
        if family == "neutral_control":
            continue
        q_spans = item.query if hasattr(item, "query") else item.get("query", [])
        c_spans = item.current if hasattr(item, "current") else item.get("current", [])
        s_spans = item.stale if hasattr(item, "stale") else item.get("stale", [])
        d_spans = item.distractor if hasattr(item, "distractor") else item.get("distractor", [])
        if not q_spans or not c_spans:
            continue
        z_seq = z[b]
        T = z_seq.size(0)
        z_q = _pool(z_seq, q_spans[0])
        z_c = _pool(z_seq, c_spans[0])
        if z_q is None or z_c is None:
            continue
        cos_c = (z_q * z_c).sum()
        cos_others = []
        if s_spans:
            z_s = _pool(z_seq, s_spans[0])
            if z_s is not None:
                cos_others.append((z_q * z_s).sum())
        if d_spans:
            z_d = _pool(z_seq, d_spans[0])
            if z_d is not None:
                cos_others.append((z_q * z_d).sum())
        z_n = _pool(z_seq, (0, min(8, T)))
        if z_n is not None:
            cos_others.append((z_q * z_n).sum())
        if not cos_others:
            continue
        max_other = torch.stack(cos_others).max()
        margins.append(cos_c - max_other)
        item_idxs.append(b)
    if margins:
        return torch.stack(margins), item_idxs
    return z.new_zeros((0,)), []


def _v4d_causal_dependence_loss(margins_normal, margins_clamped,
                                  causal_gap: float = 0.25):
    """Hinge loss that forces margin_normal − margin_clamped ≥ causal_gap.

    Both inputs are per-item margin tensors of equal length (caller must
    align item indices across normal/clamped forwards). The loss is mean of
    relu(causal_gap − (m_n − m_c)). Penalty pushes the head to MAKE its
    margin depend on v_vec: if clamping v_vec doesn't drop the margin by at
    least `causal_gap`, the head is bypassing the SL pathway.
    """
    if margins_normal.numel() == 0 or margins_clamped.numel() == 0:
        return margins_normal.sum() * 0.0
    gap_satisfied = margins_normal - margins_clamped
    deficit = causal_gap - gap_satisfied
    return F.relu(deficit).mean()


def _v4d_pool_code_by_role(code, spans):
    """Mean-pool the salience code g ([B, T, k]) over each item's current-span
    and stale-span. Returns {batch_pos: (g_current[k], g_stale[k])} for every
    anchorable item. Used by the V4e swap-consistency + role-separation losses
    (pools by ROLE span — own offsets per item, review Blocker 3)."""
    out = {}
    if code is None:
        return out
    T = code.size(1)
    for b, it in enumerate(spans):
        if (getattr(it, "family", "") == "neutral_control"
                or not it.current or not it.stale):
            continue
        cs, ce = it.current[0]
        ss, se = it.stale[0]
        if ce <= cs or se <= ss or ce > T or se > T:
            continue
        out[b] = (code[b, cs:ce].mean(0), code[b, ss:se].mean(0))
    return out


def _v4c_answer_margin_loss(logits_mem, ids, spans, target_margin: float = 1.0):
    """Hinge loss on the LM's preference for `current` over `stale` at the
    last prompt position.

    For each anchorable item:
      - last_pos = item.original_length − 1 (the position whose logits predict
        the NEXT token, i.e., what follows the prompt's "Answer:"). Uses the
        stored original_length so the function is pad-token-id-agnostic.
      - get log p(first_token_of_current) and log p(first_token_of_stale)
        from softmax(logits_mem[last_pos])
      - margin = log_p_current − log_p_stale
      - loss_i = max(0, target_margin − margin)
    Returns mean over anchorable items.

    Note: this aux loss does NOT require multi-token generation — only the
    FIRST token of current/stale is scored. Since the generator's
    current/stale words are typically 1-2 BPE tokens that differ at the
    first token (e.g., "car" vs "hat", "Berlin" vs "Bonn", "C7" vs "A12"),
    first-token margin is a clean signal.

    The gradient path: log_p depends on logits_mem ← L-stage ← h_D ←
    DStage(v_vec=v_vec) ← v_vec ← salience_bias_up. So the answer-margin
    loss directly trains the SL pathway, which is exactly the V4c+M
    hypothesis: output-side pressure makes the LM use v_vec.
    """
    B = logits_mem.size(0)
    losses = []
    n_skipped_neutral = 0
    n_skipped_no_spans = 0
    n_skipped_no_length = 0
    n_skipped_equal_token = 0
    for b in range(B):
        item = spans[b]
        family = item.family if hasattr(item, "family") else item.get("family")
        if family == "neutral_control":
            n_skipped_neutral += 1
            continue
        c_spans = item.current if hasattr(item, "current") else item.get("current", [])
        s_spans = item.stale if hasattr(item, "stale") else item.get("stale", [])
        ids_b = item.ids if hasattr(item, "ids") else item.get("ids")
        if not c_spans or not s_spans:
            n_skipped_no_spans += 1
            continue
        orig_len = (item.original_length if hasattr(item, "original_length")
                    else item.get("original_length", 0))
        if orig_len <= 0 or orig_len > logits_mem.size(1):
            n_skipped_no_length += 1
            continue
        last_pos = orig_len - 1
        cur_tok = ids_b[c_spans[0][0]]
        stale_tok = ids_b[s_spans[0][0]]
        if cur_tok == stale_tok:
            n_skipped_equal_token += 1
            continue
        logp = F.log_softmax(logits_mem[b, last_pos].float(), dim=-1)
        log_p_cur = logp[cur_tok]
        log_p_stale = logp[stale_tok]
        margin = log_p_cur - log_p_stale
        loss_i = F.relu(target_margin - margin)
        losses.append(loss_i)
    if not losses:
        out = logits_mem.sum() * 0.0
    else:
        out = torch.stack(losses).mean()
    out._v4c_margin_n_used     = len(losses)
    out._v4c_margin_n_skip_neu = n_skipped_neutral
    out._v4c_margin_n_skip_spn = n_skipped_no_spans
    out._v4c_margin_n_skip_len = n_skipped_no_length
    out._v4c_margin_n_skip_eq  = n_skipped_equal_token
    return out


def _v4c_collect_answer_quads(ids: torch.Tensor, spans):
    """V6.1 helper. Mirrors _v4c_answer_margin_loss's per-item filtering
    (skip neutral, skip no-span, skip no-length, skip equal-token), returning
    LongTensors `(b_idx, p_idx, cur_tok, stale_tok)` of length N on `ids.device`.

    The V6.1 hook uses these as fancy-indices into [B, T, ...] activations to
    avoid materializing the full [B, T, V=50257] mixture log-prob tensor.
    Empty case (zero valid items): returns four empty `[0]` LongTensors on
    `ids.device`, NOT None — caller's `if b_idx.numel() > 0:` is the gate.
    """
    B = ids.size(0)
    T = ids.size(1)
    b_list, p_list, cur_list, stale_list = [], [], [], []
    for b in range(B):
        item = spans[b]
        family = item.family if hasattr(item, "family") else item.get("family")
        if family == "neutral_control":
            continue
        c_spans = item.current if hasattr(item, "current") else item.get("current", [])
        s_spans = item.stale if hasattr(item, "stale") else item.get("stale", [])
        ids_b = item.ids if hasattr(item, "ids") else item.get("ids")
        if not c_spans or not s_spans:
            continue
        orig_len = (item.original_length if hasattr(item, "original_length")
                    else item.get("original_length", 0))
        if orig_len <= 0 or orig_len > T:
            continue
        cur_t   = int(ids_b[c_spans[0][0]])
        stale_t = int(ids_b[s_spans[0][0]])
        if cur_t == stale_t:
            continue
        b_list.append(b)
        p_list.append(orig_len - 1)
        cur_list.append(cur_t)
        stale_list.append(stale_t)
    dev = ids.device
    return (torch.tensor(b_list,     dtype=torch.long, device=dev),
            torch.tensor(p_list,     dtype=torch.long, device=dev),
            torch.tensor(cur_list,   dtype=torch.long, device=dev),
            torch.tensor(stale_list, dtype=torch.long, device=dev))


def _v4c_collect_background_quads(ids: torch.Tensor, spans, *,
                                    ans_b_idx: torch.Tensor,
                                    ans_p_idx: torch.Tensor,
                                    k_per_item: int = 8):
    """Stage 2a — deterministic background-position collector for L_bg.

    For each V4c item:
      - Valid range = [0, original_length-1); excludes the answer p_idx.
      - INCLUDES entity-span positions (current/stale/distractor) — the head
        must learn to keep λ low at "X owns a Y" entity-mention rows and
        lift only at the structural "Answer:" answer row. That's the whole
        selectivity hypothesis; do NOT exclude entity spans here.
      - DETERMINISTIC: evenly-spaced via per-item offset = b % step. Avoids
        random background rows so fixed seeds remain reproducible.
      - If valid_len < k_per_item: cycle deterministically with replacement.
      - Skip items with valid_len == 0.

    Returns (b_bg, p_bg) LongTensors on `ids.device`, dtype torch.long.
    Empty case: two `[0]` LongTensors on `ids.device` (caller's `.numel()`
    gate is the activation check).
    """
    B = ids.size(0)
    ans_by_b = {int(_b): int(_p) for _b, _p in zip(
        ans_b_idx.tolist(), ans_p_idx.tolist())}
    b_list, p_list = [], []
    for b in range(B):
        item = spans[b]
        ol = int(getattr(item, "original_length", 0)
                  if hasattr(item, "original_length")
                  else item.get("original_length", 0))
        if ol <= 1:
            continue
        skip = ans_by_b.get(b, -1)
        valid = [p for p in range(ol - 1) if p != skip]
        if not valid:
            continue
        if len(valid) >= k_per_item:
            step = max(1, len(valid) // k_per_item)
            offset = b % step
            picks = valid[offset::step][:k_per_item]
            while len(picks) < k_per_item:
                picks.append(valid[(offset + len(picks)) % len(valid)])
        else:
            picks = [valid[i % len(valid)] for i in range(k_per_item)]
        b_list.extend([b] * len(picks))
        p_list.extend(picks)
    dev = ids.device
    return (torch.tensor(b_list, dtype=torch.long, device=dev),
            torch.tensor(p_list, dtype=torch.long, device=dev))


def _pad_v4c_items(items, T_max):
    """Right-pad each item's ids to T_max (pad=0), clipping spans, in place.
    Returns the [N, T_max] LongTensor."""
    ids_padded = []
    for it in items:
        ids = it.ids
        if len(ids) > T_max:
            ids = ids[:T_max]
            for fname in ("query", "current", "stale", "distractor",
                          "neutral", "paraphrase_positives"):
                spans = getattr(it, fname)
                spans[:] = [(s, min(e, T_max)) for (s, e) in spans if s < T_max]
        pad_len = T_max - len(ids)
        if pad_len > 0:
            ids = ids + [0] * pad_len
        ids_padded.append(ids)
        it.ids = ids
    return torch.tensor(ids_padded, dtype=torch.long)


def _build_v4c_contrastive_batch(generator_fn, tokenizer, B: int, T_max: int,
                                  *, randomize_positive_mapping: bool = False,
                                  seed: int = 0, split: str = "all",
                                  hard_collision_frac: float = 0.0,
                                  swap_fn=None,
                                  family_weights=None):
    """Build one V4c synthetic mini-batch.

    Args:
      generator_fn:  callable like data_v4c_pairs.generate(tokenizer, n, seed)
      tokenizer:     GPT2TokenizerFast (or compatible)
      B:             batch size
      T_max:         pad/truncate sequences to this length
      randomize_positive_mapping: if True, rotate anchor→positive mapping
                     across items (V4c-randlabel control). For each anchor
                     `a_i`, swap its current span's positive role with item
                     a_(i+1 mod B)'s current span (with FN masking still
                     applied at loss time).
      seed:          rng seed for this batch

    Additional kwargs:
      split:               vocabulary partition ('train'/'test'/'all') passed
                           to the generator (V4e novel-entity training).
      hard_collision_frac: fraction of but_update items with stale pinned to a
                           recent current object (V4e role-by-position pressure).
      swap_fn:             if given (e.g. data_v4c_pairs.make_entity_swap), also
                           build an entity-swapped twin per item and return
                           (ids, items, swap_ids, swap_items) for the V4e
                           swap-consistency / role-separation losses. Otherwise
                           returns (ids, items, None, None).

    Returns: (ids_tensor, items, swap_ids_or_None, swap_items_or_None)
    """
    try:
        items = generator_fn(tokenizer, B, seed=seed, split=split,
                              hard_collision_frac=hard_collision_frac,
                              family_weights=family_weights)
    except TypeError:
        try:
            items = generator_fn(tokenizer, B, seed=seed, split=split,
                                  hard_collision_frac=hard_collision_frac)
        except TypeError:
            items = generator_fn(tokenizer, B, seed=seed)
    ids_tensor = _pad_v4c_items(items, T_max)
    if swap_fn is None:
        return ids_tensor, items, None, None
    import random as _random
    swap_rng = _random.Random(seed + 10007)
    twins = [swap_fn(it, tokenizer, swap_rng, split) for it in items]
    swap_tensor = _pad_v4c_items(twins, T_max)
    return ids_tensor, items, swap_tensor, twins


class V6SourceAdapter(nn.Module):
    """Small residual MLP used only by the V6 direct specialist readout."""

    def __init__(self, d_model: int):
        super().__init__()
        self.ln = nn.LayerNorm(d_model)
        self.fc1 = nn.Linear(d_model, d_model)
        self.fc2 = nn.Linear(d_model, d_model)
        nn.init.normal_(self.fc2.weight, std=1e-3)
        nn.init.zeros_(self.fc2.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.fc2(F.gelu(self.fc1(self.ln(x))))


class V6ContextAdapter(nn.Module):
    """Causal prefix adapter for the V6 direct specialist readout."""

    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.ln = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(
            d_model, n_heads, dropout=0.0, batch_first=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.dim() != 3:
            raise RuntimeError(
                "V6ContextAdapter expects [B,T,d] source states")
        T = x.size(1)
        mask = torch.ones(T, T, device=x.device, dtype=torch.bool).triu(1)
        h = self.ln(x)
        ctx, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
        return x + ctx


class CEDLTwoLoop100M(nn.Module):
    """CEDL Two-Loop at 100M scale with coupled feedback loops.

    Forward pass:
        Pass 0 (feedforward): C -> E -> D -> L (no feedback)
        Pass 1..N (feedback): L->C and D->E signals enrich re-encoding
        Final: NeuromodulatorGate applies supplementary gain modulation

    Auxiliary losses (training only):
        - CSR: contrastive support repulsion for E-stage pattern separation
        - SIGReg: Gaussian regularizer on D-stage outputs (replaces attractor contrastive)
        - Predictive latent: D(t) predicts C(t+1) — JEPA-inspired forward model
    """

    def __init__(self, cfg: CEDLConfig, vocab: int, max_seq: int):
        super().__init__()
        d = cfg.d_model
        d_expand = d * cfg.e_expand
        self.n_feedback_iters = cfg.n_feedback_iters
        self.feedback_decay = cfg.feedback_decay
        self.use_salience = cfg.use_salience
        self.salience_mode = cfg.salience_mode
        self.lex_anchor_weight = cfg.lex_anchor_weight
        self.v4c_variant = getattr(cfg, "v4c_variant", "base")
        self._v4c_cfg = cfg
        self.nll_ema_decay = 0.99
        self.sl_window_k   = 3
        self.sl_z_clip     = 4.0
        self.register_buffer('feedback_alpha', torch.tensor(0.0))

        _sb_sigma = 0.0
        if (cfg.salience_mode in ("v4c", "v4c_gate")
                and self.v4c_variant != "cold"):
            _sb_sigma = float(getattr(cfg, "v4c_warm_start_sigma", 0.05))
        _v4c_norm_cap = float(getattr(cfg, "v4c_norm_cap", 0.3))
        _v4d_w_sal_sigma = float(getattr(cfg, "v4d_w_sal_sigma", 0.02))
        _v4d_logit_cap   = float(getattr(cfg, "v4d_logit_cap",   4.0))
        _v4d_noinject    = bool(getattr(cfg,  "v4d_noinject",    False))
        _v4d_w_sal_rank  = int(getattr(cfg,   "v4d_w_sal_rank",   0))

        self.c_stage = CStageRetention(
            vocab, max_seq, d, cfg.n_heads, cfg.c_layers, cfg.ffn_dim, cfg.dropout)
        self.e_stage = EStage(d, cfg.e_expand, cfg.e_sparsity)
        self.d_stage = DStage(d, d_expand, cfg.n_heads,
                              cfg.d_refine, cfg.d_slots, cfg.dropout,
                              use_salience=cfg.use_salience,
                              salience_mode=cfg.salience_mode,
                              salience_bias_init_sigma=_sb_sigma,
                              v4c_norm_cap=_v4c_norm_cap,
                              v4d_w_sal_sigma=_v4d_w_sal_sigma,
                              v4d_logit_cap=_v4d_logit_cap,
                              v4d_noinject=_v4d_noinject,
                              v4d_w_sal_rank=_v4d_w_sal_rank)
        _role_w = float(getattr(cfg, "v4d_role_sep_weight", 0.0))
        if cfg.use_salience and cfg.salience_mode == "v4d" and _role_w > 0:
            _code_dim = _v4d_w_sal_rank if _v4d_w_sal_rank > 0 else d
            self.v4d_role_head = nn.Linear(_code_dim, 2)
        self.l_stage = LStage(d, vocab, cfg.n_heads, cfg.dropout)
        self.modulator = NeuromodulatorGate(d)

        if self.use_salience:
            with torch.random.fork_rng(devices=[]):
                self.register_buffer('sl_alpha', torch.tensor(0.0))
                self.register_buffer('v4c_bce_alpha', torch.tensor(0.0),
                                     persistent=False)
                self.register_buffer('v4c_nce_alpha', torch.tensor(0.0),
                                     persistent=False)
                self.salience = SalienceLoop(
                    d, H=8, kernel=5, gate_rank=16,
                    scalar_from_contrast=(cfg.salience_mode in (
                        "v3", "v3m", "v3m2", "v3m2_nll",
                        "v4c", "v4c_gate", "v4d")),
                    contrastive_head=(cfg.salience_mode in (
                        "v4c", "v4c_gate", "v4d")),
                    z_dim=int(getattr(cfg, "v4c_z_dim", 64)),
                    proj_hidden=int(getattr(cfg, "v4c_proj_hidden", 256)),
                )
                self.val_a = nn.Parameter(torch.tensor(1.0))
                self.val_b = nn.Parameter(torch.tensor(0.0))
                self.register_buffer('marker_token_ids', build_marker_token_ids(),
                                     persistent=False)
                _need_nll_ema = cfg.salience_mode in ("v3", "v3m", "v3m2_nll")
                if _need_nll_ema:
                    self.register_buffer('nll_ema_mean',
                                         torch.tensor(0.0))
                    self.register_buffer('nll_ema_var',
                                         torch.tensor(1.0))
                    self.register_buffer('nll_ema_initialized',
                                         torch.tensor(0.0))
                if cfg.salience_mode in ("v3m", "v3m2", "v3m2_nll"):
                    for _name in ("curve_C", "curve_E", "cosmis"):
                        self.register_buffer(f"{_name}_ema_mean",
                                             torch.tensor(0.0))
                        self.register_buffer(f"{_name}_ema_var",
                                             torch.tensor(1.0))
                        self.register_buffer(f"{_name}_ema_initialized",
                                             torch.tensor(0.0))
                if cfg.salience_mode in ("v3m2", "v3m2_nll"):
                    _anchor_neg, _label_neg = build_v3m2_lexicons()
                    self.register_buffer('anchor_negative_ids', _anchor_neg,
                                         persistent=False)
                    self.register_buffer('label_negative_ids',  _label_neg,
                                         persistent=False)

        self.l_stage.per_head.weight = self.c_stage.tok_emb.weight
        self.l_stage.mem_head.weight = self.c_stage.tok_emb.weight

        self.d_model = d
        self.v6_mixture = bool(getattr(cfg, "v6_mixture", False))
        self.v6_lambda_init = float(getattr(cfg, "v6_lambda_init", -4.0))
        self.v6_lambda_a_init = float(getattr(cfg, "v6_lambda_a_init", 1.0))
        self.v6_bank_query_source = str(
            getattr(cfg, "v6_bank_query_source", "h_d"))
        _valid_bank_sources = {"h_d", "h_e", "q_attractor", "q_mem"}
        if self.v6_bank_query_source not in _valid_bank_sources:
            raise ValueError(
                f"unknown v6_bank_query_source={self.v6_bank_query_source!r}; "
                f"expected one of {sorted(_valid_bank_sources)}")
        self.v6_bank_readout_mode = str(
            getattr(cfg, "v6_bank_readout_mode", "bank"))
        _valid_readout_modes = {"bank", "direct"}
        if self.v6_bank_readout_mode not in _valid_readout_modes:
            raise ValueError(
                f"unknown v6_bank_readout_mode={self.v6_bank_readout_mode!r}; "
                f"expected one of {sorted(_valid_readout_modes)}")
        self.use_v6_source_adapter = bool(
            getattr(cfg, "v6_source_adapter", False))
        self.use_v6_context_adapter = bool(
            getattr(cfg, "v6_context_adapter", False))
        self.bank_off = False
        self.outputs_log_probs = self.v6_mixture
        if self.v6_mixture:
            self.bank_q_proj = nn.Linear(d, d)
            nn.init.eye_(self.bank_q_proj.weight)
            nn.init.zeros_(self.bank_q_proj.bias)
            self.v6_lambda_a = nn.Parameter(torch.tensor(self.v6_lambda_a_init))
            self.v6_lambda_b = nn.Parameter(torch.tensor(self.v6_lambda_init))

            if self.use_v6_source_adapter:
                self.v6_source_adapter = V6SourceAdapter(d)
            if self.use_v6_context_adapter:
                self.v6_context_adapter = V6ContextAdapter(d, cfg.n_heads)

            self.use_v6_lambda_head = bool(
                getattr(cfg, "v6_lambda_head", False))
            if self.use_v6_lambda_head:
                _h = int(getattr(cfg, "v6_lambda_head_hidden", 160))
                _b_init = float(getattr(cfg, "v6_lambda_head_bias_init", -7.0))
                _w_std = float(getattr(cfg, "v6_lambda_head_w_init_std", 1e-3))
                self.v6_lambda_head = nn.Sequential(
                    nn.Linear(d, _h),
                    nn.GELU(),
                    nn.Linear(_h, 1),
                )
                nn.init.normal_(self.v6_lambda_head[2].weight, std=_w_std)
                nn.init.constant_(self.v6_lambda_head[2].bias, _b_init)

            self.use_v6_mem_head_bank = bool(
                getattr(cfg, "v6_mem_head_bank", False))
            if self.use_v6_mem_head_bank:
                self.mem_head_bank = nn.Linear(d, vocab, bias=True)
                with torch.no_grad():
                    self.mem_head_bank.weight.copy_(
                        self.c_stage.tok_emb.weight)
                    if self.l_stage.mem_head.bias is not None:
                        self.mem_head_bank.bias.copy_(
                            self.l_stage.mem_head.bias)
                    else:
                        nn.init.zeros_(self.mem_head_bank.bias)

    def _v6_needs_dstage_bank_sources(self) -> bool:
        return self.v6_bank_query_source in ("q_attractor", "q_mem")

    def _v6_bank_query_input(self, h_D: torch.Tensor,
                             h_E: Optional[torch.Tensor] = None,
                             b_idx: Optional[torch.Tensor] = None,
                             p_idx: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Select the source vector used by V6 bank_q_proj.

        The default h_d path preserves original V6 behavior. q_attractor/q_mem
        are populated by DStage when store_v6_bank_sources is enabled.
        """
        src = self.v6_bank_query_source
        if src == "h_d":
            base = h_D
        elif src == "h_e":
            if h_E is None:
                raise RuntimeError("v6_bank_query_source='h_e' requires h_E")
            base = h_E
        elif src == "q_attractor":
            base = getattr(self.d_stage, "_last_q_attractor_v6", None)
            if base is None:
                raise RuntimeError(
                    "q_attractor source requested but DStage did not store it")
        elif src == "q_mem":
            base = getattr(self.d_stage, "_last_q_mem_v6", None)
            if base is None:
                raise RuntimeError(
                    "q_mem source requested but DStage did not store it")
        else:
            raise RuntimeError(f"unknown V6 bank query source {src!r}")
        if b_idx is None:
            return base
        return base[b_idx, p_idx]

    def _v6_direct_readout_source(self, bank_src: torch.Tensor) -> torch.Tensor:
        if getattr(self, "use_v6_context_adapter", False):
            bank_src = self.v6_context_adapter(bank_src)
        if getattr(self, "use_v6_source_adapter", False):
            return self.v6_source_adapter(bank_src)
        return bank_src

    def _v6_bank_logits_from_source(self, bank_src: torch.Tensor,
                                    adapt_direct: bool = True) -> torch.Tensor:
        """Project the selected V6 source to vocab logits.

        `bank` is the original V6 external readout. `direct` bypasses the
        256-slot mem_keys/mem_vals compression after direct-source probing
        showed q_mem is decodable while the bank bottleneck is not.
        """
        if self.v6_bank_readout_mode == "direct":
            if adapt_direct:
                bank_src = self._v6_direct_readout_source(bank_src)
            if getattr(self, "use_v6_mem_head_bank", False):
                return self.mem_head_bank(bank_src)
            return self.l_stage.mem_head(bank_src)

        q_bank = self.bank_q_proj(bank_src)
        alpha = F.softmax(
            q_bank @ self.d_stage.mem_keys.T / math.sqrt(self.d_model),
            dim=-1,
        )
        v_bank = alpha @ self.d_stage.mem_vals
        if getattr(self, "use_v6_mem_head_bank", False):
            return self.mem_head_bank(v_bank)
        return self.l_stage.mem_head(v_bank)

    def forward(self, ids: torch.Tensor, v4c_spans=None,
                v4c_swap_ids=None, v4c_swap_spans=None):
        if v4c_spans is not None:
            assert self.use_salience and self.salience_mode in (
                "v4c", "v4c_gate", "v4d"), \
                "v4c_spans only meaningful for V4c-family models with SL enabled"
            _swap_w = float(getattr(self._v4c_cfg,
                                     "v4d_swap_consistency_weight", 0.0))
            _role_w = float(getattr(self._v4c_cfg, "v4d_role_sep_weight", 0.0))
            _v4e_active = (self.salience_mode == "v4d"
                           and (_swap_w > 0 or _role_w > 0)
                           and v4c_swap_ids is not None)
            if _v4e_active:
                self.d_stage.store_v4d_code = True
            h_C = self.c_stage(ids, feedback=None)
            h_E, h_E_sparse = self.e_stage(h_C, feedback=None)
            v_vec, v_scalar = self.salience(h_E)
            _prev_v6_store = getattr(
                self.d_stage, "store_v6_bank_sources", False)
            self.d_stage.store_v6_bank_sources = (
                self._v6_needs_dstage_bank_sources())
            try:
                h_D, _ = self.d_stage(h_E, h_E_sparse, v_vec=v_vec)
            finally:
                self.d_stage.store_v6_bank_sources = _prev_v6_store
            g_main = getattr(self.d_stage, "_last_sal_code", None) if _v4e_active else None
            if _v4e_active:
                self.d_stage.store_v4d_code = False
            z = self.salience.embed_z(h_D)
            _randlabel = (getattr(self, "v4c_variant", "base") == "randlabel")
            v4c_aux = _v4c_contrastive_loss(
                z=z, v_scalar=v_scalar, spans=v4c_spans, cfg=self._v4c_cfg,
                salience_module=self.salience,
                bce_alpha=getattr(self, "v4c_bce_alpha",
                                   torch.tensor(0.0, device=ids.device)),
                nce_alpha=getattr(self, "v4c_nce_alpha",
                                   torch.tensor(0.0, device=ids.device)),
                randomize_positive_mapping=_randlabel,
            )
            margin_weight = float(getattr(self._v4c_cfg, "v4c_margin_weight", 0.0))
            if margin_weight > 0 and not _randlabel:
                logits_mem, _loop1 = self.l_stage(h_D, h_C)
                margin_target = float(getattr(self._v4c_cfg,
                                                "v4c_margin_target", 1.0))
                L_margin = _v4c_answer_margin_loss(
                    logits_mem, ids, v4c_spans,
                    target_margin=margin_target)
                self.salience._last_v4c_margin_stats = {
                    "loss":            float(L_margin.detach().item()),
                    "n_used":          int(getattr(L_margin, "_v4c_margin_n_used", 0)),
                    "n_skip_neutral":  int(getattr(L_margin, "_v4c_margin_n_skip_neu", 0)),
                    "n_skip_no_spans": int(getattr(L_margin, "_v4c_margin_n_skip_spn", 0)),
                    "n_skip_no_len":   int(getattr(L_margin, "_v4c_margin_n_skip_len", 0)),
                    "n_skip_eq_token": int(getattr(L_margin, "_v4c_margin_n_skip_eq", 0)),
                }
                bce_alpha = getattr(self, "v4c_bce_alpha",
                                    torch.tensor(0.0, device=ids.device))
                v4c_aux = v4c_aux + bce_alpha * margin_weight * L_margin
            self.salience._last_v6_aux_stats = None
            _v6_w = float(getattr(self._v4c_cfg, "v6_aux_weight", 0.0))
            if (self.v6_mixture and _v6_w > 0 and not _randlabel
                    and self.salience_mode == "v4d"):
                if margin_weight <= 0:
                    raise RuntimeError(
                        "V6.1 (v6_aux_weight > 0) requires v4c_margin_weight "
                        "> 0 so logits_mem can be reused from the V4c+M "
                        "branch — otherwise the fallback materializes a "
                        "[B,T,V] tensor (~800 MB per V4c batch). Set both "
                        "knobs together.")
                b_idx, p_idx, cur_tok, stale_tok = _v4c_collect_answer_quads(
                    ids, v4c_spans)
                if b_idx.numel() > 0:
                    h_C_mod = self.modulator(h_C, _loop1, h_D)
                    h_D_ans       = h_D[b_idx, p_idx]
                    h_C_mod_ans   = h_C_mod[b_idx, p_idx]
                    v_scalar_ans  = v_scalar[b_idx, p_idx]
                    logits_mem_a  = logits_mem[b_idx, p_idx]
                    logits_per_a  = self.l_stage.per_head(h_C_mod_ans)
                    gate_a = self.l_stage.gate_net(
                        torch.cat([h_D_ans, h_C_mod_ans], dim=-1))
                    logits_trunk_a = (gate_a * logits_mem_a
                                      + (1 - gate_a) * logits_per_a)
                    if getattr(self, "use_v6_lambda_head", False):
                        lam_logit_a = self.v6_lambda_head(h_D_ans)
                    else:
                        lam_logit_a = (self.v6_lambda_a * v_scalar_ans.detach()
                                       + self.v6_lambda_b)
                    lam_a = torch.sigmoid(lam_logit_a).clamp(
                        1e-4, 1.0 - 1e-4)
                    if getattr(self, "use_v6_context_adapter", False):
                        bank_src_seq = self._v6_bank_query_input(h_D, h_E)
                        bank_src_seq = self._v6_direct_readout_source(
                            bank_src_seq)
                        bank_src_a = bank_src_seq[b_idx, p_idx]
                        logits_bank_a = self._v6_bank_logits_from_source(
                            bank_src_a, adapt_direct=False)
                    else:
                        bank_src_a = self._v6_bank_query_input(
                            h_D, h_E, b_idx, p_idx)
                        logits_bank_a = self._v6_bank_logits_from_source(
                            bank_src_a)
                    log_p_mix_a = torch.logaddexp(
                        torch.log1p(-lam_a) + F.log_softmax(
                            logits_trunk_a, dim=-1),
                        torch.log(lam_a) + F.log_softmax(
                            logits_bank_a, dim=-1))
                    _target = float(getattr(self._v4c_cfg,
                                            "v6_margin_target", 1.0))
                    log_p_bank_a = F.log_softmax(logits_bank_a, dim=-1)
                    lp_bk_cur   = log_p_bank_a.gather(
                        1, cur_tok.unsqueeze(1)).squeeze(1)
                    lp_bk_stale = log_p_bank_a.gather(
                        1, stale_tok.unsqueeze(1)).squeeze(1)
                    L_bank = F.relu(_target + lp_bk_stale - lp_bk_cur).mean()
                    L_bank_ce = torch.zeros((), device=ids.device)
                    _w_bank_ce = float(getattr(self._v4c_cfg,
                                                "v6_bank_ce_weight", 0.0))
                    if (_w_bank_ce > 0
                            and getattr(self, "use_v6_mem_head_bank", False)):
                        L_bank_ce = F.cross_entropy(logits_bank_a, cur_tok)
                    L_bank_pair_ce = torch.zeros((), device=ids.device)
                    _w_bank_pair_ce = float(getattr(
                        self._v4c_cfg, "v6_bank_pair_ce_weight", 0.0))
                    if (_w_bank_pair_ce > 0
                            and getattr(self, "use_v6_mem_head_bank", False)):
                        bk_cur_logit = logits_bank_a.gather(
                            1, cur_tok.unsqueeze(1)).squeeze(1)
                        bk_stale_logit = logits_bank_a.gather(
                            1, stale_tok.unsqueeze(1)).squeeze(1)
                        pair_logits = torch.stack(
                            [bk_cur_logit, bk_stale_logit], dim=1)
                        L_bank_pair_ce = F.cross_entropy(
                            pair_logits,
                            torch.zeros_like(cur_tok, dtype=torch.long))
                    lp_mix_cur   = log_p_mix_a.gather(
                        1, cur_tok.unsqueeze(1)).squeeze(1)
                    lp_mix_stale = log_p_mix_a.gather(
                        1, stale_tok.unsqueeze(1)).squeeze(1)
                    L_mix = F.relu(_target + lp_mix_stale - lp_mix_cur).mean()
                    _lam_floor = float(getattr(self._v4c_cfg,
                                                "v6_lambda_floor", 0.05))
                    _lam_floor_logit = math.log(
                        _lam_floor / (1.0 - _lam_floor))
                    L_gate = F.relu(_lam_floor_logit
                                    - lam_logit_a.view(-1)).mean()
                    L_bg = torch.zeros((), device=ids.device)
                    _bg_mean_val = 0.0
                    _bg_std_val  = 0.0
                    _bg_n_items  = 0
                    if getattr(self, "use_v6_lambda_head", False):
                        b_bg, p_bg = _v4c_collect_background_quads(
                            ids, v4c_spans,
                            ans_b_idx=b_idx, ans_p_idx=p_idx,
                            k_per_item=8)
                        if b_bg.numel() > 0:
                            h_D_bg = h_D[b_bg, p_bg]
                            lam_logit_bg = self.v6_lambda_head(h_D_bg).squeeze(-1)
                            lam_bg = torch.sigmoid(lam_logit_bg).clamp(
                                1e-4, 1.0 - 1e-4)
                            _bg_target = float(getattr(self._v4c_cfg,
                                                        "v6_bg_target", 0.01))
                            L_bg = F.relu(lam_bg - _bg_target).mean()
                            with torch.no_grad():
                                _lb = lam_bg.detach().float()
                                _bg_mean_val = float(_lb.mean().item())
                                _bg_std_val  = float(_lb.std(
                                    unbiased=False).item())
                                _bg_n_items  = int(_lb.numel())

                    _use_bce = (getattr(self, "use_v6_lambda_head", False)
                                and bool(getattr(self._v4c_cfg,
                                                 "v6_bce_objective", False)))
                    _w_mix  = float(getattr(self._v4c_cfg,
                                            "v6_mix_weight",  0.5))
                    L_sel_ans = torch.zeros((), device=ids.device)
                    L_sel_bg  = torch.zeros((), device=ids.device)
                    L_sel     = torch.zeros((), device=ids.device)
                    if _use_bce:
                        L_sel_ans = -F.logsigmoid(lam_logit_a.view(-1)).mean()
                        if (getattr(self, "use_v6_lambda_head", False)
                                and b_bg.numel() > 0):
                            L_sel_bg = -F.logsigmoid(
                                -lam_logit_bg.view(-1)).mean()
                        L_sel = 0.5 * L_sel_ans + 0.5 * L_sel_bg
                        _w_sel = float(getattr(self._v4c_cfg,
                                                "v6_sel_weight", 1.0))
                        L_v6 = (L_bank + _w_mix * L_mix + _w_sel * L_sel
                                + _w_bank_ce * L_bank_ce
                                + _w_bank_pair_ce * L_bank_pair_ce)
                    else:
                        _w_gate = float(getattr(self._v4c_cfg,
                                                "v6_gate_weight", 0.01))
                        _w_bg   = float(getattr(self._v4c_cfg,
                                                "v6_bg_weight",   1.0))
                        L_v6 = (L_bank + _w_mix * L_mix
                                + _w_gate * L_gate + _w_bg * L_bg
                                + _w_bank_ce * L_bank_ce
                                + _w_bank_pair_ce * L_bank_pair_ce)
                    bce_alpha = getattr(self, "v4c_bce_alpha",
                                        torch.tensor(0.0, device=ids.device))
                    v4c_aux = v4c_aux + bce_alpha * _v6_w * L_v6
                    with torch.no_grad():
                        _l = lam_a.detach().float().view(-1)
                        _llg = lam_logit_a.detach().float().view(-1)
                        _stats = {
                            "loss":   float(L_v6.detach().item()),
                            "L_bank": float(L_bank.detach().item()),
                            "L_bank_ce": float(L_bank_ce.detach().item()),
                            "L_bank_pair_ce": float(
                                L_bank_pair_ce.detach().item()),
                            "L_mix":  float(L_mix.detach().item()),
                            "L_gate": float(L_gate.detach().item()),
                            "lam_ans_mean":     float(_l.mean().item()),
                            "lam_ans_std":      float(_l.std(unbiased=False).item()),
                            "lam_ans_sat_low":  float((_l < 1e-3).float().mean().item()),
                            "lam_ans_sat_high": float((_l > 1.0 - 1e-3).float().mean().item()),
                            "n_items":          int(_l.numel()),
                            "v6_lambda_a":         float(self.v6_lambda_a.item()),
                            "v6_lambda_b":         float(self.v6_lambda_b.item()),
                            "lam_logit_ans_mean":  float(_llg.mean().item()),
                            "lam_logit_ans_std":   float(_llg.std(unbiased=False).item()),
                            "L_bg":         float(L_bg.detach().item()),
                            "lam_bg_mean":  _bg_mean_val,
                            "lam_bg_std":   _bg_std_val,
                            "n_bg_items":   _bg_n_items,
                            "L_sel":        float(L_sel.detach().item()),
                            "L_sel_ans":    float(L_sel_ans.detach().item()),
                            "L_sel_bg":     float(L_sel_bg.detach().item()),
                            "bce_active":   bool(_use_bce),
                        }
                        if getattr(self, "use_v6_lambda_head", False):
                            _stats["v6_lambda_head_bias"] = float(
                                self.v6_lambda_head[2].bias.item())
                        self.salience._last_v6_aux_stats = _stats
            causal_weight   = float(getattr(self._v4c_cfg, "v4d_causal_weight",   0.0))
            causal_z_weight = float(getattr(self._v4c_cfg, "v4d_causal_z_weight", 0.0))
            run_clamped_fwd = ((causal_weight > 0 or causal_z_weight > 0)
                                and not _randlabel
                                and self.salience_mode == "v4d")
            if run_clamped_fwd:
                def _zero_v_vec_hook(_mod, _inputs, output):
                    v_vec, v_scalar = output
                    return torch.zeros_like(v_vec), v_scalar
                handle = self.salience.register_forward_hook(_zero_v_vec_hook)
                try:
                    h_C2 = self.c_stage(ids, feedback=None)
                    h_E2, h_E_sparse2 = self.e_stage(h_C2, feedback=None)
                    v_vec2, v_scalar2 = self.salience(h_E2)
                    h_D2, _ = self.d_stage(h_E2, h_E_sparse2, v_vec=v_vec2)
                    logits_mem_clamped = None
                    if causal_weight > 0 and margin_weight > 0:
                        logits_mem_clamped, _ = self.l_stage(h_D2, h_C2)
                    z_clamped = self.salience.embed_z(h_D2)
                finally:
                    handle.remove()
                nce_alpha = getattr(self, "v4c_nce_alpha",
                                     torch.tensor(0.0, device=ids.device))
                if (causal_weight > 0 and margin_weight > 0
                        and logits_mem_clamped is not None):
                    causal_gap = float(getattr(self._v4c_cfg,
                                                "v4d_causal_gap", 0.25))
                    margins_n, item_idxs_n = _v4d_compute_per_item_margins(
                        logits_mem, ids, v4c_spans)
                    margins_c, item_idxs_c = _v4d_compute_per_item_margins(
                        logits_mem_clamped, ids, v4c_spans)
                    if item_idxs_n == item_idxs_c and margins_n.numel() > 0:
                        L_causal = _v4d_causal_dependence_loss(
                            margins_n, margins_c, causal_gap=causal_gap)
                        v4c_aux = v4c_aux + nce_alpha * causal_weight * L_causal
                        self.salience._last_v4d_causal_stats = {
                            "loss": float(L_causal.detach().item()),
                            "margin_normal_mean":  float(margins_n.detach().mean().item()),
                            "margin_clamped_mean": float(margins_c.detach().mean().item()),
                            "margin_gap_mean": float((margins_n.detach()
                                                    - margins_c.detach()).mean().item()),
                            "n_items": int(margins_n.numel()),
                        }
                if causal_z_weight > 0:
                    causal_z_gap = float(getattr(self._v4c_cfg,
                                                  "v4d_causal_z_gap", 0.15))
                    mz_n, item_idxs_zn = _v4d_compute_per_item_z_margins(z, v4c_spans)
                    mz_c, item_idxs_zc = _v4d_compute_per_item_z_margins(z_clamped, v4c_spans)
                    if item_idxs_zn == item_idxs_zc and mz_n.numel() > 0:
                        L_causal_z = _v4d_causal_dependence_loss(
                            mz_n, mz_c, causal_gap=causal_z_gap)
                        v4c_aux = v4c_aux + nce_alpha * causal_z_weight * L_causal_z
                        self.salience._last_v4d_causal_z_stats = {
                            "loss": float(L_causal_z.detach().item()),
                            "mz_normal_mean":  float(mz_n.detach().mean().item()),
                            "mz_clamped_mean": float(mz_c.detach().mean().item()),
                            "mz_gap_mean": float((mz_n.detach()
                                                  - mz_c.detach()).mean().item()),
                            "n_items": int(mz_n.numel()),
                        }
            if _v4e_active and g_main is not None:
                self.d_stage.store_v4d_code = True
                try:
                    h_Cs = self.c_stage(v4c_swap_ids, feedback=None)
                    h_Es, h_Es_sp = self.e_stage(h_Cs, feedback=None)
                    v_vecs, _ = self.salience(h_Es)
                    self.d_stage(h_Es, h_Es_sp, v_vec=v_vecs)
                    g_swap = getattr(self.d_stage, "_last_sal_code", None)
                finally:
                    self.d_stage.store_v4d_code = False
                nce_alpha = getattr(self, "v4c_nce_alpha",
                                     torch.tensor(0.0, device=ids.device))
                m_pool = _v4d_pool_code_by_role(g_main, v4c_spans)
                s_pool = _v4d_pool_code_by_role(g_swap, v4c_swap_spans)
                common = sorted(set(m_pool) & set(s_pool))
                if common:
                    gc_m = torch.stack([m_pool[b][0] for b in common])
                    gs_m = torch.stack([m_pool[b][1] for b in common])
                    gc_s = torch.stack([s_pool[b][0] for b in common])
                    gs_s = torch.stack([s_pool[b][1] for b in common])
                    if _swap_w > 0:
                        L_swap = ((gc_m - gc_s).pow(2).sum(-1)
                                   + (gs_m - gs_s).pow(2).sum(-1)).mean()
                        v4c_aux = v4c_aux + nce_alpha * _swap_w * L_swap
                        self.salience._last_v4d_swap_stats = {
                            "loss": float(L_swap.detach().item()),
                            "n_pairs": int(len(common)),
                        }
                    if _role_w > 0 and hasattr(self, "v4d_role_head"):
                        feats = torch.cat([gc_m, gs_m, gc_s, gs_s], dim=0)
                        n = gc_m.size(0)
                        labels = torch.cat([
                            torch.ones(n, dtype=torch.long, device=ids.device),
                            torch.zeros(n, dtype=torch.long, device=ids.device),
                            torch.ones(n, dtype=torch.long, device=ids.device),
                            torch.zeros(n, dtype=torch.long, device=ids.device)])
                        role_logits = self.v4d_role_head(feats)
                        L_role = F.cross_entropy(role_logits, labels)
                        v4c_aux = v4c_aux + nce_alpha * _role_w * L_role
                        with torch.no_grad():
                            role_acc = (role_logits.argmax(-1)
                                        == labels).float().mean().item()
                        self.salience._last_v4d_role_stats = {
                            "loss": float(L_role.detach().item()),
                            "acc": float(role_acc),
                            "n": int(4 * n),
                        }
            return None, v4c_aux

        h_C = self.c_stage(ids, feedback=None)
        h_E, h_E_sparse = self.e_stage(h_C, feedback=None)
        if self.use_salience:
            v_vec, v_scalar = self.salience(h_E)
        else:
            v_vec, v_scalar = None, None
        _prev_v6_store = getattr(self.d_stage, "store_v6_bank_sources", False)
        self.d_stage.store_v6_bank_sources = (
            self._v6_needs_dstage_bank_sources())
        try:
            h_D, loop2_signal = self.d_stage(h_E, h_E_sparse, v_vec=v_vec)
            logits_mem, loop1_fb = self.l_stage(h_D, h_C)
            v6_bank_src_p0 = (
                self._v6_bank_query_input(h_D, h_E)
                if (self.v6_mixture and not self.bank_off)
                else None
            )

            h_C_p0, h_E_p0, h_E_sparse_p0, h_D_p0 = h_C, h_E, h_E_sparse, h_D
            if self.use_salience:
                logits_mem_p0 = logits_mem
                v_scalar_p0 = v_scalar

            alpha = self.feedback_alpha.item()
            for i in range(self.n_feedback_iters):
                decay = self.feedback_decay ** i
                fb_scale = alpha * decay

                if fb_scale > 0:
                    h_C = self.c_stage(ids, feedback=fb_scale * loop1_fb)
                    h_E, h_E_sparse = self.e_stage(h_C, feedback=fb_scale * loop2_signal)
                    h_D, loop2_signal = self.d_stage(h_E, h_E_sparse, v_vec=v_vec)
                    logits_mem, loop1_fb = self.l_stage(h_D, h_C)
        finally:
            self.d_stage.store_v6_bank_sources = _prev_v6_store

        aux_loss = torch.tensor(0.0, device=ids.device)
        if self.training:
            aux_loss = aux_loss + 0.05 * self.e_stage.csr_loss(h_C_p0, h_E_sparse_p0)
            aux_loss = aux_loss + 0.1 * sigreg_loss(h_D_p0)
            aux_loss = aux_loss + 0.05 * predictive_latent_loss(h_D_p0, h_C_p0)

            if self.use_salience:
                aux_loss = aux_loss + 1e-3 * anchor_ortho_reg(self.salience.U)

                sl_a = self.sl_alpha

                if self.salience_mode in ("v0", "v2a"):
                    if sl_a.item() > 0:
                        with torch.no_grad():
                            flat_logits = logits_mem_p0[:, :-1].reshape(
                                -1, logits_mem_p0.size(-1))
                            nll_flat = F.cross_entropy(
                                flat_logits.float(), ids[:, 1:].reshape(-1),
                                reduction='none')
                            nll = nll_flat.view(ids.size(0), ids.size(1) - 1)
                            nll_z = torch.sigmoid(
                                (nll - nll.mean()) / (nll.std() + 1e-6))
                        pred = torch.sigmoid(
                            self.val_a * v_scalar_p0[:, :-1, 0] + self.val_b)
                        aux_loss = aux_loss + 0.03 * sl_a * F.mse_loss(pred, nll_z)
                        if self.marker_token_ids.numel() > 0:
                            mask = build_marker_mask(ids, self.marker_token_ids)
                            logits_mark = (self.val_a * v_scalar_p0[:, :, 0]
                                           + self.val_b)
                            aux_loss = aux_loss + 0.02 * sl_a * \
                                F.binary_cross_entropy_with_logits(
                                    logits_mark, mask)

                elif self.salience_mode in ("v3", "v3m"):
                    with torch.no_grad():
                        flat_logits = logits_mem_p0[:, :-1].reshape(
                            -1, logits_mem_p0.size(-1)).float()
                        nll = F.cross_entropy(
                            flat_logits, ids[:, 1:].reshape(-1),
                            reduction='none').view(ids.size(0), -1)
                        nll_z = _ema_zscore(
                            nll, self.nll_ema_mean, self.nll_ema_var,
                            self.nll_ema_initialized,
                            self.nll_ema_decay, self.training
                        ).clamp(-self.sl_z_clip, self.sl_z_clip)

                        if self.salience_mode == "v3":
                            raw_t = (nll_z >= 1.5).float()
                            pad = F.pad(raw_t.unsqueeze(1), (0, self.sl_window_k))
                            target = F.max_pool1d(
                                pad, kernel_size=self.sl_window_k + 1,
                                stride=1).squeeze(1)[:, :nll_z.size(1)]
                        else:
                            def _curvature(h):
                                h_n = F.normalize(h.detach().float(), dim=-1)
                                d1 = F.normalize(h_n[:, 1:-1] - h_n[:, :-2], dim=-1)
                                d2 = F.normalize(h_n[:, 2:]   - h_n[:, 1:-1], dim=-1)
                                cc = 1.0 - (d1 * d2).sum(-1)
                                return F.pad(cc, (1, 0))

                            curve_C = _curvature(h_C_p0)
                            curve_E = _curvature(h_E_p0)
                            h_C_n = F.normalize(h_C_p0.detach().float(), dim=-1)
                            h_E_n = F.normalize(h_E_p0.detach().float(), dim=-1)
                            cosmis = (1.0 - (h_C_n * h_E_n).sum(-1))[:, :-1]

                            curve_C_z = _ema_zscore(
                                curve_C, self.curve_C_ema_mean,
                                self.curve_C_ema_var,
                                self.curve_C_ema_initialized,
                                self.nll_ema_decay, self.training
                            ).clamp(-self.sl_z_clip, self.sl_z_clip)
                            curve_E_z = _ema_zscore(
                                curve_E, self.curve_E_ema_mean,
                                self.curve_E_ema_var,
                                self.curve_E_ema_initialized,
                                self.nll_ema_decay, self.training
                            ).clamp(-self.sl_z_clip, self.sl_z_clip)
                            cosmis_z = _ema_zscore(
                                cosmis, self.cosmis_ema_mean,
                                self.cosmis_ema_var,
                                self.cosmis_ema_initialized,
                                self.nll_ema_decay, self.training
                            ).clamp(-self.sl_z_clip, self.sl_z_clip)

                            target_logit = (0.7 * nll_z
                                            + 0.5 * curve_C_z
                                            + 0.5 * curve_E_z
                                            + 0.7 * cosmis_z
                                            - 2.5)
                            target_soft = torch.sigmoid(target_logit)
                            pad = F.pad(target_soft.unsqueeze(1),
                                        (0, self.sl_window_k))
                            target = F.max_pool1d(
                                pad, kernel_size=self.sl_window_k + 1,
                                stride=1).squeeze(1)[:, :nll_z.size(1)]

                            if self.training:
                                mean_target = target.mean().item()
                                frac_high   = (target > 0.5).float().mean().item()
                                pred_sig    = torch.sigmoid(
                                    v_scalar_p0[:, :-1, 0]).detach()
                                frac_v_high = (pred_sig > 0.5).float().mean().item()
                                step_idx = getattr(
                                    self, '_v3m_density_log_step', 0)
                                if step_idx % 100 == 0:
                                    print(f"  [v3m_density step={step_idx:>5d}] "
                                          f"mean_target={mean_target:.3f}  "
                                          f"frac_target>0.5={frac_high:.3f}  "
                                          f"frac_sig_v_scalar>0.5={frac_v_high:.3f}")
                                if 200 <= step_idx < 1000 and frac_v_high > 0.8:
                                    print(f"  [v3m_density EARLY-ABORT] "
                                          f"frac_sig_v_scalar>0.5={frac_v_high:.3f} "
                                          f"> 0.8 before sl_alpha starts at "
                                          f"step 1000 — scalar_head init is "
                                          f"saturated, ABORT and check "
                                          f"scalar_head.bias initialization "
                                          f"(should be ≲ -2.0 for sparse "
                                          f"salience prior)")
                                if step_idx >= 2000:
                                    if frac_high > 0.50:
                                        print(f"  [v3m_density WARNING] "
                                              f"frac_target>0.5={frac_high:.3f} "
                                              f"> 0.50 — V3m target severely broad")
                                    if frac_v_high > 0.80:
                                        print(f"  [v3m_density WARNING] "
                                              f"frac_sig_v_scalar>0.5={frac_v_high:.3f} "
                                              f"> 0.80 — v_scalar saturating")
                                self._v3m_density_log_step = step_idx + 1

                        if self.salience_mode == "v3":
                            w = 1.0 + 4.0 * torch.sigmoid(nll_z.detach())
                        else:
                            w = 1.0 + 4.0 * target.detach()
                        w = w / w.mean()

                    if sl_a.item() > 0:
                        v_logit = v_scalar_p0[:, :-1, 0]
                        aux_loss = aux_loss + 0.05 * sl_a * (
                            F.binary_cross_entropy_with_logits(
                                v_logit, target, reduction='none') * w
                        ).mean()

                        if (self.lex_anchor_weight > 0
                                and self.marker_token_ids.numel() > 0):
                            mask = build_marker_mask(ids, self.marker_token_ids)
                            v_logit_full = v_scalar_p0[:, :, 0]
                            aux_loss = aux_loss + self.lex_anchor_weight * sl_a * \
                                F.binary_cross_entropy_with_logits(
                                    v_logit_full, mask)

                elif self.salience_mode in ("v3m2", "v3m2_nll"):
                    with torch.no_grad():
                        target, w, _dbg = _build_v3m2_targets(
                            ids=ids,
                            h_C_p0=h_C_p0, h_E_p0=h_E_p0,
                            logits_mem_p0=logits_mem_p0,
                            model=self,
                            salience_mode=self.salience_mode,
                            update_ema=True,
                            with_nll=(self.salience_mode == "v3m2_nll"),
                        )

                        mean_target = target.mean().item()
                        frac_pos    = (target > 0.1).float().mean().item()
                        pred_sig    = torch.sigmoid(
                            v_scalar_p0[:, :-1, 0]).detach()
                        frac_v_high = (pred_sig > 0.5).float().mean().item()
                        step_idx = getattr(self, '_v3m2_density_log_step', 0)
                        if step_idx % 100 == 0:
                            print(f"  [v3m2_density step={step_idx:>5d}] "
                                  f"mean_target={mean_target:.3f}  "
                                  f"frac_pos={frac_pos:.3f}  "
                                  f"frac_sig_v_scalar>0.5={frac_v_high:.3f}")
                        if 200 <= step_idx < 1000 and frac_v_high > 0.8:
                            print(f"  [v3m2_density EARLY-ABORT] "
                                  f"frac_sig_v_scalar>0.5={frac_v_high:.3f} "
                                  f"> 0.8 before sl_alpha starts at step 1000 "
                                  f"— scalar_head init saturated, ABORT")
                        if step_idx >= 3000:
                            if frac_pos > 0.30:
                                print(f"  [v3m2_density WARNING] "
                                      f"frac_pos={frac_pos:.3f} > 0.30 "
                                      f"— V3m2 target too broad (anchor "
                                      f"selection should reject routine text)")
                            if frac_v_high > 0.50:
                                print(f"  [v3m2_density WARNING] "
                                      f"frac_sig_v_scalar>0.5={frac_v_high:.3f} "
                                      f"> 0.50 — v_scalar saturating")
                        self._v3m2_density_log_step = step_idx + 1

                    if sl_a.item() > 0:
                        v_logit = v_scalar_p0[:, :-1, 0]
                        aux_loss = aux_loss + 0.05 * sl_a * (
                            F.binary_cross_entropy_with_logits(
                                v_logit, target, reduction='none') * w
                        ).mean()
        del h_C_p0, h_E_p0, h_E_sparse_p0, h_D_p0
        lam = None
        if (self.v6_mixture and not self.bank_off
                and self.use_salience and v_scalar is not None):
            if getattr(self, "use_v6_lambda_head", False):
                _wt_w = float(getattr(self._v4c_cfg,
                                       "v6_wt_sparsity_weight", 0.0))
                if self.training and _wt_w > 0:
                    lam_logit_main = self.v6_lambda_head(h_D)
                    lam_with_grad = torch.sigmoid(
                        lam_logit_main).clamp(1e-4, 1.0 - 1e-4)
                    _wt_target = float(getattr(self._v4c_cfg,
                                                "v6_wt_sparsity_target", 0.0))
                    if _wt_target > 0:
                        wt_term = F.relu(
                            lam_with_grad - _wt_target).mean()
                    else:
                        wt_term = lam_with_grad.mean()
                    aux_loss = aux_loss + _wt_w * wt_term
                    lam = lam_with_grad.detach()
                else:
                    with torch.no_grad():
                        lam_logit_main = self.v6_lambda_head(h_D)
                    lam = torch.sigmoid(lam_logit_main).clamp(1e-4, 1.0 - 1e-4)
            else:
                lam = torch.sigmoid(self.v6_lambda_a * v_scalar.detach()
                                    + self.v6_lambda_b)
                lam = lam.clamp(1e-4, 1.0 - 1e-4)
            if not self.training:
                with torch.no_grad():
                    _l = lam.detach().float().view(-1)
                    _l_cpu = _l.cpu()
                    _gt = (_l > 0.7)
                    _stats = {
                        "lam_mean": float(_l.mean().item()),
                        "lam_std":  float(_l.std(unbiased=False).item()),
                        "lam_sat_low":  float((_l < 1e-3).float().mean().item()),
                        "lam_sat_high": float((_l > 1.0 - 1e-3).float().mean().item()),
                        "lam_max":  float(_l.max().item()),
                        "lam_gt_0_7_count": int(_gt.sum().item()),
                        "lam_total_count": int(_l.numel()),
                        "lam_tensor_cpu": _l_cpu,
                        "v6_lambda_a": float(self.v6_lambda_a.item()),
                        "v6_lambda_b": float(self.v6_lambda_b.item()),
                    }
                    self._last_lam_stats = _stats
        if self.use_salience:
            del logits_mem_p0, v_scalar_p0, v_vec, v_scalar

        h_E_bank = None
        del h_E_sparse
        del h_E

        h_C_mod = self.modulator(h_C, loop1_fb, h_D)
        del loop1_fb
        logits_per_mod = self.l_stage.per_head(h_C_mod)
        gate_alpha = self.l_stage.gate_net(torch.cat([h_D, h_C_mod], dim=-1))
        del h_C, h_C_mod

        logits_per_mod.mul_(1 - gate_alpha)
        logits_mem.mul_(gate_alpha)
        logits_per_mod.add_(logits_mem)
        del logits_mem

        if lam is not None:
            bank_src = v6_bank_src_p0
            if bank_src is None:
                raise RuntimeError(
                    "V6 readout requested but pass-0 bank source was not captured")
            logits_bank = self._v6_bank_logits_from_source(bank_src)
            del h_D, bank_src
            if h_E_bank is not None:
                del h_E_bank
            log_p_trunk = F.log_softmax(logits_per_mod, dim=-1)
            del logits_per_mod
            log_p_bank = F.log_softmax(logits_bank, dim=-1)
            del logits_bank
            log_p_final = torch.logaddexp(
                torch.log1p(-lam) + log_p_trunk,
                torch.log(lam) + log_p_bank)
            del log_p_trunk, log_p_bank, lam
            return log_p_final, aux_loss

        del h_D
        if h_E_bank is not None:
            del h_E_bank
        return logits_per_mod, aux_loss



class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
        super().__init__()
        self.attn = nn.MultiheadAttention(
            d_model, n_heads, dropout=dropout, batch_first=True)

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        out, _ = self.attn(x, x, x, attn_mask=mask)
        return out


class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, ffn_dim: int,
                 dropout: float = 0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads, dropout)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, ffn_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_dim, d_model),
            nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln1(x), mask)
        x = x + self.ffn(self.ln2(x))
        return x


class TransformerLM100M(nn.Module):
    """GPT-2-style causal Transformer at 100M scale."""

    def __init__(self, cfg: TransformerConfig, vocab: int, max_seq: int):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab, cfg.d_model)
        self.pos_emb = nn.Embedding(max_seq, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList([
            TransformerBlock(cfg.d_model, cfg.n_heads, cfg.ffn_dim, cfg.dropout)
            for _ in range(cfg.n_layers)
        ])
        self.ln_f = nn.LayerNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, vocab, bias=False)
        self.head.weight = self.tok_emb.weight

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        B, T = ids.shape
        pos = torch.arange(T, device=ids.device).unsqueeze(0)
        x = self.drop(self.tok_emb(ids) + self.pos_emb(pos))
        mask = torch.triu(torch.ones(T, T, device=ids.device, dtype=torch.bool), 1)
        for block in self.blocks:
            x = block(x, mask)
        return self.head(self.ln_f(x))



class TransformerXLLM100M(nn.Module):
    """Transformer-XL at 100M scale with segment-level recurrence.

    Note: Uses learned absolute position embeddings (applied only to the
    current segment) rather than the relative position encodings of the
    original Dai et al. 2019 paper. This is a simplification — memory
    tokens carry positional information from their original segment.

    The attention mask gives memory tokens full visibility while maintaining
    causal masking within the current segment.
    """

    def __init__(self, cfg: TransformerXLConfig, vocab: int, max_seq: int):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(vocab, cfg.d_model)
        self.pos_emb = nn.Embedding(max_seq, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList([
            TransformerBlock(cfg.d_model, cfg.n_heads, cfg.ffn_dim, cfg.dropout)
            for _ in range(cfg.n_layers)
        ])
        self.ln_f = nn.LayerNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, vocab, bias=False)
        self.head.weight = self.tok_emb.weight
        self.mems: List[Optional[torch.Tensor]] = [None] * cfg.n_layers

    def reset_memory(self):
        self.mems = [None] * self.cfg.n_layers

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        B, T = ids.shape
        pos = torch.arange(T, device=ids.device).unsqueeze(0)
        x = self.drop(self.tok_emb(ids) + self.pos_emb(pos))

        for i, block in enumerate(self.blocks):
            mem = self.mems[i]
            if mem is not None:
                M = mem.size(1)
                cat = torch.cat([mem, x], dim=1)
            else:
                M = 0
                cat = x
            S = M + T
            mask = torch.zeros(S, S, device=ids.device, dtype=torch.bool)
            if T > 1:
                mask[M:, M:] = torch.triu(
                    torch.ones(T, T, device=ids.device, dtype=torch.bool), 1)

            if self.training:
                out = torch.utils.checkpoint.checkpoint(
                    block, cat, mask, use_reentrant=False)
            else:
                out = block(cat, mask)
            x = out[:, -T:]

            with torch.no_grad():
                self.mems[i] = x[:, -self.cfg.mem_len:].detach()

        return self.head(self.ln_f(x))



class RetNetLM100M(nn.Module):
    """Multi-scale retention LM at 100M scale.

    Uses learned absolute position embeddings instead of the original RetNet
    paper's xPos. The decay matrix provides temporal weighting; absolute
    embeddings provide positional discrimination — a simpler but effective
    substitute for a benchmark comparison.
    """

    def __init__(self, cfg: RetNetConfig, vocab: int, max_seq: int):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab, cfg.d_model)
        self.pos_emb = nn.Embedding(max_seq, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.layers = nn.ModuleList([
            RetentionLayer(cfg.d_model, cfg.n_heads, cfg.ffn_dim, cfg.dropout)
            for _ in range(cfg.n_layers)
        ])
        self.ln_f = nn.LayerNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, vocab, bias=False)
        self.head.weight = self.tok_emb.weight

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        B, T = ids.shape
        pos = torch.arange(T, device=ids.device).unsqueeze(0)
        x = self.drop(self.tok_emb(ids) + self.pos_emb(pos))
        for layer in self.layers:
            x = layer(x)
        return self.head(self.ln_f(x))



class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization (used by Mamba)."""
    def __init__(self, d: int, eps: float = 1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight


class SelectiveSSM(nn.Module):
    """Selective state-space model (S6) — pure PyTorch implementation.

    Faithful to Gu & Dao 2023: per-channel delta via low-rank dt_proj,
    input-dependent B/C, S4D-Lin initialization for A, parallel scan.

    Core: h[t] = A_bar[t] * h[t-1] + B_bar[t] * x[t]
          y[t] = C[t] * h[t]
    """

    def __init__(self, d_inner: int, d_state: int = 16, d_conv: int = 4,
                 dt_rank: Optional[int] = None):
        super().__init__()
        self.d_inner = d_inner
        self.d_state = d_state
        if dt_rank is None:
            dt_rank = math.ceil(d_inner / 16)
        self.dt_rank = dt_rank

        self.conv1d = nn.Conv1d(d_inner, d_inner, d_conv, padding=d_conv - 1,
                                groups=d_inner)

        self.x_proj = nn.Linear(d_inner, dt_rank + d_state * 2, bias=False)
        self.dt_proj = nn.Linear(dt_rank, d_inner)
        dt = torch.exp(
            torch.rand(d_inner) * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
        )
        inv_dt = dt + torch.log(-torch.expm1(-dt))
        with torch.no_grad():
            self.dt_proj.bias.copy_(inv_dt)

        self.A_log = nn.Parameter(
            torch.log(torch.arange(1, d_state + 1, dtype=torch.float32)
                      .unsqueeze(0).expand(d_inner, -1).clone())
        )
        self.D = nn.Parameter(torch.ones(d_inner))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """x: [B, T, d_inner] -> y: [B, T, d_inner]"""
        B, T, D = x.shape
        N = self.d_state

        x_conv = self.conv1d(x.transpose(1, 2))[:, :, :T].transpose(1, 2)
        x_conv = F.silu(x_conv)

        proj = self.x_proj(x_conv)
        dt_input, B_input, C_input = proj.split(
            [self.dt_rank, N, N], dim=-1)
        delta = F.softplus(self.dt_proj(dt_input))

        A = -torch.exp(self.A_log)
        A_bar = torch.exp(
            delta.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0)
        )
        B_bar = delta.unsqueeze(-1) * B_input.unsqueeze(2)

        a = A_bar
        b = B_bar * x_conv.unsqueeze(-1)
        h_all = self._chunked_scan(a, b)

        y = (h_all * C_input[:, :, None, :]).sum(-1)
        y = y + x * self.D
        return y

    @staticmethod
    def _chunked_scan(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """Chunked sequential scan for linear recurrence h[t] = a[t]*h[t-1] + b[t].

        Sequential within chunks of 64 steps, state propagated between chunks.
        O(T) total iterations with good cache locality.

        Args:
            a: [B, T, D, N] multiplicative coefficients
            b: [B, T, D, N] additive terms
        Returns:
            h: [B, T, D, N] all hidden states
        """
        T = a.shape[1]
        CHUNK = 64
        h_list = []
        h_prev = torch.zeros_like(a[:, 0])
        for start in range(0, T, CHUNK):
            end = min(start + CHUNK, T)
            a_chunk = a[:, start:end]
            b_chunk = b[:, start:end]
            h_chunk = []
            h = h_prev
            for t in range(end - start):
                h = a_chunk[:, t] * h + b_chunk[:, t]
                h_chunk.append(h)
            h_prev = h
            h_list.append(torch.stack(h_chunk, dim=1))

        return torch.cat(h_list, dim=1)[:, :T]


class MambaBlock(nn.Module):
    """Single Mamba block: in_proj -> SSM -> out_proj with gating."""

    def __init__(self, d_model: int, d_state: int = 16, d_conv: int = 4,
                 expand: int = 2):
        super().__init__()
        d_inner = d_model * expand
        self.ln = RMSNorm(d_model)
        self.in_proj = nn.Linear(d_model, d_inner * 2, bias=False)
        self.ssm = SelectiveSSM(d_inner, d_state, d_conv)
        self.out_proj = nn.Linear(d_inner, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.ln(x)
        xz = self.in_proj(x)
        x_ssm, z = xz.chunk(2, dim=-1)
        x_ssm = self.ssm(x_ssm)
        x_ssm = x_ssm * F.silu(z)
        return residual + self.out_proj(x_ssm)


class MambaLM100M(nn.Module):
    """Mamba language model at 100M scale."""

    def __init__(self, cfg: MambaConfig, vocab: int, max_seq: int):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab, cfg.d_model)
        self.blocks = nn.ModuleList([
            MambaBlock(cfg.d_model, cfg.d_state, cfg.d_conv, cfg.expand)
            for _ in range(cfg.n_layers)
        ])
        self.ln_f = RMSNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, vocab, bias=False)
        self.head.weight = self.tok_emb.weight

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        x = self.tok_emb(ids)
        for block in self.blocks:
            x = block(x)
        return self.head(self.ln_f(x))



class LSTMLM100M(nn.Module):
    """LSTM language model at 100M scale."""

    def __init__(self, cfg: LSTMConfig, vocab: int, max_seq: int):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab, cfg.d_model)
        self.lstm = nn.LSTM(
            cfg.d_model, cfg.hidden_size, cfg.n_layers,
            batch_first=True, dropout=cfg.dropout if cfg.n_layers > 1 else 0.0)
        for name, param in self.lstm.named_parameters():
            if 'bias' in name:
                n = param.size(0) // 4
                param.data[n:2*n].fill_(1.0)
        self.ln_f = nn.LayerNorm(cfg.hidden_size)
        self.proj = nn.Linear(cfg.hidden_size, cfg.d_model, bias=False)
        self.head = nn.Linear(cfg.d_model, vocab, bias=False)
        self.head.weight = self.tok_emb.weight

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        x = self.tok_emb(ids)
        x, _ = self.lstm(x)
        x = self.proj(self.ln_f(x))
        return self.head(x)



ALL_MODELS = ["CEDL", "CEDL-V2a", "CEDL-V3", "CEDL-V3m", "CEDL-V3m2",
              "CEDL-V3m2-NLL", "CEDL-V3lex",
              "CEDL-V4c", "CEDL-V4c-G", "CEDL-V4c-cold",
              "CEDL-V4c-randlabel", "CEDL-V4c-frozen", "CEDL-V4c-M",
              "CEDL-V4d", "CEDL-V4d-cold", "CEDL-V4d-noinject",
              "CEDL-V4d-strong", "CEDL-V4d-strong-noinject",
              "CEDL-V4d-causal", "CEDL-V4d-causal-strong", "CEDL-V4d-causalz",
              "CEDL-V4d-combo", "CEDL-V4e", "CEDL-V4e-noinject",
              "CEDL-V4e-k1", "CEDL-V4e-k2", "CEDL-V4e-k4",
              "CEDL-V4e-k16", "CEDL-V4e-k32", "CEDL-V5-LM",
              "CEDL-V6",
              "CEDL-B0",
              "Transformer", "Transformer-XL", "RetNet", "Mamba"]


def build_model(tag: str, vocab: int, max_seq: int,
                *, v6_mixture: bool = False,
                v6_lambda_init: float = -4.0,
                v6_lambda_a_init: float = 1.0,
                v6_aux_weight: float = 0.0,
                v6_margin_target: float = 1.0,
                v6_mix_weight: float = 0.5,
                v6_gate_weight: float = 0.01,
                v6_lambda_floor: float = 0.05,
                v6_lambda_head: bool = False,
                v6_lambda_head_hidden: int = 160,
                v6_lambda_head_bias_init: float = -7.0,
                v6_bg_weight: float = 1.0,
                v6_bg_target: float = 0.01,
                v6_bce_objective: bool = False,
                v6_sel_weight: float = 1.0,
                v6_lambda_head_w_init_std: float = 1e-3,
                v6_wt_sparsity_weight: float = 0.0,
                v6_wt_sparsity_target: float = 0.0,
                v6_mem_head_bank: bool = False,
                v6_bank_ce_weight: float = 0.0,
                v6_bank_pair_ce_weight: float = 0.0,
                v6_bank_query_source: str = "h_d",
                v6_bank_readout_mode: str = "bank",
                v6_source_adapter: bool = False,
                v6_context_adapter: bool = False,
                v6_specialist_noinject: bool = False,
                lambda_init: float | None = None,
                lambda_a_init: float | None = None,
                aux_weight: float | None = None,
                margin_target: float | None = None,
                mix_weight: float | None = None,
                gate_weight: float | None = None,
                lambda_floor: float | None = None,
                lambda_head: bool | None = None,
                lambda_head_hidden: int | None = None,
                lambda_head_bias_init: float | None = None,
                bg_weight: float | None = None,
                bg_target: float | None = None,
                bce_objective: bool | None = None,
                sel_weight: float | None = None,
                lambda_head_w_init_std: float | None = None,
                wt_sparsity_weight: float | None = None,
                wt_sparsity_target: float | None = None,
                memory_head_enabled: bool | None = None,
                memory_ce_weight: float | None = None,
                memory_pair_ce_weight: float | None = None,
                memory_query_source: str | None = None,
                memory_readout_mode: str | None = None,
                source_adapter: bool | None = None,
                context_adapter: bool | None = None,
                specialist_noinject: bool | None = None) -> nn.Module:
    """Build a CEDL model by tag."""
    if lambda_init is not None:
        v6_lambda_init = lambda_init
    if lambda_a_init is not None:
        v6_lambda_a_init = lambda_a_init
    if aux_weight is not None:
        v6_aux_weight = aux_weight
    if margin_target is not None:
        v6_margin_target = margin_target
    if mix_weight is not None:
        v6_mix_weight = mix_weight
    if gate_weight is not None:
        v6_gate_weight = gate_weight
    if lambda_floor is not None:
        v6_lambda_floor = lambda_floor
    if lambda_head is not None:
        v6_lambda_head = lambda_head
    if lambda_head_hidden is not None:
        v6_lambda_head_hidden = lambda_head_hidden
    if lambda_head_bias_init is not None:
        v6_lambda_head_bias_init = lambda_head_bias_init
    if bg_weight is not None:
        v6_bg_weight = bg_weight
    if bg_target is not None:
        v6_bg_target = bg_target
    if bce_objective is not None:
        v6_bce_objective = bce_objective
    if sel_weight is not None:
        v6_sel_weight = sel_weight
    if lambda_head_w_init_std is not None:
        v6_lambda_head_w_init_std = lambda_head_w_init_std
    if wt_sparsity_weight is not None:
        v6_wt_sparsity_weight = wt_sparsity_weight
    if wt_sparsity_target is not None:
        v6_wt_sparsity_target = wt_sparsity_target
    if memory_head_enabled is not None:
        v6_mem_head_bank = memory_head_enabled
    if memory_ce_weight is not None:
        v6_bank_ce_weight = memory_ce_weight
    if memory_pair_ce_weight is not None:
        v6_bank_pair_ce_weight = memory_pair_ce_weight
    if memory_query_source is not None:
        v6_bank_query_source = memory_query_source
    if memory_readout_mode is not None:
        v6_bank_readout_mode = memory_readout_mode
    if source_adapter is not None:
        v6_source_adapter = source_adapter
    if context_adapter is not None:
        v6_context_adapter = context_adapter
    if specialist_noinject is not None:
        v6_specialist_noinject = specialist_noinject

    if tag == "CEDL":
        tag = "CEDL-V6"
    if tag == "CEDL-legacy":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v0"), vocab, max_seq)
    elif tag == "CEDL-V2a":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v2a"), vocab, max_seq)
    elif tag == "CEDL-V3":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v3",
                       lex_anchor_weight=0.0), vocab, max_seq)
    elif tag == "CEDL-V3m":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v3m",
                       lex_anchor_weight=0.0), vocab, max_seq)
    elif tag == "CEDL-V3m2":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v3m2",
                       lex_anchor_weight=0.0), vocab, max_seq)
    elif tag == "CEDL-V3m2-NLL":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v3m2_nll",
                       lex_anchor_weight=0.0), vocab, max_seq)
    elif tag == "CEDL-V4c":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4c",
                       v4c_variant="base", lex_anchor_weight=0.0),
            vocab, max_seq)
    elif tag == "CEDL-V4c-G":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4c_gate",
                       v4c_variant="base", lex_anchor_weight=0.0),
            vocab, max_seq)
    elif tag == "CEDL-V4c-cold":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4c",
                       v4c_variant="cold", lex_anchor_weight=0.0),
            vocab, max_seq)
    elif tag == "CEDL-V4c-randlabel":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4c",
                       v4c_variant="randlabel", lex_anchor_weight=0.0),
            vocab, max_seq)
    elif tag == "CEDL-V4c-M":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4c",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0),
            vocab, max_seq)
    elif tag == "CEDL-V4c-frozen":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4c",
                       v4c_variant="frozen", lex_anchor_weight=0.0),
            vocab, max_seq)
    elif tag == "CEDL-V4d":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4d",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0,
                       v4d_w_sal_sigma=0.02, v4d_logit_cap=4.0,
                       v4d_noinject=False),
            vocab, max_seq)
    elif tag == "CEDL-V4d-strong":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4d",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0,
                       v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0,
                       v4d_noinject=False),
            vocab, max_seq)
    elif tag == "CEDL-V4d-causalz":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4d",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0,
                       v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0,
                       v4d_noinject=False,
                       v4d_causal_weight=0.0,
                       v4d_causal_z_weight=0.10,
                       v4d_causal_z_gap=0.15),
            vocab, max_seq)
    elif tag == "CEDL-V4d-combo":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4d",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0,
                       v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0,
                       v4d_noinject=False,
                       v4d_causal_weight=0.10, v4d_causal_gap=0.25,
                       v4d_causal_z_weight=0.05, v4d_causal_z_gap=0.15),
            vocab, max_seq)
    elif tag == "CEDL-V4e" or tag == "CEDL-V4e-noinject" \
            or tag.startswith("CEDL-V4e-k"):
        _noinj = (tag == "CEDL-V4e-noinject")
        _rank = 8
        if tag.startswith("CEDL-V4e-k"):
            try:
                _rank = int(tag.rsplit("k", 1)[1])
            except ValueError:
                _rank = 8
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4d",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0,
                       v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0,
                       v4d_noinject=_noinj,
                       v4d_causal_weight=0.10, v4d_causal_gap=0.25,
                       v4d_causal_z_weight=0.05, v4d_causal_z_gap=0.15,
                       v4d_w_sal_rank=_rank,
                       v4d_swap_consistency_weight=0.05,
                       v4d_role_sep_weight=0.05),
            vocab, max_seq)
    elif tag == "CEDL-V5-LM":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4d",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0,
                       v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0,
                       v4d_noinject=False,
                       v4d_causal_weight=0.10, v4d_causal_gap=0.25,
                       v4d_causal_z_weight=0.05, v4d_causal_z_gap=0.15,
                       v4d_w_sal_rank=8,
                       v4d_swap_consistency_weight=0.05,
                       v4d_role_sep_weight=0.05),
            vocab, max_seq)
    elif tag == "CEDL-V6":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4d",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0,
                       v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0,
                       v4d_noinject=bool(v6_specialist_noinject),
                       v4d_causal_weight=0.10, v4d_causal_gap=0.25,
                       v4d_causal_z_weight=0.05, v4d_causal_z_gap=0.15,
                       v4d_w_sal_rank=8,
                       v4d_swap_consistency_weight=0.05,
                       v4d_role_sep_weight=0.05,
                       v6_mixture=True,
                       v6_lambda_init=float(v6_lambda_init),
                       v6_lambda_a_init=float(v6_lambda_a_init),
                       v6_aux_weight=float(v6_aux_weight),
                       v6_margin_target=float(v6_margin_target),
                       v6_mix_weight=float(v6_mix_weight),
                       v6_gate_weight=float(v6_gate_weight),
                       v6_lambda_floor=float(v6_lambda_floor),
                       v6_lambda_head=bool(v6_lambda_head),
                       v6_lambda_head_hidden=int(v6_lambda_head_hidden),
                       v6_lambda_head_bias_init=float(v6_lambda_head_bias_init),
                       v6_bg_weight=float(v6_bg_weight),
                       v6_bg_target=float(v6_bg_target),
                       v6_bce_objective=bool(v6_bce_objective),
                       v6_sel_weight=float(v6_sel_weight),
                       v6_lambda_head_w_init_std=float(v6_lambda_head_w_init_std),
                       v6_wt_sparsity_weight=float(v6_wt_sparsity_weight),
                       v6_wt_sparsity_target=float(v6_wt_sparsity_target),
                       v6_mem_head_bank=bool(v6_mem_head_bank),
                       v6_bank_ce_weight=float(v6_bank_ce_weight),
                       v6_bank_pair_ce_weight=float(v6_bank_pair_ce_weight),
                       v6_bank_query_source=str(v6_bank_query_source),
                       v6_bank_readout_mode=str(v6_bank_readout_mode),
                       v6_source_adapter=bool(v6_source_adapter),
                       v6_context_adapter=bool(v6_context_adapter)),
            vocab, max_seq)
    elif tag == "CEDL-V4d-causal-strong":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4d",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0,
                       v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0,
                       v4d_noinject=False,
                       v4d_causal_weight=0.10, v4d_causal_gap=0.25),
            vocab, max_seq)
    elif tag == "CEDL-V4d-causal":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4d",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0,
                       v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0,
                       v4d_noinject=False,
                       v4d_causal_weight=0.02, v4d_causal_gap=0.25),
            vocab, max_seq)
    elif tag == "CEDL-V4d-strong-noinject":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4d",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0,
                       v4d_w_sal_sigma=0.10, v4d_logit_cap=4.0,
                       v4d_noinject=True),
            vocab, max_seq)
    elif tag == "CEDL-V4d-cold":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4d",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0,
                       v4d_w_sal_sigma=0.0, v4d_logit_cap=4.0,
                       v4d_noinject=False),
            vocab, max_seq)
    elif tag == "CEDL-V4d-noinject":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v4d",
                       v4c_variant="base", lex_anchor_weight=0.0,
                       v4c_margin_weight=0.05, v4c_margin_target=1.0,
                       v4d_w_sal_sigma=0.02, v4d_logit_cap=4.0,
                       v4d_noinject=True),
            vocab, max_seq)
    elif tag == "CEDL-V3lex":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=True, salience_mode="v3",
                       lex_anchor_weight=0.005), vocab, max_seq)
    elif tag == "CEDL-B0":
        return CEDLTwoLoop100M(
            CEDLConfig(use_salience=False), vocab, max_seq)
    elif tag == "Transformer":
        return TransformerLM100M(TransformerConfig(), vocab, max_seq)
    elif tag == "Transformer-XL":
        return TransformerXLLM100M(TransformerXLConfig(), vocab, max_seq)
    elif tag == "RetNet":
        return RetNetLM100M(RetNetConfig(), vocab, max_seq)
    elif tag == "Mamba":
        return MambaLM100M(MambaConfig(), vocab, max_seq)
    elif tag == "LSTM":
        return LSTMLM100M(LSTMConfig(), vocab, max_seq)
    else:
        raise ValueError(f"Unknown model: {tag}")


def count_params(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters())


def verify_all_params():
    """Print parameter counts for all models."""
    print("=" * 60)
    print("Parameter Verification (target: ~100M)")
    print("=" * 60)
    for tag in ALL_MODELS:
        model = build_model(tag, 50257, 1024)
        n = count_params(model)
        print(f"  {tag:20s}: {n:>12,d} ({n/1e6:.1f}M)")
        del model
    print("=" * 60)



class TextChunkDataset(Dataset):
    """Pre-tokenized text split into fixed-length chunks."""

    def __init__(self, token_ids: torch.Tensor, chunk_size: int):
        self.chunk_size = chunk_size
        n_chunks = len(token_ids) // chunk_size
        self.data = token_ids[:n_chunks * chunk_size].view(n_chunks, chunk_size)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


def load_data(cfg: Config):
    """Load and tokenize dataset. Returns train/val/test DataLoaders."""
    from transformers import GPT2TokenizerFast
    tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")

    def tokenize_rows(rows):
        """Tokenize row-by-row to avoid OOM from giant string join."""
        all_ids = []
        for text in rows:
            if text.strip():
                ids = tokenizer.encode(text, add_special_tokens=False)
                all_ids.extend(ids)
        return torch.tensor(all_ids, dtype=torch.long)

    if cfg.dataset == "wikitext103":
        from datasets import load_dataset
        ds = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1")
        print("Tokenizing train split...")
        train_ids = tokenize_rows(ds["train"]["text"])
        print("Tokenizing val split...")
        val_ids = tokenize_rows(ds["validation"]["text"])
        print("Tokenizing test split...")
        test_ids = tokenize_rows(ds["test"]["text"])
        del ds
    elif cfg.dataset == "c4":
        from datasets import load_dataset
        ds_train = load_dataset("allenai/c4", "en", split="train", streaming=True)
        ds_val = load_dataset("allenai/c4", "en", split="validation", streaming=True)
        all_ids = []
        total_tokens = 0
        for example in ds_train:
            ids = tokenizer.encode(example["text"], add_special_tokens=False)
            all_ids.extend(ids)
            total_tokens += len(ids)
            if total_tokens > 400_000_000:
                break
        train_ids = torch.tensor(all_ids, dtype=torch.long)
        del all_ids
        val_ids_list = []
        val_tokens = 0
        for example in ds_val:
            ids = tokenizer.encode(example["text"], add_special_tokens=False)
            val_ids_list.extend(ids)
            val_tokens += len(ids)
            if val_tokens > 2_000_000:
                break
        val_ids = torch.tensor(val_ids_list, dtype=torch.long)
        del val_ids_list
        test_ids = val_ids
    else:
        raise ValueError(f"Unknown dataset: {cfg.dataset}")

    chunk = cfg.max_seq + 1
    train_ds = TextChunkDataset(train_ids, chunk)
    val_ds = TextChunkDataset(val_ids, chunk)
    test_ds = TextChunkDataset(test_ids, chunk)

    pin = not cfg.tpu
    train_loader = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True,
                              num_workers=4, pin_memory=pin, drop_last=True,
                              persistent_workers=True)
    val_loader = DataLoader(val_ds, batch_size=cfg.batch_size, shuffle=False,
                            num_workers=2, pin_memory=pin, drop_last=True,
                            persistent_workers=True)
    test_loader = DataLoader(test_ds, batch_size=cfg.batch_size, shuffle=False,
                             num_workers=2, pin_memory=pin, drop_last=True,
                             persistent_workers=True)

    print(f"Dataset: {cfg.dataset}")
    print(f"  Train chunks: {len(train_ds):,}")
    print(f"  Val chunks:   {len(val_ds):,}")
    print(f"  Test chunks:  {len(test_ds):,}")

    return train_loader, val_loader, test_loader, tokenizer


def make_loaders(train_ds, val_ds, test_ds, batch_size: int, tpu: bool = False,
                 shuffle_seed: Optional[int] = None):
    """Build DataLoaders from pre-tokenized datasets (fast, no re-tokenization).

    If shuffle_seed is provided, the train_loader uses a seeded torch.Generator
    so batch ordering is identical across model tags in a paired run.
    """
    pin = not tpu
    gen = None
    if shuffle_seed is not None:
        gen = torch.Generator()
        gen.manual_seed(shuffle_seed)
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
                              num_workers=4, pin_memory=pin, drop_last=True,
                              persistent_workers=True, generator=gen)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,
                            num_workers=2, pin_memory=pin, drop_last=True,
                            persistent_workers=True)
    test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False,
                             num_workers=2, pin_memory=pin, drop_last=True,
                             persistent_workers=True)
    return train_loader, val_loader, test_loader



def get_lr(step: int, cfg: Config) -> float:
    """Cosine schedule with linear warmup."""
    if step < cfg.warmup_steps:
        return cfg.lr * step / cfg.warmup_steps
    if step >= cfg.max_steps:
        return cfg.min_lr
    progress = (step - cfg.warmup_steps) / (cfg.max_steps - cfg.warmup_steps)
    return cfg.min_lr + 0.5 * (cfg.lr - cfg.min_lr) * (1 + math.cos(math.pi * progress))


def _v5_effective_v4c_every(gstep: int, cfg=None) -> int:
    """Synthetic-batch cadence by GLOBAL step (early/mid/late)."""
    e  = int(getattr(cfg, "v5_cadence_early", 8))   if cfg else 8
    m  = int(getattr(cfg, "v5_cadence_mid",   16))  if cfg else 16
    l_ = int(getattr(cfg, "v5_cadence_late",  32))  if cfg else 32
    if gstep < 10000:
        return e
    if gstep < 20000:
        return m
    return l_


def _v5_aux_scale(gstep: int, cfg=None) -> float:
    """Global aux REPLAY strength applied to the TOTAL aux loss (early/mid/late)."""
    e  = float(getattr(cfg, "v5_aux_scale_early", 0.50)) if cfg else 0.50
    m  = float(getattr(cfg, "v5_aux_scale_mid",   0.25)) if cfg else 0.25
    l_ = float(getattr(cfg, "v5_aux_scale_late",  0.10)) if cfg else 0.10
    if gstep < 10000:
        return e
    if gstep < 20000:
        return m
    return l_


_V5_SAL_PREFIXES = ("salience.", "d_stage.salience_bias_",
                    "d_stage.w_sal", "v4d_role_head.",
                    "bank_q_proj.", "v6_lambda_", "mem_head_bank.",
                    "v6_source_adapter.", "v6_context_adapter.")


def _v5_salience_param_ids(raw_model):
    """ids + names of AUX-ELIGIBLE params (V5 gradient isolation). Matches
    salience.*, d_stage.salience_bias_*, d_stage.w_sal* (incl. w_sal_A/B),
    v4d_role_head.* (parent-level), AND (V6.1) the bank-readout params
    bank_q_proj.* + v6_lambda_* — the latter are bank-readout semantically
    (not salience), but routed through the same isolation allowlist so the
    V4c-aux gradient survives the restore loop. d_stage.cross_attn.* is NOT
    matched → trunk."""
    ids, names = set(), []
    for n, p in raw_model.named_parameters():
        if any(n.startswith(pre) for pre in _V5_SAL_PREFIXES):
            ids.add(id(p))
            names.append(n)
    return ids, names


def _unwrap_model_state_dict(obj, src_label: str = "checkpoint"):
    """Return a plain model state_dict from bare, train-state, or wrapper ckpts."""
    if isinstance(obj, dict) and "model" in obj and isinstance(obj["model"], dict):
        wrapper_keys = sorted(k for k in obj.keys() if k != "model")
        wrapper_msg = f" wrapper_keys={wrapper_keys}" if wrapper_keys else ""
        print(f"  [load] unwrapping {src_label}['model'] -> model weights{wrapper_msg}")
        obj = obj["model"]
    if not isinstance(obj, dict):
        raise RuntimeError(f"{src_label} did not contain a state_dict-like object")
    if any(k.startswith("_orig_mod.") for k in obj):
        obj = {k.replace("_orig_mod.", ""): v for k, v in obj.items()}
    return obj


def _load_model_state_from_path(path: str, src_label: str = "checkpoint"):
    state = torch.load(path, map_location="cpu", weights_only=True)
    return _unwrap_model_state_dict(state, src_label=src_label)


def _is_v6_specialist_b0_missing_key(name: str) -> bool:
    """Keys present in a CEDL-V6 specialist shell but absent from B0."""
    exact = {
        "sl_alpha",
        "val_a", "val_b",
        "v6_lambda_a", "v6_lambda_b",
    }
    prefixes = (
        "salience.",
        "d_stage.salience_bias_",
        "d_stage.w_sal",
        "v4d_role_head.",
        "bank_q_proj.",
        "v6_lambda_head.",
        "mem_head_bank.",
        "v6_source_adapter.",
        "v6_context_adapter.",
    )
    return name in exact or any(name.startswith(pre) for pre in prefixes)


def _apply_v6_specialist_freeze(raw_model):
    """Freeze the B0 trunk and leave only the direct specialist trainable."""
    trainable_prefixes = ["v6_lambda_head.", "mem_head_bank."]
    if not hasattr(raw_model, "v6_lambda_head"):
        raise RuntimeError("V6 specialist freeze requires v6_lambda_head.*")
    if not hasattr(raw_model, "mem_head_bank"):
        raise RuntimeError("V6 specialist freeze requires mem_head_bank.*")
    if getattr(raw_model, "use_v6_source_adapter", False):
        if not hasattr(raw_model, "v6_source_adapter"):
            raise RuntimeError(
                "V6 specialist freeze expected v6_source_adapter.*")
        trainable_prefixes.append("v6_source_adapter.")
    if getattr(raw_model, "use_v6_context_adapter", False):
        if not hasattr(raw_model, "v6_context_adapter"):
            raise RuntimeError(
                "V6 specialist freeze expected v6_context_adapter.*")
        trainable_prefixes.append("v6_context_adapter.")
    trainable_prefixes = tuple(trainable_prefixes)

    trainable_names = []
    frozen_names = []
    for n, p in raw_model.named_parameters():
        keep_trainable = any(n.startswith(pre) for pre in trainable_prefixes)
        p.requires_grad = keep_trainable
        if keep_trainable:
            trainable_names.append(n)
        else:
            frozen_names.append(n)

    raw_model._v6_specialist_trainable_prefixes = list(trainable_prefixes)
    raw_model._v6_specialist_trainable_names = trainable_names
    n_trainable = sum(p.numel() for p in raw_model.parameters()
                      if p.requires_grad)
    n_frozen = sum(p.numel() for p in raw_model.parameters()
                   if not p.requires_grad)
    print(f"  [V6 specialist freeze] trainable params={n_trainable:,} "
          f"frozen params={n_frozen:,} "
          f"trainable_prefixes={list(trainable_prefixes)}")
    print(f"  [V6 specialist freeze] trainable names={trainable_names}")



@torch.no_grad()
def evaluate(model: nn.Module, loader: DataLoader, device: torch.device,
             max_steps: int = 200, tpu: bool = False) -> float:
    """Compute cross-entropy loss → perplexity."""
    _xm = None
    if tpu:
        import torch_xla.core.xla_model as _xm
    model.eval()
    if hasattr(model, 'reset_memory'):
        model.reset_memory()
    total_loss = 0.0
    count = 0
    _raw = model._orig_mod if hasattr(model, '_orig_mod') else model
    _use_logprobs = bool(getattr(_raw, "outputs_log_probs", False)
                          and not getattr(_raw, "bank_off", False))
    _lam_acc = {"lam_mean": 0.0, "lam_std": 0.0,
                "lam_sat_low": 0.0, "lam_sat_high": 0.0}
    _lam_n = 0
    _lam_max_global = -1.0
    _lam_gt_0_7_sum = 0
    _lam_total_sum  = 0
    _lam_tensors    = []
    for i, batch in enumerate(loader):
        if i >= max_steps:
            break
        batch = batch.to(device)
        input_ids = batch[:, :-1]
        targets = batch[:, 1:]
        logits = model(input_ids)
        if isinstance(logits, tuple):
            logits = logits[0]
        if _use_logprobs:
            loss = F.nll_loss(logits.reshape(-1, logits.size(-1)),
                              targets.reshape(-1))
            _ls = getattr(_raw, "_last_lam_stats", None)
            if _ls is not None:
                for _k in _lam_acc:
                    _lam_acc[_k] += float(_ls[_k])
                _lam_n += 1
                _lam_max_global = max(_lam_max_global, float(_ls["lam_max"]))
                _lam_gt_0_7_sum += int(_ls["lam_gt_0_7_count"])
                _lam_total_sum  += int(_ls["lam_total_count"])
                _lam_tensors.append(_ls["lam_tensor_cpu"])
        else:
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)),
                                   targets.reshape(-1))
        total_loss += loss.item()
        del logits, loss
        count += 1
        if _xm is not None:
            _xm.mark_step()
    if _lam_n > 0:
        _raw._last_eval_lam_avg = {k: v / _lam_n for k, v in _lam_acc.items()}
        _raw._last_eval_lam_avg["n_batches"] = _lam_n
        _raw._last_eval_lam_avg["lam_max"] = _lam_max_global
        _raw._last_eval_lam_avg["lam_frac_gt_0_7"] = (
            _lam_gt_0_7_sum / _lam_total_sum if _lam_total_sum > 0 else 0.0)
        try:
            _pooled = torch.cat(_lam_tensors).numpy()
            _raw._last_eval_lam_avg["lam_p99"] = float(
                np.percentile(_pooled, 99))
        except Exception as _pe:
            print(f"  [V6 λ] WARN: lam_p99 computation failed ({_pe})")
            _raw._last_eval_lam_avg["lam_p99"] = float("nan")
        _raw._last_eval_lam_avg["v6_lambda_a"] = float(_raw.v6_lambda_a.item())
        _raw._last_eval_lam_avg["v6_lambda_b"] = float(_raw.v6_lambda_b.item())
        _use_head = bool(getattr(_raw, "use_v6_lambda_head", False))
        if _use_head:
            _raw._last_eval_lam_avg["v6_lambda_head_bias"] = float(
                _raw.v6_lambda_head[2].bias.item())
        _avg = _raw._last_eval_lam_avg
        print(f"  [V6 λ] mean={_avg['lam_mean']:.4f}  "
              f"std={_avg['lam_std']:.4f}  "
              f"sat_low={_avg['lam_sat_low']:.3f}  "
              f"sat_high={_avg['lam_sat_high']:.3f}  "
              f"(over {_lam_n} eval batches)")
        if _use_head:
            print(f"  [V6 λ tail] max={_avg['lam_max']:.4f}  "
                  f"p99={_avg['lam_p99']:.4f}  "
                  f"frac>0.7={_avg['lam_frac_gt_0_7']:.4f}  "
                  f"(head_bias={_avg['v6_lambda_head_bias']:+.3f})")
        else:
            print(f"  [V6 λ tail] max={_avg['lam_max']:.4f}  "
                  f"p99={_avg['lam_p99']:.4f}  "
                  f"frac>0.7={_avg['lam_frac_gt_0_7']:.4f}  "
                  f"(a={_avg['v6_lambda_a']:+.3f}, b={_avg['v6_lambda_b']:+.3f})")
    elif _use_logprobs:
        print("  [V6 λ] WARN: outputs_log_probs=True but no _last_lam_stats "
              "recorded — forward block may not have executed the λ path.")
    model.train()
    avg_loss = total_loss / max(count, 1)
    return math.exp(avg_loss)


def train(model: nn.Module, tag: str, cfg: Config, device: torch.device,
          train_loader: DataLoader, val_loader: DataLoader):
    """Main training loop with gradient accumulation and cosine LR."""
    xm = None
    if cfg.tpu:
        import torch_xla.core.xla_model as xm

    model.train()

    raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
    is_cedl = isinstance(raw_model, CEDLTwoLoop100M)
    _resume_ckpt = getattr(cfg, "resume_checkpoint", None)

    if is_cedl and getattr(raw_model, "v4c_variant", "base") == "frozen":
        for n, p in raw_model.named_parameters():
            in_sl_mod = (n.startswith("salience.")
                         or n.startswith("d_stage.salience_bias_"))
            p.requires_grad = bool(in_sl_mod)
        n_trainable = sum(1 for p in raw_model.parameters() if p.requires_grad)
        n_frozen = sum(1 for p in raw_model.parameters() if not p.requires_grad)
        print(f"  [V4c-frozen] {n_trainable} trainable params, "
              f"{n_frozen} frozen (LM backbone)")

    _specialist_from_b0 = (
        is_cedl and bool(getattr(cfg, "v6_specialist_from_b0", False)))
    if _specialist_from_b0:
        if not _resume_ckpt:
            raise RuntimeError(
                "v6_specialist_from_b0=True but cfg.resume_checkpoint is empty")
        if bool(getattr(cfg, "v6_specialist_noinject", False)):
            _noinj_live = bool(getattr(raw_model._v4c_cfg,
                                       "v4d_noinject", False))
            if not _noinj_live:
                raise RuntimeError(
                    "v6_specialist_noinject=True but model was built with "
                    "v4d_noinject=False")
        print(f"  [V6 specialist load] source B0 checkpoint: {_resume_ckpt}")
        _b0_sd = _load_model_state_from_path(_resume_ckpt, src_label="B0")
        _res = raw_model.load_state_dict(_b0_sd, strict=False)
        _miss = set(_res.missing_keys)
        _unexp = set(_res.unexpected_keys)
        _bad_miss = sorted(k for k in _miss
                           if not _is_v6_specialist_b0_missing_key(k))
        if _bad_miss or _unexp:
            raise RuntimeError(
                "[V6 specialist load] B0->V6 audit failed: "
                f"unexpected_missing={_bad_miss}, "
                f"unexpected={sorted(_unexp)}")
        _fresh = sorted(_miss)
        raw_model._v6_specialist_loaded_missing = _fresh
        raw_model._v6_specialist_base_checkpoint = _resume_ckpt
        print(f"  [V6 specialist load] matched B0 trunk; freshly initialized "
              f"{len(_fresh)} specialist keys")
        if _fresh:
            print(f"  [V6 specialist load] fresh keys={_fresh}")
        if int(getattr(cfg, "resume_start_step", 0) or 0) != 0:
            print("  [V6 specialist load] WARN: ignoring --start-step for "
                  "fresh B0-specialist bootstrap")
            cfg.resume_start_step = 0
        _resume_ckpt = None

    if _specialist_from_b0 and getattr(cfg, "v6_specialist_freeze", "none") == "trunk":
        _apply_v6_specialist_freeze(raw_model)

    if is_cedl:
        c_param_ids = {id(p) for p in raw_model.c_stage.parameters()
                        if p.requires_grad}
        c_params = [p for p in raw_model.c_stage.parameters()
                    if p.requires_grad]
        _bank_head_lr = float(getattr(cfg, "v6_bank_head_lr", 0.0))
        _use_bank_lr_group = (_bank_head_lr > 0
                              and bool(getattr(cfg, "v6_mem_head_bank", False)))
        bank_head_param_ids = set()
        if _use_bank_lr_group and hasattr(raw_model, "mem_head_bank"):
            bank_head_param_ids = {
                id(p) for p in raw_model.mem_head_bank.parameters()
                if p.requires_grad}
            if hasattr(raw_model, "v6_source_adapter"):
                bank_head_param_ids.update(
                    id(p) for p in raw_model.v6_source_adapter.parameters()
                    if p.requires_grad)
            if hasattr(raw_model, "v6_context_adapter"):
                bank_head_param_ids.update(
                    id(p) for p in raw_model.v6_context_adapter.parameters()
                    if p.requires_grad)
        edl_params = [p for p in raw_model.parameters()
                      if p.requires_grad and id(p) not in c_param_ids
                      and id(p) not in bank_head_param_ids]
        groups = []
        if c_params:
            groups.append({"params": c_params, "lr": cfg.lr})
        if edl_params:
            groups.append({"params": edl_params, "lr": cfg.lr * 1.5})
        if _use_bank_lr_group:
            bank_head_params = [p for p in raw_model.mem_head_bank.parameters()
                                if p.requires_grad]
            if hasattr(raw_model, "v6_source_adapter"):
                bank_head_params.extend(
                    p for p in raw_model.v6_source_adapter.parameters()
                    if p.requires_grad)
            if hasattr(raw_model, "v6_context_adapter"):
                bank_head_params.extend(
                    p for p in raw_model.v6_context_adapter.parameters()
                    if p.requires_grad)
            if bank_head_params:
                groups.append({
                    "params": bank_head_params,
                    "lr": _bank_head_lr,
                    "bank_head_lr_constant": True,
                })
                print(f"  [Stage2b-lr] direct readout params in dedicated AdamW group "
                      f"with CONSTANT lr={_bank_head_lr} "
                      f"({len(bank_head_params)} params, "
                      f"{sum(p.numel() for p in bank_head_params):,} elements) "
                      f"— cosine schedule applies to trunk/gate only.")
        optimizer = torch.optim.AdamW(
            groups, weight_decay=cfg.weight_decay, betas=(0.9, 0.95))
    else:
        optimizer = torch.optim.AdamW(
            model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay,
            betas=(0.9, 0.95))

    scaler = None

    step = int(getattr(cfg, "resume_start_step", 0) or 0)
    _resumed_best_ppl = None
    _resumed_run_id = None
    _arch_fork = False

    _v6_allowed_missing = {"bank_q_proj.weight", "bank_q_proj.bias",
                            "v6_lambda_a", "v6_lambda_b"}
    if bool(getattr(cfg, "v6_lambda_head", False)) and is_cedl:
        _v6_allowed_missing |= {
            "v6_lambda_head.0.weight", "v6_lambda_head.0.bias",
            "v6_lambda_head.2.weight", "v6_lambda_head.2.bias",
        }
    if bool(getattr(cfg, "v6_mem_head_bank", False)) and is_cedl:
        _v6_allowed_missing |= {
            "mem_head_bank.weight", "mem_head_bank.bias",
        }
    _v6_mixture_on = bool(getattr(cfg, "v6_mixture", False)) and is_cedl

    def _audited_load(sd, src_label):
        result = raw_model.load_state_dict(sd, strict=False)
        miss = set(result.missing_keys)
        unexp = set(result.unexpected_keys)
        if _v6_mixture_on:
            extra_miss = miss - _v6_allowed_missing
            if extra_miss or unexp:
                raise RuntimeError(
                    f"[V6 resume {src_label}] strict audit failed: "
                    f"unexpected_missing={sorted(extra_miss)}, "
                    f"unexpected={sorted(unexp)}")
            new_init = sorted(miss & _v6_allowed_missing)
            if new_init:
                print(f"  [V6 resume {src_label}] freshly initialized: "
                      f"{new_init}")
        else:
            if miss or unexp:
                raise RuntimeError(
                    f"[resume {src_label}] strict load failed: "
                    f"missing={sorted(miss)}, unexpected={sorted(unexp)}")

    if _resume_ckpt:
        _rs = torch.load(_resume_ckpt, map_location="cpu", weights_only=True)
        if isinstance(_rs, dict) and "model" in _rs and "optimizer" in _rs:
            _msd = _rs["model"]
            if any(k.startswith("_orig_mod.") for k in _msd):
                _msd = {k.replace("_orig_mod.", ""): v for k, v in _msd.items()}
            _audited_load(_msd, "trainstate")
            _src_tag = _rs.get("model_tag")
            _arch_fork = (_src_tag is not None and _src_tag != tag)
            try:
                optimizer.load_state_dict(_rs["optimizer"])
                print("  [resume] restored optimizer state from train-state ckpt")
            except Exception as _e:
                print(f"  [resume] WARN: optimizer state not restored ({_e}); "
                      f"continuing with fresh AdamW")
            step = int(_rs.get("global_step", step))
            if _arch_fork:
                print(f"  [resume] ARCH-FORK detected: trainstate tag "
                      f"{_src_tag!r} != current tag {tag!r}. Resetting "
                      f"best_val_ppl=inf, best_step={step}, new run_started_at.")
                _resumed_best_ppl = None
                _resumed_run_id = None
            else:
                if _rs.get("best_val_ppl") is not None:
                    _resumed_best_ppl = float(_rs["best_val_ppl"])
                _resumed_run_id = _rs.get("run_started_at")
                print(f"  [resume] restored best_val_ppl={_resumed_best_ppl} "
                      f"run_id={_resumed_run_id}")
        else:
            if any(k.startswith("_orig_mod.") for k in _rs):
                _rs = {k.replace("_orig_mod.", ""): v for k, v in _rs.items()}
            _audited_load(_rs, "bare")
            print("  [resume] loaded bare state_dict; using FRESH AdamW "
                  "(no optimizer state in source checkpoint)")
        print(f"  [resume] {_resume_ckpt}  global start_step={step}  "
              f"target total_steps={cfg.max_steps}")

    _v5_iso = bool(getattr(cfg, "v5_grad_isolation", False)) and is_cedl
    _v5_trunk_frac = float(getattr(cfg, "v5_trunk_aux_frac", 0.0))
    _sal_param_ids = set()
    if _v5_iso:
        _sal_param_ids, _sal_names = _v5_salience_param_ids(raw_model)
        _n_trunk = sum(1 for p in raw_model.parameters()
                       if id(p) not in _sal_param_ids)
        _has_v4d = any(n.startswith("d_stage.w_sal")
                       or n.startswith("v4d_role_head.") for n in _sal_names)
        if not _has_v4d:
            raise RuntimeError(
                "v5_grad_isolation=True but NO params matched d_stage.w_sal* / "
                "v4d_role_head.* — the salience set is wrong (a V5 model must "
                "have these). Aborting to avoid an invalid isolation run.")
        print(f"  [V5 isolation] salience params={len(_sal_param_ids)}  "
              f"trunk params={_n_trunk}  trunk_aux_frac={_v5_trunk_frac}  "
              f"(cross_attn stays trunk; aux-only trunk grads discarded)")

    if is_cedl and bool(getattr(cfg, "v6_mixture", False)):
        _v6_param_names = [n for n, _ in raw_model.named_parameters()
                           if n.startswith("bank_q_proj.")
                           or n in ("v6_lambda_a", "v6_lambda_b")]
        if len(_v6_param_names) == 0:
            raise RuntimeError(
                "cfg.v6_mixture=True but the model has NO bank_q_proj.* / "
                "v6_lambda_* parameters — V6 forward will be a silent no-op. "
                "Build the model with v6_mixture=True (use the CEDL-V6 tag, "
                "or pass v6_mixture=True to build_model for other CEDL tags).")
        _a_init = float(getattr(cfg, "v6_lambda_a_init", 1.0))
        _a_live = float(raw_model.v6_lambda_a.item())
        print(f"  [V6 mixture] params={_v6_param_names}  "
              f"v6_lambda_a_init={_a_init:+.3f} (live={_a_live:+.3f})  "
              f"λ_b_init=sigmoid({raw_model.v6_lambda_b.item():.3f})="
              f"{torch.sigmoid(raw_model.v6_lambda_b).item():.4f}  "
              f"(AUX-eligible via V5 allowlist; V4c bank-aux can reach)")
        for _k in ("v6_aux_weight", "v6_bg_weight", "v6_bg_target",
                   "v6_lambda_head_hidden", "v6_lambda_head_bias_init",
                   "v6_bce_objective", "v6_sel_weight",
                   "v6_lambda_head_w_init_std",
                   "v6_wt_sparsity_weight",
                   "v6_wt_sparsity_target",
                  "v6_mem_head_bank", "v6_bank_ce_weight",
                   "v6_bank_pair_ce_weight",
                   "v6_bank_query_source", "v6_bank_readout_mode",
                   "v6_source_adapter", "v6_context_adapter"):
            _cv = getattr(cfg, _k, None)
            _mv = getattr(raw_model._v4c_cfg, _k, None)
            if isinstance(_cv, float) and isinstance(_mv, float):
                if abs(_cv - _mv) > 1e-9:
                    raise RuntimeError(
                        f"V6 propagation bug: cfg.{_k}={_cv} but "
                        f"raw_model._v4c_cfg.{_k}={_mv}. The V4c hook reads "
                        f"_v4c_cfg; without matching values the supervision "
                        f"silently falls back to the default.")
            elif _cv != _mv:
                raise RuntimeError(
                    f"V6 propagation bug: cfg.{_k}={_cv} but "
                    f"raw_model._v4c_cfg.{_k}={_mv}.")
        _cfg_w  = float(getattr(cfg, "v6_aux_weight", 0.0))
        _mdl_w  = float(getattr(raw_model._v4c_cfg, "v6_aux_weight", 0.0))
        print(f"  [V6.1] cfg.v6_aux_weight={_cfg_w}  "
              f"_v4c_cfg.v6_aux_weight={_mdl_w}  "
              f"v6_margin_target={getattr(raw_model._v4c_cfg, 'v6_margin_target', 1.0)}  "
              f"v6_mix_weight={getattr(raw_model._v4c_cfg, 'v6_mix_weight', 0.5)}  "
              f"v6_gate_weight={getattr(raw_model._v4c_cfg, 'v6_gate_weight', 0.01)}  "
              f"v6_lambda_floor={getattr(raw_model._v4c_cfg, 'v6_lambda_floor', 0.05)}  "
              f"(0 = V6 baseline; >0 = V6.1 aux-through-mixture active)")
        if getattr(raw_model, "use_v6_lambda_head", False):
            _head_bias = float(raw_model.v6_lambda_head[2].bias.item())
            _bce_on = bool(getattr(raw_model._v4c_cfg,
                                    "v6_bce_objective", False))
            _w_std = float(getattr(raw_model._v4c_cfg,
                                    "v6_lambda_head_w_init_std", 1e-3))
            _w_live = float(raw_model.v6_lambda_head[2].weight.std().item())
            print(f"  [Stage2a] v6_lambda_head ON  hidden="
                  f"{getattr(raw_model._v4c_cfg, 'v6_lambda_head_hidden', 160)}  "
                  f"head_bias={_head_bias:+.3f}  "
                  f"w_init_std={_w_std:.4f} (live_std={_w_live:.4f})  "
                  f"bg_weight={getattr(raw_model._v4c_cfg, 'v6_bg_weight', 1.0)}  "
                  f"bg_target={getattr(raw_model._v4c_cfg, 'v6_bg_target', 0.01)}  "
                  f"(aux-only via torch.no_grad in main forward + V5 allowlist)")
            if _bce_on:
                print(f"  [Stage2a-bce] BCE objective ON  "
                      f"sel_weight={getattr(raw_model._v4c_cfg, 'v6_sel_weight', 1.0)}  "
                      f"(L_gate + L_bg are telemetry-only; v4c_aux uses "
                      f"L_bank + w_mix·L_mix + w_sel·L_sel where "
                      f"L_sel = 0.5·BCE_ans + 0.5·BCE_bg)")
            _wt_w = float(getattr(raw_model._v4c_cfg,
                                    "v6_wt_sparsity_weight", 0.0))
            if _wt_w > 0:
                _wt_t = float(getattr(raw_model._v4c_cfg,
                                       "v6_wt_sparsity_target", 0.0))
                if _wt_t > 0:
                    _wt_form = (f"hinge form aux_loss += {_wt_w}·relu(λ_wt"
                                f" - {_wt_t}).mean() — free zone [0,"
                                f"{_wt_t}]")
                else:
                    _wt_form = f"mean-λ form aux_loss += {_wt_w}·E[λ_wt]"
                print(f"  [Stage2a-wts] WT sparsity loss ON  "
                      f"wt_sparsity_weight={_wt_w}  "
                      f"wt_sparsity_target={_wt_t}  ({_wt_form}; head "
                      f"receives gradient via parallel non-no_grad path; "
                      f"mixture lam stays detached → LM-CE still cannot "
                      f"drive head)")
        if getattr(raw_model, "use_v6_mem_head_bank", False):
            _bce_w = float(getattr(raw_model._v4c_cfg,
                                    "v6_bank_ce_weight", 0.0))
            _pair_w = float(getattr(raw_model._v4c_cfg,
                                     "v6_bank_pair_ce_weight", 0.0))
            _mhb_w_std = float(raw_model.mem_head_bank.weight.std().item())
            _mhb_b_mean = float(raw_model.mem_head_bank.bias.mean().item())
            _tied = (raw_model.mem_head_bank.weight is
                     raw_model.c_stage.tok_emb.weight)
            _bhlr = float(getattr(cfg, "v6_bank_head_lr", 0.0))
            print(f"  [Stage2b] mem_head_bank ON  "
                  f"weight.std={_mhb_w_std:.4f}  "
                  f"bias.mean={_mhb_b_mean:+.3f}  "
                  f"tied_to_tok_emb={_tied} (expect False)  "
                  f"bank_ce_weight={_bce_w}  "
                  f"bank_pair_ce_weight={_pair_w}  "
                  f"bank_query_source={getattr(raw_model, 'v6_bank_query_source', 'h_d')}  "
                  f"bank_readout_mode={getattr(raw_model, 'v6_bank_readout_mode', 'bank')}  "
                  f"source_adapter={getattr(raw_model, 'use_v6_source_adapter', False)}  "
                  f"context_adapter={getattr(raw_model, 'use_v6_context_adapter', False)}  "
                  f"bank_head_lr={_bhlr}{' (CONSTANT; separate AdamW group)' if _bhlr > 0 else ' (= main lr; cosine schedule)'}  "
                  f"(separate trainable bank vocab projection; aux-eligible "
                  f"via V5 allowlist; L_v6 += {_bce_w}·CE(logits_bank_a, "
                  f"cur_tok) at V4c answer rows)")
            assert not _tied, (
                "Stage 2b: mem_head_bank.weight must be a separate "
                "nn.Parameter, not aliased to tok_emb.weight. "
                "Check __init__ — use .copy_() not assignment.")

    accum_loss = 0.0
    best_val_ppl = _resumed_best_ppl if _resumed_best_ppl is not None else float('inf')
    best_step = step
    os.makedirs(cfg.save_dir, exist_ok=True)
    _run_started_at = _resumed_run_id or time.strftime("%Y%m%dT%H%M%S")
    _run_start_step = step

    def _save_trainstate(_gstep):
        """Resume-only train-state with best-so-far and run identity."""
        torch.save({"model": model.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "global_step": _gstep,
                    "best_val_ppl": float(best_val_ppl),
                    "best_step": int(best_step),
                    "run_started_at": _run_started_at,
                    "model_tag": tag},
                   os.path.join(cfg.save_dir, f"{tag}_trainstate_latest.pt"))

    print(f"\nTraining {tag} ({count_params(model)/1e6:.1f}M params)")
    print(f"  Device: {device}")
    print(f"  Effective batch: {cfg.batch_size * cfg.grad_accum} seqs "
          f"= {cfg.batch_size * cfg.grad_accum * cfg.max_seq / 1e3:.0f}K tokens")
    print(f"  Schedule total (max_steps): {cfg.max_steps:,}"
          + (f"  | run_until_step: {cfg.run_until_step:,}"
             if getattr(cfg, "run_until_step", None) else ""))

    t0 = time.time()
    data_iter = iter(train_loader)
    cedl_cfg = CEDLConfig() if is_cedl else None

    _v3m2_diagnostic_active = (
        is_cedl and getattr(raw_model, 'use_salience', False)
        and getattr(raw_model, 'salience_mode', 'v0') in ("v3m2", "v3m2_nll")
    )
    _v3m2_eval_batch = None
    if _v3m2_diagnostic_active:
        try:
            _v3m2_eval_batch = next(iter(val_loader)).to(device)[:16, :512]
            print(f"  [v_vec_pathway] diagnostic ACTIVE  "
                  f"(eval batch: {tuple(_v3m2_eval_batch.shape)})")
        except Exception as _e:
            print(f"  [v_vec_pathway] WARN: could not grab eval batch ({_e}); "
                  f"diagnostic DISABLED")
            _v3m2_diagnostic_active = False

    _v4c_active = (
        is_cedl and getattr(raw_model, 'use_salience', False)
        and getattr(raw_model, 'salience_mode', 'v0') in (
            "v4c", "v4c_gate", "v4d")
    )
    _v4c_pair_gen = None
    _v4c_tok = None
    _v4c_log_idx = 0
    _v4c_last_lm_grad_norm = None
    _v4c_swap_fn = None
    _v4c_split = "all"
    _v4c_collision = 0.0
    if _v4c_active:
        try:
            import data_v4c_pairs as _v4c_module
            from transformers import GPT2TokenizerFast as _V4cTok
            _v4c_pair_gen = _v4c_module.generate
            _v4c_tok = _V4cTok.from_pretrained("gpt2")
            v4c_variant = getattr(raw_model, "v4c_variant", "base")
            print(f"  [V4c] ACTIVE  variant={v4c_variant}  "
                  f"every={raw_model._v4c_cfg.v4c_every} steps  "
                  f"B={raw_model._v4c_cfg.v4c_batch_size}  "
                  f"T_max={raw_model._v4c_cfg.v4c_max_seq}")
            _v4e_on = (getattr(raw_model, "salience_mode", "v0") == "v4d"
                       and (float(getattr(raw_model._v4c_cfg,
                                          "v4d_swap_consistency_weight", 0.0)) > 0
                            or float(getattr(raw_model._v4c_cfg,
                                             "v4d_role_sep_weight", 0.0)) > 0))
            if _v4e_on:
                _v4c_swap_fn = _v4c_module.make_entity_swap
                _v4c_split = "train"
                _v4c_collision = 0.2
                print(f"  [V4e] swap/role losses ON  split=train  "
                      f"hard_collision=0.2  "
                      f"rank={getattr(raw_model._v4c_cfg, 'v4d_w_sal_rank', 0)}")
        except Exception as _e:
            raise RuntimeError(
                f"V4c setup failed for tag {tag!r} (model is "
                f"salience_mode={getattr(raw_model, 'salience_mode', '?')}): "
                f"{_e}. Check that data_v4c_pairs.py is on PYTHONPATH and "
                f"that the GPT2 tokenizer is available.") from _e

    _stop_step = int(cfg.run_until_step) if getattr(cfg, "run_until_step", None) \
        else cfg.max_steps
    while step < _stop_step:
        optimizer.zero_grad()

        for micro in range(cfg.grad_accum):
            try:
                batch = next(data_iter)
            except StopIteration:
                if hasattr(model, 'reset_memory'):
                    model.reset_memory()
                data_iter = iter(train_loader)
                batch = next(data_iter)

            batch = batch.to(device)
            input_ids = batch[:, :-1]
            targets = batch[:, 1:]

            autocast_device = 'xla' if cfg.tpu else device.type
            with torch.amp.autocast(autocast_device, dtype=torch.bfloat16,
                                    enabled=cfg.bfloat16):
                output = model(input_ids)
                if isinstance(output, tuple):
                    logits, aux_loss = output[0], output[1]
                else:
                    logits, aux_loss = output, torch.tensor(0.0, device=device)
                if bool(getattr(raw_model, "outputs_log_probs", False)
                        and not getattr(raw_model, "bank_off", False)):
                    loss = F.nll_loss(logits.reshape(-1, logits.size(-1)),
                                      targets.reshape(-1))
                else:
                    loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)),
                                           targets.reshape(-1))
                loss = (loss + aux_loss) / cfg.grad_accum

            loss.backward()
            accum_loss += loss.item()
            del logits, loss
            if xm is not None:
                xm.mark_step()

        _v4c_eff_every = (_v5_effective_v4c_every(step, cfg) if _v5_iso
                          else raw_model._v4c_cfg.v4c_every) if _v4c_active else 1
        if _v4c_active and (step + 1) % _v4c_eff_every == 0:
            v4c_cfg_obj = raw_model._v4c_cfg
            _lm_grad_snapshot = {
                id(p): p.grad.detach().clone()
                for p in model.parameters()
                if p.grad is not None
            }
            _lm_grad_norm_sq = 0.0
            for g in _lm_grad_snapshot.values():
                _lm_grad_norm_sq += float((g * g).sum())
            _v4c_last_lm_grad_norm = _lm_grad_norm_sq ** 0.5
            try:
                _v4c_family_weights = (
                    {"but_update": 0.50, "however_revision": 0.20,
                     "temporal_update": 0.15, "paraphrased_equiv": 0.05,
                     "neutral_control": 0.10}
                    if getattr(cfg, "v5_family_reweight", False) else None)
                _v4c_ids, _v4c_items, _v4c_swap_ids, _v4c_swap_items = \
                    _build_v4c_contrastive_batch(
                        _v4c_pair_gen, _v4c_tok,
                        B=v4c_cfg_obj.v4c_batch_size,
                        T_max=v4c_cfg_obj.v4c_max_seq,
                        randomize_positive_mapping=(
                            getattr(raw_model, "v4c_variant", "base") == "randlabel"),
                        seed=step * 7919 + 13,
                        split=_v4c_split,
                        hard_collision_frac=_v4c_collision,
                        swap_fn=_v4c_swap_fn,
                        family_weights=_v4c_family_weights,
                    )
                _v4c_ids = _v4c_ids.to(device)
                if _v4c_swap_ids is not None:
                    _v4c_swap_ids = _v4c_swap_ids.to(device)
                _log_this_step = (_v4c_log_idx % 50 == 0)
                _v4d_mode = (getattr(raw_model, "salience_mode", "v0") == "v4d")
                if _v4d_mode and _log_this_step:
                    raw_model.d_stage.store_v4d_attn = True
                _, _v4c_aux = model(_v4c_ids, v4c_spans=_v4c_items,
                                     v4c_swap_ids=_v4c_swap_ids,
                                     v4c_swap_spans=_v4c_swap_items)
                _v5_scale = _v5_aux_scale(step, cfg) if _v5_iso else 1.0
                if _v5_iso:
                    _v4c_aux = _v4c_aux * _v5_scale
                _v4c_aux.backward()
                _aux_grad_norm_sq = 0.0
                for p in model.parameters():
                    if p.grad is None:
                        continue
                    g_lm = _lm_grad_snapshot.get(id(p))
                    if g_lm is None:
                        _aux_grad_norm_sq += float((p.grad * p.grad).sum())
                    else:
                        diff = p.grad - g_lm
                        _aux_grad_norm_sq += float((diff * diff).sum())
                _aux_gn_true = _aux_grad_norm_sq ** 0.5
                if _v5_iso:
                    for p in model.parameters():
                        if p.grad is None or id(p) in _sal_param_ids:
                            continue
                        snap = _lm_grad_snapshot.get(id(p))
                        if snap is None:
                            p.grad = None
                        elif _v5_trunk_frac > 0.0:
                            p.grad = snap + _v5_trunk_frac * (p.grad - snap)
                        else:
                            p.grad = snap
                if _v4c_log_idx % 50 == 0:
                    _ratio = (_aux_gn_true / _v4c_last_lm_grad_norm
                              if _v4c_last_lm_grad_norm > 1e-6 else float('inf'))
                    _msg = (f"  [V4c step={step + 1:>5d}] "
                            f"aux={_v4c_aux.item():.4f}  "
                            f"LM_grad={_v4c_last_lm_grad_norm:.3f}  "
                            f"V4c_grad={_aux_gn_true:.3f}  "
                            f"ratio={_ratio:.3f}  "
                            f"bce_α={raw_model.v4c_bce_alpha.item():.2f}  "
                            f"nce_α={raw_model.v4c_nce_alpha.item():.2f}")
                    if _v5_iso:
                        _msg += (f"  [V5 iso every={_v4c_eff_every} "
                                  f"aux_scale={_v5_scale:.2f} "
                                  f"trunk_frac={_v5_trunk_frac:.2f}]")
                    _mstats = getattr(raw_model.salience,
                                       "_last_v4c_margin_stats", None)
                    if _mstats is not None:
                        _msg += (f"  margin={_mstats['loss']:.3f}  "
                                  f"used={_mstats['n_used']}  "
                                  f"skip={_mstats['n_skip_neutral']}/n,"
                                  f"{_mstats['n_skip_eq_token']}/eq")
                    _cstats = getattr(raw_model.salience,
                                       "_last_v4d_causal_stats", None)
                    if _cstats is not None:
                        _msg += (f"  causal={_cstats['loss']:.3f}  "
                                  f"gap_mean={_cstats['margin_gap_mean']:+.3f}  "
                                  f"m_n={_cstats['margin_normal_mean']:+.2f}  "
                                  f"m_c={_cstats['margin_clamped_mean']:+.2f}")
                    _czstats = getattr(raw_model.salience,
                                       "_last_v4d_causal_z_stats", None)
                    if _czstats is not None:
                        _msg += (f"  causalZ={_czstats['loss']:.4f}  "
                                  f"mz_gap={_czstats['mz_gap_mean']:+.4f}  "
                                  f"mz_n={_czstats['mz_normal_mean']:+.3f}  "
                                  f"mz_c={_czstats['mz_clamped_mean']:+.3f}")
                    _swstats = getattr(raw_model.salience,
                                       "_last_v4d_swap_stats", None)
                    if _swstats is not None:
                        _msg += (f"  swap={_swstats['loss']:.4f}"
                                  f"(n={_swstats['n_pairs']})")
                    _rlstats = getattr(raw_model.salience,
                                       "_last_v4d_role_stats", None)
                    if _rlstats is not None:
                        _msg += (f"  role={_rlstats['loss']:.4f}"
                                  f"(acc={_rlstats['acc']:.2f})")
                    _v6stats = getattr(raw_model.salience,
                                       "_last_v6_aux_stats", None)
                    if _v6stats is not None:
                        if bool(_v6stats.get("bce_active", False)):
                            _msg += (f"  v6={_v6stats['loss']:.3f}"
                                     f"(B={_v6stats['L_bank']:.3f},"
                                     f"P={_v6stats.get('L_bank_pair_ce', 0.0):.3f},"
                                     f"M={_v6stats['L_mix']:.3f},"
                                     f"Sel={_v6stats['L_sel']:.3f}"
                                     f"[ans={_v6stats['L_sel_ans']:.3f},"
                                     f"bg={_v6stats['L_sel_bg']:.3f}])  "
                                     f"lam_ans={_v6stats['lam_ans_mean']:.3f}"
                                     f"±{_v6stats['lam_ans_std']:.3f}"
                                     f"(n={_v6stats['n_items']})")
                        else:
                            _msg += (f"  v6={_v6stats['loss']:.3f}"
                                     f"(B={_v6stats['L_bank']:.3f},"
                                     f"P={_v6stats.get('L_bank_pair_ce', 0.0):.3f},"
                                     f"M={_v6stats['L_mix']:.3f},"
                                     f"G={_v6stats['L_gate']:.3f}"
                                     f",Bg={_v6stats.get('L_bg', 0.0):.3f})  "
                                     f"lam_ans={_v6stats['lam_ans_mean']:.3f}"
                                     f"±{_v6stats['lam_ans_std']:.3f}"
                                     f"(sat_lo={_v6stats['lam_ans_sat_low']:.2f},"
                                     f"hi={_v6stats['lam_ans_sat_high']:.2f},"
                                     f"n={_v6stats['n_items']})")
                        if int(_v6stats.get("n_bg_items", 0)) > 0:
                            _msg += (f"  lam_bg={_v6stats['lam_bg_mean']:.3f}"
                                     f"±{_v6stats['lam_bg_std']:.3f}"
                                     f"(n={_v6stats['n_bg_items']})")
                            if "v6_lambda_head_bias" in _v6stats:
                                _msg += (f"  head_bias="
                                         f"{_v6stats['v6_lambda_head_bias']:+.3f}")
                    if _v4d_mode:
                        _xattn = raw_model.d_stage.cross_attn
                        _attn_t = getattr(_xattn, "_last_attn", None)
                        if _attn_t is not None:
                            with torch.no_grad():
                                _ent = -(_attn_t.clamp(min=1e-12)
                                          * _attn_t.clamp(min=1e-12).log()
                                          ).sum(-1).mean().item()
                            _msg += (f"  attn_ent={_ent:.2f}  "
                                      f"sal_std={_xattn._last_sal_logits_std:.3f}  "
                                      f"base_std={_xattn._last_base_logits_std:.3f}")
                            if _ent < 1.5:
                                print(_msg)
                                raise RuntimeError(
                                    f"V4d attention saturation: mean attn entropy "
                                    f"{_ent:.2f} < 1.5 nats at step {step + 1}. "
                                    f"Attention collapsed onto a single slot. "
                                    f"Lower v4d_logit_cap (4.0 → 2.0) or zero-init W_sal.")
                    print(_msg)
                if _v4d_mode and _log_this_step:
                    raw_model.d_stage.store_v4d_attn = False
                _v4c_log_idx += 1
                del _lm_grad_snapshot
            except Exception as _e:
                raise RuntimeError(
                    f"V4c aux forward failed at step {step + 1} "
                    f"(tag={tag!r}, variant={getattr(raw_model, 'v4c_variant', '?')}): "
                    f"{_e}") from _e

        torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)

        lr = get_lr(step, cfg)
        for i, pg in enumerate(optimizer.param_groups):
            if pg.get("bank_head_lr_constant", False):
                continue
            multiplier = 1.0 if (not is_cedl or i == 0) else 1.5
            pg['lr'] = lr * multiplier

        if xm is not None:
            xm.optimizer_step(optimizer)
        else:
            optimizer.step()

        step += 1

        if is_cedl:
            if step < cedl_cfg.feedback_warmup_start:
                raw_model.feedback_alpha.fill_(0.0)
            elif step < cedl_cfg.feedback_warmup_end:
                frac = (step - cedl_cfg.feedback_warmup_start) / (
                    cedl_cfg.feedback_warmup_end - cedl_cfg.feedback_warmup_start)
                raw_model.feedback_alpha.fill_(frac)
            else:
                raw_model.feedback_alpha.fill_(1.0)
            if getattr(raw_model, 'use_salience', False):
                if step < 1000:
                    raw_model.sl_alpha.fill_(0.0)
                elif step < 2000:
                    raw_model.sl_alpha.fill_((step - 1000) / 1000.0)
                else:
                    raw_model.sl_alpha.fill_(1.0)
                if getattr(raw_model, "salience_mode", "v0") in (
                        "v4c", "v4c_gate", "v4d"):
                    if step < 500:
                        raw_model.v4c_bce_alpha.fill_(0.0)
                    elif step < 1500:
                        raw_model.v4c_bce_alpha.fill_((step - 500) / 1000.0)
                    else:
                        raw_model.v4c_bce_alpha.fill_(1.0)
                    if step < 1500:
                        raw_model.v4c_nce_alpha.fill_(0.0)
                    elif step < 3000:
                        raw_model.v4c_nce_alpha.fill_((step - 1500) / 1500.0)
                    else:
                        raw_model.v4c_nce_alpha.fill_(1.0)
            anneal_step = int(cfg.max_steps * cedl_cfg.sparsity_anneal_frac)
            if step >= anneal_step and cfg.max_steps > anneal_step:
                frac = (step - anneal_step) / (cfg.max_steps - anneal_step)
                sparsity = cedl_cfg.e_sparsity + frac * (
                    cedl_cfg.sparsity_final - cedl_cfg.e_sparsity)
                raw_model.e_stage.set_sparsity(sparsity)

        if step % 100 == 0:
            elapsed = time.time() - t0
            tokens_per_sec = (step * cfg.batch_size * cfg.grad_accum * cfg.max_seq) / elapsed
            print(f"  [{step:>6d}/{cfg.max_steps}] loss={accum_loss/100:.4f} "
                  f"lr={lr:.2e} tok/s={tokens_per_sec:.0f}")
            accum_loss = 0.0

        if step % cfg.eval_interval == 0:
            val_ppl = evaluate(model, val_loader, device, cfg.eval_steps, tpu=cfg.tpu)
            print(f"  >>> Val PPL: {val_ppl:.2f}")
            _is_best = val_ppl < best_val_ppl
            if _is_best:
                best_val_ppl = val_ppl
                best_step = step
                ckpt_path = os.path.join(cfg.save_dir, f"{tag}_best.pt")
                if xm is not None:
                    xm.save(model.state_dict(), ckpt_path)
                else:
                    torch.save(model.state_dict(), ckpt_path)
                print(f"  >>> Saved best checkpoint: {ckpt_path}")
                if is_cedl and getattr(raw_model, "salience_mode", "v0") in (
                        "v4c", "v4c_gate", "v4d"):
                    import json as _json
                    cfg_path = os.path.join(cfg.save_dir, f"{tag}_config.json")
                    with open(cfg_path, "w") as _f:
                        _json.dump({
                            "tag": tag,
                            "salience_mode": raw_model.salience_mode,
                            "v4c_variant": getattr(
                                raw_model, "v4c_variant", "base"),
                            "v4c_margin_weight": float(getattr(
                                raw_model._v4c_cfg, "v4c_margin_weight", 0.0)),
                            "v4c_margin_target": float(getattr(
                                raw_model._v4c_cfg, "v4c_margin_target", 1.0)),
                            "v4c_warm_start_sigma": float(getattr(
                                raw_model._v4c_cfg, "v4c_warm_start_sigma", 0.05)),
                            "v4c_norm_cap": float(getattr(
                                raw_model._v4c_cfg, "v4c_norm_cap", 0.3)),
                            "v4d_w_sal_sigma": float(getattr(
                                raw_model._v4c_cfg, "v4d_w_sal_sigma", 0.02)),
                            "v4d_logit_cap": float(getattr(
                                raw_model._v4c_cfg, "v4d_logit_cap", 4.0)),
                            "v4d_noinject": bool(getattr(
                                raw_model._v4c_cfg, "v4d_noinject", False)),
                            "v4d_causal_weight": float(getattr(
                                raw_model._v4c_cfg, "v4d_causal_weight", 0.0)),
                            "v4d_causal_gap": float(getattr(
                                raw_model._v4c_cfg, "v4d_causal_gap", 0.25)),
                            "v4d_causal_z_weight": float(getattr(
                                raw_model._v4c_cfg, "v4d_causal_z_weight", 0.0)),
                            "v4d_causal_z_gap": float(getattr(
                                raw_model._v4c_cfg, "v4d_causal_z_gap", 0.15)),
                            "v4d_w_sal_rank": int(getattr(
                                raw_model._v4c_cfg, "v4d_w_sal_rank", 0)),
                            "v4d_swap_consistency_weight": float(getattr(
                                raw_model._v4c_cfg, "v4d_swap_consistency_weight", 0.0)),
                            "v4d_role_sep_weight": float(getattr(
                                raw_model._v4c_cfg, "v4d_role_sep_weight", 0.0)),
                            "v6_mixture": bool(getattr(
                                raw_model, "v6_mixture", False)),
                            "v6_lambda_init": float(getattr(
                                raw_model, "v6_lambda_init", -4.0)),
                            "v6_lambda_a_init": float(getattr(
                                raw_model, "v6_lambda_a_init", 1.0)),
                            "v6_lambda_head": bool(getattr(
                                raw_model, "use_v6_lambda_head", False)),
                            "v6_lambda_head_hidden": int(getattr(
                                raw_model._v4c_cfg, "v6_lambda_head_hidden", 160)),
                            "v6_lambda_head_bias_init": float(getattr(
                                raw_model._v4c_cfg, "v6_lambda_head_bias_init", -7.0)),
                            "v6_bg_weight": float(getattr(
                                raw_model._v4c_cfg, "v6_bg_weight", 1.0)),
                            "v6_bg_target": float(getattr(
                                raw_model._v4c_cfg, "v6_bg_target", 0.01)),
                            "v6_bce_objective": bool(getattr(
                                raw_model._v4c_cfg, "v6_bce_objective", False)),
                            "v6_sel_weight": float(getattr(
                                raw_model._v4c_cfg, "v6_sel_weight", 1.0)),
                            "v6_lambda_head_w_init_std": float(getattr(
                                raw_model._v4c_cfg, "v6_lambda_head_w_init_std", 1e-3)),
                            "v6_wt_sparsity_weight": float(getattr(
                                raw_model._v4c_cfg, "v6_wt_sparsity_weight", 0.0)),
                            "v6_wt_sparsity_target": float(getattr(
                                raw_model._v4c_cfg, "v6_wt_sparsity_target", 0.0)),
                            "v6_mem_head_bank": bool(getattr(
                                raw_model._v4c_cfg, "v6_mem_head_bank", False)),
                            "v6_bank_ce_weight": float(getattr(
                                raw_model._v4c_cfg, "v6_bank_ce_weight", 0.0)),
                            "v6_bank_pair_ce_weight": float(getattr(
                                raw_model._v4c_cfg, "v6_bank_pair_ce_weight", 0.0)),
                            "v6_bank_query_source": str(getattr(
                                raw_model._v4c_cfg, "v6_bank_query_source", "h_d")),
                            "v6_bank_readout_mode": str(getattr(
                                raw_model._v4c_cfg, "v6_bank_readout_mode", "bank")),
                            "v6_source_adapter": bool(getattr(
                                raw_model._v4c_cfg, "v6_source_adapter", False)),
                            "v6_context_adapter": bool(getattr(
                                raw_model._v4c_cfg, "v6_context_adapter", False)),
                            "v6_bank_head_lr": float(getattr(
                                cfg, "v6_bank_head_lr", 0.0)),
                            "v6_specialist_from_b0": bool(getattr(
                                cfg, "v6_specialist_from_b0", False)),
                            "v6_specialist_freeze": str(getattr(
                                cfg, "v6_specialist_freeze", "none")),
                            "v6_specialist_noinject": bool(getattr(
                                cfg, "v6_specialist_noinject", False)),
                            "v6_specialist_base_checkpoint": getattr(
                                raw_model, "_v6_specialist_base_checkpoint",
                                getattr(cfg, "resume_checkpoint", None)),
                            "v6_specialist_trainable_prefixes": getattr(
                                raw_model, "_v6_specialist_trainable_prefixes",
                                None),
                            "v6_specialist_loaded_missing": getattr(
                                raw_model, "_v6_specialist_loaded_missing",
                                None),
                            "v6_lambda_head_bias_live": (
                                float(raw_model.v6_lambda_head[2].bias.item())
                                if getattr(raw_model, "use_v6_lambda_head", False)
                                else None),
                            "v6_lambda_trajectory": getattr(
                                raw_model, "_last_eval_lam_avg", None),
                            "v6_aux_weight": float(getattr(
                                raw_model._v4c_cfg, "v6_aux_weight", 0.0)),
                            "v6_margin_target": float(getattr(
                                raw_model._v4c_cfg, "v6_margin_target", 1.0)),
                            "v6_mix_weight": float(getattr(
                                raw_model._v4c_cfg, "v6_mix_weight", 0.5)),
                            "v6_gate_weight": float(getattr(
                                raw_model._v4c_cfg, "v6_gate_weight", 0.01)),
                            "v6_lambda_floor": float(getattr(
                                raw_model._v4c_cfg, "v6_lambda_floor", 0.05)),
                            "v6_aux_trajectory": getattr(
                                getattr(raw_model, "salience", None),
                                "_last_v6_aux_stats", None),
                            "v5_training": {
                                "v5_grad_isolation": bool(getattr(cfg, "v5_grad_isolation", False)),
                                "v5_trunk_aux_frac": float(getattr(cfg, "v5_trunk_aux_frac", 0.0)),
                                "resume_checkpoint": getattr(cfg, "resume_checkpoint", None),
                                "resume_start_step": int(getattr(cfg, "resume_start_step", 0) or 0),
                                "schedule_total_steps": int(cfg.max_steps),
                                "aux_scale_schedule": (
                                    f"5-10k:{getattr(cfg,'v5_aux_scale_early',0.50)}, "
                                    f"10-20k:{getattr(cfg,'v5_aux_scale_mid',0.25)}, "
                                    f"20k+:{getattr(cfg,'v5_aux_scale_late',0.10)}"),
                                "v4c_every_schedule": (
                                    f"5-10k:{getattr(cfg,'v5_cadence_early',8)}, "
                                    f"10-20k:{getattr(cfg,'v5_cadence_mid',16)}, "
                                    f"20k+:{getattr(cfg,'v5_cadence_late',32)}"),
                            },
                            "step": step,
                            "val_ppl": float(val_ppl),
                        }, _f, indent=2)

        if step % cfg.eval_interval == 0:
            try:
                import json as _sj
                _sc_path = os.path.join(cfg.save_dir, f"{tag}_trainscore.jsonl")
                with open(_sc_path, "a") as _scf:
                    _row = {
                        "step": step, "global_step": step, "model_tag": tag,
                        "ppl": float(val_ppl), "is_best": bool(_is_best),
                        "best_ppl": float(best_val_ppl),
                        "val_loss": float(math.log(max(val_ppl, 1e-9))),
                        "checkpoint_path": os.path.join(cfg.save_dir, f"{tag}_best.pt"),
                        "run_started_at": _run_started_at,
                        "run_start_step": _run_start_step,
                    }
                    _lam_avg = getattr(raw_model, "_last_eval_lam_avg", None)
                    if _lam_avg is not None:
                        _row["v6_lambda"] = _lam_avg
                    _sal = getattr(raw_model, "salience", None)
                    _v6_aux = getattr(_sal, "_last_v6_aux_stats", None) if _sal else None
                    if _v6_aux is not None:
                        _row["v6_aux"] = _v6_aux
                    _scf.write(_sj.dumps(_row) + "\n")
            except Exception as _e:
                print(f"  [scorecard] WARN: could not append train-score ({_e})")
            if getattr(cfg, "save_trainstate", False) and xm is None:
                _save_trainstate(step)

        if step % cfg.save_interval == 0:
            ckpt_path = os.path.join(cfg.save_dir, f"{tag}_step{step}.pt")
            if xm is not None:
                xm.save(model.state_dict(), ckpt_path)
            else:
                torch.save(model.state_dict(), ckpt_path)
            if getattr(cfg, "save_trainstate", False) and xm is None:
                _save_trainstate(step)

        if _v3m2_diagnostic_active and step > 0 and step % 250 == 0:
            _run_vvec_pathway_diagnostic(
                raw_model, _v3m2_eval_batch, step, device)

    final_path = os.path.join(cfg.save_dir, f"{tag}_final.pt")
    if xm is not None:
        xm.save(model.state_dict(), final_path)
    else:
        torch.save(model.state_dict(), final_path)
    print(f"  >>> Saved final checkpoint: {final_path}")
    if getattr(cfg, "save_trainstate", False) and xm is None:
        _save_trainstate(step)
    if is_cedl and getattr(raw_model, "salience_mode", "v0") in (
            "v4c", "v4c_gate", "v4d"):
        import json as _json
        cfg_path = os.path.join(cfg.save_dir, f"{tag}_config.json")
        if not os.path.exists(cfg_path):
            with open(cfg_path, "w") as _f:
                _json.dump({
                    "tag": tag,
                    "salience_mode": raw_model.salience_mode,
                    "v4c_variant": getattr(raw_model, "v4c_variant", "base"),
                    "v4c_margin_weight": float(getattr(
                        raw_model._v4c_cfg, "v4c_margin_weight", 0.0)),
                    "v4c_margin_target": float(getattr(
                        raw_model._v4c_cfg, "v4c_margin_target", 1.0)),
                    "v4c_warm_start_sigma": float(getattr(
                        raw_model._v4c_cfg, "v4c_warm_start_sigma", 0.05)),
                    "v4c_norm_cap": float(getattr(
                        raw_model._v4c_cfg, "v4c_norm_cap", 0.3)),
                    "v4d_w_sal_sigma": float(getattr(
                        raw_model._v4c_cfg, "v4d_w_sal_sigma", 0.02)),
                    "v4d_logit_cap": float(getattr(
                        raw_model._v4c_cfg, "v4d_logit_cap", 4.0)),
                    "v4d_noinject": bool(getattr(
                        raw_model._v4c_cfg, "v4d_noinject", False)),
                    "v4d_causal_weight": float(getattr(
                        raw_model._v4c_cfg, "v4d_causal_weight", 0.0)),
                    "v4d_causal_gap": float(getattr(
                        raw_model._v4c_cfg, "v4d_causal_gap", 0.25)),
                    "v4d_causal_z_weight": float(getattr(
                        raw_model._v4c_cfg, "v4d_causal_z_weight", 0.0)),
                    "v4d_causal_z_gap": float(getattr(
                        raw_model._v4c_cfg, "v4d_causal_z_gap", 0.15)),
                    "v4d_w_sal_rank": int(getattr(
                        raw_model._v4c_cfg, "v4d_w_sal_rank", 0)),
                    "v4d_swap_consistency_weight": float(getattr(
                        raw_model._v4c_cfg, "v4d_swap_consistency_weight", 0.0)),
                    "v4d_role_sep_weight": float(getattr(
                        raw_model._v4c_cfg, "v4d_role_sep_weight", 0.0)),
                    "v6_mixture": bool(getattr(
                        raw_model, "v6_mixture", False)),
                    "v6_lambda_init": float(getattr(
                        raw_model, "v6_lambda_init", -4.0)),
                    "v6_lambda_a_init": float(getattr(
                        raw_model, "v6_lambda_a_init", 1.0)),
                    "v6_lambda_head": bool(getattr(
                        raw_model, "use_v6_lambda_head", False)),
                    "v6_lambda_head_hidden": int(getattr(
                        raw_model._v4c_cfg, "v6_lambda_head_hidden", 160)),
                    "v6_lambda_head_bias_init": float(getattr(
                        raw_model._v4c_cfg, "v6_lambda_head_bias_init", -7.0)),
                    "v6_bg_weight": float(getattr(
                        raw_model._v4c_cfg, "v6_bg_weight", 1.0)),
                    "v6_bg_target": float(getattr(
                        raw_model._v4c_cfg, "v6_bg_target", 0.01)),
                    "v6_bce_objective": bool(getattr(
                        raw_model._v4c_cfg, "v6_bce_objective", False)),
                    "v6_sel_weight": float(getattr(
                        raw_model._v4c_cfg, "v6_sel_weight", 1.0)),
                    "v6_lambda_head_w_init_std": float(getattr(
                        raw_model._v4c_cfg, "v6_lambda_head_w_init_std", 1e-3)),
                    "v6_wt_sparsity_weight": float(getattr(
                        raw_model._v4c_cfg, "v6_wt_sparsity_weight", 0.0)),
                    "v6_wt_sparsity_target": float(getattr(
                        raw_model._v4c_cfg, "v6_wt_sparsity_target", 0.0)),
                    "v6_mem_head_bank": bool(getattr(
                        raw_model._v4c_cfg, "v6_mem_head_bank", False)),
                    "v6_bank_ce_weight": float(getattr(
                        raw_model._v4c_cfg, "v6_bank_ce_weight", 0.0)),
                    "v6_bank_pair_ce_weight": float(getattr(
                        raw_model._v4c_cfg, "v6_bank_pair_ce_weight", 0.0)),
                    "v6_bank_query_source": str(getattr(
                        raw_model._v4c_cfg, "v6_bank_query_source", "h_d")),
                    "v6_bank_readout_mode": str(getattr(
                        raw_model._v4c_cfg, "v6_bank_readout_mode", "bank")),
                    "v6_source_adapter": bool(getattr(
                        raw_model._v4c_cfg, "v6_source_adapter", False)),
                    "v6_context_adapter": bool(getattr(
                        raw_model._v4c_cfg, "v6_context_adapter", False)),
                    "v6_bank_head_lr": float(getattr(
                        cfg, "v6_bank_head_lr", 0.0)),
                    "v6_specialist_from_b0": bool(getattr(
                        cfg, "v6_specialist_from_b0", False)),
                    "v6_specialist_freeze": str(getattr(
                        cfg, "v6_specialist_freeze", "none")),
                    "v6_specialist_noinject": bool(getattr(
                        cfg, "v6_specialist_noinject", False)),
                    "v6_specialist_base_checkpoint": getattr(
                        raw_model, "_v6_specialist_base_checkpoint",
                        getattr(cfg, "resume_checkpoint", None)),
                    "v6_specialist_trainable_prefixes": getattr(
                        raw_model, "_v6_specialist_trainable_prefixes",
                        None),
                    "v6_specialist_loaded_missing": getattr(
                        raw_model, "_v6_specialist_loaded_missing",
                        None),
                    "v6_lambda_head_bias_live": (
                        float(raw_model.v6_lambda_head[2].bias.item())
                        if getattr(raw_model, "use_v6_lambda_head", False)
                        else None),
                    "v6_lambda_trajectory": getattr(
                        raw_model, "_last_eval_lam_avg", None),
                    "v6_aux_weight": float(getattr(
                        raw_model._v4c_cfg, "v6_aux_weight", 0.0)),
                    "v6_margin_target": float(getattr(
                        raw_model._v4c_cfg, "v6_margin_target", 1.0)),
                    "v6_mix_weight": float(getattr(
                        raw_model._v4c_cfg, "v6_mix_weight", 0.5)),
                    "v6_gate_weight": float(getattr(
                        raw_model._v4c_cfg, "v6_gate_weight", 0.01)),
                    "v6_lambda_floor": float(getattr(
                        raw_model._v4c_cfg, "v6_lambda_floor", 0.05)),
                    "v6_aux_trajectory": getattr(
                        getattr(raw_model, "salience", None),
                        "_last_v6_aux_stats", None),
                    "step": step,
                    "val_ppl": float(best_val_ppl) if best_val_ppl != float('inf')
                                else None,
                }, _f, indent=2)

    return best_val_ppl


@torch.no_grad()
def _run_vvec_pathway_diagnostic(raw_model, eval_batch, step, device):
    """V3m2 mid-training diagnostic. Runs the FIXED eval batch through the
    model twice (once normal, once with v_vec clamped to zero via a hook),
    measures the logit Effective-Contribution-Ratio (EVR), and prints
    salience_bias_up.norm. Switches to eval mode for the entire diagnostic
    so EMA / density / sl_alpha-gated branches don't fire."""
    was_training = raw_model.training
    raw_model.eval()
    try:
        sb_up_norm = raw_model.d_stage.salience_bias_up.weight.norm().item()

        out_normal = raw_model(eval_batch)
        logits_normal = out_normal[0] if isinstance(out_normal, tuple) else out_normal

        def _zero_v_vec_hook(_module, _inputs, output):
            v_vec, v_scalar = output
            return torch.zeros_like(v_vec), v_scalar
        handle = raw_model.salience.register_forward_hook(_zero_v_vec_hook)
        try:
            out_clamped = raw_model(eval_batch)
            logits_clamped = (out_clamped[0] if isinstance(out_clamped, tuple)
                              else out_clamped)
        finally:
            handle.remove()

        delta = (logits_normal - logits_clamped).norm().item()
        denom = logits_normal.norm().item() + 1e-9
        evr = delta / denom
        print(f"  [v_vec_pathway step={step:>5d}] "
              f"sb_up.norm={sb_up_norm:.4f}  EVR={evr:.4f}")
        if step >= 3000 and evr < 0.05:
            print(f"  [v_vec_pathway WARN] EVR<0.05 at step {step}: "
                  f"pathway may be too weak; calibrate against V0 baseline "
                  f"(load V0 ckpt + run diagnostic) before declaring failure.")
    finally:
        raw_model.train(was_training)



def run_structured_benchmark(model, tokenizer, device, max_seq=1024, tpu=False):
    """Zero-shot structured benchmark: 12 tasks x 4 difficulty levels."""
    import random as _rng
    _rng.seed(42)

    xm_mod = None
    if tpu:
        try:
            import torch_xla.core.xla_model as _xm
            xm_mod = _xm
        except ImportError:
            pass

    N_INSTANCES = 200

    NAMES = ["Alice", "Bob", "Carol", "Dave", "Eve", "Frank", "Grace",
             "Hank", "Iris", "Jack", "Kate", "Leo", "Mia", "Nick",
             "Olivia", "Paul", "Quinn", "Rosa", "Sam", "Tina"]
    COLORS = ["red", "blue", "green", "pink", "gray", "brown", "white",
              "black", "gold", "silver", "orange", "purple", "yellow"]
    OBJECTS = ["car", "hat", "bag", "pen", "cup", "box", "ring",
               "lamp", "book", "ball", "shoe", "coat", "desk", "bell",
               "fork", "drum", "fish", "kite", "coin", "vase"]
    TRAITS = ["tall", "short", "fast", "slow", "kind", "bold",
              "calm", "warm", "cold", "rich", "poor", "wise",
              "young", "old", "brave"]
    ANIMALS = ["dog", "cat", "bird", "fish", "frog", "bear", "deer",
               "wolf", "duck", "fox", "pig", "cow", "ant", "bee", "rat"]
    FRUITS = ["apple", "banana", "grape", "lemon", "mango", "peach",
              "plum", "cherry", "melon", "berry"]

    def gen_A1(level):
        n = [2, 4, 6, 8][level]
        names = _rng.sample(NAMES, n)
        colors = [_rng.choice(COLORS) for _ in range(n)]
        objs = [_rng.choice(OBJECTS) for _ in range(n)]
        facts = ". ".join(f"{names[i]} has a {colors[i]} {objs[i]}" for i in range(n))
        qi = _rng.randint(0, n - 1)
        return f"{facts}. Who has a {colors[qi]} {objs[qi]}? Answer:", f" {names[qi]}"

    def gen_A2(level):
        n = [2, 4, 6, 8][level]
        names = _rng.sample(NAMES, n)
        traits = [_rng.choice(TRAITS) for _ in range(n)]
        facts = ". ".join(f"{names[i]} is {traits[i]}" for i in range(n))
        qi = _rng.randint(0, n - 1)
        is_true = _rng.random() < 0.5
        ct = traits[qi] if is_true else _rng.choice([t for t in TRAITS if t != traits[qi]])
        return f"{facts}. Claim: {names[qi]} is {ct}. True or false? Answer:", " true" if is_true else " false"

    def gen_A3(level):
        n_hops = [1, 2, 3, 3][level]
        n_dist = [0, 0, 0, 3][level]
        all_n = _rng.sample(NAMES, n_hops + 1 + n_dist * 2)
        chain = all_n[:n_hops + 1]
        rels = [f"{chain[i]} is parent of {chain[i+1]}" for i in range(n_hops)]
        for i in range(n_dist):
            rels.append(f"{all_n[n_hops+1+2*i]} is friend of {all_n[n_hops+2+2*i]}")
        _rng.shuffle(rels)
        label = {1: "child", 2: "grandchild", 3: "great grandchild"}
        return f"{'. '.join(rels)}. Who is {chain[0]}'s {label[n_hops]}? Answer:", f" {chain[-1]}"

    def gen_A4(level):
        n_pairs, noise = 4, [0, 2, 4, 8][level]
        keys = _rng.sample(FRUITS, n_pairs)
        vals = _rng.sample(COLORS, n_pairs)
        parts = []
        for i in range(n_pairs):
            parts.append(f"key: {keys[i]} value: {vals[i]}.")
            if noise > 0:
                parts.append(" ".join(_rng.choice(ANIMALS) for _ in range(noise)) + ".")
        qi = _rng.randint(0, n_pairs - 1)
        return f"{' '.join(parts)} What is the value of {keys[qi]}? Answer:", f" {vals[qi]}"

    def gen_B1(level):
        if level == 0:
            a, b = _rng.randint(1, 9), _rng.randint(1, 9); r = a + b
            return f"What is {a} + {b}? Answer:", f" {r}"
        elif level == 1:
            a, b = _rng.randint(10, 49), _rng.randint(10, 49); r = a + b
            return f"What is {a} + {b}? Answer:", f" {r}"
        elif level == 2:
            a, b = _rng.randint(1, 30), _rng.randint(1, 30)
            op = _rng.choice(["+", "-"])
            if op == "-": a, b = max(a, b), min(a, b)
            r = (a + b) if op == "+" else (a - b)
            return f"What is {a} {op} {b}? Answer:", f" {r}"
        else:
            a, b, c = _rng.randint(1, 20), _rng.randint(1, 20), _rng.randint(1, 15)
            s1 = a + b; op2 = _rng.choice(["+", "-"])
            if op2 == "-" and s1 < c: c = _rng.randint(1, s1)
            r = (s1 + c) if op2 == "+" else (s1 - c)
            return f"What is {a} + {b} {op2} {c}? Answer:", f" {r}"

    def gen_B2(level):
        if level == 0:
            a, b = _rng.randint(1, 30), _rng.randint(1, 30)
            return f"If x = {a} + {b}, what is x? Answer:", f" {a+b}"
        elif level == 1:
            a, b, c = _rng.randint(1, 15), _rng.randint(1, 15), _rng.randint(1, 15)
            return f"If x = {a} + {b}, and y = x + {c}, what is y? Answer:", f" {a+b+c}"
        elif level == 2:
            a, b, c = _rng.randint(2, 9), _rng.randint(2, 9), _rng.randint(1, 15)
            return f"If x = {a} * {b} + {c}, what is x? Answer:", f" {a*b+c}"
        else:
            a, b, c, d = _rng.randint(1, 12), _rng.randint(1, 12), _rng.randint(1, 12), _rng.randint(1, 12)
            return f"If x = {a} + {b}, and y = {c} + {d}, what is x + y? Answer:", f" {a+b+c+d}"

    def gen_B3(level):
        if level == 0:
            a, b = _rng.randint(1, 99), _rng.randint(1, 99)
            while a == b: b = _rng.randint(1, 99)
            return f"Which is larger, {a} or {b}? Answer:", f" {max(a, b)}"
        elif level == 1:
            vals = _rng.sample(range(1, 99), 3)
            return f"Which is largest, {vals[0]}, {vals[1]}, or {vals[2]}? Answer:", f" {max(vals)}"
        elif level == 2:
            a, b, c, d = _rng.randint(1, 40), _rng.randint(1, 40), _rng.randint(1, 40), _rng.randint(1, 40)
            while a + b == c + d: d = _rng.randint(1, 40)
            return f"Which is larger, {a} + {b} or {c} + {d}? Answer:", f" {max(a+b, c+d)}"
        else:
            b = _rng.choice([3, 4, 5, 6, 7, 8, 9]); a = _rng.randint(10, 99)
            return f"What is {a} mod {b}? Answer:", f" {a % b}"

    def gen_B4(level):
        if level <= 2:
            target = _rng.choice(ANIMALS[:5])
            others = [x for x in ANIMALS[:8] if x != target]
            seq_len = [_rng.randint(4, 8), _rng.randint(6, 10), _rng.randint(8, 15)][level]
            seq = [_rng.choice([target] + others) for _ in range(seq_len)]
            r = sum(1 for s in seq if s == target)
            return f"Count the number of times '{target}' appears: {' '.join(seq)}. Answer:", f" {r}"
        else:
            t1, t2 = _rng.sample(ANIMALS[:6], 2)
            others = [x for x in ANIMALS[:10] if x not in (t1, t2)]
            seq = [_rng.choice([t1, t2] + others) for _ in range(_rng.randint(6, 12))]
            r = sum(1 for s in seq if s in (t1, t2))
            return f"Count the total times '{t1}' or '{t2}' appear: {' '.join(seq)}. Answer:", f" {r}"

    def gen_A5(level):
        """Memory update: present facts, then update some, query latest value."""
        n_facts, n_updates = [(2, 1), (4, 2), (6, 3), (8, 4)][level]
        names = _rng.sample(NAMES, n_facts)
        attrs = _rng.sample(OBJECTS, n_facts)
        facts = ". ".join(f"{names[i]} owns a {attrs[i]}" for i in range(n_facts))
        update_idx = set(_rng.sample(range(n_facts), n_updates))
        new_attrs = {}
        for ui in update_idx:
            remaining = [o for o in OBJECTS if o != attrs[ui]]
            new_attrs[ui] = _rng.choice(remaining)
        updates = ". ".join(f"{names[ui]} now owns a {new_attrs[ui]}"
                            for ui in sorted(update_idx))
        query_updated = _rng.random() < 0.5
        updated_list = list(update_idx)
        unchanged_list = [i for i in range(n_facts) if i not in update_idx]
        if query_updated and updated_list:
            qi = _rng.choice(updated_list)
        elif unchanged_list:
            qi = _rng.choice(unchanged_list)
        else:
            qi = _rng.choice(updated_list)
        answer = new_attrs[qi] if qi in update_idx else attrs[qi]
        return f"{facts}. Update: {updates}. What does {names[qi]} own now? Answer:", f" {answer}"


    def gen_C1(level):
        """Periodic pattern detection: predict next element in a repeating sequence.
        Probes C-stage grid-cell periodic retention."""
        period, noise_per, length = [
            (2, 0, 8), (3, 0, 15), (4, 2, 20), (6, 3, 30)][level]
        pattern_items = _rng.sample(OBJECTS[:10], period)
        sequence = []
        for i in range(length):
            sequence.append(pattern_items[i % period])
            if noise_per > 0 and _rng.random() < 0.3:
                for _ in range(min(noise_per, 2)):
                    sequence.append(_rng.choice(ANIMALS[:6]))
        next_item = pattern_items[len([s for s in sequence if s in pattern_items]) % period]
        prompt = " ".join(sequence)
        return f"Repeating pattern: {prompt}. Next item:", f" {next_item}"

    def gen_C2(level):
        """Near-miss disambiguation: distinguish entities with overlapping attributes.
        Probes E-stage AHSD/CSR pattern separation."""
        n_entities, n_shared = [(2, 1), (2, 2), (3, 2), (4, 3)][level]
        names = _rng.sample(NAMES, n_entities)
        shared_colors = _rng.sample(COLORS, n_shared)
        shared_objs = _rng.sample(OBJECTS[:10], n_shared)
        unique_objs = _rng.sample(OBJECTS[10:], n_entities)

        facts = []
        for i, name in enumerate(names):
            for j in range(n_shared):
                facts.append(f"{name} has a {shared_colors[j]} {shared_objs[j]}")
            facts.append(f"{name} has a {unique_objs[i]}")
        _rng.shuffle(facts)

        qi = _rng.randint(0, n_entities - 1)
        return (f"{'. '.join(facts)}. Who has a {unique_objs[qi]}? Answer:",
                f" {names[qi]}")

    def gen_C3(level):
        """Partial cue completion: complete a degraded fact from stored patterns.
        Probes D-stage attractor pattern completion."""
        n_facts, n_missing = [(2, 1), (3, 1), (4, 2), (5, 2)][level]
        names = _rng.sample(NAMES, n_facts)
        colors = [_rng.choice(COLORS) for _ in range(n_facts)]
        objs = _rng.sample(OBJECTS, n_facts)
        traits = _rng.sample(TRAITS, n_facts)

        facts = [f"{names[i]} is {traits[i]} and has a {colors[i]} {objs[i]}"
                 for i in range(n_facts)]
        context = ". ".join(facts)

        qi = _rng.randint(0, n_facts - 1)
        cue_parts = []
        query_attr = None
        attrs = [("trait", traits[qi]), ("color", colors[qi]), ("object", objs[qi])]
        mask_indices = _rng.sample(range(len(attrs)), min(n_missing, len(attrs)))
        query_idx = _rng.choice(mask_indices)
        query_attr_name, query_attr_val = attrs[query_idx]

        cue = f"{names[qi]}"
        for j, (aname, aval) in enumerate(attrs):
            if j in mask_indices:
                cue += f" ??? {aname}"
            else:
                cue += f" {aval}"

        if query_attr_name == "color":
            return (f"{context}. Partial cue: {cue}. What is {names[qi]}'s color? Answer:",
                    f" {query_attr_val}")
        elif query_attr_name == "object":
            return (f"{context}. Partial cue: {cue}. What does {names[qi]} have? Answer:",
                    f" {query_attr_val}")
        else:
            return (f"{context}. Partial cue: {cue}. What trait does {names[qi]} have? Answer:",
                    f" {query_attr_val}")

    TASKS = {
        "A1_fact_retrieval": gen_A1, "A2_consistency": gen_A2,
        "A3_multihop": gen_A3, "A4_noisy_retrieval": gen_A4,
        "A5_memory_update": gen_A5,
        "B1_arithmetic": gen_B1, "B2_algebra": gen_B2,
        "B3_comparison": gen_B3, "B4_counting": gen_B4,
        "C1_periodic_pattern": gen_C1, "C2_near_miss": gen_C2,
        "C3_partial_completion": gen_C3,
    }

    CANDIDATES = {
        "A1_fact_retrieval": NAMES,
        "A2_consistency": ["true", "false"],
        "A3_multihop": NAMES,
        "A4_noisy_retrieval": COLORS,
        "A5_memory_update": OBJECTS,
        "B1_arithmetic": [str(i) for i in range(-50, 200)],
        "B2_algebra": [str(i) for i in range(-50, 200)],
        "B3_comparison": [str(i) for i in range(0, 200)],
        "B4_counting": [str(i) for i in range(0, 30)],
        "C1_periodic_pattern": OBJECTS[:10],
        "C2_near_miss": NAMES,
        "C3_partial_completion": COLORS + OBJECTS + TRAITS,
    }

    cand_ids_cache = {}
    for task_name, cands in CANDIDATES.items():
        cand_ids_cache[task_name] = []
        for c in cands:
            ids = tokenizer.encode(f" {c}", add_special_tokens=False)
            cand_ids_cache[task_name].append((c, ids))

    model.eval()
    if hasattr(model, 'reset_memory'):
        model.reset_memory()
    results = {}
    for task_name, gen_fn in TASKS.items():
        results[task_name] = {}
        cand_list = cand_ids_cache[task_name]
        for level in range(4):
            correct = 0
            for _ in range(N_INSTANCES):
                if hasattr(model, 'reset_memory'):
                    model.reset_memory()
                prompt_str, answer_str = gen_fn(level)
                answer_str_stripped = answer_str.strip()
                prompt_ids = tokenizer.encode(prompt_str)
                if len(prompt_ids) >= max_seq:
                    prompt_ids = prompt_ids[-(max_seq - 1):]
                inp = torch.tensor([prompt_ids], dtype=torch.long, device=device)
                with torch.no_grad():
                    logits = model(inp)
                    if isinstance(logits, tuple):
                        logits = logits[0]
                    log_probs = F.log_softmax(logits[0, -1], dim=-1)
                    del logits
                    best_score = float('-inf')
                    best_cand = None
                    for cand_str, cand_tok_ids in cand_list:
                        if not cand_tok_ids:
                            continue
                        score = log_probs[cand_tok_ids[0]].item()
                        if score > best_score:
                            best_score = score
                            best_cand = cand_str
                    if best_cand == answer_str_stripped:
                        correct += 1
                if xm_mod is not None:
                    xm_mod.mark_step()
            results[task_name][level] = correct / N_INSTANCES

    print(f"\n{'='*72}")
    print("Structured Benchmark Results (accuracy per task per level)")
    print(f"{'='*72}")
    header = f"{'Task':<22s}" + "".join(f" {'L'+str(i):>7s}" for i in range(4)) + f" {'Avg':>7s}"
    print(header)
    print("-" * len(header))
    a_accs, b_accs, c_accs = [], [], []
    for tn in TASKS:
        row = f"{tn:<22s}"
        avgs = []
        for lv in range(4):
            acc = results[tn][lv]
            row += f" {acc:>6.1%}"; avgs.append(acc)
        row += f" {sum(avgs)/4:>6.1%}"
        print(row)
        if tn.startswith("A"):
            a_accs.extend(avgs)
        elif tn.startswith("B"):
            b_accs.extend(avgs)
        else:
            c_accs.extend(avgs)
    print("-" * len(header))
    print(f"{'Content (A1-A5)':<22s}{' ':>31s} {sum(a_accs)/len(a_accs):>6.1%}")
    print(f"{'Math (B1-B4)':<22s}{' ':>31s} {sum(b_accs)/len(b_accs):>6.1%}")
    print(f"{'CEDL-probe (C1-C3)':<22s}{' ':>31s} {sum(c_accs)/len(c_accs):>6.1%}")
    all_accs = a_accs + b_accs + c_accs
    print(f"{'Overall':<22s}{' ':>31s} {sum(all_accs)/len(all_accs):>6.1%}")
    print(f"{'='*72}\n")
    return results



def run_downstream_eval(model: nn.Module, tag: str, device: torch.device):
    """Zero-shot evaluation on LAMBADA, HellaSwag, ARC-Easy."""
    try:
        from lm_eval import evaluator
        from lm_eval.api.model import LM
        from transformers import GPT2TokenizerFast
    except ImportError:
        print("  lm-eval not installed. Skipping downstream eval.")
        print("  Install: pip install lm-eval transformers")
        return {}

    tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")

    class WrappedModel(LM):
        def __init__(self, model, tokenizer, device):
            super().__init__()
            self._model = model
            self._tokenizer = tokenizer
            self._device = device

        @property
        def eot_token_id(self):
            return self._tokenizer.eos_token_id

        @property
        def max_length(self):
            return 1024

        @property
        def max_gen_toks(self):
            return 256

        @property
        def batch_size(self):
            return 8

        @property
        def device(self):
            return self._device

        def tok_encode(self, string):
            return self._tokenizer.encode(string, add_special_tokens=False)

        def tok_decode(self, tokens):
            return self._tokenizer.decode(tokens)

        def loglikelihood(self, requests):
            results = []
            for req in requests:
                if hasattr(req, 'args'):
                    ctx, cont = req.args
                else:
                    ctx, cont = req
                ctx_ids = self.tok_encode(ctx)
                cont_ids = self.tok_encode(cont)
                all_ids = ctx_ids + cont_ids
                if len(all_ids) > 1024:
                    excess = len(all_ids) - 1024
                    all_ids = all_ids[excess:]
                    ctx_ids = ctx_ids[excess:]
                input_ids = torch.tensor([all_ids], device=self._device)
                with torch.no_grad():
                    logits = self._model(input_ids)
                    if isinstance(logits, tuple):
                        logits = logits[0]
                start = len(ctx_ids)
                log_probs = F.log_softmax(logits[0], dim=-1)
                cont_log_prob = 0.0
                is_greedy = True
                for j in range(start, len(all_ids)):
                    tok = all_ids[j]
                    lp = log_probs[j - 1, tok].item()
                    cont_log_prob += lp
                    if logits[0, j - 1].argmax().item() != tok:
                        is_greedy = False
                results.append((cont_log_prob, is_greedy))
            return results

        def loglikelihood_rolling(self, requests):
            results = []
            for req in requests:
                if hasattr(req, 'args'):
                    (text,) = req.args
                else:
                    text = req[0] if isinstance(req, tuple) else req
                all_ids = self.tok_encode(text)
                max_len = 1024
                stride = 512
                total_lp = 0.0
                scored = set()
                for start in range(0, max(1, len(all_ids) - 1), stride):
                    end = min(start + max_len, len(all_ids))
                    chunk = all_ids[start:end]
                    if len(chunk) < 2:
                        continue
                    input_t = torch.tensor([chunk], device=self._device)
                    with torch.no_grad():
                        logits = self._model(input_t)
                        if isinstance(logits, tuple):
                            logits = logits[0]
                    log_probs = F.log_softmax(logits[0], dim=-1)
                    score_start = 1 if start == 0 else (end - start - stride)
                    for j in range(max(1, score_start), len(chunk)):
                        abs_pos = start + j
                        if abs_pos not in scored:
                            total_lp += log_probs[j - 1, chunk[j]].item()
                            scored.add(abs_pos)
                    if end == len(all_ids):
                        break
                results.append((total_lp,))
            return results

        def generate_until(self, requests):
            results = []
            for req in requests:
                if hasattr(req, 'args'):
                    ctx, kwargs = req.args
                else:
                    ctx, kwargs = req
                input_ids = self.tok_encode(ctx)[-900:]
                input_t = torch.tensor([input_ids], device=self._device)
                max_gen = kwargs.get("max_gen_toks", 64)
                stop = kwargs.get("until", [])
                for _ in range(max_gen):
                    with torch.no_grad():
                        logits = self._model(input_t[:, -1024:])
                        if isinstance(logits, tuple):
                            logits = logits[0]
                    next_tok = logits[0, -1].argmax().item()
                    input_t = torch.cat([input_t,
                        torch.tensor([[next_tok]], device=self._device)], dim=1)
                    gen_text = self._tokenizer.decode(
                        input_t[0, len(input_ids):].tolist())
                    if any(s in gen_text for s in stop):
                        break
                results.append(gen_text)
            return results

    wrapped = WrappedModel(model, tokenizer, device)
    tasks = ["lambada_openai", "hellaswag", "arc_easy"]

    try:
        results = evaluator.simple_evaluate(
            model=wrapped, tasks=tasks, batch_size=8)
        print(f"\n  Downstream Results ({tag}):")
        for task_name, task_res in results.get("results", {}).items():
            acc = (task_res.get("acc,none",
                   task_res.get("acc_norm,none",
                   task_res.get("acc",
                   task_res.get("acc_norm",
                   task_res.get("perplexity,none",
                   task_res.get("perplexity", "N/A")))))))
            print(f"    {task_name}: {acc}")
        return results
    except Exception as e:
        print(f"  Downstream eval failed: {e}")
        return {}


def _downstream_metric(
    downstream: Dict[str, Any],
    task_names: Tuple[str, ...],
    metric_names: Tuple[str, ...],
) -> Optional[float]:
    """Extract a scalar metric across lm-eval task/key naming variants."""
    results = downstream.get("results", {}) if isinstance(downstream, dict) else {}
    for task_name in task_names:
        task_results = results.get(task_name, {})
        if not isinstance(task_results, dict):
            continue
        for metric_name in metric_names:
            value = task_results.get(metric_name)
            if isinstance(value, (int, float)):
                return float(value)
    return None


def _fmt_pct(value: Optional[float]) -> str:
    return "N/A" if value is None else f"{value:.1%}"


def _structured_group_avg(
    structured: Dict[str, Dict[int, float]],
    prefix: str,
) -> Optional[float]:
    vals = [
        v
        for task_name, levels in structured.items()
        if task_name.startswith(prefix) and isinstance(levels, dict)
        for v in levels.values()
        if isinstance(v, (int, float))
    ]
    return sum(vals) / len(vals) if vals else None


def print_final_results(results_table: Dict[str, Dict[str, Any]]) -> None:
    """Print only metrics that were actually run; skipped suites show as N/A."""
    downstream_cols = [
        ("LAMBADA", ("lambada_openai", "lambada"), ("acc,none", "acc")),
        ("HellaSwag", ("hellaswag",), ("acc_norm,none", "acc_norm", "acc,none", "acc")),
        ("ARC-Easy", ("arc_easy",), ("acc_norm,none", "acc_norm", "acc,none", "acc")),
    ]
    has_downstream = any(
        bool((res.get("downstream") or {}).get("results"))
        for res in results_table.values()
    )
    has_structured = any(bool(res.get("structured")) for res in results_table.values())

    print(f"\n{'='*70}")
    print("FINAL RESULTS")
    table_width = 90 if has_downstream else 46
    print(f"{'='*table_width}")

    header = f"{'Model':20s} {'Params':>10s} {'Test PPL':>10s}"
    if has_downstream:
        for label, _, _ in downstream_cols:
            header += f" {label:>10s}"
    print(header)
    print("-" * len(header))

    for tag, res in results_table.items():
        row = f"{tag:20s} {res['params']/1e6:>8.1f}M {res['test_ppl']:>10.2f}"
        if has_downstream:
            downstream = res.get("downstream", {})
            for _, task_names, metric_names in downstream_cols:
                metric = _downstream_metric(downstream, task_names, metric_names)
                row += f" {_fmt_pct(metric):>10s}"
        print(row)

    if not has_downstream:
        print("\nDownstream eval: N/A (lm-eval not installed or downstream eval failed).")

    if not has_structured:
        print("Structured benchmark: N/A (skipped for 100M path; use PPL/downstream metrics).")
        return

    print(f"\n{'='*70}")
    print("STRUCTURED RESULTS")
    print(f"{'Model':20s} {'Struct A':>9s} {'Struct B':>9s} {'Struct C':>9s} {'Overall':>9s}")
    print("-" * 62)
    for tag, res in results_table.items():
        sr = res.get("structured", {})
        a_avg = _structured_group_avg(sr, "A")
        b_avg = _structured_group_avg(sr, "B")
        c_avg = _structured_group_avg(sr, "C")
        all_vals = [
            v
            for levels in sr.values()
            if isinstance(levels, dict)
            for v in levels.values()
            if isinstance(v, (int, float))
        ]
        overall = sum(all_vals) / len(all_vals) if all_vals else None
        print(
            f"{tag:20s} {_fmt_pct(a_avg):>9s} {_fmt_pct(b_avg):>9s} "
            f"{_fmt_pct(c_avg):>9s} {_fmt_pct(overall):>9s}"
        )



def _run_v4c_pair_preview(args):
    """V4c synthetic-pair preview + audit. No model build, no training.

    Generates 500 pairs via data_v4c_pairs.generate(), runs the family-aware
    audit (hard-fail on span-validity violations), reports the duplicate-
    collision rate (soft-fail at ≥5%), and prints the first 10 decoded items.
    Exits 0 on PASS, prints failures and exits nonzero on hard-fail."""
    try:
        import data_v4c_pairs as v4c_module
        from transformers import GPT2TokenizerFast
    except Exception as e:
        print(f"FATAL: V4c preview imports failed: {e}")
        raise SystemExit(1)
    print("=" * 70)
    print(f"V4c SYNTHETIC-PAIR PREVIEW — seed={args.seed}")
    print("=" * 70)
    tok = GPT2TokenizerFast.from_pretrained("gpt2")
    items = v4c_module.generate(tok, n=500, seed=args.seed)
    report = v4c_module.audit(items, tok, verbose=True)
    print()
    print(f"  Total items:     {report['total']}")
    print(f"  By family:       {report['by_family']}")
    print(f"  Hard-fails:      {len(report['hard_fails'])}")
    print(f"  Anchorable:      {report['n_anchorable']}")
    print(f"  Dup collisions:  {report['duplicate_collisions']}")
    print()
    if report["hard_fails"]:
        print("  STATUS: HARD-FAIL — fix data_v4c_pairs.py before training")
        raise SystemExit(1)
    print(f"  STATUS: span-validity PASS")

    n_update_items = report["total"] - report["by_family"].get(
        "neutral_control", 0)
    dup_rate = report["duplicate_collisions"] / max(1, n_update_items)
    print(f"\n  --- Duplicate-collision audit ---")
    print(f"  Overall: {report['duplicate_collisions']} / {n_update_items} "
          f"update-family items = {dup_rate:.2%}")
    fams_over = []
    for fam, dup in report.get("duplicate_collisions_by_family", {}).items():
        fam_total = report["by_family"].get(fam, 0)
        fam_rate = dup / max(1, fam_total)
        flag = "FAIL" if fam_rate >= 0.10 else "warn" if fam_rate >= 0.05 else "ok"
        print(f"    {fam:<22s}  {dup:>3d}/{fam_total:>3d} = {fam_rate:>6.2%}  [{flag}]")
        if fam_rate >= 0.10:
            fams_over.append((fam, fam_rate))
    if report["top_repeated_pairs"]:
        print(f"  Top repeated (family, query, current) triples:")
        for entry in report["top_repeated_pairs"][:8]:
            print(f"    {entry['family']:<22s}  q={entry['query']!r:<14s}  "
                  f"c={entry['current']!r:<14s}  ×{entry['count']}")
    if fams_over:
        print(f"\n  STATUS: HARD-FAIL — {len(fams_over)} families ≥10% dup rate:")
        for fam, rate in fams_over:
            print(f"    {fam}: {rate:.2%}")
        print(f"  Expand the corresponding pools in data_v4c_pairs.py "
              f"(NAMES/OBJECTS/CAPITALS_OLD/GATES/PARAPHRASE_*_CITIES) "
              f"before training. Reviewer-blocker: at this rate, InfoNCE "
              f"in-batch FN masking is doing major work every batch.")
        raise SystemExit(1)
    if dup_rate >= 0.05:
        print(f"  WARN: overall rate {dup_rate:.2%} ≥ 5%; under per-family "
              f"10% hard-fail bar, but worth monitoring.")
    print()
    print("  --- First 10 decoded pairs ---")
    for i, it in enumerate(items[:10]):
        print(f"\n  Item {i}  family={it.family}")
        text = tok.decode(it.ids).replace("\n", " ").strip()
        print(f"    text:    {text!r}")
        decoded = v4c_module.decode_spans(it, tok)
        for k, v in decoded.items():
            if v:
                print(f"    {k:>21s}: {v!r}")


def _run_v3m2_preview(args):
    """Build the V3m2 model, load shared backbone weights from an existing
    checkpoint (strict=False, with audit), run ONE forward in eval mode on a
    16-seq WikiText sample, and print:
      (a) decoded anchor_negative_ids / label_negative_ids audit
      (b) belief-revision-lexicon eligibility check (must all be True)
      (c) per-sentence anchor + generated-label inspection for first 5 seqs.

    Aborts BEFORE training if (b) fails or if labels are still firing mostly
    on periods / function words."""
    if args.model not in ("CEDL-V3m2", "CEDL-V3m2-NLL"):
        print(f"ERROR: --preview-only requires --model CEDL-V3m2 or "
              f"CEDL-V3m2-NLL, got {args.model!r}")
        return
    if not args.preview_shared_checkpoint:
        print("ERROR: --preview-only requires --preview-shared-checkpoint "
              "(path to V0/V3/V3m checkpoint to load shared backbone "
              "weights from). Random-init preview is meaningless.")
        return

    print("=" * 70)
    print(f"V3m2 TARGET-PREVIEW — model={args.model}")
    print(f"  shared-backbone checkpoint: {args.preview_shared_checkpoint}")
    print("=" * 70)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    model = build_model(args.model, vocab=50257, max_seq=1024).to(device)
    model.eval()

    state = torch.load(args.preview_shared_checkpoint,
                       map_location="cpu", weights_only=True)
    if any(k.startswith("_orig_mod.") for k in state):
        state = {k.replace("_orig_mod.", ""): v for k, v in state.items()}
    res = model.load_state_dict(state, strict=False)
    print(f"\n  Loaded shared weights with strict=False:")
    print(f"    missing keys    (V3m2-new):     {len(res.missing_keys)}")
    for k in res.missing_keys[:12]:
        print(f"      - {k}")
    if len(res.missing_keys) > 12:
        print(f"      ... ({len(res.missing_keys) - 12} more)")
    print(f"    unexpected keys (V3/V3m-only):  {len(res.unexpected_keys)}")
    for k in res.unexpected_keys[:12]:
        print(f"      - {k}")
    if len(res.unexpected_keys) > 12:
        print(f"      ... ({len(res.unexpected_keys) - 12} more)")

    raw = model
    if not hasattr(raw, "anchor_negative_ids"):
        print("\nFATAL: model has no anchor_negative_ids buffer; lexicon "
              "construction failed at build time. Cannot preview.")
        return

    try:
        from transformers import GPT2TokenizerFast
        tok = GPT2TokenizerFast.from_pretrained("gpt2")
    except Exception as e:
        print(f"\nWARN: transformers unavailable ({e}); decoded mask audit "
              f"skipped.")
        tok = None

    print(f"\n  --- (a) Decoded mask audit ---")
    print(f"  anchor_negative_ids (N={raw.anchor_negative_ids.numel()}):")
    if tok is not None and raw.anchor_negative_ids.numel() > 0:
        decoded = [repr(tok.decode([int(i)]))
                   for i in raw.anchor_negative_ids[:8].tolist()]
        print(f"    sample: {decoded}")
    print(f"  label_negative_ids  (N={raw.label_negative_ids.numel()}):")
    if tok is not None and raw.label_negative_ids.numel() > 0:
        decoded = [repr(tok.decode([int(i)]))
                   for i in raw.label_negative_ids[:30].tolist()]
        print(f"    first 30: {decoded}")

    print(f"\n  --- (b) Belief-revision eligibility check ---")
    abort = False
    if tok is not None:
        belief_words = ["not", "no", "never", "but", "however", "although",
                        "now", "currently", "previously", "instead",
                        "all", "every", "some", "any"]
        label_neg_set = set(raw.label_negative_ids.tolist())
        for w in belief_words:
            for form in (w, " " + w):
                ids = tok.encode(form, add_special_tokens=False)
                if not ids:
                    continue
                tid = ids[0]
                eligible = tid not in label_neg_set
                marker = "OK  " if eligible else "FAIL"
                print(f"    [{marker}] {form!r:>14s}  -> token_id={tid:>5d}  "
                      f"label_eligible={eligible}")
                if not eligible:
                    abort = True
        if abort:
            print("\n  ABORT: a belief-revision marker is in label_negative_ids. "
                  "Fix V3M2_LABEL_NEGATIVE_WORDS in CEDL.py and re-run.")
            return

    print(f"\n  --- (c) Per-sentence anchor + label inspection ---")
    try:
        from datasets import load_dataset
        ds = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split="validation")
        texts = [t for t in ds["text"] if len(t.strip()) > 200][:128]
        if tok is None:
            print("    skipping (tokenizer unavailable)")
            return
        seqs = []
        for t in texts:
            ids = tok.encode(t, add_special_tokens=False)[:512]
            if len(ids) >= 64:
                seqs.append(ids)
            if len(seqs) >= 64:
                break
        if not seqs:
            print("    no eligible sequences found; aborting preview")
            return
        T = min(len(s) for s in seqs)
        ids_tensor = torch.tensor([s[:T] for s in seqs],
                                  dtype=torch.long, device=device)
    except Exception as e:
        print(f"    WARN: could not load WikiText ({e}); aborting preview")
        return

    calib_n = min(8, max(1, ids_tensor.size(0) // 4))
    calib_batches = [ids_tensor[i * 4:(i + 1) * 4]
                     for i in range(calib_n) if (i + 1) * 4 <= ids_tensor.size(0)]
    needs_calib = (raw.curve_C_ema_initialized.item() == 0.0)
    if needs_calib:
        print(f"    EMA calibration: manifold buffers are UNINITIALIZED "
              f"(shared checkpoint was likely V0/V3, not V3m/V3m2). "
              f"Running {len(calib_batches)} warmup batches to populate "
              f"curve_C/curve_E/cosmis EMA before preview...")
        with torch.no_grad():
            for cb in calib_batches:
                hC = raw.c_stage(cb, feedback=None)
                hE, _ = raw.e_stage(hC, feedback=None)
                _build_v3m2_targets(
                    ids=cb, h_C_p0=hC, h_E_p0=hE,
                    logits_mem_p0=None,
                    model=raw, salience_mode=raw.salience_mode,
                    update_ema=True, with_nll=False)
        cinit = raw.curve_C_ema_initialized.item()
        cvar  = raw.curve_C_ema_var.item()
        evar  = raw.curve_E_ema_var.item()
        cosvar = raw.cosmis_ema_var.item()
        print(f"    after calib: curve_C_init={cinit:.4f}  "
              f"curve_C_var={cvar:.4f}  curve_E_var={evar:.4f}  "
              f"cosmis_var={cosvar:.4f}")
        if cinit < 0.05 or cvar < 1e-6 or evar < 1e-6 or cosvar < 1e-6:
            print("    ABORT: EMA buffers did not move after warmup. "
                  "Anchor threshold cannot be calibrated. Use a V3m/V3m2 "
                  "checkpoint instead, or extend calibration batches.")
            return

    with torch.no_grad():
        h_C = raw.c_stage(ids_tensor, feedback=None)
        h_E, _ = raw.e_stage(h_C, feedback=None)
        target, w, dbg = _build_v3m2_targets(
            ids=ids_tensor,
            h_C_p0=h_C, h_E_p0=h_E,
            logits_mem_p0=None,
            model=raw,
            salience_mode=raw.salience_mode,
            update_ema=False,
            with_nll=False,
        )

    anchors = dbg["anchors"]
    strength = dbg["strength"]
    n_show = min(5, ids_tensor.size(0))
    label_neg_set = set(raw.label_negative_ids.tolist())
    sent_end_ids = set()
    if tok is not None:
        for s in (".", "!", "?", " .", " !", " ?"):
            try:
                ids_s = tok.encode(s, add_special_tokens=False)
                if ids_s:
                    sent_end_ids.add(ids_s[0])
            except Exception:
                pass
    n_pos = 0
    n_pos_in_label_neg = 0
    n_pos_sent_initial = 0
    pos_mask = target > 0.05
    for b in range(ids_tensor.size(0)):
        for p in pos_mask[b].nonzero(as_tuple=True)[0].tolist():
            n_pos += 1
            tid = int(ids_tensor[b, p].item())
            if tid in label_neg_set:
                n_pos_in_label_neg += 1
            prev_id = int(ids_tensor[b, p - 1].item()) if p > 0 else -1
            if prev_id in sent_end_ids:
                n_pos_sent_initial += 1
    for b in range(n_show):
        anchor_pos = anchors[b].nonzero(as_tuple=True)[0].tolist()
        print(f"\n    seq {b}  T={ids_tensor.size(1)}  "
              f"#anchors={len(anchor_pos)}")
        for p in anchor_pos[:5]:
            tok_id = int(ids_tensor[b, p].item())
            decoded = tok.decode([tok_id])
            s = float(strength[b, p].item())
            print(f"      anchor @ pos={p:>3d}  token={decoded!r:>12s}  "
                  f"strength={s:.3f}")
        pos_targets = (target[b] > 0.05).nonzero(as_tuple=True)[0].tolist()
        labels_show = pos_targets[:10]
        if labels_show:
            print(f"      positive-label tokens:")
            for p in labels_show:
                tok_id = int(ids_tensor[b, p].item())
                decoded = tok.decode([tok_id])
                v = float(target[b, p].item())
                print(f"        pos={p:>3d}  token={decoded!r:>12s}  "
                      f"target={v:.3f}")

    print(f"\n  --- preview hard checks (N={n_pos} positive labels total) ---")
    if n_pos == 0:
        print("    WARN: no positive labels produced — anchor threshold may "
              "be too strict on uncalibrated EMA, or input is genuinely flat.")
    else:
        frac_neg = n_pos_in_label_neg / n_pos
        frac_si  = n_pos_sent_initial / n_pos
        flag_neg = "PASS" if frac_neg <= 0.05 else "FAIL"
        flag_si  = "PASS" if frac_si  <= 0.50 else "FAIL"
        print(f"    positive labels IN label_negative_ids: {frac_neg:>5.3f}  "
              f"(criterion 3 target <= 0.05)  [{flag_neg}]")
        print(f"    positive labels at sentence-initial:   {frac_si:>5.3f}  "
              f"(preview soft target <= 0.50)  [{flag_si}]")
        if flag_neg == "FAIL":
            print("    ABORT: label mask is letting punctuation/stopwords "
                  "through. Fix V3M2_LABEL_NEGATIVE_WORDS in CEDL.py.")

    print(f"\n  --- (d) Curated discourse-prompts preview ---")
    curated = [
        ("Alice wanted to leave early, however she stayed until the end.",
         "explicit"),
        ("The forecast called for rain, but the day turned out sunny.",
         "explicit"),
        ("Although Tom had studied hard, the exam was very difficult.",
         "explicit"),
        ("John loves coffee. He never drinks it.",            "implicit"),
        ("Mary said she would attend. She stayed home.",      "implicit"),
        ("The cake was delicious. Nobody ate it.",            "implicit"),
        ("Alice used to live in Paris. She now lives in Berlin.",
         "belief"),
        ("The capital was Bonn until 1990. Today it is Berlin.",
         "belief"),
        ("Previously the gate was A12. Currently it is C7.",  "belief"),
        ("Alice owns a hat. Bob owns a pen. Carol owns a cup. "
         "But Alice now owns a car. What does Alice own now? Answer:",
         "stale_current"),
        ("Bob owns a ring. Dave owns a lamp. Eve owns a book. "
         "But Bob now owns a kite. What does Bob own now? Answer:",
         "stale_current"),
        ("Carol owns a bag. Frank owns a fork. Grace owns a desk. "
         "But Carol now owns a fish. What does Carol own now? Answer:",
         "stale_current"),
    ]

    MARKER_FORMS = ["however", "but", "although", "yet", "instead",
                    "nevertheless", "though", "never", "no", "not",
                    "nobody", "nothing", "now", "currently", "previously",
                    "today", "originally"]
    marker_ids = set()
    for w in MARKER_FORMS:
        for f in (w, " " + w, w.capitalize(), " " + w.capitalize()):
            try:
                ids = tok.encode(f, add_special_tokens=False)
                if ids:
                    marker_ids.add(ids[0])
            except Exception:
                pass

    n_on_marker = 0
    n_total_anchors = 0
    for text, cat in curated:
        ids = tok.encode(text, add_special_tokens=False)
        ids_t = torch.tensor([ids], dtype=torch.long, device=device)
        with torch.no_grad():
            hC = raw.c_stage(ids_t, feedback=None)
            hE, _ = raw.e_stage(hC, feedback=None)
            tgt, _w, dbg = _build_v3m2_targets(
                ids=ids_t, h_C_p0=hC, h_E_p0=hE,
                logits_mem_p0=None, model=raw,
                salience_mode=raw.salience_mode,
                update_ema=False, with_nll=False)
        apos = dbg["anchors"][0].nonzero(as_tuple=True)[0].tolist()
        lpos = (tgt[0] > 0.05).nonzero(as_tuple=True)[0].tolist()
        marker_hit = any(int(ids_t[0, p].item()) in marker_ids for p in apos)
        n_total_anchors += len(apos)
        n_on_marker += sum(1 for p in apos
                           if int(ids_t[0, p].item()) in marker_ids)
        anc_str = ", ".join(
            f"{p}={repr(tok.decode([int(ids_t[0, p].item())]))}"
            for p in apos) or "(none)"
        lbl_str = ", ".join(
            f"{p}={repr(tok.decode([int(ids_t[0, p].item())]))}"
            f":{tgt[0, p].item():.2f}"
            for p in lpos[:8]) or "(none)"
        flag = "  [MARKER]" if marker_hit else "  [no-marker]"
        print(f"    [{cat:>13s}] {text[:64]}...")
        print(f"      anchors: {anc_str}{flag}")
        print(f"      labels:  {lbl_str}")
    frac_marker = (n_on_marker / max(1, n_total_anchors))
    flag_m = "PASS" if frac_marker >= 0.40 else "FAIL"
    print(f"\n    anchor-on-discourse-marker fraction: {frac_marker:>5.3f}  "
          f"(curated target >= 0.40)  [{flag_m}]")
    print(f"    (anchors here should land on discourse markers — "
          f"however/but/now/currently/never/etc. — and labels should "
          f"land on the contrasting clause that follows.)")

    print(f"\n  Preview complete. Review printed anchors/labels above for "
          f"qualitative confirmation.")



def main():
    parser = argparse.ArgumentParser(description="CEDL 100M Benchmark")
    parser.add_argument("--model", type=str, default="CEDL",
                        help="Model tag, comma-separated list, or 'all'. "
                             "Examples: CEDL  |  CEDL-B0  |  CEDL,CEDL-B0  |  all")
    parser.add_argument("--dataset", type=str, default="wikitext103")
    parser.add_argument("--tpu", action="store_true")
    parser.add_argument("--eval-only", action="store_true")
    parser.add_argument("--checkpoint", type=str, default=None)
    parser.add_argument("--verify-params", action="store_true",
                        help="Print param counts and exit")
    parser.add_argument("--max-steps", type=int, default=30_000)
    parser.add_argument("--eval-interval", type=int, default=None,
                        help="Eval cadence in steps (Config default 1000). "
                             "Lower for short Stage-0 runs so the λ trajectory "
                             "/ PPL trend has more datapoints.")
    parser.add_argument("--save-interval", type=int, default=None,
                        help="Periodic checkpoint cadence in steps (Config "
                             "default 5000).")
    parser.add_argument("--total-steps", type=int, default=None,
                        help="V5-LM: GLOBAL target step (alias for --max-steps; "
                             "schedules key off the global step). If set, "
                             "overrides --max-steps.")
    parser.add_argument("--resume-checkpoint", type=str, default=None,
                        help="V5-LM: resume from a checkpoint. Bare state_dict "
                             "(e.g. CEDL-V4e_best.pt) → fresh AdamW; train-state "
                             "dict {model,optimizer,global_step} → restore all.")
    parser.add_argument("--start-step", type=int, default=0,
                        help="V5-LM: global step to resume the loop counter at "
                             "(e.g. 5000). Schedules continue from here.")
    parser.add_argument("--run-until-step", type=int, default=None,
                        help="V5-LM: STOP the loop at this global step WITHOUT "
                             "compressing the schedule (which uses --total-steps "
                             "as its denominator). Stage 0: --total-steps 30000 "
                             "--run-until-step 6000.")
    parser.add_argument("--save-trainstate", action="store_true",
                        help="V5-LM/B0-long: also write resume-only "
                             "{model,optimizer,global_step} checkpoints for "
                             "multi-session continuation.")
    parser.add_argument("--save-dir", type=str, default=None,
                        help="Checkpoint/scorecard output dir. For Colab "
                             "multi-session V5, point at Drive (e.g. "
                             "/content/drive/MyDrive/cedl_ckpt) so train-state "
                             "survives session resets. Use a SEPARATE dir per "
                             "run to avoid mixing scorecard logs across runs.")
    parser.add_argument("--v5-grad-isolation", action="store_true",
                        help="V5-LM: route the V4c aux gradient to the salience "
                             "pathway only (trunk trains on LM-CE). Auto-enabled "
                             "for CEDL-V5-LM* tags.")
    parser.add_argument("--v5-trunk-aux-frac", type=float, default=0.0,
                        help="V5-LM: fraction of aux gradient still allowed into "
                             "the trunk (0.0 = full isolation; Stage-0 may use "
                             "0.05).")
    parser.add_argument("--v5-aux-scale-early", type=float, default=None,
                        help="V5-LM: aux scale for global steps 5K-10K (default 0.50)")
    parser.add_argument("--v5-aux-scale-mid", type=float, default=None,
                        help="V5-LM: aux scale for global steps 10K-20K (default 0.25)")
    parser.add_argument("--v5-aux-scale-late", type=float, default=None,
                        help="V5-LM: aux scale floor for steps 20K+ (default 0.10)")
    parser.add_argument("--v5-cadence-early", type=int, default=None,
                        help="V5-LM: v4c_every for steps 5K-10K (default 8)")
    parser.add_argument("--v5-cadence-mid", type=int, default=None,
                        help="V5-LM: v4c_every for steps 10K-20K (default 16)")
    parser.add_argument("--v5-cadence-late", type=int, default=None,
                        help="V5-LM: v4c_every for steps 20K+ (default 32)")
    parser.add_argument("--v5-family-reweight", action="store_true",
                        help="V5-LM: oversample but_update (0.30→0.50) and "
                             "drop neutral (0.20→0.10) in V4c batches — "
                             "Path B for B1 (but_update is the weakest B2 "
                             "family). No test-format contamination.")
    parser.add_argument("--v6-mixture", action="store_true",
                        help="V6: documentary flag. The mixture is wired by "
                             "the CEDL-V6 tag itself (build_model only "
                             "registers bank_q_proj + v6_lambda_* for that "
                             "tag), so this flag is auto-implied for CEDL-V6 "
                             "and a NO-OP on every other tag — passing it "
                             "with e.g. CEDL-V5-LM logs a warning and is "
                             "ignored. Use --v6-lambda-init to override the "
                             "initial λ logit.")
    parser.add_argument("--v6-lambda-init", type=float, default=-4.0,
                        help="V6: BIAS LOGIT for the λ sigmoid gate "
                             "(λ = sigmoid(v6_lambda_a · v_scalar + v6_lambda_b); "
                             "v6_lambda_b inits to this value). Actual λ̄ also "
                             "depends on v_scalar's distribution from the V5 "
                             "trainstate — observed range at init: -4.0 → "
                             "λ̄≈0.001, -2.0 → λ̄≈0.06 (empirical, NOT just "
                             "the bias sigmoid). Try -2.0 if λ stays "
                             "collapsed in Stage 0.")
    parser.add_argument("--v6-lambda-a-init", type=float, default=1.0,
                        help="V6.1b: slope (a) of the λ affine over v_scalar. "
                             "Default +1.0 matches V6.1 v1/v2/v3. V6.1b "
                             "sign-correction sets -1.0 (V4c-ans has LOWER "
                             "v_scalar than WikiText per d=-0.85 separability, "
                             "so negative slope is needed for V4c-ans to get "
                             "HIGHER λ). Only meaningful for the CEDL-V6 tag; "
                             "ignored elsewhere with a WARN.")
    parser.add_argument("--v6-aux-weight", type=float, default=0.0,
                        help="V6.1: global scale on the V4c-routed bank-aux "
                             "loss (L_bank + v6_mix_weight·L_mix + "
                             "v6_gate_weight·L_gate). 0.0 = V6 baseline "
                             "(mixture-only, falsified). 0.05 = V6.1 main run.")
    parser.add_argument("--v6-margin-target", type=float, default=1.0,
                        help="V6.1: hinge target (nats) for L_bank and L_mix.")
    parser.add_argument("--v6-mix-weight", type=float, default=0.5,
                        help="V6.1: weight on the mixture-output margin L_mix. "
                             "Mostly 0 in practice (trunk margins ≈ 30 nats > "
                             "target=1), but trains v6_lambda_a/b when bank "
                             "beats trunk on a few examples.")
    parser.add_argument("--v6-gate-weight", type=float, default=0.01,
                        help="V6.1: weight on the LOGIT-SPACE λ floor "
                             "regularizer L_gate (anti-collapse insurance). "
                             "Sigmoid-space floor has vanishing gradient when "
                             "λ collapses; logit-space gives constant gradient.")
    parser.add_argument("--v6-lambda-floor", type=float, default=0.05,
                        help="V6.1: λ floor at answer positions (default 0.05; "
                             "logit ≈ -2.94). L_gate fires when λ < this.")
    parser.add_argument("--v6-lambda-head", action="store_true",
                        help="Stage 2a: replace the v_scalar-coupled affine "
                             "gate with a separate λ head MLP(h_D). Gradient "
                             "is aux-only (torch.no_grad() in main forward + "
                             "V5 allowlist for the head's params), so LM-CE "
                             "cannot drive λ→0 the way it did in V6/V6.1/V6.1b.")
    parser.add_argument("--v6-lambda-head-hidden", type=int, default=160,
                        help="Stage 2a: hidden dim of the λ head MLP (d/4=160 "
                             "default at d=640). ~102K total head params.")
    parser.add_argument("--v6-lambda-head-bias-init", type=float, default=-7.0,
                        help="Stage 2a: final-layer bias init (sigmoid(-7)≈9e-4 "
                             "so λ starts near zero everywhere; the head LEARNS "
                             "h_D→λ selectivity from L_gate + L_bg supervision).")
    parser.add_argument("--v6-bg-weight", type=float, default=1.0,
                        help="Stage 2a: weight on L_bg, the background-row "
                             "sparsity hinge. Asymmetric counterpart to "
                             "L_gate — pushes λ DOWN at sampled non-answer V4c "
                             "rows so smoke tests measure selectivity rather "
                             "than bias lift.")
    parser.add_argument("--v6-bg-target", type=float, default=0.01,
                        help="Stage 2a: λ target at background V4c rows; "
                             "L_bg = relu(λ_bg - target).mean().")
    parser.add_argument("--v6-bce-objective", action="store_true",
                        help="Stage 2a-bce: replace hinge L_gate + L_bg with "
                             "class-balanced BCE (target=1 at ans, target=0 "
                             "at bg, 0.5/0.5 mean). Strong gradient at the "
                             "floor exactly proportional to how wrong each "
                             "row is — fixes the stalemate where the hinge "
                             "objective could not lift λ_ans above floor.")
    parser.add_argument("--v6-sel-weight", type=float, default=1.0,
                        help="Stage 2a-bce: weight on L_sel (subsumes "
                             "gate_weight + bg_weight under BCE). Default 1.0 "
                             "with v6_aux_weight=0.05 gives effective scale "
                             "0.05 — 20× the failed gate=0.05 hinge run.")
    parser.add_argument("--v6-lambda-head-w-init-std", type=float, default=1e-3,
                        help="Stage 2a-bce: final-layer W init std. Default "
                             "1e-3 keeps prior Stage 2a runs bit-reproducible. "
                             "Try 0.05 to let the head's output vary across "
                             "positions from step 0 (BCE shapes existing "
                             "differentiation instead of building it). Tail "
                             "at bias=-3, std=0.05 → λ_max ≈ 0.25 at 3σ "
                             "outliers, still well below 0.7 hard kill.")
    parser.add_argument("--v6-wt-sparsity-weight", type=float, default=0.0,
                        help="Stage 2a-wts: WikiText-side sparsity loss weight. "
                             "BCE alone supervises V4c-ans vs V4c-bg; the head "
                             "learns ANY separating axis, which can over-fire "
                             "on WikiText (empirically: w_init scan showed "
                             "this is geometric, not optimization-driven). "
                             "This term adds mean-λ-on-WikiText to aux_loss "
                             "in the main forward, with gradient flowing to "
                             "the head via a parallel non-no_grad path. "
                             "Mixture lam stays detached → LM-CE still cannot "
                             "drive the head. Try 0.05 (parity with v6_aux).")
    parser.add_argument("--v6-wt-sparsity-target", type=float, default=0.0,
                        help="Stage 2a-wts: hinge target. 0 = mean-λ form "
                             "(empirically crushes lam_ans because Adam "
                             "normalizes consistent-sign gradient). >0 = "
                             "hinge form `relu(λ-target).mean()`: only "
                             "suppresses positions ABOVE target, leaving a "
                             "free zone in [0,target] where BCE can lift "
                             "lam_ans without WT pressure dragging it back. "
                             "Try 0.05 (parity with v6_lambda_floor).")
    parser.add_argument("--v6-mem-head-bank", action="store_true",
                        help="Stage 2b: enable untied bank vocab head. The "
                             "existing bank readout uses l_stage.mem_head, "
                             "whose .weight is WEIGHT-TIED to tok_emb. B2 "
                             "probe confirmed the resulting bank distribution "
                             "is near-uniform (entropy 9.15 / 10.83). This "
                             "flag adds self.mem_head_bank = nn.Linear(d, V) "
                             "as a SEPARATE trainable param initialized from "
                             "tok_emb (copy, not alias) + l_stage.mem_head."
                             "bias. Bank branch (V4c hook + main mixture) "
                             "uses it when set. Aux-only via V5 allowlist.")
    parser.add_argument("--v6-bank-ce-weight", type=float, default=0.0,
                        help="Stage 2b: weight on direct cross-entropy of "
                             "logits_bank_a at V4c-answer rows vs cur_tok. "
                             "The existing L_bank margin hinge is satisfiable "
                             "through near-uniform logits with cur_lp slightly "
                             "above stale_lp (B2 confirmed). Adding CE anchors "
                             "the bank distribution to put high mass on the "
                             "correct token. Try 1.0 (CE is on natural log "
                             "scale; values <0.5 will be dominated by other "
                             "losses).")
    parser.add_argument("--v6-bank-pair-ce-weight", type=float, default=0.0,
                        help="Stage 2f: weight on pairwise CE over "
                             "[bank_logit(cur), bank_logit(stale)] at V4c "
                             "answer rows. This directly trains current-vs-"
                             "stale role binding; use when full-vocab bank CE "
                             "learns semantic candidates but leaves stale tied "
                             "or above current. Try 2.0.")
    parser.add_argument("--v6-bank-query-source", type=str, default="h_d",
                        choices=("h_d", "h_e", "q_attractor", "q_mem"),
                        help="Stage 2c: source vector passed through "
                             "bank_q_proj for the external V6 bank. h_d is "
                             "the original V6 path; h_e uses the dense "
                             "E-stage state; q_attractor/q_mem use DStage "
                             "internal memory states.")
    parser.add_argument("--v6-bank-readout-mode", type=str, default="bank",
                        choices=("bank", "direct"),
                        help="Stage 2d: V6 external readout mode. bank keeps "
                             "source→bank_q_proj→mem_vals→mem_head_bank. "
                             "direct bypasses the 256-slot bank bottleneck "
                             "and uses source→mem_head_bank.")
    parser.add_argument("--v6-source-adapter", action="store_true",
                        help="Stage 2e: add a small nonlinear residual source "
                             "adapter before the direct V6 readout head. Use "
                             "with B0 specialist mode when direct-source probes "
                             "rank current/stale candidates but cannot separate "
                             "current from stale.")
    parser.add_argument("--v6-context-adapter", action="store_true",
                        help="Stage 2g: add a causal prefix self-attention "
                             "adapter before the direct V6 readout head. Use "
                             "when answer-row source vectors rank plausible "
                             "candidates but fail current-vs-stale role binding.")
    parser.add_argument("--v6-bank-head-lr", type=float, default=0.0,
                        help="Stage 2b-lr: dedicated CONSTANT lr for "
                             "mem_head_bank.* (separate AdamW group, no "
                             "cosine schedule). The 50K-output head needs "
                             "~10⁴+ CE examples to calibrate but V4c aux "
                             "fires sparsely; without a dedicated higher "
                             "lr the bank head doesn't move (Stage 2b "
                             "smoke at main lr=1e-5 confirmed L_bank_ce "
                             "stuck at log(V)=10.83 over 500 steps). "
                             "Default 0.0 = use the main schedule lr. "
                             "Try 3e-4 (~30× the tail lr).")
    parser.add_argument("--v6-specialist-from-b0", action="store_true",
                        help="B0-preserving V6 specialist mode: build CEDL-V6, "
                             "partially load matching weights from "
                             "--resume-checkpoint, freshly initialize only the "
                             "V6/salience specialist keys, and use a fresh "
                             "optimizer. Requires CEDL-V6, direct q_mem readout, "
                             "--v6-lambda-head, --v6-mem-head-bank, and "
                             "--v6-specialist-noinject.")
    parser.add_argument("--v6-specialist-freeze", type=str, default="none",
                        choices=("none", "trunk"),
                        help="Freeze policy for --v6-specialist-from-b0. "
                             "'trunk' freezes the B0-shared LM path and trains "
                             "only v6_lambda_head.* and mem_head_bank.*.")
    parser.add_argument("--v6-specialist-noinject", action="store_true",
                        help="For CEDL-V6 specialist mode, build the V4d shell "
                             "with v4d_noinject=True so random/fresh salience "
                             "injection cannot perturb the B0 trunk.")
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--grad-accum", type=int, default=4)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--preview-only", action="store_true",
                        help="V3m2 only: build the model, load shared "
                             "backbone weights via --preview-shared-checkpoint "
                             "(strict=False), run ONE forward in eval mode on "
                             "a 16-seq WikiText sample, print the decoded "
                             "mask audit + first-5-sequence anchor / label "
                             "inspection, then exit. No training, no save.")
    parser.add_argument("--preview-shared-checkpoint", type=str, default=None,
                        help="Path to an existing checkpoint (V0/V3/V3m) to "
                             "load shared backbone weights from for "
                             "--preview-only mode. SL-specific keys (scalar_"
                             "head, EMA buffers, lexicon buffers) will be "
                             "missing — strict=False is used, and missing/"
                             "unexpected keys are printed for audit.")
    parser.add_argument("--preview-v4c-pairs", action="store_true",
                        help="V4c only: generate 500 synthetic contrastive "
                             "pairs, run the family-aware audit (hard-fail on "
                             "span-validity violations), and print the first "
                             "10 pairs with decoded spans. No model build, "
                             "no training. Use BEFORE training to verify the "
                             "pair generator is well-formed.")
    args = parser.parse_args()

    if args.verify_params:
        verify_all_params()
        return

    if args.preview_only:
        return _run_v3m2_preview(args)

    if args.preview_v4c_pairs:
        return _run_v4c_pair_preview(args)

    cfg = Config(
        dataset=args.dataset,
        tpu=args.tpu,
        max_steps=(args.total_steps if args.total_steps is not None
                   else args.max_steps),
        batch_size=args.batch_size,
        grad_accum=args.grad_accum,
        seed=args.seed,
        resume_checkpoint=args.resume_checkpoint,
        resume_start_step=args.start_step,
        save_trainstate=args.save_trainstate,
        run_until_step=args.run_until_step,
        **({"save_dir": args.save_dir} if args.save_dir else {}),
        **({"eval_interval": args.eval_interval} if args.eval_interval is not None else {}),
        **({"save_interval": args.save_interval} if args.save_interval is not None else {}),
        **({"v5_aux_scale_early": args.v5_aux_scale_early} if args.v5_aux_scale_early is not None else {}),
        **({"v5_aux_scale_mid":   args.v5_aux_scale_mid}   if args.v5_aux_scale_mid   is not None else {}),
        **({"v5_aux_scale_late":  args.v5_aux_scale_late}  if args.v5_aux_scale_late  is not None else {}),
        **({"v5_cadence_early":   args.v5_cadence_early}   if args.v5_cadence_early   is not None else {}),
        **({"v5_cadence_mid":     args.v5_cadence_mid}     if args.v5_cadence_mid     is not None else {}),
        **({"v5_cadence_late":    args.v5_cadence_late}    if args.v5_cadence_late    is not None else {}),
        **({"v5_family_reweight": True} if args.v5_family_reweight else {}),
        **({"v6_specialist_from_b0": True} if args.v6_specialist_from_b0 else {}),
        **({"v6_specialist_freeze": args.v6_specialist_freeze}
           if args.v6_specialist_freeze != "none" else {}),
        **({"v6_specialist_noinject": True} if args.v6_specialist_noinject else {}),
    )

    torch.manual_seed(cfg.seed)
    np.random.seed(cfg.seed)

    torch.set_float32_matmul_precision('medium')
    torch.backends.cudnn.benchmark = True

    if cfg.tpu:
        import torch_xla.core.xla_model as xm
        device = xm.xla_device()
        print(f"Using TPU: {device}")
    elif torch.cuda.is_available():
        device = torch.device("cuda")
        print(f"Using GPU: {torch.cuda.get_device_name()}")
    else:
        device = torch.device("cpu")
        print("Using CPU (will be slow!)")

    train_loader, val_loader, test_loader, tokenizer = load_data(cfg)
    _train_ds = train_loader.dataset
    _val_ds = val_loader.dataset
    _test_ds = test_loader.dataset

    if args.model == "all":
        models_to_run = ALL_MODELS
    elif "," in args.model:
        models_to_run = [t.strip() for t in args.model.split(",") if t.strip()]
    else:
        models_to_run = [args.model]
    print(f"Models to run ({len(models_to_run)}): {models_to_run}")

    results_table = {}

    for tag in models_to_run:
        torch.manual_seed(cfg.seed)
        np.random.seed(cfg.seed)
        if device.type == 'cuda':
            torch.cuda.manual_seed_all(cfg.seed)

        if tag in ("CEDL", "CEDL-V2a", "CEDL-V3", "CEDL-V3m",
                    "CEDL-V3m2", "CEDL-V3m2-NLL",
                    "CEDL-V4c", "CEDL-V4c-G", "CEDL-V4c-cold",
                    "CEDL-V4c-randlabel", "CEDL-V4c-frozen", "CEDL-V4c-M",
                    "CEDL-V4d", "CEDL-V4d-cold", "CEDL-V4d-noinject",
                    "CEDL-V4d-strong", "CEDL-V4d-strong-noinject",
                    "CEDL-V4d-causal", "CEDL-V4d-causal-strong",
                    "CEDL-V4d-causalz", "CEDL-V4d-combo",
                    "CEDL-V4e", "CEDL-V4e-noinject",
                    "CEDL-V5-LM", "CEDL-V6",
                    "CEDL-V3lex", "CEDL-B0") or tag.startswith("CEDL-V4e-k"):
            cfg.batch_size = min(args.batch_size, 4)
        else:
            cfg.batch_size = args.batch_size
        cfg.v5_grad_isolation = bool(args.v5_grad_isolation
                                      or tag.startswith("CEDL-V5-LM")
                                      or tag == "CEDL-V6")
        cfg.v5_trunk_aux_frac = float(args.v5_trunk_aux_frac)
        if args.v6_mixture and tag != "CEDL-V6":
            print(f"  [V6] WARN: --v6-mixture is a no-op on tag {tag!r} "
                  f"(only CEDL-V6 wires the mixture). Ignoring.")
        cfg.v6_mixture = (tag == "CEDL-V6")
        cfg.v6_lambda_init = float(args.v6_lambda_init)
        if args.v6_lambda_a_init != 1.0 and tag != "CEDL-V6":
            print(f"  [V6.1b] WARN: --v6-lambda-a-init={args.v6_lambda_a_init} "
                  f"is a no-op on tag {tag!r} (only CEDL-V6 wires the "
                  f"mixture). Ignoring.")
        cfg.v6_lambda_a_init = float(args.v6_lambda_a_init) if tag == "CEDL-V6" else 1.0
        if args.v6_aux_weight > 0 and tag != "CEDL-V6":
            print(f"  [V6.1] WARN: --v6-aux-weight={args.v6_aux_weight} is a "
                  f"no-op on tag {tag!r} (only CEDL-V6 wires the mixture). "
                  f"Ignoring.")
        cfg.v6_aux_weight    = float(args.v6_aux_weight) if tag == "CEDL-V6" else 0.0
        cfg.v6_margin_target = float(args.v6_margin_target)
        cfg.v6_mix_weight    = float(args.v6_mix_weight)
        cfg.v6_gate_weight   = float(args.v6_gate_weight)
        cfg.v6_lambda_floor  = float(args.v6_lambda_floor)
        if args.v6_lambda_head and tag != "CEDL-V6":
            print(f"  [Stage2a] WARN: --v6-lambda-head is a no-op on tag "
                  f"{tag!r} (only CEDL-V6 wires the mixture). Ignoring.")
        cfg.v6_lambda_head           = bool(args.v6_lambda_head) and tag == "CEDL-V6"
        cfg.v6_lambda_head_hidden    = int(args.v6_lambda_head_hidden)
        cfg.v6_lambda_head_bias_init = float(args.v6_lambda_head_bias_init)
        cfg.v6_bg_weight             = float(args.v6_bg_weight) if tag == "CEDL-V6" else 1.0
        cfg.v6_bg_target             = float(args.v6_bg_target)
        if args.v6_bce_objective and tag != "CEDL-V6":
            print(f"  [Stage2a-bce] WARN: --v6-bce-objective is a no-op on "
                  f"tag {tag!r} (only CEDL-V6 wires the mixture). Ignoring.")
        cfg.v6_bce_objective = bool(args.v6_bce_objective) and tag == "CEDL-V6"
        cfg.v6_sel_weight    = float(args.v6_sel_weight) if tag == "CEDL-V6" else 1.0
        cfg.v6_lambda_head_w_init_std = float(args.v6_lambda_head_w_init_std)
        if args.v6_wt_sparsity_weight > 0 and tag != "CEDL-V6":
            print(f"  [Stage2a-wts] WARN: --v6-wt-sparsity-weight="
                  f"{args.v6_wt_sparsity_weight} is a no-op on tag {tag!r} "
                  f"(only CEDL-V6 wires the mixture). Ignoring.")
        cfg.v6_wt_sparsity_weight = (
            float(args.v6_wt_sparsity_weight) if tag == "CEDL-V6" else 0.0)
        cfg.v6_wt_sparsity_target = float(args.v6_wt_sparsity_target)
        if args.v6_mem_head_bank and tag != "CEDL-V6":
            print(f"  [Stage2b] WARN: --v6-mem-head-bank is a no-op on tag "
                  f"{tag!r} (only CEDL-V6 wires the mixture). Ignoring.")
        cfg.v6_mem_head_bank = bool(args.v6_mem_head_bank) and tag == "CEDL-V6"
        if args.v6_bank_ce_weight > 0 and tag != "CEDL-V6":
            print(f"  [Stage2b] WARN: --v6-bank-ce-weight="
                  f"{args.v6_bank_ce_weight} is a no-op on tag {tag!r} "
                  f"(only CEDL-V6 wires the mixture). Ignoring.")
        cfg.v6_bank_ce_weight = (
            float(args.v6_bank_ce_weight) if tag == "CEDL-V6" else 0.0)
        if args.v6_bank_pair_ce_weight > 0 and tag != "CEDL-V6":
            print(f"  [Stage2f] WARN: --v6-bank-pair-ce-weight="
                  f"{args.v6_bank_pair_ce_weight} is a no-op on tag {tag!r} "
                  f"(only CEDL-V6 wires the mixture). Ignoring.")
        cfg.v6_bank_pair_ce_weight = (
            float(args.v6_bank_pair_ce_weight) if tag == "CEDL-V6" else 0.0)
        if args.v6_bank_query_source != "h_d" and tag != "CEDL-V6":
            print(f"  [Stage2c] WARN: --v6-bank-query-source="
                  f"{args.v6_bank_query_source!r} is a no-op on tag {tag!r} "
                  f"(only CEDL-V6 wires the mixture). Ignoring.")
        cfg.v6_bank_query_source = (
            str(args.v6_bank_query_source) if tag == "CEDL-V6" else "h_d")
        if args.v6_bank_readout_mode != "bank" and tag != "CEDL-V6":
            print(f"  [Stage2d] WARN: --v6-bank-readout-mode="
                  f"{args.v6_bank_readout_mode!r} is a no-op on tag {tag!r} "
                  f"(only CEDL-V6 wires the mixture). Ignoring.")
        cfg.v6_bank_readout_mode = (
            str(args.v6_bank_readout_mode) if tag == "CEDL-V6" else "bank")
        if args.v6_source_adapter and tag != "CEDL-V6":
            print(f"  [Stage2e] WARN: --v6-source-adapter is a no-op on tag "
                  f"{tag!r} (only CEDL-V6 wires the mixture). Ignoring.")
        cfg.v6_source_adapter = bool(args.v6_source_adapter) and tag == "CEDL-V6"
        if args.v6_context_adapter and tag != "CEDL-V6":
            print(f"  [Stage2g] WARN: --v6-context-adapter is a no-op on tag "
                  f"{tag!r} (only CEDL-V6 wires the mixture). Ignoring.")
        cfg.v6_context_adapter = (
            bool(args.v6_context_adapter) and tag == "CEDL-V6")
        if args.v6_bank_head_lr > 0 and tag != "CEDL-V6":
            print(f"  [Stage2b-lr] WARN: --v6-bank-head-lr="
                  f"{args.v6_bank_head_lr} is a no-op on tag {tag!r} "
                  f"(only CEDL-V6 wires the mixture). Ignoring.")
        cfg.v6_bank_head_lr = (
            float(args.v6_bank_head_lr) if tag == "CEDL-V6" else 0.0)
        if args.v6_specialist_from_b0 and tag != "CEDL-V6":
            raise RuntimeError(
                "--v6-specialist-from-b0 is only valid with --model CEDL-V6")
        cfg.v6_specialist_from_b0 = (
            bool(args.v6_specialist_from_b0) and tag == "CEDL-V6")
        cfg.v6_specialist_freeze = (
            str(args.v6_specialist_freeze)
            if cfg.v6_specialist_from_b0 else "none")
        cfg.v6_specialist_noinject = (
            bool(args.v6_specialist_noinject) and tag == "CEDL-V6")
        if cfg.v6_specialist_from_b0:
            if not cfg.resume_checkpoint:
                raise RuntimeError(
                    "--v6-specialist-from-b0 requires --resume-checkpoint "
                    "pointing at the B0 source checkpoint.")
            if cfg.v6_specialist_freeze != "trunk":
                raise RuntimeError(
                    "--v6-specialist-from-b0 requires "
                    "--v6-specialist-freeze trunk.")
            if not cfg.v6_specialist_noinject:
                raise RuntimeError(
                    "--v6-specialist-from-b0 requires "
                    "--v6-specialist-noinject.")
            if not cfg.v6_lambda_head:
                raise RuntimeError(
                    "--v6-specialist-from-b0 requires --v6-lambda-head.")
            if not cfg.v6_mem_head_bank:
                raise RuntimeError(
                    "--v6-specialist-from-b0 requires --v6-mem-head-bank.")
            if cfg.v6_bank_query_source != "q_mem":
                raise RuntimeError(
                    "--v6-specialist-from-b0 requires "
                    "--v6-bank-query-source q_mem.")
            if cfg.v6_bank_readout_mode != "direct":
                raise RuntimeError(
                    "--v6-specialist-from-b0 requires "
                    "--v6-bank-readout-mode direct.")
            print("  [V6 specialist] B0 bootstrap mode ON  "
                  "freeze=trunk  source=q_mem  readout=direct  noinject=True")
        train_loader, val_loader, test_loader = make_loaders(
            _train_ds, _val_ds, _test_ds, cfg.batch_size, cfg.tpu,
            shuffle_seed=cfg.seed)

        print(f"\n{'='*60}")
        print(f"  Model: {tag} (batch_size={cfg.batch_size})")
        print(f"{'='*60}")

        model = build_model(
            tag, cfg.vocab_size, cfg.max_seq,
            v6_mixture=bool(getattr(cfg, "v6_mixture", False)),
            v6_lambda_init=float(getattr(cfg, "v6_lambda_init", -4.0)),
            v6_lambda_a_init=float(getattr(cfg, "v6_lambda_a_init", 1.0)),
            v6_aux_weight=float(getattr(cfg, "v6_aux_weight", 0.0)),
            v6_margin_target=float(getattr(cfg, "v6_margin_target", 1.0)),
            v6_mix_weight=float(getattr(cfg, "v6_mix_weight", 0.5)),
            v6_gate_weight=float(getattr(cfg, "v6_gate_weight", 0.01)),
            v6_lambda_floor=float(getattr(cfg, "v6_lambda_floor", 0.05)),
            v6_lambda_head=bool(getattr(cfg, "v6_lambda_head", False)),
            v6_lambda_head_hidden=int(getattr(cfg, "v6_lambda_head_hidden", 160)),
            v6_lambda_head_bias_init=float(getattr(cfg, "v6_lambda_head_bias_init", -7.0)),
            v6_bg_weight=float(getattr(cfg, "v6_bg_weight", 1.0)),
            v6_bg_target=float(getattr(cfg, "v6_bg_target", 0.01)),
            v6_bce_objective=bool(getattr(cfg, "v6_bce_objective", False)),
            v6_sel_weight=float(getattr(cfg, "v6_sel_weight", 1.0)),
            v6_lambda_head_w_init_std=float(getattr(cfg, "v6_lambda_head_w_init_std", 1e-3)),
            v6_wt_sparsity_weight=float(getattr(cfg, "v6_wt_sparsity_weight", 0.0)),
            v6_wt_sparsity_target=float(getattr(cfg, "v6_wt_sparsity_target", 0.0)),
            v6_mem_head_bank=bool(getattr(cfg, "v6_mem_head_bank", False)),
            v6_bank_ce_weight=float(getattr(cfg, "v6_bank_ce_weight", 0.0)),
            v6_bank_pair_ce_weight=float(getattr(
                cfg, "v6_bank_pair_ce_weight", 0.0)),
            v6_bank_query_source=str(getattr(cfg, "v6_bank_query_source", "h_d")),
            v6_bank_readout_mode=str(getattr(cfg, "v6_bank_readout_mode", "bank")),
            v6_source_adapter=bool(getattr(cfg, "v6_source_adapter", False)),
            v6_context_adapter=bool(getattr(cfg, "v6_context_adapter", False)),
            v6_specialist_noinject=bool(getattr(cfg, "v6_specialist_noinject", False)),
        ).to(device)
        n_params = count_params(model)
        print(f"  Parameters: {n_params:,} ({n_params/1e6:.1f}M)")

        if device.type == 'cuda' and not args.eval_only and tag not in (
                "Transformer-XL", "CEDL", "CEDL-V2a", "CEDL-V3", "CEDL-V3m",
                "CEDL-V3m2", "CEDL-V3m2-NLL",
                "CEDL-V4c", "CEDL-V4c-G", "CEDL-V4c-cold",
                "CEDL-V4c-randlabel", "CEDL-V4c-frozen", "CEDL-V4c-M",
                "CEDL-V4d", "CEDL-V4d-cold", "CEDL-V4d-noinject",
                "CEDL-V4d-strong", "CEDL-V4d-strong-noinject",
                "CEDL-V4d-causal", "CEDL-V4d-causal-strong",
                "CEDL-V4d-causalz", "CEDL-V4d-combo",
                "CEDL-V4e", "CEDL-V4e-noinject",
                "CEDL-V5-LM", "CEDL-V6",
                "CEDL-V3lex", "CEDL-B0", "Mamba") \
                or tag.startswith("CEDL-V4e-k"):
            try:
                model = torch.compile(model)
                print("  torch.compile: enabled")
            except Exception as e:
                print(f"  torch.compile: skipped ({e})")

        if args.eval_only:
            ckpt = args.checkpoint or os.path.join(cfg.save_dir, f"{tag}_best.pt")
            if os.path.exists(ckpt):
                state = torch.load(ckpt, map_location='cpu', weights_only=True)
                if isinstance(state, dict) and "model" in state and isinstance(state["model"], dict):
                    wrapper_keys = sorted(k for k in state.keys() if k != "model")
                    wrapper_msg = f" (wrapper keys={wrapper_keys})" if wrapper_keys else ""
                    print(f"  [eval] unwrapping checkpoint['model'] -> model weights{wrapper_msg}")
                    state = state["model"]
                if any(k.startswith('_orig_mod.') for k in state):
                    state = {k.replace('_orig_mod.', ''): v for k, v in state.items()}
                model.load_state_dict(state)
                print(f"  Loaded checkpoint: {ckpt}")
            else:
                print(f"  WARNING: no checkpoint found at {ckpt}, skipping {tag}")
                del model
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                continue
        else:
            best_val_ppl = train(model, tag, cfg, device, train_loader, val_loader)
            best_path = os.path.join(cfg.save_dir, f"{tag}_best.pt")
            if os.path.exists(best_path):
                model.load_state_dict(torch.load(best_path, map_location='cpu', weights_only=True))

        raw_m = model._orig_mod if hasattr(model, '_orig_mod') else model
        if isinstance(raw_m, CEDLTwoLoop100M):
            raw_m.feedback_alpha.fill_(1.0)

        test_ppl = evaluate(model, test_loader, device, max_steps=500, tpu=cfg.tpu)
        print(f"\n  Test Perplexity: {test_ppl:.2f}")

        structured_reason = "skipped for 100M path; use PPL/downstream metrics"
        print(f"  Structured benchmark: {structured_reason}.")
        struct_results = {}

        downstream = run_downstream_eval(model, tag, device)

        results_table[tag] = {
            "params": n_params,
            "test_ppl": test_ppl,
            "structured": struct_results,
            "structured_reason": structured_reason,
            "downstream": downstream,
        }

        del model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    print_final_results(results_table)


if __name__ == "__main__":
    main()