File size: 5,586 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 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """Actual H3 projection dimensions, M5 layouts, and new quantization formats.
Run as a module from the repository root. Each result is persisted immediately.
The inputs are synthetic; weights come from the original checkpoint.
"""
from __future__ import annotations
import argparse
import gc
import json
from pathlib import Path
import mlx.core as mx
from h3_mlx.benchmark import error, measure
from h3_mlx.ops import swiglu
from h3_mlx.runtime import configure_memory, environment
from h3_mlx.weights import Linear, WeightStore
def run(args):
configure_memory()
mx.random.seed(731)
dest = Path(args.output)
dest.parent.mkdir(parents=True, exist_ok=True)
report = {
"environment": environment(),
"warmup": 2,
"repeats": args.repeats,
"scope": "Original block-0 weights; synthetic BF16 inputs. Operator latency only, excludes conversion and loading.",
"trials": [],
}
def record(op, variant, fn, ref, **extra):
try:
timing = measure(fn, repeats=args.repeats)
row = {
"operator": op,
"variant": variant,
**extra,
**timing,
**error(ref, fn()),
}
except Exception as exc:
row = {
"operator": op,
"variant": variant,
**extra,
"error": f"{type(exc).__name__}: {exc}",
}
report["trials"].append(row)
dest.write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps(row), flush=True)
store = WeightStore(Path(args.model) / "transformer")
for name in ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2"):
layer = Linear.load(
store,
"blocks.0." + name,
grouped_heads=(56, 128) if name == "attn.qkv_proj" else None,
)
w, bias = layer.weight, layer.bias
x = mx.random.normal((args.tokens, w.shape[1])).astype(mx.bfloat16)
mx.eval(x, w)
ref = layer(x)
mx.eval(ref)
shape = [args.tokens, w.shape[1], w.shape[0]]
record(name, "bf16_nt", lambda layer=layer, x=x: layer(x), ref, shape=shape)
wt = mx.contiguous(w.T)
mx.eval(wt)
fn = (
(lambda x=x, wt=wt: x @ wt)
if bias is None
else (lambda bias=bias, x=x, wt=wt: mx.addmm(bias, x, wt))
)
record(name, "bf16_nn", fn, ref, shape=shape)
del wt, fn
for pad in (64, 128, 256):
n = ((args.tokens + pad - 1) // pad) * pad
padded = mx.pad(x, ((0, n - args.tokens), (0, 0)))
mx.eval(padded)
record(
name,
f"bf16_pad{pad}",
lambda layer=layer, padded=padded: layer(padded)[: args.tokens],
ref,
shape=shape,
padded_tokens=n,
)
del padded
half, halfw = x.astype(mx.float16), w.astype(mx.float16)
halfb = None if bias is None else bias.astype(mx.float16)
mx.eval(half, halfw, halfb)
fn = (
(lambda half=half, halfw=halfw: (half @ halfw.T).astype(mx.bfloat16))
if bias is None
else (
lambda halfb=halfb, half=half, halfw=halfw: mx.addmm(
halfb, half, halfw.T
).astype(mx.bfloat16)
)
)
record(name, "fp16_round_bf16", fn, ref, shape=shape)
del half, halfw, halfb, fn
for mode, bits, group in (
("mxfp8", 8, 32),
("mxfp4", 4, 32),
("nvfp4", 4, 16),
("affine", 8, 32),
("affine", 6, 32),
("affine", 4, 32),
):
try:
packed = mx.quantize(w, mode=mode, bits=bits, group_size=group)
mx.eval(packed)
def quant(
x=x, packed=packed, mode=mode, bits=bits, group=group, bias=bias
):
y = mx.quantized_matmul(
x, *packed, mode=mode, bits=bits, group_size=group
)
return y if bias is None else y + bias
record(
name,
f"{mode}{bits}_g{group}",
quant,
ref,
shape=shape,
weight_gib=sum(v.nbytes for v in packed) / 1024**3,
)
del packed, quant
except Exception as exc:
record(
name,
f"{mode}{bits}_g{group}",
lambda exc=exc: (_ for _ in ()).throw(exc),
ref,
)
del x, w, bias, layer, ref
gc.collect()
mx.clear_cache()
a = Linear.load(store, "blocks.0.mlp.fc1")
b = Linear.load(store, "blocks.0.mlp.fc2")
x = mx.random.normal((args.tokens, 5376)).astype(mx.bfloat16)
def mlp(x):
return b(swiglu(a(x)))
ref = mlp(x)
mx.eval(ref)
record("mlp", "eager", lambda: mlp(x), ref)
compiled = mx.compile(mlp)
record("mlp", "compiled", lambda: compiled(x), ref)
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/operators.json")
run(p.parse_args())
|