# 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)