Add verified causal-v4 experiment source

recurrent_ctx_interaction_v4_causal.py

@@ -0,0 +1,308 @@
+"""
+Recurrent Depth × Context Length Interaction at Tiny Scale (v2 - memory fix)
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import time
+import json
+
+# --- Config ---
+D_MODEL = 256
+N_HEADS = 4
+VOCAB = 256
+BATCH = 32
+STEPS = 2000
+LR = 3e-4
+WARMUP = 200
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+SEED = 42
+
+# --- Data: pre-generate tensors (memory efficient) ---
+def load_tinystories_bytes(max_train=50_000_000, max_val=2_000_000):
+    with open('/work/data/tinystories.txt', 'rb') as f:
+        raw = f.read()
+    # Split into train/val: 95/5
+    split = int(len(raw) * 0.95)
+    train_bytes = raw[:split][:max_train]
+    val_bytes = raw[split:][:max_val]
+    
+    # Pre-generate all (x, y) pairs as contiguous tensors
+    def make_pairs(data, seq_len=256):
+        n = len(data) // seq_len
+        arr = torch.tensor(list(data[:n * seq_len + 1]), dtype=torch.long)
+        x = arr[:-1].view(n, seq_len)
+        y = arr[1:].view(n, seq_len)
+        return x, y
+    
+    train_x, train_y = make_pairs(train_bytes)
+    val_x, val_y = make_pairs(val_bytes)
+    print(f'  Train: {train_x.shape[0]} sequences, Val: {val_x.shape[0]} sequences')
+    return train_x, train_y, val_x, val_y
+
+# --- Model ---
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, n_heads, d_ff):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(d_model)
+        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
+        self.ln2 = nn.LayerNorm(d_model)
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, d_ff),
+            nn.GELU(),
+            nn.Linear(d_ff, d_model)
+        )
+    
+    def forward(self, x):
+        h = self.ln1(x)
+        causal_mask = torch.ones(h.size(1), h.size(1), device=h.device, dtype=torch.bool).triu(1)
+        attn_out, _ = self.attn(h, h, h, attn_mask=causal_mask, need_weights=False)
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        return x
+
+class TinyGPT(nn.Module):
+    def __init__(self, d_model, n_heads, d_ff, n_blocks, recurrent=False, ctx_len=128):
+        super().__init__()
+        self.d_model = d_model
+        self.ctx_len = ctx_len
+        self.recurrent = recurrent
+        self.n_blocks = n_blocks
+        
+        self.tok_emb = nn.Embedding(VOCAB, d_model)
+        self.pos_emb = nn.Embedding(ctx_len, d_model)
+        
+        if recurrent:
+            self.block = TransformerBlock(d_model, n_heads, d_ff)
+        else:
+            self.blocks = nn.ModuleList([
+                TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_blocks)
+            ])
+        
+        self.ln_f = nn.LayerNorm(d_model)
+        self.head = nn.Linear(d_model, VOCAB, bias=False)
+        self.head.weight = self.tok_emb.weight
+    
+    def forward(self, idx, targets=None):
+        B, T = idx.shape
+        T = min(T, self.ctx_len)
+        idx = idx[:, :T]
+        if targets is not None:
+            targets = targets[:, :T]
+        
+        pos = torch.arange(T, device=idx.device).unsqueeze(0)
+        x = self.tok_emb(idx) + self.pos_emb(pos)
+        
+        if self.recurrent:
+            for _ in range(self.n_blocks):
+                x = self.block(x)
+        else:
+            for block in self.blocks:
+                x = block(x)
+        
+        x = self.ln_f(x)
+        logits = self.head(x)
+        
+        if targets is not None:
+            loss = F.cross_entropy(logits.reshape(-1, VOCAB), targets.reshape(-1))
+            return loss
+        return logits
+
+def count_params(model):
+    return sum(p.numel() for p in model.parameters())
+
+def train(model, train_x, train_y, val_x, val_y, steps=STEPS, batch=BATCH, lr=LR, warmup=WARMUP):
+    model = model.to(DEVICE)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.01)
+    
+    def get_lr(step):
+        if step < warmup:
+            return lr * step / max(warmup, 1)
+        return lr * (0.5 + 0.5 * math.cos(math.pi * (step - warmup) / max(steps - warmup, 1)))
+    
+    n_train = train_x.shape[0]
+    model.train()
+    losses = []
+    val_losses = []
+    best_val = float("inf")
+    
+    for step in range(steps):
+        optimizer.param_groups[0]["lr"] = get_lr(step)
+        
+        idx = torch.randint(0, n_train, (batch,))
+        x = train_x[idx].to(DEVICE)
+        y = train_y[idx].to(DEVICE)
+        
+        loss = model(x, y)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+        optimizer.zero_grad()
+        
+        losses.append(loss.item())
+        
+        if (step + 1) % 200 == 0:
+            model.eval()
+            val_loss = 0.0
+            val_count = 0
+            with torch.no_grad():
+                for i in range(0, val_x.shape[0], batch):
+                    vx = val_x[i:i+batch].to(DEVICE)
+                    vy = val_y[i:i+batch].to(DEVICE)
+                    vl = model(vx, vy)
+                    val_loss += vl.item()
+                    val_count += 1
+            avg_val = val_loss / max(val_count, 1)
+            val_losses.append((step + 1, avg_val))
+            if avg_val < best_val:
+                best_val = avg_val
+            model.train()
+    
+    # Final validation
+    model.eval()
+    val_loss = 0.0
+    val_count = 0
+    with torch.no_grad():
+        for i in range(0, val_x.shape[0], batch):
+            vx = val_x[i:i+batch].to(DEVICE)
+            vy = val_y[i:i+batch].to(DEVICE)
+            vl = model(vx, vy)
+            val_loss += vl.item()
+            val_count += 1
+    final_val = val_loss / max(val_count, 1)
+    
+    return {
+        "final_train_loss": losses[-1],
+        "avg_train_loss_last_100": sum(losses[-100:]) / 100,
+        "best_val_loss": best_val,
+        "final_val_loss": final_val,
+        "final_val_ppl": math.exp(final_val),
+        "val_curve": val_losses,
+    }
+
+def main():
+    torch.manual_seed(SEED)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(SEED)
+    
+    print("=" * 60)
+    print("Recurrent Depth × Context Length Interaction")
+    print("=" * 60)
+    print(f"Device: {DEVICE}")
+    if torch.cuda.is_available():
+        free, total = torch.cuda.mem_get_info()
+        print(f"GPU: {free/1e9:.1f} GB free / {total/1e9:.1f} GB total")
+    
+    print("\nLoading data...")
+    train_x, train_y, val_x, val_y = load_tinystories_bytes()
+    
+    configs = [
+        ("recurrent_ctx64",  True,  4, 2048, 64),
+        ("recurrent_ctx256", True,  4, 2048, 256),
+        ("stacked_ctx64",    False, 4, 128,  64),
+        ("stacked_ctx256",   False, 4, 128,  256),
+    ]
+    
+    results = {}
+    for name, recurrent, n_blocks, d_ff, ctx_len in configs:
+        print(f"\n{'='*40}")
+        print(f"Config: {name}")
+        print(f"  recurrent={recurrent}, blocks={n_blocks}, d_ff={d_ff}, ctx={ctx_len}")
+        
+        model = TinyGPT(D_MODEL, N_HEADS, d_ff, n_blocks, recurrent=recurrent, ctx_len=ctx_len)
+        n_params = count_params(model)
+        print(f"  Params: {n_params:,}")
+        
+        t0 = time.time()
+        res = train(model, train_x, train_y, val_x, val_y)
+        elapsed = time.time() - t0
+        res["params"] = n_params
+        res["train_time_s"] = elapsed
+        results[name] = res
+        
+        print(f"  Train loss (final): {res['final_train_loss']:.4f}")
+        print(f"  Val PPL: {res['final_val_ppl']:.2f}")
+        print(f"  Time: {elapsed:.1f}s")
+        
+        # Free GPU memory between configs
+        del model
+        torch.cuda.empty_cache()
+    
+    # Summary
+    print("\n" + "=" * 60)
+    print("SUMMARY")
+    print("=" * 60)
+    print(f"\n{'Config':<25} {'Params':>10} {'Val PPL':>10} {'Time':>8}")
+    print("-" * 60)
+    for name, res in results.items():
+        print(f"{name:<25} {res['params']:>10,} {res['final_val_ppl']:>10.2f} {res['train_time_s']:>7.1f}s")
+    
+    # Interaction analysis
+    print("\n--- Interaction Analysis ---")
+    rec_64 = results["recurrent_ctx64"]["final_val_ppl"]
+    rec_256 = results["recurrent_ctx256"]["final_val_ppl"]
+    stk_64 = results["stacked_ctx64"]["final_val_ppl"]
+    stk_256 = results["stacked_ctx256"]["final_val_ppl"]
+    
+    rec_gain = rec_64 - rec_256
+    stk_gain = stk_64 - stk_256
+    
+    print(f"Recurrent: ctx64={rec_64:.2f} → ctx256={rec_256:.2f} (Δ={rec_gain:+.2f})")
+    print(f"Stacked:   ctx64={stk_64:.2f} → ctx256={stk_256:.2f} (Δ={stk_gain:+.2f})")
+    print(f"\nRecurrent advantage at ctx64:  {stk_64 - rec_64:+.2f}")
+    print(f"Recurrent advantage at ctx256: {stk_256 - rec_256:+.2f}")
+    
+    if rec_gain > 0.5 and stk_gain < 0.3:
+        print("\n>>> FINDING: Longer context disproportionately benefits recurrent models!")
+    elif abs(rec_gain) < 0.3 and abs(stk_gain) < 0.3:
+        print("\n>>> FINDING: Context length has minimal effect on either architecture.")
+    elif rec_gain < 0 and stk_gain < 0:
+        print("\n>>> FINDING: Longer context HURTS both architectures (attention dilution).")
+    else:
+        print(f"\n>>> Mixed: rec_gain={rec_gain:+.2f}, stk_gain={stk_gain:+.2f}")
+    
+    out = {
+        "experiment": "recurrent_ctx_interaction",
+        "question": "Does recurrent depth amplify the benefit of longer context at tiny scale?",
+        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "config": {
+            "d_model": D_MODEL,
+            "n_heads": N_HEADS,
+            "vocab": VOCAB,
+            "batch": BATCH,
+            "steps": STEPS,
+            "lr": LR,
+            "device": DEVICE,
+        },
+        "results": results,
+        "analysis": {
+            "recurrent_ctx_gain": rec_gain,
+            "stacked_ctx_gain": stk_gain,
+            "recurrent_advantage_ctx64": stk_64 - rec_64,
+            "recurrent_advantage_ctx256": stk_256 - rec_256,
+        }
+    }
+    
+    out_path = "/work/work/recurrent_ctx_interaction_results.json"
+    with open(out_path, "w") as f:
+        json.dump(out, f, indent=2)
+    print(f"\nSaved to {out_path}")
+
+if __name__ == "__main__":
+    main()
+
+def preflight_case():
+    """CPU-only correctness fixture using the real TinyGPT and shifted TinyStories bytes."""
+    from pathlib import Path
+    import torch
+    source = Path('/work/data/tinystories.txt')
+    with source.open('rb') as f:
+        raw = f.read(130)
+    if len(raw) != 130:
+        raise ValueError('Need at least 130 real training bytes for causal preflight')
+    tokens = torch.tensor(list(raw), dtype=torch.long)
+    x = torch.stack((tokens[0:64], tokens[65:129]))
+    shifted_y = torch.stack((tokens[1:65], tokens[66:130]))
+    model = TinyGPT(d_model=32, n_heads=4, d_ff=64, n_blocks=2, recurrent=True, ctx_len=64).cpu()
+    return model, x, shifted_y, []  # No dedicated memory-write modules; recurrence reuses TransformerBlock weights.