363 lines
11 KiB
Python
363 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
PyTorch MPS Mixture of Experts (MoE) Benchmark - Comparison baseline for RustyTorch.
|
|
|
|
This script benchmarks MoE routing and expert computation on Apple MPS backend
|
|
and outputs JSON results for automated comparison with RustyTorch MoE.
|
|
|
|
Usage:
|
|
python3 benchmarks/metal/bench_moe_mps.py
|
|
python3 benchmarks/metal/bench_moe_mps.py --json > benchmarks/reports/pytorch_moe.json
|
|
|
|
Run Rust benchmark with:
|
|
cargo bench -p rtx-transformers --bench metal_moe_bench
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import time
|
|
import sys
|
|
import json
|
|
import argparse
|
|
import platform
|
|
import statistics
|
|
from dataclasses import dataclass, asdict
|
|
from typing import List, Optional
|
|
|
|
|
|
@dataclass
|
|
class MoEBenchmarkResult:
|
|
"""Single MoE benchmark result."""
|
|
name: str
|
|
batch_size: int
|
|
seq_len: int
|
|
hidden_size: int
|
|
num_experts: int
|
|
top_k: int
|
|
routing_time_ms: float
|
|
expert_compute_time_ms: float
|
|
total_time_ms: float
|
|
std_time_ms: float
|
|
throughput_tokens_per_sec: float
|
|
expert_utilization: float
|
|
memory_mb: Optional[float] = None
|
|
|
|
|
|
@dataclass
|
|
class MoEBenchmarkReport:
|
|
"""Complete MoE benchmark report."""
|
|
framework: str
|
|
backend: str
|
|
pytorch_version: str
|
|
python_version: str
|
|
macos_version: str
|
|
chip: str
|
|
timestamp: str
|
|
iterations: int
|
|
warmup: int
|
|
results: List[dict]
|
|
|
|
|
|
class TopKRouter(nn.Module):
|
|
"""Top-K expert router with softmax gating."""
|
|
|
|
def __init__(self, hidden_size: int, num_experts: int, top_k: int = 2):
|
|
super().__init__()
|
|
self.top_k = top_k
|
|
self.num_experts = num_experts
|
|
self.gate = nn.Linear(hidden_size, num_experts, bias=False)
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
"""
|
|
Route tokens to experts.
|
|
|
|
Args:
|
|
x: Input tensor [batch, seq, hidden]
|
|
|
|
Returns:
|
|
expert_indices: [batch, seq, top_k]
|
|
expert_weights: [batch, seq, top_k]
|
|
"""
|
|
# Compute gating scores
|
|
scores = self.gate(x) # [batch, seq, num_experts]
|
|
probs = F.softmax(scores, dim=-1)
|
|
|
|
# Select top-k experts
|
|
weights, indices = torch.topk(probs, self.top_k, dim=-1)
|
|
weights = weights / weights.sum(dim=-1, keepdim=True) # Renormalize
|
|
|
|
return indices, weights
|
|
|
|
|
|
class ExpertFFN(nn.Module):
|
|
"""Single expert FFN."""
|
|
|
|
def __init__(self, hidden_size: int, intermediate_size: int):
|
|
super().__init__()
|
|
self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
|
|
self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
|
|
self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|
|
|
|
|
|
class SimpleMoE(nn.Module):
|
|
"""Simple Mixture of Experts layer for benchmarking."""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
num_experts: int,
|
|
top_k: int = 2,
|
|
):
|
|
super().__init__()
|
|
self.num_experts = num_experts
|
|
self.top_k = top_k
|
|
self.router = TopKRouter(hidden_size, num_experts, top_k)
|
|
self.experts = nn.ModuleList([
|
|
ExpertFFN(hidden_size, intermediate_size) for _ in range(num_experts)
|
|
])
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
MoE forward pass.
|
|
|
|
Args:
|
|
x: Input tensor [batch, seq, hidden]
|
|
|
|
Returns:
|
|
Output tensor [batch, seq, hidden]
|
|
"""
|
|
batch_size, seq_len, hidden_size = x.shape
|
|
|
|
# Get routing decisions
|
|
expert_indices, expert_weights = self.router(x)
|
|
|
|
# Compute expert outputs (simplified - not optimized for production)
|
|
output = torch.zeros_like(x)
|
|
|
|
for k in range(self.top_k):
|
|
for e in range(self.num_experts):
|
|
# Create mask for tokens routed to this expert
|
|
mask = expert_indices[:, :, k] == e # [batch, seq]
|
|
if mask.any():
|
|
expert_input = x[mask] # [num_tokens, hidden]
|
|
expert_output = self.experts[e](expert_input)
|
|
weight = expert_weights[:, :, k][mask].unsqueeze(-1)
|
|
output[mask] += weight * expert_output
|
|
|
|
return output
|
|
|
|
|
|
def check_mps() -> bool:
|
|
"""Check if MPS backend is available."""
|
|
if not torch.backends.mps.is_available():
|
|
print("ERROR: MPS not available on this system", file=sys.stderr)
|
|
return False
|
|
return True
|
|
|
|
|
|
def benchmark_moe(
|
|
batch_size: int,
|
|
seq_len: int,
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
num_experts: int,
|
|
top_k: int = 2,
|
|
iterations: int = 500,
|
|
warmup: int = 50,
|
|
) -> MoEBenchmarkResult:
|
|
"""Run MoE benchmark."""
|
|
device = torch.device("mps")
|
|
|
|
# Create model
|
|
model = SimpleMoE(
|
|
hidden_size=hidden_size,
|
|
intermediate_size=intermediate_size,
|
|
num_experts=num_experts,
|
|
top_k=top_k,
|
|
).to(device)
|
|
model.eval()
|
|
|
|
# Create input
|
|
x = torch.randn(batch_size, seq_len, hidden_size, device=device)
|
|
|
|
# Warmup
|
|
with torch.no_grad():
|
|
for _ in range(warmup):
|
|
_ = model(x)
|
|
torch.mps.synchronize()
|
|
|
|
# Benchmark
|
|
total_times = []
|
|
routing_times = []
|
|
expert_times = []
|
|
|
|
with torch.no_grad():
|
|
for _ in range(iterations):
|
|
# Time routing
|
|
start_routing = time.perf_counter()
|
|
expert_indices, expert_weights = model.router(x)
|
|
torch.mps.synchronize()
|
|
end_routing = time.perf_counter()
|
|
|
|
# Time expert computation
|
|
start_expert = time.perf_counter()
|
|
_ = model(x)
|
|
torch.mps.synchronize()
|
|
end_expert = time.perf_counter()
|
|
|
|
routing_times.append((end_routing - start_routing) * 1000)
|
|
expert_times.append((end_expert - start_expert) * 1000 - routing_times[-1])
|
|
total_times.append((end_expert - start_routing) * 1000)
|
|
|
|
# Calculate statistics
|
|
avg_routing = statistics.mean(routing_times)
|
|
avg_expert = statistics.mean(expert_times)
|
|
avg_total = statistics.mean(total_times)
|
|
std_total = statistics.stdev(total_times) if len(total_times) > 1 else 0.0
|
|
|
|
# Calculate throughput
|
|
total_tokens = batch_size * seq_len
|
|
throughput = total_tokens / (avg_total / 1000)
|
|
|
|
# Estimate expert utilization (uniform ideal = 100%)
|
|
expert_utilization = 100.0 # Simplified - would need actual load tracking
|
|
|
|
# Memory estimate
|
|
memory_mb = None
|
|
try:
|
|
params = sum(p.numel() * p.element_size() for p in model.parameters())
|
|
activations = x.numel() * x.element_size()
|
|
memory_mb = (params + activations) / (1024 * 1024)
|
|
except Exception:
|
|
pass
|
|
|
|
name = f"MoE_E{num_experts}_K{top_k}_BS{batch_size}"
|
|
|
|
return MoEBenchmarkResult(
|
|
name=name,
|
|
batch_size=batch_size,
|
|
seq_len=seq_len,
|
|
hidden_size=hidden_size,
|
|
num_experts=num_experts,
|
|
top_k=top_k,
|
|
routing_time_ms=avg_routing,
|
|
expert_compute_time_ms=avg_expert,
|
|
total_time_ms=avg_total,
|
|
std_time_ms=std_total,
|
|
throughput_tokens_per_sec=throughput,
|
|
expert_utilization=expert_utilization,
|
|
memory_mb=memory_mb,
|
|
)
|
|
|
|
|
|
def run_benchmark_suite(iterations: int = 500, warmup: int = 50) -> MoEBenchmarkReport:
|
|
"""Run the complete MoE benchmark suite."""
|
|
import datetime
|
|
|
|
results = []
|
|
|
|
# MoE configurations
|
|
# (batch_size, seq_len, hidden_size, intermediate_size, num_experts, top_k)
|
|
scenarios = [
|
|
# Small model (GPT-2 like)
|
|
(1, 128, 768, 3072, 4, 2),
|
|
(8, 128, 768, 3072, 4, 2),
|
|
(32, 128, 768, 3072, 4, 2),
|
|
|
|
# Medium model with more experts
|
|
(1, 128, 1024, 4096, 8, 2),
|
|
(8, 128, 1024, 4096, 8, 2),
|
|
(16, 128, 1024, 4096, 8, 2),
|
|
|
|
# Large model (LLaMA-like)
|
|
(1, 128, 2048, 5504, 8, 2),
|
|
(4, 128, 2048, 5504, 8, 2),
|
|
|
|
# DeepSeek-style (many experts)
|
|
(1, 128, 1024, 2816, 16, 2),
|
|
(4, 128, 1024, 2816, 16, 2),
|
|
]
|
|
|
|
for batch_size, seq_len, hidden_size, intermediate_size, num_experts, top_k in scenarios:
|
|
try:
|
|
result = benchmark_moe(
|
|
batch_size=batch_size,
|
|
seq_len=seq_len,
|
|
hidden_size=hidden_size,
|
|
intermediate_size=intermediate_size,
|
|
num_experts=num_experts,
|
|
top_k=top_k,
|
|
iterations=iterations,
|
|
warmup=warmup,
|
|
)
|
|
results.append(asdict(result))
|
|
except Exception as e:
|
|
print(f"Warning: Benchmark failed: {e}", file=sys.stderr)
|
|
|
|
return MoEBenchmarkReport(
|
|
framework="PyTorch",
|
|
backend="MPS",
|
|
pytorch_version=torch.__version__,
|
|
python_version=platform.python_version(),
|
|
macos_version=platform.mac_ver()[0],
|
|
chip=platform.processor() or "Apple Silicon",
|
|
timestamp=datetime.datetime.now().isoformat(),
|
|
iterations=iterations,
|
|
warmup=warmup,
|
|
results=results,
|
|
)
|
|
|
|
|
|
def print_table(report: MoEBenchmarkReport):
|
|
"""Print results as formatted table."""
|
|
print("=" * 100)
|
|
print("PyTorch MPS Mixture of Experts (MoE) Benchmark")
|
|
print("=" * 100)
|
|
print(f"PyTorch version: {report.pytorch_version}")
|
|
print(f"Backend: {report.backend}")
|
|
print(f"macOS version: {report.macos_version}")
|
|
print(f"Chip: {report.chip}")
|
|
print(f"Iterations: {report.iterations}")
|
|
print()
|
|
|
|
print("-" * 100)
|
|
print(f"| {'Scenario':<25} | {'Routing':>10} | {'Expert':>10} | {'Total':>10} | {'Throughput':>18} |")
|
|
print(f"| {'':<25} | {'(ms)':>10} | {'(ms)':>10} | {'(ms)':>10} | {'(tokens/s)':>18} |")
|
|
print("-" * 100)
|
|
|
|
for r in report.results:
|
|
throughput_str = f"{r['throughput_tokens_per_sec']:.0f}"
|
|
print(f"| {r['name']:<25} | {r['routing_time_ms']:>10.4f} | {r['expert_compute_time_ms']:>10.4f} | {r['total_time_ms']:>10.4f} | {throughput_str:>18} |")
|
|
|
|
print("-" * 100)
|
|
print()
|
|
print("To compare with RustyTorch MoE:")
|
|
print(" cargo bench -p rtx-transformers --bench metal_moe_bench")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="PyTorch MPS MoE Benchmark")
|
|
parser.add_argument("--json", action="store_true", help="Output JSON format")
|
|
parser.add_argument("--iterations", type=int, default=500, help="Number of iterations")
|
|
parser.add_argument("--warmup", type=int, default=50, help="Number of warmup iterations")
|
|
args = parser.parse_args()
|
|
|
|
if not check_mps():
|
|
sys.exit(1)
|
|
|
|
report = run_benchmark_suite(iterations=args.iterations, warmup=args.warmup)
|
|
|
|
if args.json:
|
|
print(json.dumps(asdict(report), indent=2))
|
|
else:
|
|
print_table(report)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|