File size: 4,664 Bytes
3da88e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """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()
# A single FP32 CPU reference distinguishes accumulation precision.
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())
|