File size: 2,474 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 | """Compare the original tensor reader with lazy native MLX safetensors I/O."""
from __future__ import annotations
import argparse
import gc
import hashlib
import json
import statistics
import time
from pathlib import Path
import mlx.core as mx
import numpy as np
from h3_mlx.runtime import configure_memory, environment
from h3_mlx.weights import WeightStore
def main(a):
configure_memory(wired_gb=4)
store = WeightStore(Path(a.model) / "transformer", reader="memmap")
report = {
"environment": environment(),
"scope": "One original tensor at a time. OS page cache is not flushed; includes read/copy and Metal synchronization.",
"trials": [],
}
for key in (
"blocks.0.attn.qkv_proj.weight",
"blocks.25.mlp.fc1.weight",
"blocks.49.mlp.fc2.weight",
):
native = mx.load(str(store.index[key]))
for mode in ("memmap_copy", "mlx_lazy_load"):
samples = []
for _ in range(a.repeats):
gc.collect()
mx.clear_cache()
mx.synchronize()
tick = time.perf_counter()
if mode == "memmap_copy":
value = store.get(key)
else:
# Reload dictionary so every sample has an unevaluated Load node.
value = mx.load(str(store.index[key]))[key]
mx.eval(value)
mx.synchronize()
samples.append(time.perf_counter() - tick)
bits = hashlib.sha256(
np.asarray(value.view(mx.uint16)).tobytes()
).hexdigest()
del value
report["trials"].append(
{
"key": key,
"mode": mode,
"median_s": statistics.median(samples),
"samples_s": samples,
"uint16_sha256": bits,
}
)
print(json.dumps(report["trials"][-1]), flush=True)
del native
dest = Path(a.output)
dest.parent.mkdir(parents=True, exist_ok=True)
dest.write_text(json.dumps(report, indent=2) + "\n")
if __name__ == "__main__":
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--model", required=True)
p.add_argument("--repeats", type=int, default=3)
p.add_argument("--output", default="reports/optimization-v2/loading.json")
main(p.parse_args())
|