| """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): |
| |
| 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() |
|
|
