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