v0.6.0: stream MoE experts from the SSD (LRU cache, direct I/O) when the model does not fit; expert memory setting
3a7bf50 verified | """SSD read speed for expert streaming: random experts from experts.bin, direct (unbuffered) reads, 1 to 16 threads. | |
| usage: python bench_ssd.py MODEL_DIR [SECONDS_PER_SETTING]""" | |
| import json, os, random, sys, time | |
| from concurrent.futures import ThreadPoolExecutor | |
| import numpy as np | |
| from onw.xio import DirectReader | |
| def main(): | |
| d = sys.argv[1] | |
| secs = float(sys.argv[2]) if len(sys.argv) > 2 else 4 | |
| idx = json.load(open(os.path.join(d, "experts.bin.json"))) | |
| ents = [(w[0], w[1] + s[1]) for ex in idx["layers"].values() for (w, s) in ex.values()] # w and s are adjacent | |
| size = max(n for _, n in ents) | |
| r = DirectReader(os.path.join(d, "experts.bin"), size) | |
| print(f"{len(ents)} experts, {size / 2**20:.2f} MB each, direct I/O: {r.direct}") | |
| bufs = {} | |
| def one(e): | |
| import threading | |
| b = bufs.setdefault(threading.get_ident(), np.empty(size, np.uint8)) | |
| r.read_into(b[:e[1]], e[0], e[1]) | |
| return e[1] | |
| for threads in (1, 2, 4, 8, 16): | |
| with ThreadPoolExecutor(threads) as pool: | |
| t0, total, lat = time.time(), 0, [] | |
| while time.time() - t0 < secs: | |
| batch = random.sample(ents, threads * 4) | |
| t1 = time.perf_counter() | |
| total += sum(pool.map(one, batch)) | |
| lat.append((time.perf_counter() - t1) / 4) | |
| dt = time.time() - t0 | |
| print(f"threads {threads:2d}: {total / dt / 2**30:5.2f} GB/s, {total / dt / size:6.0f} experts/s, " | |
| f"{np.median(lat) * 1000:5.2f} ms per batch of {threads}", flush=True) | |
| if __name__ == "__main__": | |
| main() | |