File size: 3,311 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 | """Check schedule caching against individual GEMVs and PyTorch on actual weights."""
from __future__ import annotations
import argparse
import json
import math
from pathlib import Path
import mlx.core as mx
import numpy as np
import torch
import torch.nn.functional as F
from h3_mlx.benchmark import error
from h3_mlx.config import load_sigmas
from h3_mlx.ops import silu
from h3_mlx.runtime import configure_memory
from h3_mlx.weights import Linear, WeightStore
def tensor(x):
# Bit-preserve BF16 rather than round-tripping it through FP32.
if x.dtype == mx.bfloat16:
return torch.from_numpy(np.asarray(x.view(mx.uint16)).copy()).view(
torch.bfloat16
)
return torch.from_numpy(np.asarray(x).copy())
def main():
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--model", required=True)
p.add_argument("--layers", type=int, nargs="+", default=[0, 25, 49])
p.add_argument("--output", default="reports/modulation-audit.json")
a = p.parse_args()
configure_memory()
torch.set_num_threads(6)
path = Path(a.model)
store = WeightStore(path / "transformer")
sigmas = load_sigmas(path)
times = mx.array([1 - s for s in sigmas[:-1]], mx.float32)
freq = mx.exp(-math.log(10000) * mx.arange(128, dtype=mx.float32) / 128)
args = times[:, None] * freq
features = mx.concatenate((mx.cos(args), mx.sin(args)), axis=-1)
first = Linear.load(store, "time_embedder.proj_in")
last = Linear.load(store, "time_embedder.proj_out")
t = silu(last(silu(first(features)))).astype(mx.bfloat16)
mx.eval(t)
tt = torch.tensor([1 - s for s in sigmas[:-1]], dtype=torch.float32)
tf = torch.exp(-math.log(10000) * torch.arange(128, dtype=torch.float32) / 128)
tf = tt[:, None] * tf
tf = torch.cat((tf.cos(), tf.sin()), -1)
tf = F.silu(
F.linear(
F.silu(F.linear(tf, tensor(first.weight), tensor(first.bias))),
tensor(last.weight),
tensor(last.bias),
)
).bfloat16()
records = []
for index in a.layers:
linear = Linear.load(store, f"blocks.{index}.adaln_proj.linear")
batched = linear(t)
individual = mx.concatenate([linear(t[i : i + 1]) for i in range(6)])
mx.eval(batched, individual)
tw, tb = tensor(linear.weight), tensor(linear.bias)
with torch.inference_mode():
ref = torch.cat([F.linear(tf[i : i + 1], tw, tb) for i in range(6)])
reference = mx.array(ref.float().numpy()).astype(mx.bfloat16)
row = {
"layer": index,
"batched_vs_individual_mlx": error(individual, batched),
"mlx_vs_native_torch_cpu": error(reference, batched),
}
print(json.dumps(row), flush=True)
records.append(row)
del linear, tw, tb, batched, individual, reference, ref
mx.clear_cache()
Path(a.output).write_text(
json.dumps(
{
"source": store.fingerprint(),
"layers": records,
"scope": "Actual original weights; original one-row-per-step BF16 projections, FP32 timestep MLP; CPU PyTorch is not a CUDA image reference",
},
indent=2,
)
+ "\n"
)
if __name__ == "__main__":
main()
|