Add experiments/joint_test.py

experiments/joint_test.py

@@ -0,0 +1,234 @@
+#!/usr/bin/env python3
+"""
+Joint AR+SAT training - what AGILLM-3 actually does
+Test which attention mechanism works best for BOTH modes simultaneously
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import time
+import math
+import argparse
+
+DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+VOCAB = 128256
+
+def get_mask(n, mode, block_size=4):
+    if mode == "nar":
+        return None
+    elif mode == "ar":
+        return torch.triu(torch.full((n, n), float("-inf"), device=DEV), 1)
+    elif mode == "sat":
+        idx = torch.arange(n, device=DEV)
+        block_idx = idx // block_size
+        mask = torch.where(
+            block_idx.unsqueeze(0) <= block_idx.unsqueeze(1),
+            torch.tensor(0.0, device=DEV),
+            torch.tensor(float("-inf"), device=DEV)
+        )
+        return mask
+
+def alibi_bias(h, n):
+    def slopes(n):
+        start = 2 ** (-2 ** -(math.log2(n) - 3))
+        return [start * (start ** i) for i in range(n)]
+    s = slopes(h) if h > 0 and math.log2(h).is_integer() else slopes(2 ** math.floor(math.log2(max(1,h))))[:h]
+    s = torch.tensor(s, device=DEV).view(1, h, 1, 1)
+    i = torch.arange(n, device=DEV).view(1, 1, n, 1)
+    j = torch.arange(n, device=DEV).view(1, 1, 1, n)
+    return -s * (j - i).clamp_min(0).float()
+
+
+class StandardAttn(nn.Module):
+    def __init__(self, d, h):
+        super().__init__()
+        self.h, self.dk = h, d // h
+        self.qkv = nn.Linear(d, 3*d, bias=False)
+        self.proj = nn.Linear(d, d, bias=False)
+    
+    def forward(self, x, mask=None):
+        B, N, _ = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.h, self.dk).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk) + alibi_bias(self.h, N)
+        if mask is not None:
+            att = att + mask.unsqueeze(0).unsqueeze(0)
+        return self.proj((att.softmax(-1) @ v).transpose(1, 2).reshape(B, N, -1))
+
+
+class MQAAttn(nn.Module):
+    def __init__(self, d, h):
+        super().__init__()
+        self.h, self.dk = h, d // h
+        self.q = nn.Linear(d, d, bias=False)
+        self.k = nn.Linear(d, self.dk, bias=False)
+        self.v = nn.Linear(d, self.dk, bias=False)
+        self.proj = nn.Linear(d, d, bias=False)
+    
+    def forward(self, x, mask=None):
+        B, N, _ = x.shape
+        q = self.q(x).view(B, N, self.h, self.dk).transpose(1, 2)
+        k = self.k(x).view(B, N, 1, self.dk).transpose(1, 2)
+        v = self.v(x).view(B, N, 1, self.dk).transpose(1, 2)
+        att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk) + alibi_bias(self.h, N)
+        if mask is not None:
+            att = att + mask.unsqueeze(0).unsqueeze(0)
+        return self.proj((att.softmax(-1) @ v).transpose(1, 2).reshape(B, N, -1))
+
+
+class GQAAttn(nn.Module):
+    def __init__(self, d, h, kv_heads=2):
+        super().__init__()
+        self.h, self.dk, self.kv_heads = h, d // h, kv_heads
+        self.q = nn.Linear(d, d, bias=False)
+        self.k = nn.Linear(d, kv_heads * self.dk, bias=False)
+        self.v = nn.Linear(d, kv_heads * self.dk, bias=False)
+        self.proj = nn.Linear(d, d, bias=False)
+    
+    def forward(self, x, mask=None):
+        B, N, _ = x.shape
+        q = self.q(x).view(B, N, self.h, self.dk).transpose(1, 2)
+        k = self.k(x).view(B, N, self.kv_heads, self.dk).transpose(1, 2)
+        v = self.v(x).view(B, N, self.kv_heads, self.dk).transpose(1, 2)
+        k = k.repeat_interleave(self.h // self.kv_heads, dim=1)
+        v = v.repeat_interleave(self.h // self.kv_heads, dim=1)
+        att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk) + alibi_bias(self.h, N)
+        if mask is not None:
+            att = att + mask.unsqueeze(0).unsqueeze(0)
+        return self.proj((att.softmax(-1) @ v).transpose(1, 2).reshape(B, N, -1))
+
+
+class Block(nn.Module):
+    def __init__(self, d, h, attn_type):
+        super().__init__()
+        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
+        if attn_type == "standard":
+            self.attn = StandardAttn(d, h)
+        elif attn_type == "mqa":
+            self.attn = MQAAttn(d, h)
+        elif attn_type == "gqa":
+            self.attn = GQAAttn(d, h, kv_heads=2)
+        elif attn_type == "gqa4":
+            self.attn = GQAAttn(d, h, kv_heads=4)
+        self.ff = nn.Sequential(nn.Linear(d, 4*d), nn.GELU(), nn.Linear(4*d, d))
+    
+    def forward(self, x, mask=None):
+        x = x + self.attn(self.ln1(x), mask)
+        return x + self.ff(self.ln2(x))
+
+
+class Model(nn.Module):
+    def __init__(self, d, layers, h, attn_type):
+        super().__init__()
+        self.emb = nn.Embedding(VOCAB, d)
+        self.blocks = nn.ModuleList([Block(d, h, attn_type) for _ in range(layers)])
+        self.ln = nn.LayerNorm(d)
+        self.head = nn.Linear(d, VOCAB, bias=False)
+        self.head.weight = self.emb.weight
+    
+    def forward(self, x, mask=None):
+        x = self.emb(x)
+        for b in self.blocks:
+            x = b(x, mask)
+        return self.head(self.ln(x))
+
+
+def train_joint(attn_type, d, layers, h, batch, seq, steps, ar_weight=0.5, block_size=4):
+    """Train with mixed AR and SAT objectives"""
+    print(f"\n{'='*60}")
+    print(f"JOINT AR+SAT: {attn_type.upper()} (AR weight={ar_weight})")
+    print(f"{'='*60}")
+    
+    model = Model(d, layers, h, attn_type).to(DEV)
+    params = sum(p.numel() for p in model.parameters())
+    print(f"Parameters: {params:,}")
+    
+    opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
+    
+    ar_mask = get_mask(seq - 1, "ar")
+    sat_mask = get_mask(seq - 1, "sat", block_size)
+    
+    ar_losses, sat_losses, times = [], [], []
+    
+    for step in range(steps):
+        ids = torch.randint(0, VOCAB, (batch, seq), device=DEV)
+        target = ids[:, 1:]
+        input_ids = ids[:, :-1]
+        
+        start = time.time()
+        opt.zero_grad()
+        
+        # AR forward
+        ar_logits = model(input_ids, ar_mask)
+        ar_loss = F.cross_entropy(ar_logits.view(-1, VOCAB), target.reshape(-1))
+        
+        # SAT forward (same input, different mask)
+        sat_logits = model(input_ids, sat_mask)
+        sat_loss = F.cross_entropy(sat_logits.view(-1, VOCAB), target.reshape(-1))
+        
+        # Combined loss
+        loss = ar_weight * ar_loss + (1 - ar_weight) * sat_loss
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step()
+        
+        elapsed = time.time() - start
+        ar_losses.append(ar_loss.item())
+        sat_losses.append(sat_loss.item())
+        times.append(elapsed)
+        
+        if step % 50 == 0 or step == steps - 1:
+            tok_s = batch * seq / elapsed
+            print(f"Step {step:3d} | AR: {ar_loss.item():.2f} | SAT: {sat_loss.item():.2f} | {tok_s:.0f} tok/s")
+    
+    avg_ar = sum(ar_losses[-20:]) / 20
+    avg_sat = sum(sat_losses[-20:]) / 20
+    avg_tok = batch * seq / (sum(times[-20:]) / 20)
+    
+    return {"type": attn_type, "ar_loss": avg_ar, "sat_loss": avg_sat, "tok_s": avg_tok, "params": params}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--d", type=int, default=256)
+    parser.add_argument("--layers", type=int, default=4)
+    parser.add_argument("--heads", type=int, default=8)
+    parser.add_argument("--batch", type=int, default=16)
+    parser.add_argument("--seq", type=int, default=128)
+    parser.add_argument("--steps", type=int, default=200)
+    parser.add_argument("--block_size", type=int, default=4)
+    args = parser.parse_args()
+    
+    print(f"Device: {DEV}")
+    if torch.cuda.is_available():
+        print(f"GPU: {torch.cuda.get_device_name()}")
+    
+    print(f"\nJoint AR+SAT Training (block_size={args.block_size})")
+    
+    results = []
+    for attn_type in ["standard", "mqa", "gqa", "gqa4"]:
+        r = train_joint(attn_type, args.d, args.layers, args.heads,
+                       args.batch, args.seq, args.steps, 
+                       ar_weight=0.5, block_size=args.block_size)
+        results.append(r)
+        torch.cuda.empty_cache()
+    
+    print(f"\n{'='*60}")
+    print("JOINT AR+SAT RESULTS")
+    print(f"{'='*60}")
+    
+    std = next(r for r in results if r['type'] == 'standard')
+    for r in sorted(results, key=lambda x: x['ar_loss'] + x['sat_loss']):
+        combined = r['ar_loss'] + r['sat_loss']
+        std_combined = std['ar_loss'] + std['sat_loss']
+        diff = (std_combined - combined) / std_combined * 100
+        
+        kv_ratio = {"standard": "1.00", "mqa": "0.12", "gqa": "0.25", "gqa4": "0.50"}[r['type']]
+        
+        print(f"{r['type']:10s} | AR: {r['ar_loss']:.2f} | SAT: {r['sat_loss']:.2f} | "
+              f"Combined: {combined:.2f} ({diff:+.1f}%) | {r['tok_s']:.0f} tok/s | KV: {kv_ratio}x")
+
+
+if __name__ == "__main__":
+    main()