File size: 2,051 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 | """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"
# Same number of padded/effective token rows as product, with duplicates,
# padding, and vocabulary endpoints. No tokenizer/download dependency.
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())
|