diff --git a/prain/policy.py b/prain/policy.py new file mode 100644 index 0000000..c823156 --- /dev/null +++ b/prain/policy.py @@ -0,0 +1,245 @@ +"""Per-layer VRAM slot cache: which slot holds which expert, what to evict, what to copy +first, how scans and prefetched experts are admitted. Pure Python, no torch.""" + +import json +import math +import os +from bisect import bisect_right + +from prain.config import Policy + + +class Lru: + """Evict the expert used longest ago.""" + + def __init__(self, num_experts, policy, cost): + self.last = [0.0] * num_experts + + def score(self, e, now): + return self.last[e] + + def hit(self, e, now, heat): + self.last[e] = now + + def admit(self, e, now, heat, low): + self.last[e] = -math.inf if low else now + + def evicted(self, e): + pass + + +class Lfu: + """Evict the expert with the lowest exponentially decayed heat. Scores survive eviction, + so an expert that comes back keeps its history.""" + + def __init__(self, num_experts, policy, cost): + self.decay = 0.5 ** (1 / policy.lfu_half_life) + self.value = [0.0] * num_experts + self.when = [0] * num_experts + + def score(self, e, now): + return self.value[e] * self.decay ** (now - self.when[e]) + + def hit(self, e, now, heat): + self.value[e] = self.score(e, now) + heat + self.when[e] = now + + def admit(self, e, now, heat, low): + if not low: # a low-priority admit keeps its history: never-seen experts score 0 + self.hit(e, now, heat) + + def evicted(self, e): + pass + + +class Gdsf: + """GreedyDual-Size-Frequency: score = floor + frequency x refetch cost. The floor rises + to each victim's score, so experts that stopped being used age out. cost[e] is 1 for a + RAM expert and policy.disk_cost for a disk one, so slow-to-refetch experts stay longer.""" + + def __init__(self, num_experts, policy, cost): + self.cost = cost + self.floor = 0.0 + self.freq = [0.0] * num_experts + self.value = [0.0] * num_experts + + def score(self, e, now): + return self.value[e] + + def hit(self, e, now, heat): + self.freq[e] += heat + self.value[e] = self.floor + self.freq[e] * self.cost[e] + + def admit(self, e, now, heat, low): + self.freq[e] = 0.0 if low else heat + self.value[e] = self.floor + self.freq[e] * self.cost[e] + + def evicted(self, e): + self.floor = max(self.floor, self.value[e]) + + +class Belady: + """The offline optimum, for the simulator only: evict the expert used furthest in the + future. future[e] lists the forward numbers that route to e, ascending.""" + + def __init__(self, future): + self.future = future + + def score(self, e, now): + uses = self.future.get(e, ()) + i = bisect_right(uses, now) + return -uses[i] if i < len(uses) else -math.inf + + def hit(self, e, now, heat): + pass + + def admit(self, e, now, heat, low): + pass + + def evicted(self, e): + pass + + +POLICIES = {"lru": Lru, "lfu": Lfu, "gdsf": Gdsf} + + +class SlotCache: + """Slots [0, len(pinned)) hold the pinned experts for the whole run. The cache_slots + slots after them hold experts from RAM or disk, managed by an eviction policy.""" + + def __init__(self, pinned, num_experts, cache_slots, policy=Policy(), cost=None, evictor=None): + self.num_experts = num_experts + self.cache_slots = cache_slots + self.policy = policy + self.pinned = {e: slot for slot, e in enumerate(pinned)} + self.slots = {} # cached expert -> slot + first = len(self.pinned) + self.free = list(range(first + cache_slots - 1, first - 1, -1)) # pop() -> lowest slot + cost = cost or [1.0] * num_experts + self.evictor = evictor or POLICIES[policy.eviction](num_experts, policy, cost) + self.now = 0 # forwards planned so far + self.recent = set() # experts routed to in the last forward + self.prefetched = set() # prefetched experts that were not used yet + self.uses = [0] * num_experts # forwards that routed to each expert + self.hits = 0 # routed expert already in VRAM + self.misses = 0 # routed expert copied in + self.cpu = 0 # routed expert computed on the CPU instead + self.prefetches = 0 # experts copied in ahead of time (prefetch or admit-behind) + self.prefetch_hits = 0 # ... that were used before being evicted + self.prefetch_wasted = 0 # ... that were evicted unused + + def choose_cpu(self, counts, miss): + """Experts the CPU should compute: counts maps expert -> routed tokens.""" + if miss.mode == "copy": + return set() + away = (e for e in counts if e not in self.pinned and e not in self.slots) + if miss.mode == "cpu": + return set(away) + return {e for e in away if counts[e] <= miss.cpu_max_tokens} + + def plan(self, heat, tokens=1, cpu=()): + """Split one forward pass into rounds. + + heat: {expert: heat > 0} for every expert routed to in this pass. + tokens: tokens in the pass; a pass with >= policy.scan_tokens tokens is a scan. + cpu: routed experts the CPU computes; they get no slot. + Returns [(slot_of, loads), ...]: slot_of maps each expert computed in that round to + its slot; loads lists the (expert, slot) copies to do before the round. Experts in + VRAM run in round 0. The others follow coldest first, so the hottest are copied last + and are the ones left in VRAM afterwards. + """ + self.now += 1 + now, scan = self.now, tokens >= self.policy.scan_tokens + norm = {e: h / tokens for e, h in heat.items()} # per token, so scans don't dominate + for e in heat: + self.uses[e] += 1 + self.recent = set(heat) + pinned = [e for e in heat if e in self.pinned] + resident = [e for e in heat if e in self.slots] + missing = sorted((e for e in heat if e not in self.pinned and e not in self.slots + and e not in cpu), key=lambda e: (heat[e], e)) + self.hits += len(pinned) + len(resident) + self.misses += len(missing) + self.cpu += len(cpu) + for e in resident: + self.evictor.hit(e, now, norm[e]) + if e in self.prefetched: + self.prefetched.discard(e) + self.prefetch_hits += 1 + if missing and not self.cache_slots: + raise RuntimeError(f"experts {missing} are not in VRAM and there are no cache slots") + protected = set() + if scan and self.policy.scan_keep: + ranked = sorted(self.slots, key=lambda e: (self.evictor.score(e, now), e)) + protected = set(ranked[-self.policy.scan_keep:]) + busy = set(resident) # experts computed in the round being built + slot_of = {e: self.pinned[e] for e in pinned} | {e: self.slots[e] for e in resident} + rounds, i = [], 0 + while True: + room = len(self.free) + sum(1 for e in self.slots if e not in busy and e not in protected) + group = missing[i:i + room] + if i < len(missing) and not group and not busy: + raise RuntimeError("no evictable cache slot: lower policy.scan_keep") + i += len(group) + loads = [] + for e in group: + slot = self._take_slot(now, busy | protected) + self.slots[e] = slot + self.evictor.admit(e, now, norm[e], low=scan) + busy.add(e) + loads.append((e, slot)) + slot_of.update(loads) + rounds.append((slot_of, loads)) + if i >= len(missing): + break + slot_of, busy = {}, set() + if scan and self.policy.promote: + for e in missing[-self.policy.promote:]: + if e in self.slots: + self.evictor.hit(e, now, norm[e]) + return rounds + + def prefetch(self, predicted, budget): + """Copy predicted experts in ahead of time, at low priority: they are evicted first + unless a forward uses them. Never evicts an expert routed to in the last forward. + predicted: [(expert, probability)], most likely first. Returns the (expert, slot) loads.""" + loads, keep = [], set(self.recent) + for e, _ in predicted: + if len(loads) >= budget: + break + if e in self.pinned or e in self.slots: + continue + if not self.free and all(r in keep for r in self.slots): + break + slot = self._take_slot(self.now, keep) + self.slots[e] = slot + self.evictor.admit(e, self.now, 0.0, low=True) + self.prefetched.add(e) + keep.add(e) + loads.append((e, slot)) + self.prefetches += len(loads) + return loads + + def _take_slot(self, now, keep): + if self.free: + return self.free.pop() + victim = min((e for e in self.slots if e not in keep), + key=lambda e: (self.evictor.score(e, now), e)) + self.evictor.evicted(victim) + if victim in self.prefetched: + self.prefetched.discard(victim) + self.prefetch_wasted += 1 + return self.slots.pop(victim) + + +COUNTERS = ("hits", "misses", "cpu", "prefetches", "prefetch_hits", "prefetch_wasted") + + +def write_stats(path, caches): + """Atomically write {"layers": {"": {"uses": [...], "hits": n, ...}}}.""" + data = {"layers": {str(layer): {"uses": c.uses, **{k: getattr(c, k) for k in COUNTERS}} + for layer, c in sorted(caches.items())}} + tmp = f"{path}.tmp" + with open(tmp, "w") as f: + json.dump(data, f) + os.replace(tmp, path) diff --git a/tests/test_policy.py b/tests/test_policy.py new file mode 100644 index 0000000..414ff4b --- /dev/null +++ b/tests/test_policy.py @@ -0,0 +1,121 @@ +import json +import math +import os +import random +import unittest + +from prain.config import Miss, Policy +from prain.policy import SlotCache, write_stats +from tests.helpers import TmpCase + + +def lru(**kw): + return Policy(eviction="lru", **kw) + + +class Eviction(unittest.TestCase): + def test_pinned_experts_stay_and_lru_evicts_the_oldest(self): + cache = SlotCache(pinned=[3, 5], num_experts=8, cache_slots=2, policy=lru()) + self.assertEqual(cache.plan({0: 1, 3: 1}), [({0: 2, 3: 0}, [(0, 2)])]) + self.assertEqual(cache.plan({0: 1, 1: 1}), [({0: 2, 1: 3}, [(1, 3)])]) + self.assertEqual(cache.plan({1: 1, 2: 1}), [({1: 3, 2: 2}, [(2, 2)])]) # 0 was oldest + self.assertEqual((cache.hits, cache.misses), (3, 3)) + self.assertEqual(cache.uses, [2, 2, 1, 1, 0, 0, 0, 0]) + + def test_lfu_keeps_the_frequent_expert_lru_the_recent_one(self): + for policy, survivor in ((lru(), 1), (Policy(eviction="lfu"), 0)): + cache = SlotCache([], 3, 2, policy) + for e in (0, 0, 0, 1, 2): + cache.plan({e: 1}) + self.assertIn(survivor, cache.slots, policy.eviction) + + def test_gdsf_keeps_the_expensive_disk_expert(self): + cost = [1.0, 4.0, 1.0] # expert 1 lives on disk + for eviction, survivor in (("lfu", 0), ("gdsf", 1)): + cache = SlotCache([], 3, 2, Policy(eviction=eviction), cost) + for e in (0, 0, 1, 2): + cache.plan({e: 1}) + self.assertIn(survivor, cache.slots, eviction) + + +class Rounds(unittest.TestCase): + def test_hottest_experts_are_copied_last_and_stay(self): + cache = SlotCache([], 10, 2, lru()) + rounds = cache.plan({0: 1, 1: 9, 2: 3, 3: 8, 4: 2, 5: 1}) + self.assertEqual([sorted(slot_of) for slot_of, _ in rounds], [[0, 5], [2, 4], [1, 3]]) + self.assertEqual(set(cache.slots), {1, 3}) + + def test_scan_keeps_the_protected_expert_and_promotes_the_hottest(self): + cache = SlotCache([], 6, 3, lru(scan_tokens=4, scan_keep=1, promote=1)) + cache.plan({0: 1}) + cache.plan({1: 1}) + rounds = cache.plan({2: 1, 3: 5, 4: 2, 5: 3}, tokens=8) + self.assertEqual(rounds, [({2: 2, 4: 0}, [(2, 2), (4, 0)]), + ({5: 2, 3: 0}, [(5, 2), (3, 0)])]) + self.assertEqual(set(cache.slots), {1, 3, 5}) # 1 was protected + self.assertEqual(cache.evictor.last[3], cache.now) # promoted + self.assertEqual(cache.evictor.last[5], -math.inf) # admitted at low priority + + def test_cpu_experts_get_no_slot(self): + cache = SlotCache([0], 6, 2, lru()) + cache.plan({1: 1}) + counts = {0: 1, 1: 1, 2: 1, 3: 4} + self.assertEqual(cache.choose_cpu(counts, Miss()), set()) + self.assertEqual(cache.choose_cpu(counts, Miss(mode="cpu")), {2, 3}) + cpu = cache.choose_cpu(counts, Miss(mode="auto", cpu_max_tokens=2)) + self.assertEqual(cpu, {2}) + self.assertEqual(cache.plan(counts, 4, cpu), [({0: 0, 1: 1, 3: 2}, [(3, 2)])]) + self.assertEqual((cache.cpu, cache.misses), (1, 2)) + + def test_random_plans_always_compute_every_expert_from_the_right_slot(self): + rng = random.Random(1) + for _ in range(300): + n = rng.randint(2, 12) + pinned = rng.sample(range(n), rng.randint(0, n // 2)) + slots = rng.randint(1, 4) + policy = Policy(eviction=rng.choice(["lru", "lfu", "gdsf"]), + scan_tokens=rng.randint(1, 6), scan_keep=rng.randint(0, slots - 1), + promote=rng.randint(0, 3)) + cache = SlotCache(pinned, n, slots, policy, [rng.choice([1.0, 4.0]) for _ in range(n)]) + content = {slot: e for e, slot in cache.pinned.items()} # what each slot really holds + for _ in range(12): + if rng.random() < 0.3: + guesses = [(e, 0.5) for e in rng.sample(range(n), rng.randint(1, n))] + recent = set(cache.recent) + for e, slot in cache.prefetch(guesses, rng.randint(0, 3)): + self.assertNotIn(content.get(slot), recent - {None}) + content[slot] = e + routed = rng.sample(range(n), rng.randint(1, n)) + heat = {e: rng.randint(1, 9) for e in routed} + cpu = cache.choose_cpu(heat, Miss(mode=rng.choice(["copy", "auto"]), cpu_max_tokens=3)) + computed = [] + for slot_of, loads in cache.plan(heat, rng.randint(1, 8), cpu): + self.assertEqual(len(set(slot_of.values())), len(slot_of)) + for e, slot in loads: + self.assertEqual(slot_of[e], slot) + content[slot] = e + for e, slot in slot_of.items(): + self.assertEqual(content[slot], e) + computed += slot_of + self.assertEqual(sorted(computed + sorted(cpu)), sorted(routed)) + held = list(cache.slots.values()) + self.assertEqual(len(set(held)), len(held)) + self.assertEqual(len(held) + len(cache.free), slots) + + +class Prefetch(TmpCase): + def test_prefetch_is_low_priority_and_never_evicts_the_last_forward(self): + cache = SlotCache([], 4, 2, Policy(eviction="lfu")) + cache.plan({0: 1}) + self.assertEqual(cache.prefetch([(1, 0.5), (2, 0.4)], budget=2), [(1, 1)]) # 0 is recent + cache.plan({1: 1}) # the prefetch paid off + self.assertEqual(cache.prefetch([(3, 0.9)], budget=1), [(3, 0)]) # evicts 0, not 1 + cache.plan({2: 1}) # 3 was never used: it goes first + self.assertEqual((cache.prefetches, cache.prefetch_hits, cache.prefetch_wasted), (2, 1, 1)) + self.assertEqual(set(cache.slots), {1, 2}) + path = os.path.join(self.dir, "stats.json") + write_stats(path, {7: cache}) + with open(path) as f: + stats = json.load(f)["layers"]["7"] + self.assertEqual((stats["uses"], stats["prefetch_wasted"], stats["misses"]), + ([1, 1, 1, 0], 1, 2))