| """Fresh-process embedding-row experiment; never initialize Torch/MPS.""" |
|
|
| import argparse |
| import hashlib |
| import json |
| import time |
| from pathlib import Path |
|
|
| import mlx.core as mx |
| import numpy as np |
|
|
| from h3_mlx.memory import FootprintMonitor |
| from h3_mlx.runtime import configure_memory, environment |
| from h3_mlx.weights import WeightStore |
|
|
|
|
| def main(a): |
| configure_memory() |
| store = WeightStore(a.encoder) |
| key = "model.language_model.embed_tokens.weight" |
| |
| |
| ids = ([0, 1, 3, store.headers[key]["shape"][0] - 1, 3, 0, 100, 999] * 11)[:88] |
| if max(ids) >= store.headers[key]["shape"][0]: |
| raise ValueError("Unexpected vocabulary") |
| tick = time.perf_counter() |
| with FootprintMonitor(interval=0.001) as physical: |
| if a.mode == "full": |
| embedding = store.get(key) |
| value = embedding[mx.array(ids, mx.int32)] |
| mx.eval(value) |
| del embedding |
| else: |
| value = store.get(key, rows=ids) |
| mx.eval(value) |
| elapsed = time.perf_counter() - tick |
| bits = np.asarray(value.view(mx.uint16)) |
| row = { |
| "mode": a.mode, |
| "seconds": elapsed, |
| "physical": physical.result(), |
| "mlx_peak_bytes": mx.get_peak_memory(), |
| "output_sha256": hashlib.sha256(bits.tobytes()).hexdigest(), |
| "shape": list(value.shape), |
| "environment": environment(), |
| "scope": "Cold process, page-cache warm tensor read + gather only. Excludes tokenizer and full 50-layer encoder.", |
| } |
| dest = Path(a.output) |
| dest.write_text(json.dumps(row, indent=2) + "\n") |
| print(json.dumps(row), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--encoder", required=True) |
| p.add_argument("--mode", choices=["full", "rows"], required=True) |
| p.add_argument("--output", required=True) |
| main(p.parse_args()) |
|
|
