| """Apple GPU/CPU alternatives through PyTorch, including MPS Inductor fusion.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import json |
| import platform |
| import statistics |
| import time |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
| from safetensors import safe_open |
|
|
|
|
| def load(folder, key): |
| for p in sorted(Path(folder).glob("*.safetensors")): |
| with safe_open(p, framework="pt", device="cpu") as f: |
| if key in f.keys(): |
| return f.get_tensor(key) |
| raise KeyError(key) |
|
|
|
|
| def measure(fn, device, repeats): |
| def sync(): |
| if device == "mps": |
| torch.mps.synchronize() |
|
|
| for _ in range(2): |
| fn() |
| sync() |
| samples = [] |
| for _ in range(repeats): |
| sync() |
| tick = time.perf_counter() |
| fn() |
| sync() |
| samples.append(1000 * (time.perf_counter() - tick)) |
| return {"median_ms": statistics.median(samples), "samples_ms": samples} |
|
|
|
|
| def main(a): |
| torch.set_num_threads(10) |
| torch.manual_seed(731) |
| root = Path(a.model) / "transformer" |
| out = Path(a.output) |
| out.parent.mkdir(parents=True, exist_ok=True) |
| report = { |
| "platform": platform.platform(), |
| "torch": torch.__version__, |
| "threads": 10, |
| "repeats": a.repeats, |
| "warmup": 2, |
| "scope": "Actual H3 block-0 weights, synthetic activations; model loading and dtype conversion excluded.", |
| "trials": [], |
| } |
| with torch.inference_mode(): |
| for name in ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2"): |
| w = load(root, "blocks.0." + name + ".weight") |
| x = torch.randn(a.tokens, w.shape[1]).bfloat16() |
| |
| ref = F.linear(x.float(), w.float()) |
| for device, dtype in ( |
| ("cpu", torch.bfloat16), |
| ("cpu", torch.float32), |
| ("mps", torch.bfloat16), |
| ("mps", torch.float16), |
| ("mps", torch.float32), |
| ): |
| row = { |
| "operator": name, |
| "device": device, |
| "dtype": str(dtype), |
| "shape": [a.tokens, w.shape[1], w.shape[0]], |
| } |
| try: |
| xx, ww = x.to(device, dtype), w.to(device, dtype) |
| row.update( |
| measure( |
| lambda xx=xx, ww=ww: F.linear(xx, ww), device, a.repeats |
| ) |
| ) |
| y = F.linear(xx, ww).cpu().float() |
| row["relative_rmse"] = torch.sqrt( |
| torch.mean((ref - y) ** 2) / torch.mean(ref**2) |
| ).item() |
| del xx, ww, y |
| if device == "mps": |
| torch.mps.empty_cache() |
| except Exception as exc: |
| row["error"] = f"{type(exc).__name__}: {exc}" |
| report["trials"].append(row) |
| out.write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps(row), flush=True) |
| del x, w, ref |
| gc.collect() |
| w1 = load(root, "blocks.0.mlp.fc1.weight").to("mps") |
| w2 = load(root, "blocks.0.mlp.fc2.weight").to("mps") |
| x = torch.randn(a.tokens, 5376, device="mps").bfloat16() |
|
|
| def mlp(x): |
| gate, up = F.linear(x, w1).chunk(2, dim=-1) |
| return F.linear(F.silu(gate) * up, w2) |
|
|
| reference = mlp(x) |
| for compiled in (False, True): |
| row = { |
| "operator": "mlp", |
| "compiled": compiled, |
| "device": "mps", |
| "dtype": "torch.bfloat16", |
| } |
| try: |
| fn = torch.compile(mlp) if compiled else mlp |
| row.update(measure(lambda fn=fn: fn(x), "mps", a.repeats)) |
| y = fn(x) |
| row["equal_fraction"] = (reference == y).float().mean().item() |
| except Exception as exc: |
| row["error"] = f"{type(exc).__name__}: {exc}" |
| report["trials"].append(row) |
| out.write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps(row), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--model", required=True) |
| p.add_argument("--tokens", type=int, default=4181) |
| p.add_argument("--repeats", type=int, default=5) |
| p.add_argument("--output", default="reports/optimization-v2/torch.json") |
| main(p.parse_args()) |
|
|
