prain: traces, simulator and tuner
This commit is contained in:
+311
@@ -0,0 +1,311 @@
|
||||
"""Routing traces, a synthetic trace generator, a simulator with a simple cost model, and a
|
||||
tuner that searches policy variants and per-layer slot counts. Pure Python, no torch.
|
||||
|
||||
A trace is gzip JSON lines, one per MoE layer per forward, in execution order:
|
||||
{"l": layer, "t": tokens, "e": [experts], "c": [tokens per expert], "w": [router weight
|
||||
mass], "r": [recency-weighted mass], "p": {"1": [[expert, prob], ...], "2": [...]}}.
|
||||
"p" holds predictions for layers l+1 and l+2 (pregate) and is optional. Routing does not
|
||||
depend on placement, so one trace can evaluate any placement and policy.
|
||||
"""
|
||||
|
||||
import gzip
|
||||
import json
|
||||
import random
|
||||
from dataclasses import dataclass, field, replace
|
||||
|
||||
from prain.config import GIB, LayerPlacement, Miss, Policy, Prefetch
|
||||
from prain.policy import COUNTERS, Belady, SlotCache
|
||||
|
||||
HEAT_FIELD = {"count": "c", "weight": "w", "recency": "r"}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Cost:
|
||||
"""Rough RX 6600 XT numbers. scripts/measure.py writes measured ones as JSON."""
|
||||
expert_mib: float = 11.34 # one expert
|
||||
h2d_gbps: float = 12.0 # pinned RAM -> VRAM, PCIe 4.0 x8
|
||||
disk_gbps: float = 3.0 # pack file -> RAM
|
||||
round_ms: float = 0.12 # one fused-MoE kernel launch
|
||||
route_ms: float = 0.0012 # GPU work per (token, expert) route
|
||||
other_ms: float = 0.6 # attention + dense MLP of one layer per forward
|
||||
sync_ms: float = 0.05 # host sync per layer
|
||||
cpu_ms: float = 0.5 # CPU work per (token, expert) route
|
||||
|
||||
def copy_ms(self, disk):
|
||||
mb = self.expert_mib * 1.048576
|
||||
return mb / self.h2d_gbps + (mb / self.disk_gbps if disk else 0.0)
|
||||
|
||||
|
||||
def load_cost(path, expert_mib):
|
||||
values = {"expert_mib": expert_mib}
|
||||
if path:
|
||||
with open(path) as f:
|
||||
values |= json.load(f)
|
||||
values["expert_mib"] = expert_mib
|
||||
return Cost(**values)
|
||||
|
||||
|
||||
def write_trace(path, records, mode="wt"):
|
||||
with gzip.open(path, mode) as f:
|
||||
for r in records:
|
||||
f.write(json.dumps(r, separators=(",", ":")) + "\n")
|
||||
|
||||
|
||||
def read_forwards(path):
|
||||
"""[{layer: record}] per forward, in execution order."""
|
||||
forwards, last = [], None
|
||||
with gzip.open(path, "rt") as f:
|
||||
for line in f:
|
||||
if not line.strip():
|
||||
continue
|
||||
r = json.loads(line)
|
||||
if last is None or r["l"] <= last:
|
||||
forwards.append({})
|
||||
forwards[-1][r["l"]] = r
|
||||
last = r["l"]
|
||||
return forwards
|
||||
|
||||
|
||||
def synthetic_trace(layers=30, experts=128, top_k=8, forwards=300, prompt_every=100,
|
||||
prompt_tokens=128, zipf=1.1, stay=0.5, accuracy=0.7, seed=0):
|
||||
"""Zipf-popular experts per layer, decode tokens that keep some of the previous token's
|
||||
experts, a prompt every prompt_every forwards, and pregate predictions that are right
|
||||
with probability accuracy (accuracy**2 two layers ahead)."""
|
||||
rng = random.Random(seed)
|
||||
rank = [rng.sample(range(experts), experts) for _ in range(layers)]
|
||||
weights = [1 / (i + 1) ** zipf for i in range(experts)]
|
||||
previous = [[] for _ in range(layers)]
|
||||
records = []
|
||||
for f in range(forwards):
|
||||
tokens = prompt_tokens if f % prompt_every == 0 else 1
|
||||
routed = []
|
||||
for layer in range(layers):
|
||||
count, recency = {}, {}
|
||||
for t in range(tokens):
|
||||
chosen = {e for e in previous[layer] if rng.random() < stay} if tokens == 1 else set()
|
||||
while len(chosen) < top_k:
|
||||
chosen.add(rank[layer][rng.choices(range(experts), weights)[0]])
|
||||
decay = 0.5 ** ((tokens - 1 - t) / 32)
|
||||
for e in chosen:
|
||||
count[e] = count.get(e, 0) + 1
|
||||
recency[e] = recency.get(e, 0.0) + decay / top_k
|
||||
if tokens == 1:
|
||||
previous[layer] = sorted(chosen)
|
||||
routed.append((count, recency))
|
||||
for layer, (count, recency) in enumerate(routed):
|
||||
ids = sorted(count)
|
||||
record = {"l": layer, "t": tokens, "e": ids, "c": [count[e] for e in ids],
|
||||
"w": [round(count[e] / top_k, 4) for e in ids],
|
||||
"r": [round(recency[e], 4) for e in ids]}
|
||||
predictions = {}
|
||||
for d in (1, 2):
|
||||
if layer + d < layers:
|
||||
nxt = routed[layer + d][0]
|
||||
guess = []
|
||||
for e in sorted(nxt, key=lambda e: (-nxt[e], e))[:top_k]:
|
||||
g = e if rng.random() < accuracy ** d else rng.randrange(experts)
|
||||
if g not in guess:
|
||||
guess.append(g)
|
||||
predictions[str(d)] = [[g, round(1 / top_k, 4)] for g in guess]
|
||||
if predictions:
|
||||
record["p"] = predictions
|
||||
records.append(record)
|
||||
return records
|
||||
|
||||
|
||||
@dataclass
|
||||
class Result:
|
||||
total_ms: float = 0.0
|
||||
decode_ms: float = 0.0 # summed over forwards with fewer than scan_tokens tokens
|
||||
decode_forwards: int = 0
|
||||
prefill_ms: float = 0.0
|
||||
prefill_tokens: int = 0
|
||||
counters: dict = field(default_factory=dict)
|
||||
layer_ms: dict = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def hit_rate(self):
|
||||
c = self.counters
|
||||
return c["hits"] / max(1, c["hits"] + c["misses"] + c["cpu"])
|
||||
|
||||
@property
|
||||
def ms_per_decode(self):
|
||||
return self.decode_ms / max(1, self.decode_forwards)
|
||||
|
||||
|
||||
def _future(forwards, layer):
|
||||
future, n = {}, 0
|
||||
for fwd in forwards:
|
||||
if layer in fwd:
|
||||
n += 1
|
||||
for e in fwd[layer]["e"]:
|
||||
future.setdefault(e, []).append(n)
|
||||
return future
|
||||
|
||||
|
||||
def simulate(forwards, placements, policy, prefetch, miss, cost, belady=False):
|
||||
"""Replay a trace. Model: one PCIe link serves copies in issue order; demand copies of
|
||||
a round wait for the previous round's kernel; async copies (prefetch, admit-behind) are
|
||||
issued after the layer's demand copies; the CPU path runs alongside the GPU."""
|
||||
if prefetch.mode == "pregate" and not any("p" in r for fwd in forwards for r in fwd.values()):
|
||||
raise ValueError("prefetch.mode = pregate needs a trace with predictions ('p')")
|
||||
caches, disk = {}, {}
|
||||
for layer, p in placements.items():
|
||||
n = len(p.vram) + len(p.ram) + len(p.disk)
|
||||
disk[layer] = set(p.disk)
|
||||
cost_of = [policy.disk_cost if e in disk[layer] else 1.0 for e in range(n)]
|
||||
evictor = Belady(_future(forwards, layer)) if belady else None
|
||||
caches[layer] = SlotCache(p.vram, n, p.cache_slots, policy, cost_of, evictor)
|
||||
result = Result(layer_ms={layer: 0.0 for layer in caches})
|
||||
key, t, link, ready = HEAT_FIELD[policy.heat], 0.0, 0.0, {}
|
||||
for fwd in forwards:
|
||||
start, tokens = t, 0
|
||||
for layer in sorted(fwd):
|
||||
if layer not in caches:
|
||||
continue
|
||||
rec, cache = fwd[layer], caches[layer]
|
||||
tokens = rec["t"]
|
||||
counts = dict(zip(rec["e"], rec["c"]))
|
||||
heat = dict(zip(rec["e"], rec[key]))
|
||||
cpu = cache.choose_cpu(counts, miss)
|
||||
gpu = t + cost.sync_ms
|
||||
for slot_of, loads in cache.plan(heat, tokens, cpu):
|
||||
if loads:
|
||||
link = max(link, gpu) + sum(cost.copy_ms(e in disk[layer]) for e, _ in loads)
|
||||
gpu = link
|
||||
land = max((ready.pop((layer, e), 0.0) for e in slot_of), default=0.0)
|
||||
gpu = max(gpu, land) + cost.round_ms + cost.route_ms * sum(counts[e] for e in slot_of)
|
||||
cpu_end = t + cost.sync_ms + cost.cpu_ms * sum(counts[e] for e in cpu)
|
||||
later = []
|
||||
if cpu and miss.admit:
|
||||
hot = sorted((e for e in cpu if e not in disk[layer]), key=lambda e: (-heat[e], e))
|
||||
later += [(layer, e) for e, _ in cache.prefetch([(e, 1.0) for e in hot], miss.admit)]
|
||||
target = layer + prefetch.distance
|
||||
if prefetch.mode == "pregate" and target in caches:
|
||||
guesses = [(e, p) for e, p in rec.get("p", {}).get(str(prefetch.distance), [])
|
||||
if p >= prefetch.min_prob and e not in disk[target]]
|
||||
later += [(target, e) for e, _ in caches[target].prefetch(guesses, prefetch.budget)]
|
||||
for where in later:
|
||||
link = max(link, t + cost.sync_ms) + cost.copy_ms(False)
|
||||
ready[where] = link
|
||||
end = max(gpu, cpu_end) + cost.other_ms
|
||||
result.layer_ms[layer] += end - t
|
||||
t = end
|
||||
if tokens < policy.scan_tokens:
|
||||
result.decode_ms += t - start
|
||||
result.decode_forwards += 1
|
||||
else:
|
||||
result.prefill_ms += t - start
|
||||
result.prefill_tokens += tokens
|
||||
result.total_ms = t
|
||||
result.counters = {k: sum(getattr(c, k) for c in caches.values()) for k in COUNTERS}
|
||||
return result
|
||||
|
||||
|
||||
def describe(result, bound, cost):
|
||||
c = result.counters
|
||||
copied = (c["misses"] + c["prefetches"]) * cost.expert_mib / 1024
|
||||
rate = 1000 / result.ms_per_decode if result.decode_forwards else 0.0
|
||||
return [
|
||||
f"forwards {result.decode_forwards} decode + prefill of {result.prefill_tokens} tokens",
|
||||
f"VRAM hit rate {100 * result.hit_rate:.1f}% "
|
||||
f"(eviction alone at best: {100 * bound.hit_rate:.1f}%)",
|
||||
f"copies {c['misses']} + prefetched {c['prefetches']} (used {c['prefetch_hits']}, "
|
||||
f"wasted {c['prefetch_wasted']}) = {copied:.1f} GiB, CPU-run {c['cpu']}",
|
||||
f"decode {result.ms_per_decode:.1f} ms/forward ({rate:.1f} tok/s), "
|
||||
f"prefill {result.prefill_ms / max(1, result.prefill_tokens):.2f} ms/token",
|
||||
f"simulated total {result.total_ms / 1000:.2f} s",
|
||||
]
|
||||
|
||||
|
||||
DIMENSIONS = (
|
||||
("policy", "eviction", ("lru", "lfu", "gdsf")),
|
||||
("policy", "heat", ("count", "weight", "recency")),
|
||||
("policy", "scan_keep", (0, 1, 2)),
|
||||
("policy", "promote", (0, 2, 4)),
|
||||
("prefetch", "mode", ("none", "pregate")),
|
||||
("prefetch", "distance", (1, 2)),
|
||||
("prefetch", "budget", (1, 2, 4)),
|
||||
("miss", "mode", ("copy", "auto", "cpu")),
|
||||
("miss", "admit", (0, 1, 2)),
|
||||
)
|
||||
_PART = {"policy": 0, "prefetch": 1, "miss": 2}
|
||||
|
||||
|
||||
def _with(variant, table, name, value):
|
||||
parts = list(variant)
|
||||
parts[_PART[table]] = replace(parts[_PART[table]], **{name: value})
|
||||
return tuple(parts)
|
||||
|
||||
|
||||
def ram_tier(forwards, layers, num_experts, rows):
|
||||
"""The `rows` most-used (layer, expert) pairs of the trace -> {layer: set of experts}."""
|
||||
uses = {(layer, e): 0 for layer in layers for e in range(num_experts)}
|
||||
for fwd in forwards:
|
||||
for layer, rec in fwd.items():
|
||||
for e in rec["e"]:
|
||||
uses[(layer, e)] += 1
|
||||
ranked = sorted(uses, key=lambda le: (-uses[le], le))[:rows]
|
||||
ram = {layer: set() for layer in layers}
|
||||
for layer, e in ranked:
|
||||
ram[layer].add(e)
|
||||
return ram
|
||||
|
||||
|
||||
def tune(forwards, num_experts, expert_vram_gib, ram_gib, cost, base=(Policy(), Prefetch(), Miss()),
|
||||
passes=2, window=3):
|
||||
"""Coordinate descent over DIMENSIONS with uniform slots, then per-layer slot counts by
|
||||
greedy marginal gain. Returns (variant, ram, slots_per_layer, result, bound, rows)."""
|
||||
layers = sorted({layer for fwd in forwards for layer in fwd})
|
||||
expert_size = cost.expert_mib * 2**20
|
||||
slots = int(expert_vram_gib * GIB // expert_size) // len(layers)
|
||||
if slots < 1:
|
||||
raise ValueError("expert VRAM budget is less than one slot per layer")
|
||||
ram = ram_tier(forwards, layers, num_experts, int(ram_gib * GIB // expert_size))
|
||||
|
||||
def placements(per_layer):
|
||||
return {layer: LayerPlacement((), tuple(sorted(ram[layer])),
|
||||
tuple(e for e in range(num_experts) if e not in ram[layer]),
|
||||
per_layer[layer]) for layer in layers}
|
||||
|
||||
uniform = placements({layer: slots for layer in layers})
|
||||
predicted = any("p" in r for fwd in forwards for r in fwd.values())
|
||||
results = {}
|
||||
|
||||
def run(variant):
|
||||
if variant not in results:
|
||||
ok = variant[0].scan_keep < slots
|
||||
results[variant] = simulate(forwards, uniform, *variant, cost) if ok else None
|
||||
return results[variant]
|
||||
|
||||
best = base
|
||||
if best[1].mode == "pregate" and not predicted:
|
||||
best = _with(best, "prefetch", "mode", "none")
|
||||
for _ in range(passes):
|
||||
for table, name, values in DIMENSIONS:
|
||||
if table == "prefetch" and (not predicted or (name != "mode" and best[1].mode == "none")):
|
||||
continue
|
||||
if (table, name) == ("miss", "admit") and best[2].mode == "copy":
|
||||
continue
|
||||
scored = [(r.total_ms, i, v) for i, v in enumerate(_with(best, table, name, x) for x in values)
|
||||
if (r := run(v)) is not None]
|
||||
best = min(scored)[2]
|
||||
lo, hi = max(1, slots - window, best[0].scan_keep + 1), slots + window
|
||||
curves = {}
|
||||
for s in range(lo, hi + 1):
|
||||
r = simulate(forwards, placements({layer: s for layer in layers}), *best, cost)
|
||||
curves |= {(layer, s): ms for layer, ms in r.layer_ms.items()}
|
||||
per = {layer: lo for layer in layers}
|
||||
for _ in range((slots - lo) * len(layers)):
|
||||
gains = [(curves[(layer, per[layer])] - curves[(layer, per[layer] + 1)], -layer)
|
||||
for layer in layers if per[layer] < hi]
|
||||
if not gains:
|
||||
break
|
||||
per[-max(gains)[1]] += 1
|
||||
final = simulate(forwards, placements(per), *best, cost)
|
||||
if final.total_ms > results[best].total_ms: # greedy on approximate curves can lose
|
||||
per, final = {layer: slots for layer in layers}, results[best]
|
||||
bound = simulate(forwards, placements(per), Policy(scan_tokens=best[0].scan_tokens),
|
||||
Prefetch(), Miss(), cost, belady=True)
|
||||
rows = sorted((r.total_ms, str(v), v, r) for v, r in results.items() if r is not None)
|
||||
return best, ram, per, final, bound, [(v, r) for _, _, v, r in rows]
|
||||
@@ -0,0 +1,83 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user