# -*- coding: utf-8 -*- """31_teste_v21.py — TESTE REAL v21 (KHMAMBAJEPA, doc 33 §10). Requisitos do usuário atendidos aqui: • "treinar com datasets reais verificando progresso" — corpus REAL (cache_dados/corpus_v16.jsonl — 1015 registros coletados na v19 de LeRobot/GSM8K/MATH-500/hendrycks/OpenR1/IFM; coleta fresca opcional); • A/B PAREADO EXATO (Teo 31.7): mesma semente, re-seed APÓS construção, mesmos lotes — UMA variável por braço; • valida as 3 melhorias v21 ISOLADAS: normalização espectral do NS, hot-paths Cython (equivalência Δ=0 + wallclock), tabela v17 ativa. Braços (UMA variável por processo): A) baseline_v20 — NS Frobenius (patch v20) + hot-paths OFF + v17 OFF B) v21_ns — NS espectral (v21) + hot-paths OFF C) v21_full — NS espectral + hot-paths Cython ON (equivalência Δ=0 esperada no TREINO; wallclock comparado) D) v17_ativo — v21_full + treino.autorreg_v17.ativo=true Uso (RAM 4 GiB / processos morrem entre chamadas — retomada por partes): python3 scripts/31_teste_v21.py --variante baseline_v20 --parte 12 python3 scripts/31_teste_v21.py --variante v21_ns --parte 12 python3 scripts/31_teste_v21.py --variante v21_full --parte 12 python3 scripts/31_teste_v21.py --variante v17_ativo --parte 12 python3 scripts/31_teste_v21.py --relatorio """ from __future__ import annotations import argparse import base64 import json import os import random import sys import time RAIZ = os.environ.get("KHTST_RAIZ", os.path.dirname( os.path.dirname(os.path.abspath(__file__)))) sys.path.insert(0, os.path.join(RAIZ, "src")) sys.path.insert(0, os.path.join(RAIZ, "src", "khtst", "acelerado")) CORPUS = os.path.join(RAIZ, "cache_dados", "corpus_v16.jsonl") TOKENIZADOR = os.path.join(RAIZ, "cache_dados", "tokenizador_v16.json") TOK_ANTIGO = os.path.join(RAIZ, "cache_dados", "tokenizador.json") TELEMETRIA = os.path.join(RAIZ, "telemetria_out") SEED = 2026 PASSOS_TOTAIS = 48 def carregar_registros(caminho: str) -> list[dict]: registros = [] if not os.path.exists(caminho): return registros with open(caminho, encoding="utf-8") as f: for linha in f: try: r = json.loads(linha) except Exception: continue if isinstance(r.get("imagem"), str): try: r["imagem"] = base64.b64decode(r["imagem"]) except Exception: r["imagem"] = None if isinstance(r.get("audio"), str): try: r["audio"] = base64.b64decode(r["audio"]) except Exception: r["audio"] = None registros.append(r) return registros def est_padrao(seed: int) -> dict: import numpy as np import torch torch.manual_seed(seed) random.seed(seed) np.random.seed(seed) return {"torch": torch.get_rng_state(), "random": random.getstate(), "numpy": np.random.get_state()} def _restaurar_rng(est: dict) -> None: import numpy as np import torch torch.set_rng_state(est["torch"]) random.setstate(est["random"]) np.random.set_state(est["numpy"]) # v20 exato — braço A (replicação VERBATIM da normalização Frobenius; # a única diferença da v21 é o bloco de normalização — doc 33 §3) def _ns_frob_v20(G, iters=5, limiar_f32=1e-3, teto_espectral=None): import torch as _t a, b, c = 3.4445, -4.7750, 2.0315 X = G.float() if float(G.float().norm()) < limiar_f32 else G.bfloat16() if G.size(-2) > G.size(-1): X = X.mT X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) for _ in range(iters): A = X @ X.mT B = b * A + c * (A @ A) X = a * X + B @ X if G.size(-2) > G.size(-1): X = X.mT return X def rodar_variante(nome: str, parte: int) -> dict: import gc import numpy as np import torch if not (torch.cuda.is_available() or torch.xpu.is_available()): torch.set_num_threads(int(os.environ.get("OMP_NUM_THREADS", "2"))) import khtst.treino.otimizadores as O from khtst.config import Config from khtst.dados.tokenizador import TokenizadorKHTST from khtst.memoria.orquestrador import OrquestradorSOM from khtst.nucleo.modelo import KHTSTModel from khtst.telemetria.hub import TelemetryHub from khtst.treino.treinador import TreinadorExtensao from khtst.mamba3 import kernel as k3 from khtst.percepcao import moe as moe_mod cfg = Config() cfg.dados.semente = SEED cfg.treino.warmup = 8 cfg.treino.lote = 4 cfg.treino.lote_efetivo = 1 cfg.treino.qat_ultimos_passos = 0 cfg.treino.ensemble["a_cada"] = 24 cfg.modelo.escalacao["ativo"] = False # ---- A VARIÁVEL DO BRAÇO ------------------------------------------------ if nome == "baseline_v20": O.newton_schulz5 = _ns_frob_v20 # NS Frobenius (v20 exato) k3.set_conv_cython(False) moe_mod.set_topk_cython(False) cfg.treino.autorreg_v17 = {"ativo": False} elif nome == "v21_ns": pass # NS v21 default; resto OFF elif nome == "v21_full": pass # defaults v21 (NS + Cython) elif nome == "v17_ativo": cfg.treino.autorreg_v17 = {"ativo": True, "a_cada": 5} else: raise ValueError(nome) os.makedirs(TELEMETRIA, exist_ok=True) hub = TelemetryHub(os.path.join(TELEMETRIA, f"treino_v21_{nome}.jsonl")) tk = TokenizadorKHTST(TOKENIZADOR if os.path.exists(TOKENIZADOR) else TOK_ANTIGO) registros = carregar_registros(CORPUS) modelo = KHTSTModel(cfg, usar_multimodal=True) orquestrador = OrquestradorSOM(cfg.modelo.d_modelo, cfg.som, hub=hub, cfg_atencao=None) treinar = TreinadorExtensao(cfg, modelo, tk, hub, registros, orquestrador_som=orquestrador) # PAREAMENTO EXATO (Teo 31.7): re-seed APÓS a construção _restaurar_rng(est_padrao(SEED)) ckpt = os.path.join(TELEMETRIA, f"ckpt_v21_{nome}.pt") perdas: list[float] = [] passo_inicial = 0 t_acum = 0.0 if parte and os.path.exists(ckpt): est = torch.load(ckpt, map_location="cpu", weights_only=False) modelo.load_state_dict(est["modelo"]) treinar.otimizador.load_state_dict(est["otimizador"]) treinar.passo_global = est["passo_global"] perdas = est["perdas"] t_acum = est.get("t_acum", 0.0) torch.set_rng_state(est["rng_torch"]) random.setstate(est["rng_py"]) np.random.set_state(est["rng_np"]) passo_inicial = est["passo_global"] print(f"[retomada] passo={passo_inicial} perdas={len(perdas)}", flush=True) alvo = min(PASSOS_TOTAIS, passo_inicial + parte) t0 = time.time() while treinar.passo_global < alvo: perda = treinar._passo() if perda == perda: perdas.append(perda) if treinar.passo_global % 12 == 0: print(f"[{nome} passo {treinar.passo_global} perda {perda:.3f}]", flush=True) t_acum += time.time() - t0 resultado = {"nome": nome, "passos": treinar.passo_global, "perdas": perdas, "t_acum_s": round(t_acum, 2), "autorreg_v17": treinar.autorreg_v17 is not None} if parte and treinar.passo_global < PASSOS_TOTAIS: torch.save({"modelo": modelo.state_dict(), "otimizador": treinar.otimizador.state_dict(), "passo_global": treinar.passo_global, "perdas": perdas, "t_acum": t_acum, "rng_torch": torch.get_rng_state(), "rng_py": random.getstate(), "rng_np": np.random.get_state()}, ckpt) elif os.path.exists(ckpt): os.remove(ckpt) # braço completo — ckpt fora gc.collect() return resultado def relatorio() -> None: res = {} for nome in ("baseline_v20", "v21_ns", "v21_full", "v17_ativo"): cam = os.path.join(TELEMETRIA, f"resultado_v21_{nome}.json") if os.path.exists(cam): res[nome] = json.load(open(cam, encoding="utf-8")) if len(res) < 2: print("resultados insuficientes — rode os 4 braços primeiro") return print("=" * 74) print(f"{'braço':14s} {'perda inicial':>13s} {'perda final':>11s} " f"{'ganho ×':>8s} {'ΔCE pareado':>12s} {'s/ passo':>9s}") base = res.get("baseline_v20") linhas = {} for nome, r in res.items(): p = r["perdas"] ini = sum(p[:4]) / 4 fim = sum(p[-4:]) / 4 ganho = ini / max(fim, 1e-9) dce = (sum(p[-4:]) - sum(base["perdas"][-4:])) / 4 \ if base and nome != "baseline_v20" else 0.0 spp = r["t_acum_s"] / max(1, len(p)) print(f"{nome:14s} {ini:13.3f} {fim:11.3f} {ganho:8.3f} " f"{dce:+12.4f} {spp:9.2f}") linhas[nome] = {"ini": round(ini, 4), "fim": round(fim, 4), "ganho": round(ganho, 4), "dce": round(dce, 4), "s_por_passo": round(spp, 2)} print("=" * 74) # progresso por quarto (verificação de progresso contínuo) for nome, r in res.items(): p = r["perdas"] q = max(1, len(p) // 4) quarts = [round(sum(p[i * q:(i + 1) * q]) / q, 3) for i in range(4) if p[i * q:(i + 1) * q]] print(f"{nome}: {' → '.join(map(str, quarts))}") out = os.path.join(RAIZ, "resultados_v21.json") json.dump({"braços": linhas}, open(out, "w", encoding="utf-8"), ensure_ascii=False, indent=1) print(f"salvo em {out}") def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--variante", choices=["baseline_v20", "v21_ns", "v21_full", "v17_ativo"]) ap.add_argument("--parte", type=int, default=0, help="passos NESTA chamada (retomada por ckpt)") ap.add_argument("--relatorio", action="store_true") args = ap.parse_args() if args.relatorio: relatorio() return r = rodar_variante(args.variante, args.parte) cam = os.path.join(TELEMETRIA, f"resultado_v21_{args.variante}.json") if args.parte == 0 or r["passos"] >= PASSOS_TOTAIS: json.dump(r, open(cam, "w", encoding="utf-8"), ensure_ascii=False) print(f"[{args.variante}] passos={r['passos']} " f"perda={r['perdas'][-1] if r['perdas'] else None}") if __name__ == "__main__": main()