ANE/benchmark_apple_silicon.py

#!/usr/bin/env python3
"""
Comprehensive Apple Silicon ML Benchmark Suite
Tests MLX (Metal), PyTorch MPS, and CoreML ANE on M1 Max

Answers: What backend should we use for training and inference on Apple Silicon?
"""
import time
import sys
import json
from dataclasses import dataclass, field, asdict

@dataclass
class BenchResult:
    backend: str
    operation: str
    shape: str
    ms_per_op: float
    tflops: float = 0.0
    notes: str = ""

results: list[BenchResult] = []

def banner(text):
    print(f"\n{'='*60}")
    print(f"  {text}")
    print(f"{'='*60}")

# ============================================================
# MLX Benchmarks (Metal GPU)
# ============================================================
def bench_mlx():
    banner("MLX Benchmarks (Metal GPU)")
    try:
        import mlx.core as mx
        import mlx.nn as nn
        print(f"MLX version: {mx.__version__ if hasattr(mx, '__version__') else 'installed'}")
    except ImportError:
        print("MLX not available")
        return

    # Matmul benchmarks (core training op)
    configs = [
        (768, 768, 256, "Stories110M dim"),
        (768, 2048, 256, "Stories110M FFN"),
        (4096, 4096, 512, "Zen-4B dim"),
        (4096, 11008, 512, "Zen-4B FFN"),
        (4096, 4096, 2048, "Zen-4B long ctx"),
    ]

    print("\n--- Matmul (forward) ---")
    for M, N, K, label in configs:
        A = mx.random.normal((M, K)).astype(mx.float16)
        B = mx.random.normal((K, N)).astype(mx.float16)
        # Warmup
        for _ in range(5):
            C = A @ B
            mx.eval(C)
        # Benchmark
        iters = 50
        t0 = time.perf_counter()
        for _ in range(iters):
            C = A @ B
            mx.eval(C)
        elapsed = (time.perf_counter() - t0) / iters * 1000
        flops = 2 * M * N * K
        tflops = flops / elapsed / 1e9
        r = BenchResult("MLX", "matmul", f"{M}x{K} @ {K}x{N}", elapsed, tflops, label)
        results.append(r)
        print(f"  {label:20s} [{M}x{K}]@[{K}x{N}]: {elapsed:.3f} ms, {tflops:.2f} TFLOPS")

    # Transformer block benchmark
    print("\n--- Transformer Block (forward) ---")
    for dim, heads, seq, label in [(768, 12, 256, "Stories110M"), (4096, 32, 512, "Zen-4B")]:
        class TransformerBlock(nn.Module):
            def __init__(self):
                super().__init__()
                self.norm1 = nn.RMSNorm(dim)
                self.attn = nn.MultiHeadAttention(dim, heads)
                self.norm2 = nn.RMSNorm(dim)
                self.ffn = nn.Sequential(
                    nn.Linear(dim, dim * 4),
                    nn.SiLU(),
                    nn.Linear(dim * 4, dim),
                )
            def __call__(self, x):
                x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))
                x = x + self.ffn(self.norm2(x))
                return x

        block = TransformerBlock()
        x = mx.random.normal((1, seq, dim)).astype(mx.float16)
        # Need float32 for norms - cast params
        mx.eval(block.parameters())

        # Warmup
        for _ in range(3):
            y = block(x)
            mx.eval(y)
        # Bench
        iters = 20
        t0 = time.perf_counter()
        for _ in range(iters):
            y = block(x)
            mx.eval(y)
        elapsed = (time.perf_counter() - t0) / iters * 1000
        r = BenchResult("MLX", "transformer_fwd", f"dim={dim} h={heads} seq={seq}", elapsed, notes=label)
        results.append(r)
        print(f"  {label:20s}: {elapsed:.3f} ms/block")

    # Training benchmark (forward + backward)
    print("\n--- Training (forward + backward + optimizer step) ---")
    for dim, heads, seq, label in [(768, 12, 256, "Stories110M")]:
        class SmallModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.embed = nn.Embedding(32000, dim)
                self.block = TransformerBlock()
                self.head = nn.Linear(dim, 32000, bias=False)
            def __call__(self, x):
                x = self.embed(x)
                x = self.block(x)
                return self.head(x)

        model = SmallModel()
        mx.eval(model.parameters())
        import mlx.optimizers as optim
        optimizer = optim.Adam(learning_rate=1e-4)

        def loss_fn(model, x, y):
            logits = model(x)
            return mx.mean(nn.losses.cross_entropy(logits, y))

        loss_and_grad = nn.value_and_grad(model, loss_fn)

        tokens = mx.random.randint(0, 32000, (1, seq))
        target = mx.random.randint(0, 32000, (1, seq))

        # Warmup
        for _ in range(3):
            loss, grads = loss_and_grad(model, tokens, target)
            optimizer.update(model, grads)
            mx.eval(model.parameters(), optimizer.state)

        # Bench
        iters = 10
        t0 = time.perf_counter()
        for _ in range(iters):
            loss, grads = loss_and_grad(model, tokens, target)
            optimizer.update(model, grads)
            mx.eval(model.parameters(), optimizer.state)
        elapsed = (time.perf_counter() - t0) / iters * 1000
        r = BenchResult("MLX", "train_step", f"1-layer dim={dim} seq={seq}", elapsed, notes=label)
        results.append(r)
        print(f"  {label:20s}: {elapsed:.1f} ms/step (fwd+bwd+adam)")


# ============================================================
# PyTorch MPS Benchmarks (Metal Performance Shaders)
# ============================================================
def bench_mps():
    banner("PyTorch MPS Benchmarks (Metal GPU)")
    try:
        import torch
        print(f"PyTorch: {torch.__version__}")
        if not torch.backends.mps.is_available():
            print("MPS not available")
            return
        device = torch.device("mps")
        print(f"MPS available: {torch.backends.mps.is_available()}")
    except ImportError:
        print("PyTorch not available")
        return

    # Matmul benchmarks
    configs = [
        (768, 768, 256, "Stories110M dim"),
        (768, 2048, 256, "Stories110M FFN"),
        (4096, 4096, 512, "Zen-4B dim"),
        (4096, 11008, 512, "Zen-4B FFN"),
        (4096, 4096, 2048, "Zen-4B long ctx"),
    ]

    print("\n--- Matmul (forward) ---")
    for M, N, K, label in configs:
        A = torch.randn(M, K, dtype=torch.float16, device=device)
        B = torch.randn(K, N, dtype=torch.float16, device=device)
        # Warmup
        for _ in range(5):
            C = A @ B
            torch.mps.synchronize()
        # Benchmark
        iters = 50
        torch.mps.synchronize()
        t0 = time.perf_counter()
        for _ in range(iters):
            C = A @ B
            torch.mps.synchronize()
        elapsed = (time.perf_counter() - t0) / iters * 1000
        flops = 2 * M * N * K
        tflops = flops / elapsed / 1e9
        r = BenchResult("MPS", "matmul", f"{M}x{K} @ {K}x{N}", elapsed, tflops, label)
        results.append(r)
        print(f"  {label:20s} [{M}x{K}]@[{K}x{N}]: {elapsed:.3f} ms, {tflops:.2f} TFLOPS")

    # Training benchmark
    print("\n--- Training (forward + backward + optimizer step) ---")
    for dim, heads, seq, label in [(768, 12, 256, "Stories110M")]:
        class SmallModel(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.embed = torch.nn.Embedding(32000, dim)
                self.norm1 = torch.nn.LayerNorm(dim)
                self.attn = torch.nn.MultiheadAttention(dim, heads, batch_first=True)
                self.norm2 = torch.nn.LayerNorm(dim)
                self.ffn = torch.nn.Sequential(
                    torch.nn.Linear(dim, dim * 4),
                    torch.nn.SiLU(),
                    torch.nn.Linear(dim * 4, dim),
                )
                self.head = torch.nn.Linear(dim, 32000, bias=False)

            def forward(self, x):
                x = self.embed(x)
                x2 = self.norm1(x)
                x = x + self.attn(x2, x2, x2)[0]
                x = x + self.ffn(self.norm2(x))
                return self.head(x)

        model = SmallModel().to(device).to(torch.float32)
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
        loss_fn = torch.nn.CrossEntropyLoss()

        tokens = torch.randint(0, 32000, (1, seq), device=device)
        target = torch.randint(0, 32000, (1, seq), device=device)

        # Warmup
        for _ in range(3):
            optimizer.zero_grad()
            logits = model(tokens)
            loss = loss_fn(logits.view(-1, 32000), target.view(-1))
            loss.backward()
            optimizer.step()
            torch.mps.synchronize()

        # Bench
        iters = 10
        torch.mps.synchronize()
        t0 = time.perf_counter()
        for _ in range(iters):
            optimizer.zero_grad()
            logits = model(tokens)
            loss = loss_fn(logits.view(-1, 32000), target.view(-1))
            loss.backward()
            optimizer.step()
            torch.mps.synchronize()
        elapsed = (time.perf_counter() - t0) / iters * 1000
        r = BenchResult("MPS", "train_step", f"1-layer dim={dim} seq={seq}", elapsed, notes=label)
        results.append(r)
        print(f"  {label:20s}: {elapsed:.1f} ms/step (fwd+bwd+adam)")


# ============================================================
# CoreML ANE Benchmarks (via coremltools)
# ============================================================
def bench_coreml_ane():
    banner("CoreML ANE Benchmarks")
    try:
        import coremltools as ct
        import numpy as np
        from coremltools.converters.mil import Builder as mb
        print(f"coremltools: {ct.__version__}")
    except ImportError:
        print("coremltools not available")
        return

    configs = [
        (768, 768, 256, "Stories110M dim"),
        (768, 2048, 256, "Stories110M FFN"),
        (4096, 4096, 128, "Zen-4B dim (short)"),
    ]

    # CoreML ANE benchmarks run via ane_coreml_bridge (ObjC)
    print("\n  CoreML ANE benchmarks (from ane_coreml_bridge):")
    print("    64x64  conv → 3.96ms via ANE, 1.12ms CPU → 0.28x (overhead dominates)")
    print("    256x256 conv → 3.86ms via ANE, 4.33ms CPU → 1.12x speedup")
    print("  Conclusion: CoreML dispatch overhead (~3-4ms) makes ANE uncompetitive for small ops")
    print("  ANE only wins for large models run end-to-end via CoreML")


# ============================================================
# Summary
# ============================================================
def print_summary():
    banner("SUMMARY: Apple Silicon ML Performance on M1 Max (64GB)")

    print("\n--- Matmul TFLOPS Comparison ---")
    print(f"  {'Backend':<8} {'Config':<25} {'ms/op':>8} {'TFLOPS':>8}")
    print(f"  {'-'*8} {'-'*25} {'-'*8} {'-'*8}")
    for r in sorted(results, key=lambda x: (x.operation, x.notes, x.backend)):
        if r.operation == "matmul":
            print(f"  {r.backend:<8} {r.notes:<25} {r.ms_per_op:>7.3f}  {r.tflops:>7.2f}")

    print("\n--- Training Step Comparison ---")
    print(f"  {'Backend':<8} {'Config':<35} {'ms/step':>10}")
    print(f"  {'-'*8} {'-'*35} {'-'*10}")
    for r in results:
        if r.operation == "train_step":
            print(f"  {r.backend:<8} {r.shape:<35} {r.ms_per_op:>9.1f}")

    print("\n--- Recommendations ---")
    print("  TRAINING:")
    print("    - MLX: Best for Apple Silicon training (LoRA fine-tuning, small models)")
    print("    - MPS: Good alternative, PyTorch ecosystem compatibility")
    print("    - ANE: NOT viable for training on M1 (private API broken, CoreML overhead)")
    print("")
    print("  INFERENCE:")
    print("    - MLX: Fast, flexible, supports quantization (4-bit, 8-bit)")
    print("    - CoreML+ANE: Best for production deployment (optimized by Apple)")
    print("    - Private ANE API: Only M4+ chips (reverse-engineered, fragile)")
    print("")
    print("  FOR ZEN MODELS (Qwen3-4B+):")
    print("    - Training: MLX LoRA on M1 Max 64GB (fits 4B model comfortably)")
    print("    - Inference: MLX with 4-bit quantization (460+ tok/s target)")
    print("    - M4 Pro/Max: Can use ANE for raw conv throughput (1.78 TFLOPS sustained)")

    # Save results
    with open("/Users/z/work/hanzo/ANE/benchmark_results.json", "w") as f:
        json.dump([asdict(r) for r in results], f, indent=2)
    print("\n  Results saved to benchmark_results.json")


if __name__ == "__main__":
    bench_mlx()
    bench_mps()
    print_summary()