File size: 4,489 Bytes
53e24ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | import argparse
import json
import statistics
import time
from pathlib import Path
from gemma_rlcd import Choice, DecisionEngine, Noul, Score, State
from gemma_rlcd.cached_backend import CachedMLXBackend
from gemma_rlcd.core import ScoringRequest, decision_prompt, softmax
def main():
parser = argparse.ArgumentParser(description="Profile cache, decoder, and projection stages")
parser.add_argument("--model", required=True)
parser.add_argument("--media", required=True, type=Path)
parser.add_argument("--report", required=True, type=Path)
args = parser.parse_args()
b = CachedMLXBackend(args.model, branch_batch_size=4)
mx = b.mx
qs = [
Choice(
"Which animal is mentioned or visible?",
{"cat": "A cat", "dog": "A dog", "other": "Neither"},
),
Noul("Is a dog mentioned?"),
Score("How much red is visible?", ["No red", "Some red", "Mostly red"]),
Noul(
"Does the supplied evidence contain any mention of a sofa? Evaluate the complete supplied state."
),
]
rs = []
for q in qs:
c = DecisionEngine._criteria(q)
sy = tuple(b.symbols(len(c)))
rs.append(ScoringRequest(decision_prompt(q.instructions, c, sy), sy))
report = []
for name, state in [
("text", State(text="A cat sleeps on a sofa. No dogs are present.")),
("video_speech", State(videos=(str((args.media / "red-blue-speech.mp4").resolve()),))),
]:
p = b.prepare(state, rs)
start = time.perf_counter()
cache = b.prefill(p)
prefill = time.perf_counter() - start
lengths = [len(s) for s in p.suffixes]
width = max(lengths)
keep = width - min(lengths) + 1
tokens = mx.array([s + [b.tokenizer.pad_token_id] * (width - len(s)) for s in p.suffixes])
mx.eval(tokens)
def run(optimized):
timings = {}
start = time.perf_counter()
fork = b.fork_cache(cache, len(rs))
mx.eval(*[a for c in fork for a in (c.keys, c.values)])
timings["fork"] = time.perf_counter() - start
start = time.perf_counter()
h = b.model.language_model.model(
inputs=tokens, cache=fork, logits_to_keep=keep if optimized else None
)
mx.eval(h)
timings["decoder"] = time.perf_counter() - start
indices = mx.array(lengths) - 1 - (width - h.shape[1])
last = h[mx.arange(len(rs)), indices, :]
start = time.perf_counter()
logits = b.model.language_model.logits_from_hidden(last[:, None, :])[:, 0, :].astype(
mx.float32
)
mx.eval(logits)
timings["vocab_projection"] = time.perf_counter() - start
start = time.perf_counter()
scores = b._extract(logits, rs, [p.prefix_tokens + n for n in lengths])
timings["extract"] = time.perf_counter() - start
return timings, scores
_, ref = run(False)
_, opt = run(True)
data = {"full": [], "trimmed": []}
for iteration in range(5):
for mode in ["full", "trimmed"] if iteration % 2 == 0 else ["trimmed", "full"]:
t, s = run(mode == "trimmed")
data[mode].append(t)
med = {
mode: {stage: statistics.median(t[stage] for t in samples) for stage in samples[0]}
for mode, samples in data.items()
}
delta = max(
abs(a - z)
for x, y in zip(ref, opt)
for a, z in zip(softmax(x.logits), softmax(y.logits))
)
result = {
"state": name,
"prefix_tokens": p.prefix_tokens,
"suffix_lengths": lengths,
"tail_tokens": keep,
"prefill_first_seconds": prefill,
"prefix_kv_logical_bytes": sum(c.keys.nbytes + c.values.nbytes for c in cache),
"median_seconds": med,
"raw_seconds": data,
"max_probability_delta": delta,
"winners_match": all(
max(range(len(x.logits)), key=x.logits.__getitem__)
== max(range(len(y.logits)), key=y.logits.__getitem__)
for x, y in zip(ref, opt)
),
}
report.append(result)
print(json.dumps(result), flush=True)
args.report.write_text(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|