Add tests/test_ultron.py

tests/test_ultron.py

@@ -0,0 +1,279 @@
+"""Test suite for Ultron model — verifies forward pass, generation, stability, and all variants."""
+
+import sys
+import torch
+import torch.nn.functional as F
+import time
+
+sys.path.insert(0, "/app")
+
+from ultron.model import Ultron, UltronConfig
+from ultron.variants import ultron_small, ultron_base, ultron_medium, ultron_medium_moe
+
+
+def test_basic_forward():
+    """Test forward pass with minimal config."""
+    print("=" * 60)
+    print("TEST: Basic forward pass")
+    cfg = UltronConfig(
+        vocab_size=1000, dim=128, n_heads=4, n_kv_heads=2,
+        max_seq_len=128, prelude_layers=1, coda_layers=1,
+        recurrent_layers=2, max_loop_iters=4, lora_rank=4,
+        attn_type="gqa", use_moe=False,
+    )
+    model = Ultron(cfg)
+    total_params = model.get_num_params(non_embedding=False)
+    print(f"  Config: dim={cfg.dim}, heads={cfg.n_heads}, recurrent_layers={cfg.recurrent_layers}, loops={cfg.max_loop_iters}")
+    print(f"  Parameters: {total_params:,}")
+
+    ids = torch.randint(0, cfg.vocab_size, (2, 32))
+    logits = model(ids)
+    assert logits.shape == (2, 32, cfg.vocab_size), f"Wrong shape: {logits.shape}"
+    print(f"  Logits shape: {logits.shape} ✓")
+
+    # Check stability
+    rho = model.get_spectral_radius()
+    assert rho < 1.0, f"Spectral radius {rho} >= 1!"
+    print(f"  Spectral radius ρ(A) = {rho:.6f} (< 1 ✓)")
+    print("  PASSED ✓\n")
+
+
+def test_mla_attention():
+    """Test with Multi-Latent Attention."""
+    print("=" * 60)
+    print("TEST: MLA attention")
+    cfg = UltronConfig(
+        vocab_size=1000, dim=128, n_heads=4, n_kv_heads=4,
+        max_seq_len=128, prelude_layers=1, coda_layers=1,
+        recurrent_layers=2, max_loop_iters=4, lora_rank=4,
+        attn_type="mla",
+        kv_lora_rank=32, q_lora_rank=64,
+        qk_rope_head_dim=16, qk_nope_head_dim=16, v_head_dim=16,
+    )
+    model = Ultron(cfg)
+    ids = torch.randint(0, cfg.vocab_size, (2, 32))
+    logits = model(ids)
+    assert logits.shape == (2, 32, cfg.vocab_size)
+    print(f"  Logits shape: {logits.shape} ✓")
+    print(f"  Parameters: {model.get_num_params():,}")
+    print("  PASSED ✓\n")
+
+
+def test_moe():
+    """Test with MoE FFN in recurrent block."""
+    print("=" * 60)
+    print("TEST: MoE FFN")
+    cfg = UltronConfig(
+        vocab_size=1000, dim=128, n_heads=4, n_kv_heads=2,
+        max_seq_len=128, prelude_layers=1, coda_layers=1,
+        recurrent_layers=2, max_loop_iters=4, lora_rank=4,
+        attn_type="gqa",
+        use_moe=True, n_experts=4, n_shared_experts=1,
+        n_experts_per_tok=2, expert_dim=64,
+    )
+    model = Ultron(cfg)
+    ids = torch.randint(0, cfg.vocab_size, (2, 16))
+    logits = model(ids)
+    assert logits.shape == (2, 16, cfg.vocab_size)
+    print(f"  Logits shape: {logits.shape} ✓")
+    print(f"  Parameters: {model.get_num_params():,}")
+    print("  PASSED ✓\n")
+
+
+def test_generation():
+    """Test autoregressive generation with KV caching."""
+    print("=" * 60)
+    print("TEST: Autoregressive generation")
+    cfg = UltronConfig(
+        vocab_size=1000, dim=128, n_heads=4, n_kv_heads=2,
+        max_seq_len=256, prelude_layers=1, coda_layers=1,
+        recurrent_layers=2, max_loop_iters=4, lora_rank=4,
+    )
+    model = Ultron(cfg).eval()
+    prompt = torch.randint(0, cfg.vocab_size, (1, 8))
+    output = model.generate(prompt, max_new_tokens=16, n_loops=4, temperature=1.0, top_k=10)
+    assert output.shape == (1, 24), f"Expected (1, 24), got {output.shape}"
+    print(f"  Generated shape: {output.shape} ✓")
+    print(f"  Prompt: {prompt[0].tolist()[:8]}")
+    print(f"  Generated: {output[0, 8:].tolist()}")
+    print("  PASSED ✓\n")
+
+
+def test_depth_extrapolation():
+    """Test that model works with more loops at inference than training default."""
+    print("=" * 60)
+    print("TEST: Depth extrapolation")
+    cfg = UltronConfig(
+        vocab_size=1000, dim=128, n_heads=4, n_kv_heads=2,
+        max_seq_len=128, prelude_layers=1, coda_layers=1,
+        recurrent_layers=2, max_loop_iters=4, lora_rank=4,
+    )
+    model = Ultron(cfg).eval()
+    ids = torch.randint(0, cfg.vocab_size, (1, 16))
+
+    # Normal depth
+    logits_4 = model(ids, n_loops=4)
+    # Extrapolated depth (2x training loops)
+    logits_8 = model(ids, n_loops=8)
+    # Even deeper
+    logits_16 = model(ids, n_loops=16)
+
+    assert logits_4.shape == logits_8.shape == logits_16.shape
+    # Results should differ (different loop counts = different outputs)
+    assert not torch.allclose(logits_4, logits_8, atol=1e-4)
+    print(f"  4 loops → logit mean: {logits_4.mean():.4f}")
+    print(f"  8 loops → logit mean: {logits_8.mean():.4f}")
+    print(f"  16 loops → logit mean: {logits_16.mean():.4f}")
+    print("  PASSED ✓\n")
+
+
+def test_act_halting():
+    """Verify ACT halting stops early when positions converge."""
+    print("=" * 60)
+    print("TEST: ACT halting behavior")
+    cfg = UltronConfig(
+        vocab_size=1000, dim=128, n_heads=4, n_kv_heads=2,
+        max_seq_len=128, prelude_layers=1, coda_layers=1,
+        recurrent_layers=2, max_loop_iters=16,
+        lora_rank=4, act_threshold=0.99,
+    )
+    model = Ultron(cfg).eval()
+    ids = torch.randint(0, cfg.vocab_size, (1, 16))
+    logits = model(ids, n_loops=16)
+    print(f"  Logits shape: {logits.shape} ✓")
+    print(f"  ACT threshold: {cfg.act_threshold}")
+    print("  PASSED ✓\n")
+
+
+def test_backward():
+    """Test that gradients flow correctly through the looped model."""
+    print("=" * 60)
+    print("TEST: Backward pass / gradient flow")
+    cfg = UltronConfig(
+        vocab_size=1000, dim=128, n_heads=4, n_kv_heads=2,
+        max_seq_len=128, prelude_layers=1, coda_layers=1,
+        recurrent_layers=2, max_loop_iters=4, lora_rank=4,
+    )
+    model = Ultron(cfg)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
+
+    ids = torch.randint(0, cfg.vocab_size, (2, 32))
+    labels = torch.randint(0, cfg.vocab_size, (2, 32))
+
+    logits = model(ids)
+    loss = torch.nn.functional.cross_entropy(
+        logits.view(-1, cfg.vocab_size), labels.view(-1)
+    )
+    loss.backward()
+    optimizer.step()
+
+    print(f"  Loss: {loss.item():.4f}")
+
+    # Check key gradients exist
+    assert model.recurrent.injection.log_A.grad is not None, "No gradient on log_A!"
+    assert model.recurrent.injection.B.grad is not None, "No gradient on B!"
+    assert model.recurrent.injection.C.grad is not None, "No gradient on C!"
+    if model.recurrent.lora is not None:
+        assert model.recurrent.lora.B.grad is not None, "No gradient on LoRA B!"
+    print(f"  log_A grad norm: {model.recurrent.injection.log_A.grad.norm():.6f}")
+    print(f"  B grad norm: {model.recurrent.injection.B.grad.norm():.6f}")
+    print(f"  C grad norm: {model.recurrent.injection.C.grad.norm():.6f}")
+    print("  PASSED ✓\n")
+
+
+def test_gradient_checkpointing():
+    """Test that gradient checkpointing works without errors."""
+    print("=" * 60)
+    print("TEST: Gradient checkpointing")
+    cfg = UltronConfig(
+        vocab_size=1000, dim=128, n_heads=4, n_kv_heads=2,
+        max_seq_len=128, prelude_layers=1, coda_layers=1,
+        recurrent_layers=2, max_loop_iters=4, lora_rank=4,
+        gradient_checkpointing=True,
+    )
+    model = Ultron(cfg)
+
+    ids = torch.randint(0, cfg.vocab_size, (2, 16))
+    labels = torch.randint(0, cfg.vocab_size, (2, 16))
+
+    logits = model(ids)
+    loss = torch.nn.functional.cross_entropy(logits.view(-1, cfg.vocab_size), labels.view(-1))
+    loss.backward()
+
+    print(f"  Loss: {loss.item():.4f}")
+    print(f"  Grad checkpointing enabled: {cfg.gradient_checkpointing}")
+    print("  PASSED ✓\n")
+
+
+def test_variant_param_counts():
+    """Verify parameter counts for all variants."""
+    print("=" * 60)
+    print("TEST: Variant parameter counts")
+
+    variants = {
+        "ultron_small": ultron_small(),
+        "ultron_base": ultron_base(),
+    }
+
+    for name, cfg in variants.items():
+        model = Ultron(cfg)
+        total = model.get_num_params(non_embedding=False)
+        non_emb = model.get_num_params(non_embedding=True)
+        rho = model.get_spectral_radius()
+        print(f"  {name}: {total:>12,} total | {non_emb:>12,} non-emb | ρ(A)={rho:.6f}")
+        assert rho < 1.0
+
+    print("  PASSED ✓\n")
+
+
+def test_training_loop():
+    """Run a mini training loop to verify end-to-end training works."""
+    print("=" * 60)
+    print("TEST: Mini training loop (10 steps)")
+    cfg = UltronConfig(
+        vocab_size=1000, dim=128, n_heads=4, n_kv_heads=2,
+        max_seq_len=128, prelude_layers=1, coda_layers=1,
+        recurrent_layers=2, max_loop_iters=4, lora_rank=4,
+    )
+    model = Ultron(cfg)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
+
+    losses = []
+    for step in range(10):
+        ids = torch.randint(0, cfg.vocab_size, (4, 64))
+        targets = ids[:, 1:]
+        logits = model(ids[:, :-1])
+        loss = F.cross_entropy(logits.reshape(-1, cfg.vocab_size), targets.reshape(-1))
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        losses.append(loss.item())
+
+    rho = model.get_spectral_radius()
+    print(f"  Step 0 loss: {losses[0]:.4f}")
+    print(f"  Step 9 loss: {losses[-1]:.4f}")
+    print(f"  ρ(A) after training: {rho:.6f} (< 1 ✓)")
+    assert rho < 1.0, f"Spectral radius exploded: {rho}"
+    assert losses[-1] < losses[0], "Loss didn't decrease!"
+    print("  PASSED ✓\n")
+
+
+if __name__ == "__main__":
+    print("\n🤖 ULTRON TEST SUITE\n")
+    start = time.time()
+
+    test_basic_forward()
+    test_mla_attention()
+    test_moe()
+    test_generation()
+    test_depth_extrapolation()
+    test_act_halting()
+    test_backward()
+    test_gradient_checkpointing()
+    test_variant_param_counts()
+    test_training_loop()
+
+    elapsed = time.time() - start
+    print("=" * 60)
+    print(f"🎉 ALL TESTS PASSED in {elapsed:.1f}s")
+    print("=" * 60)