prain: slot cache with eviction, scan, prefetch and CPU policies
This commit is contained in:
+245
@@ -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": {"<layer>": {"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)
|
||||
@@ -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))
|
||||
Reference in New Issue
Block a user