File size: 5,893 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 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """Compare cached GPU batches to identical uncached multimodal token sequences."""
import argparse
import json
import statistics
import time
from dataclasses import replace
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 compare(left, right):
logit_delta = max(
abs(a - b)
for x, y in zip(left, right, strict=True)
for a, b in zip(x.logits, y.logits, strict=True)
)
probability_delta = max(
abs(a - b)
for x, y in zip(left, right, strict=True)
for a, b in zip(softmax(x.logits), softmax(y.logits), strict=True)
)
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(left, right, strict=True)
)
return {
"max_logit_delta": logit_delta,
"max_probability_delta": probability_delta,
"all_winners_match": winners_match,
}
def main():
parser = argparse.ArgumentParser()
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()
backend = CachedMLXBackend(args.model, branch_batch_size=4)
mx = backend.mx
media = args.media.resolve()
questions = [
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."
),
]
requests = []
for question in questions:
criteria = DecisionEngine._criteria(question)
symbols = tuple(backend.symbols(len(criteria)))
requests.append(
ScoringRequest(decision_prompt(question.instructions, criteria, symbols), symbols)
)
states = {
"text": State(text="A cat sleeps on a sofa. No dogs are present."),
"image": State(images=(str(media / "red.png"),)),
"speech": State(audio=(str(media / "dog.wav"),)),
"video": State(videos=(str(media / "red-blue.mp4"),)),
"video_speech": State(videos=(str(media / "red-blue-speech.mp4"),)),
"long_text": State(
text=("Background: the room has a table, a window, and a lamp. " * 80)
+ "A cat sleeps on a sofa. No dogs are present."
),
}
report = {
"model": args.model,
"compute_dtype": "float32",
"status": "cache_numerics_and_timing_probe_not_quality_benchmark",
"results": [],
}
for name, state in states.items():
try:
prepared = backend.prepare(state, requests)
prefix = backend.prefill(prepared)
before = [(c.offset, mx.array(c.keys), mx.array(c.values)) for c in prefix]
mx.eval(*[a for _, k, v in before for a in (k, v)])
reference = backend.uncached(prepared, requests)
backend.branch_batch_size = 1
serial = backend.branches(prepared, prefix, requests)
backend.branch_batch_size = 4
batched = backend.branches(prepared, prefix, requests)
reverse_prepared = replace(prepared, suffixes=list(reversed(prepared.suffixes)))
reversed_scores = list(
reversed(backend.branches(reverse_prepared, prefix, list(reversed(requests))))
)
unchanged = all(
c.offset == offset
and bool(mx.array_equal(c.keys, k).item())
and bool(mx.array_equal(c.values, v).item())
for c, (offset, k, v) in zip(prefix, before, strict=True)
)
timing = {"uncached": [], "cached_batched": []}
# All shapes have been exercised. Alternate order to reduce order bias.
for repeat in range(3):
for method in (
("uncached", "cached_batched")
if repeat % 2 == 0
else ("cached_batched", "uncached")
):
start = time.perf_counter()
if method == "uncached":
backend.uncached(prepared, requests)
else:
current_prefix = backend.prefill(prepared)
backend.branches(prepared, current_prefix, requests)
timing[method].append(time.perf_counter() - start)
medians = {key: statistics.median(values) for key, values in timing.items()}
record = {
"state": name,
"prefix_tokens": prepared.prefix_tokens,
"suffix_lengths": [len(x) for x in prepared.suffixes],
"cached_serial_vs_uncached": compare(reference, serial),
"cached_batch_vs_uncached": compare(reference, batched),
"reversed_batch_vs_original": compare(batched, reversed_scores),
"prefix_unchanged": unchanged,
"branch_batch_sizes": backend.last_stats["branch_batch_sizes"],
"warm_seconds": timing,
"warm_median_seconds": medians,
"speedup": medians["uncached"] / medians["cached_batched"],
}
except Exception as exc:
record = {"state": name, "error": f"{type(exc).__name__}: {exc}"}
report["results"].append(record)
args.report.write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps(record), flush=True)
if __name__ == "__main__":
main()
|