tests/test_cte_core.py

# tests/test_cte_core.py
"""Tests du noyau CTE natif (torch optionnel : SKIP sans torch).

Fidélité au doc de référence docs/heritage/cte-architecture.md (T1, source de vérité) :
- layout des clés EXACT : 8 top-level + n_layers×27 par bloc (440 sur le réel) ;
- chemin tick_chunk EXCLUSIF ; états porteurs batchés (thought_state, attn_S/z) ;
- recette d'inférence : reset complet → UN chunk prompt → boucle chunks de 1,
  greedy = argmax(logits[0, -1]).
"""
from pathlib import Path

import pytest

torch = pytest.importorskip("torch", reason="noyau torch optionnel (substrat sans torch)")

from fractus_vorax.model.cte_core import CteCore, CteCoreConfig

_REPO = Path(__file__).parent.parent
# Le cerveau de naissance d'abord (canonique, tâche 2 du plan du relais) ;
# repli sur l'héritage vorax (checkpoints/checkpoints/) pour compat.
REAL_CKPT_CANDIDATES = (
    _REPO / "brain" / "FRACTUS_BIRTH.pt",
    _REPO / "checkpoints" / "checkpoints" / "FRACTUS_1B_PHASE2_FROZEN_MERGED.pt",
)
REAL_CKPT = next((p for p in REAL_CKPT_CANDIDATES if p.exists()), REAL_CKPT_CANDIDATES[0])


def _tiny(**overrides) -> CteCoreConfig:
    return CteCoreConfig(vocab_size=128, **overrides)


def test_forward_logits_shape():
    torch.manual_seed(0)
    core = CteCore(_tiny())
    ids = torch.randint(0, 128, (1, 8))
    logits = core(ids)
    assert logits.shape == (1, 8, 128)
    assert torch.isfinite(logits).all()


def test_forward_supports_batched_states():
    """Les états porteurs sont batchés : forward (B, L) avec B ≤ config.batch.

    Un forward plus petit tranche les états à B (ils deviennent B-sized).
    """
    torch.manual_seed(0)
    core = CteCore(_tiny(batch=2))
    assert core.thought_state.shape == (2, 1, 64)
    logits = core(torch.randint(0, 128, (2, 8)))
    assert logits.shape == (2, 8, 128)
    assert core.thought_state.shape == (2, 1, 64)
    logits1 = core(torch.randint(0, 128, (1, 5)))
    assert logits1.shape == (1, 5, 128)
    assert core.thought_state.shape == (1, 1, 64)


def test_forward_rejects_batch_above_states():
    core = CteCore(_tiny())  # batch=1
    with pytest.raises(ValueError):
        core(torch.zeros(2, 4, dtype=torch.long))


def test_generate_greedy_deterministic():
    torch.manual_seed(0)
    core = CteCore(_tiny())
    ids = torch.tensor([[1, 2, 3]])
    a = core.generate(ids, max_new_tokens=8)
    b = core.generate(ids, max_new_tokens=8)
    assert torch.equal(a, b)
    assert a.shape == (1, 11)
    # greedy = argmax des logits du prompt chunk (recette : logits[0, -1])
    core.reset_states(1)
    with torch.no_grad():
        prompt_logits = core(ids)
    assert int(a[0, 3].item()) == int(prompt_logits[0, -1].argmax().item())
    # température > 0 : même seed → même trajectoire (Generator seedé)
    c = core.generate(ids, max_new_tokens=8, temperature=1.0, seed=7)
    d = core.generate(ids, max_new_tokens=8, temperature=1.0, seed=7)
    assert torch.equal(c, d)
    # greedy : indépendant du seed (argmax, aucun aléa)
    e = core.generate(ids, max_new_tokens=8, seed=123)
    assert torch.equal(a, e)


def test_generate_stops_at_eos():
    """Append-then-break (sémantique ratifiée P4) : l'EOS clôt, rien après lui."""
    torch.manual_seed(0)
    core = CteCore(_tiny())
    ids = torch.tensor([[1, 2, 3]])
    core.reset_states(1)
    with torch.no_grad():
        first = int(core(ids)[0, -1].argmax().item())
    out = core.generate(ids, max_new_tokens=8, eos_token_id=first)
    assert out.shape == (1, 4)  # prompt + 1 token, puis arrêt
    assert int(out[0, -1].item()) == first


def test_states_update_and_reset():
    torch.manual_seed(0)
    core = CteCore(_tiny())
    blk = core.blocks[0]
    assert torch.count_nonzero(blk.attn_S) == 0
    ids = torch.randint(0, 128, (1, 8))
    with torch.no_grad():
        l1 = core(ids).clone()
    # les états porteurs ont bougé
    assert torch.count_nonzero(blk.attn_S) > 0
    assert torch.count_nonzero(blk.attn_z) > 0
    assert torch.count_nonzero(core.thought_state) > 0
    # kuramoto_phases : write-only, shape du dernier forward (B, C, N)
    assert blk.kuramoto_phases.shape == (1, 8, core.config.n_oscillators)
    # carry : un second forward sur le MÊME prompt donne d'autres logits
    with torch.no_grad():
        l2 = core(ids).clone()
    assert not torch.equal(l1, l2)
    # reset → zéros + replay déterministe
    core.reset_states()
    assert torch.count_nonzero(blk.attn_S) == 0
    assert torch.count_nonzero(blk.attn_z) == 0
    assert torch.count_nonzero(core.thought_state) == 0
    with torch.no_grad():
        l3 = core(ids).clone()
    assert torch.equal(l1, l3)


def test_state_dict_key_layout():
    """8 clés top + n_layers×27 par bloc — le motif EXACT des 440 clés réelles
    (doc §Table des clés : 8 + 16×27 = 440)."""
    cfg = _tiny()
    core = CteCore(cfg)
    keys = set(core.state_dict().keys())

    expected = {
        "observe.weight",
        "output_head.weight",
        "confidence_head.weight", "confidence_head.bias",
        "salience_head.weight", "salience_head.bias",
        "last_lb_loss",
        "thought_state",
    }
    for i in range(cfg.n_layers):
        expected |= {
            f"blocks.{i}.attn_S", f"blocks.{i}.attn_z", f"blocks.{i}.kuramoto_phases",
            f"blocks.{i}.attn.w_qkv", f"blocks.{i}.attn.b_qkv",
            f"blocks.{i}.attn.w_out", f"blocks.{i}.attn.b_out",
            f"blocks.{i}.attn.level_logits", f"blocks.{i}.attn.level_offsets",
            f"blocks.{i}.norm_attn.weight", f"blocks.{i}.norm_attn.bias",
            f"blocks.{i}.kuramoto.omega", f"blocks.{i}.kuramoto.coupling_u",
            f"blocks.{i}.kuramoto.coupling_lambda",
            f"blocks.{i}.norm_kur.weight", f"blocks.{i}.norm_kur.bias",
            f"blocks.{i}.moe.U1", f"blocks.{i}.moe.V1",
            f"blocks.{i}.moe.U2", f"blocks.{i}.moe.V2",
            f"blocks.{i}.moe.scale1", f"blocks.{i}.moe.scale2",
            f"blocks.{i}.moe.b1", f"blocks.{i}.moe.b2",
            f"blocks.{i}.moe.expert_phases",
            f"blocks.{i}.norm_moe.weight", f"blocks.{i}.norm_moe.bias",
        }

    assert keys == expected, (
        f"manquantes: {sorted(expected - keys)[:5]} / "
        f"inattendues: {sorted(keys - expected)[:5]}"
    )
    assert len(keys) == 8 + cfg.n_layers * 27
    # tête liée (doc : les deux clés existent, poids bit-identiques)
    assert core.output_head.weight is core.observe.weight
    # shapes des états porteurs (batch = config.batch, doc §États porteurs)
    assert core.thought_state.shape == (cfg.batch, 1, cfg.d_model)
    assert core.blocks[0].attn_S.shape == (cfg.batch, cfg.d_model, cfg.d_model)
    assert core.blocks[0].attn_z.shape == (cfg.batch, cfg.d_model)
    assert core.blocks[0].kuramoto_phases.shape == (
        cfg.batch, cfg.train_seq_len, cfg.n_oscillators
    )
    assert core.last_lb_loss.shape == ()


def test_checkpoint_roundtrip(tmp_path):
    """Sauvegarde {"model_state": sd} → from_checkpoint strict → même génération."""
    torch.manual_seed(0)
    cfg = _tiny()
    core = CteCore(cfg)
    torch.save({"model_state": core.state_dict()}, tmp_path / "tiny_cte.pt")
    loaded = CteCore.from_checkpoint(tmp_path / "tiny_cte.pt")
    assert loaded.config.d_model == cfg.d_model
    # le split n_heads×d_head n'est PAS encodé dans les shapes (w_qkv cubique) —
    # l'inférence par défaut doit retomber sur le split tiny 4×16.
    assert loaded.config.n_heads == cfg.n_heads
    assert loaded.config.d_head == cfg.d_head
    # Identité bit à bit PAR CLÉ (revue T2) : le chemin mmap+assign ne doit
    # altérer AUCUN octet — y compris les clés sans effet forward que la
    # comparaison de génération ne couvrait pas (confidence/salience_head,
    # buffers porteurs). AVANT generate : la recette mut les états porteurs.
    for k, v in core.state_dict().items():
        assert torch.equal(loaded.state_dict()[k], v), k
    assert torch.equal(
        loaded.generate(torch.tensor([[1, 2, 3]]), 4),
        core.generate(torch.tensor([[1, 2, 3]]), 4),
    )


@pytest.mark.slow
def test_real_checkpoint_strict_load():
    """Le VRAI 4.66 GB : strict 440/440, UN forward B=batch, logits finis.

    Lit brain/FRACTUS_BIRTH.pt (le cerveau de naissance, canonique), sinon
    l'héritage checkpoints/checkpoints/. Skip si aucun des deux n'est présent.
    ~0.5 GB via mmap (fallback: copie complète) — venv torch uniquement.
    """
    if not REAL_CKPT.exists():
        pytest.skip("checkpoint CTE 1B FROZEN_MERGED non présent")
    core = CteCore.from_checkpoint(REAL_CKPT)
    assert len(core.state_dict()) == 440
    assert core.config.d_model == 1280
    assert core.config.n_heads == 20 and core.config.d_head == 64
    assert core.config.n_layers == 16
    assert core.config.n_experts == 128
    assert core.config.vocab_size == 50257
    assert core.config.batch == 2
    ids = torch.zeros(2, 4, dtype=torch.long)
    with torch.no_grad():
        logits = core(ids)
    assert logits.shape == (2, 4, 50257)
    assert torch.isfinite(logits).all()
    # Le chemin generate du 1B, lui aussi (revue T2 : il n'était pas
    # exercé — un bug dans la recette d'inférence batch 1 passerait inaperçu).
    # Recette complète : reset_states(1), UN chunk prompt, boucle chunks de 1.
    # CPU ~10-30 s — couvert par le marqueur slow.
    with torch.no_grad():
        gen_ids = torch.zeros(1, 2, dtype=torch.long)
        out = core.generate(gen_ids, 2)
    assert out.shape == (1, 4)