94 lines
4.2 KiB
Python
94 lines
4.2 KiB
Python
import os
|
|
import unittest
|
|
|
|
from prain.config import LayerPlacement, Miss, Policy, Prefetch
|
|
from prain.sim import (Cost, read_forwards, simulate, synthetic_trace, tune, write_trace)
|
|
from tests.helpers import TmpCase
|
|
|
|
SMALL = dict(layers=4, experts=32, top_k=4, forwards=120, prompt_every=60, prompt_tokens=32)
|
|
|
|
|
|
def forwards_of(records):
|
|
forwards, last = [], None
|
|
for r in records:
|
|
if last is None or r["l"] <= last:
|
|
forwards.append({})
|
|
forwards[-1][r["l"]] = r
|
|
last = r["l"]
|
|
return forwards
|
|
|
|
|
|
def all_ram(slots, layers=4, experts=32):
|
|
return {layer: LayerPlacement((), tuple(range(experts)), (), slots) for layer in range(layers)}
|
|
|
|
|
|
class Traces(TmpCase):
|
|
def test_round_trip_and_determinism(self):
|
|
records = synthetic_trace(**SMALL)
|
|
self.assertEqual(records, synthetic_trace(**SMALL))
|
|
path = os.path.join(self.dir, "t.jsonl.gz")
|
|
write_trace(path, records[:8])
|
|
write_trace(path, records[8:], mode="at") # runs append to the same trace
|
|
forwards = read_forwards(path)
|
|
self.assertEqual(len(forwards), 120)
|
|
self.assertEqual(forwards[0][0], records[0])
|
|
self.assertEqual(forwards[0][0]["t"], 32)
|
|
self.assertEqual(sorted(forwards[1][1]["p"]), ["1", "2"])
|
|
self.assertNotIn("p", forwards[1][3]) # nothing after the last layer
|
|
|
|
def test_copy_cost(self):
|
|
self.assertAlmostEqual(Cost().copy_ms(False), 0.991, places=3)
|
|
self.assertAlmostEqual(Cost().copy_ms(True), 4.955, places=3)
|
|
|
|
|
|
class Simulate(unittest.TestCase):
|
|
def setUp(self):
|
|
self.forwards = forwards_of(synthetic_trace(**SMALL))
|
|
|
|
def run_sim(self, policy=Policy(), prefetch=Prefetch(), miss=Miss(), slots=3, **kw):
|
|
return simulate(self.forwards, all_ram(slots), policy, prefetch, miss, Cost(), **kw)
|
|
|
|
def test_belady_bounds_every_policy(self):
|
|
bound = self.run_sim(belady=True).hit_rate
|
|
for eviction in ("lru", "lfu", "gdsf"):
|
|
self.assertLessEqual(self.run_sim(Policy(eviction=eviction)).hit_rate, bound)
|
|
|
|
def test_cpu_path_beats_copies_when_the_cpu_is_faster_than_pcie(self):
|
|
copy = self.run_sim().ms_per_decode
|
|
self.assertLess(self.run_sim(miss=Miss(mode="auto")).ms_per_decode, copy)
|
|
|
|
def test_prefetch_helps_when_predictions_are_right(self):
|
|
none = self.run_sim().hit_rate
|
|
good = self.run_sim(prefetch=Prefetch(mode="pregate", budget=4))
|
|
self.assertGreater(good.hit_rate, none)
|
|
self.assertGreater(good.counters["prefetch_hits"], good.counters["prefetch_wasted"])
|
|
|
|
def test_pregate_needs_predictions(self):
|
|
stripped = [{layer: {k: v for k, v in r.items() if k != "p"} for layer, r in f.items()}
|
|
for f in self.forwards]
|
|
with self.assertRaisesRegex(ValueError, "needs a trace with predictions"):
|
|
simulate(stripped, all_ram(3), Policy(), Prefetch(mode="pregate"), Miss(), Cost())
|
|
|
|
|
|
class Tune(unittest.TestCase):
|
|
def test_tune_beats_the_defaults_within_the_budget(self):
|
|
forwards = forwards_of(synthetic_trace(**SMALL))
|
|
cost = Cost(expert_mib=64.0) # 1 GiB = 16 slots = 4 per layer
|
|
best, ram, per, final, bound, rows = tune(forwards, 32, 1.0, 1.0, cost)
|
|
self.assertLessEqual(sum(per.values()), 16)
|
|
self.assertEqual(sum(len(r) for r in ram.values()), 16) # 1 GiB of RAM tier
|
|
default = next(r for v, r in rows if v == (Policy(), Prefetch(), Miss()))
|
|
self.assertLessEqual(final.total_ms, default.total_ms)
|
|
self.assertEqual(rows[0][1].total_ms, min(r.total_ms for _, r in rows))
|
|
self.assertGreaterEqual(bound.hit_rate, 0.0)
|
|
|
|
def test_tune_searches_the_auto_threshold_and_keeps_the_cpu_threads(self):
|
|
forwards = forwards_of(synthetic_trace(**SMALL))
|
|
cost = Cost(expert_mib=64.0, h2d_gbps=2.5, cpu_ms=0.5) # eGPU: a copy is ~50 CPU routes
|
|
base = (Policy(), Prefetch(), Miss(cpu_threads=16))
|
|
best, ram, per, final, bound, rows = tune(forwards, 32, 1.0, 1.0, cost, base=base)
|
|
self.assertEqual(best[2].cpu_threads, 16)
|
|
self.assertNotEqual(best[2].mode, "copy")
|
|
tried = {(v[2].mode, v[2].cpu_max_tokens) for v, _ in rows}
|
|
self.assertLessEqual({("copy", 1), ("cpu", 1), ("auto", 1), ("auto", 8), ("auto", 16)}, tried)
|