topoexploit 0.1.0__py3-none-any.whl

topogpt3/vanilla_control.py

@@ -0,0 +1,240 @@
+#!/usr/bin/env python3
+"""Vanilla transformer control for TopoGPT3 Hodge-CM ablation.
+Matched size: D=384, L=8, H=8, SwiGLU FFN, RoPE, RMSNorm, tied embeddings.
+Target: ~33M to compare against TopoGPT3 37.4M (which includes TorusBrain+MoE+spectral).
+Same vocab (50257), full curriculum interface (codealpaca, code_feedback,
+magicoder_evol, tiny_the_stack) with train/val/holdout disjuntos — nada
+recortado. Usa tokens reales del tier cuando existen; sintetico solo con
+--allow-synthetic explícito.
+"""
+import argparse
+import math
+from pathlib import Path
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+
+
+TIER_NAMES = ("codealpaca", "code_feedback", "magicoder_evol", "tiny_the_stack")
+
+
+def resolve_device(req: str) -> torch.device:
+    if req == "cuda":
+        if not torch.cuda.is_available():
+            print("WARNING: --device cuda pedido pero torch.cuda.is_available()=False. "
+                  "RTX 2060 detectada pero este torch +cu130 exige driver >=580 "
+                  "(tienes 550/CUDA 12.4). Reinstala cu124: pip install torch "
+                  "--index-url https://download.pytorch.org/whl/cu124. Fallo a cpu.")
+            return torch.device("cpu")
+        return torch.device("cuda")
+    if req == "cpu":
+        return torch.device("cpu")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    print("NOTA: cuda no disponible para este build torch. Uso cpu. "
+          "Para tu RTX 2060: pip install torch --index-url "
+          "https://download.pytorch.org/whl/cu124")
+    return torch.device("cpu")
+
+
+def find_tier_tokens(tier: str):
+    """Busca train real tokenizado del tier. Devuelve lista de ints o None.
+    Nada silencioso: si no hay, el llamador decide (error o sintetico explícito)."""
+    repo = Path(__file__).resolve().parents[1]
+    cands = [
+        repo / "data/tokenized" / tier / "train.bin",
+        repo / "data/tokenized" / tier / "train.npy",
+        repo / "data" / tier / "train.bin",
+        repo / "data" / tier / "train.npy",
+    ]
+    for p in cands:
+        if p.exists():
+            import numpy as np
+            if p.suffix == ".npy":
+                return np.load(str(p)).astype("int64").tolist()
+            return np.fromfile(str(p), dtype=np.uint16).astype("int64").tolist()
+    return None
+
+
+@dataclass
+class VanillaConfig:
+    VOCAB_SIZE: int = 50257
+    D_MODEL: int = 384
+    N_HEADS: int = 8
+    N_LAYERS: int = 8
+    MAX_SEQ_LEN: int = 512
+    DROPOUT: float = 0.1
+    LR: float = 3e-4
+    WEIGHT_DECAY: float = 0.1
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d))
+
+    def forward(self, x):
+        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
+        return x / rms * self.weight
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, d_head: int, max_seq_len: int = 2048, base: int = 10000):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, d_head, 2).float() / d_head))
+        self.register_buffer("inv_freq", inv_freq)
+        self._build_cache(max_seq_len)
+
+    def _build_cache(self, seq_len: int):
+        t = torch.arange(seq_len, device=self.inv_freq.device).float()
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        self.register_buffer("_cos", emb.cos(), persistent=False)
+        self.register_buffer("_sin", emb.sin(), persistent=False)
+
+    @staticmethod
+    def _rot_half(x):
+        x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+        return torch.cat([-x2, x1], dim=-1)
+
+    def forward(self, q, k):
+        sq, sk = q.shape[2], k.shape[2]
+        if max(sq, sk) > self._cos.shape[0]:
+            self._build_cache(max(sq, sk) * 2)
+        cq = self._cos[:sq].unsqueeze(0).unsqueeze(0)
+        sq_ = self._sin[:sq].unsqueeze(0).unsqueeze(0)
+        ck = self._cos[:sk].unsqueeze(0).unsqueeze(0)
+        sk_ = self._sin[:sk].unsqueeze(0).unsqueeze(0)
+        return q * cq + self._rot_half(q) * sq_, k * ck + self._rot_half(k) * sk_
+
+
+class CausalAttention(nn.Module):
+    def __init__(self, cfg: VanillaConfig):
+        super().__init__()
+        assert cfg.D_MODEL % cfg.N_HEADS == 0
+        self.n_heads = cfg.N_HEADS
+        self.d_head = cfg.D_MODEL // cfg.N_HEADS
+        self.qkv = nn.Linear(cfg.D_MODEL, 3 * cfg.D_MODEL, bias=False)
+        self.out = nn.Linear(cfg.D_MODEL, cfg.D_MODEL, bias=False)
+        self.rope = RotaryEmbedding(self.d_head, cfg.MAX_SEQ_LEN)
+        self.drop = cfg.DROPOUT
+
+    def forward(self, x):
+        B, S, D = x.shape
+        qkv = self.qkv(x).view(B, S, 3, self.n_heads, self.d_head)
+        q, k, v = qkv[:, :, 0].transpose(1, 2), qkv[:, :, 1].transpose(1, 2), qkv[:, :, 2].transpose(1, 2)
+        q, k = self.rope(q, k)
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True,
+                                           dropout_p=self.drop if self.training else 0.0)
+        return self.out(y.transpose(1, 2).contiguous().view(B, S, D))
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, d: int):
+        super().__init__()
+        inner = (int(d * 8 / 3) + 7) // 8 * 8
+        self.gate = nn.Linear(d, inner, bias=False)
+        self.up = nn.Linear(d, inner, bias=False)
+        self.down = nn.Linear(inner, d, bias=False)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class VanillaLayer(nn.Module):
+    def __init__(self, cfg: VanillaConfig):
+        super().__init__()
+        self.n1 = RMSNorm(cfg.D_MODEL)
+        self.n2 = RMSNorm(cfg.D_MODEL)
+        self.attn = CausalAttention(cfg)
+        self.ffn = SwiGLU(cfg.D_MODEL)
+        self.drop = nn.Dropout(cfg.DROPOUT)
+
+    def forward(self, x):
+        x = x + self.drop(self.attn(self.n1(x)))
+        x = x + self.drop(self.ffn(self.n2(x)))
+        return x
+
+
+class VanillaTransformer(nn.Module):
+    """Tied-embedding decoder-only transformer. Same training interface as TopoGPT2."""
+
+    def __init__(self, cfg: VanillaConfig = VanillaConfig()):
+        super().__init__()
+        self.cfg = cfg
+        self.embed = nn.Embedding(cfg.VOCAB_SIZE, cfg.D_MODEL)
+        nn.init.normal_(self.embed.weight, std=0.02)
+        self.layers = nn.ModuleList([VanillaLayer(cfg) for _ in range(cfg.N_LAYERS)])
+        self.norm = RMSNorm(cfg.D_MODEL)
+        self.lm_head = nn.Linear(cfg.D_MODEL, cfg.VOCAB_SIZE, bias=False)
+        self.lm_head.weight = self.embed.weight
+        self.apply(self._init)
+
+    @staticmethod
+    def _init(m):
+        if isinstance(m, nn.Linear) and not hasattr(m, "_tied"):
+            nn.init.normal_(m.weight, std=0.02)
+
+    def forward(self, ids):
+        x = self.embed(ids)
+        for layer in self.layers:
+            x = layer(x)
+        return self.lm_head(self.norm(x))
+
+    def count_params(self):
+        total = sum(p.numel() for p in self.parameters())
+        return {"total": total, "trainable": total}
+
+
+def build_optimizer(model: VanillaTransformer, cfg: VanillaConfig):
+    return torch.optim.AdamW(model.parameters(), lr=cfg.LR, weight_decay=cfg.WEIGHT_DECAY)
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"])
+    ap.add_argument("--tier", default="codealpaca", choices=list(TIER_NAMES),
+                    help="tier curriculum de donde tomar tokens reales")
+    ap.add_argument("--seq-len", type=int, default=32)
+    ap.add_argument("--batch", type=int, default=2)
+    ap.add_argument("--allow-synthetic", action="store_true",
+                    help="permite tokens randint solo como smoke test; sin esto, falla si no hay data real")
+    ap.add_argument("--train-step", action="store_true",
+                    help="corre 1 step AdamW + grad norm (smoke de entrenamiento en GPU si disponible)")
+    args = ap.parse_args()
+    device = resolve_device(args.device)
+    print(f"device={device} cuda_avail={torch.cuda.is_available()} "
+          f"{torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'sin-gpu-para-este-build'}")
+    cfg = VanillaConfig()
+    m = VanillaTransformer(cfg).to(device)
+    print(m.count_params())
+    toks = find_tier_tokens(args.tier)
+    if toks is not None and len(toks) >= args.batch * args.seq_len:
+        import torch as _t
+        ids = _t.tensor(toks[:args.batch * args.seq_len], dtype=_t.long).view(args.batch, args.seq_len).to(device)
+        print(f"data=tier-real:{args.tier} ntok={len(toks)}")
+    elif args.allow_synthetic:
+        print(f"WARNING: sin tokens reales de {args.tier} "
+              f"(corre topogpt3/train.py --prepare-data para los 4 tiers). Uso sintetico SOLO smoke.")
+        ids = torch.randint(0, cfg.VOCAB_SIZE, (args.batch, args.seq_len), device=device)
+    else:
+        raise SystemExit(
+            f"Sin tokens reales de tier {args.tier}. Corre: python3 -m topogpt3.train --prepare-data "
+            f"(descarga CodeAlpaca + CodeFeedback + Magicoder + Stack, splits train/val/holdout). "
+            f"O repite con --allow-synthetic para smoke. No uso sintetico silencioso.")
+    use_amp = device.type == "cuda"
+    with torch.autocast(device_type="cuda", enabled=use_amp):
+        logits = m(ids)
+    print("forward:", tuple(logits.shape))
+    if args.train_step:
+        opt = build_optimizer(m, cfg)
+        tgt = torch.roll(ids, -1, dims=1)
+        with torch.autocast(device_type="cuda", enabled=use_amp):
+            loss = F.cross_entropy(logits[:, :-1].reshape(-1, cfg.VOCAB_SIZE), tgt[:, :-1].reshape(-1))
+        opt.zero_grad()
+        loss.backward()
+        gn = torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
+        opt.step()
+        print(f"train-step loss={loss.item():.4f} grad_norm={gn:.4f}")