File size: 4,717 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 | """Separate shared-video reuse from GPU batching for two independent fields."""
import argparse
import json
import statistics
import time
from pathlib import Path
from check_cache import compare
from gemma_rlcd import State
from gemma_rlcd.cached_backend import CachedMLXBackend
from gemma_rlcd.core import ScoringRequest, decision_prompt, softmax
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)
parser.add_argument("--repeats", type=int, default=6)
args = parser.parse_args()
if args.repeats < 1:
parser.error("repeats must be positive")
backend = CachedMLXBackend(args.model)
requests = []
for animal in ("dogs", "cats"):
instruction = f"What proportion of sampled video frames contain visible {animal}?"
criteria = {
"0": f"No sampled frames contain visible {animal}",
"1": f"More than zero and at most 25 percent contain visible {animal}",
"2": f"More than 25 and at most 50 percent contain visible {animal}",
"3": f"More than 50 and at most 75 percent contain visible {animal}",
"4": f"More than 75 percent contain visible {animal}",
}
symbols = tuple(backend.symbols(len(criteria)))
requests.append(
ScoringRequest(
decision_prompt(instruction, criteria, symbols),
symbols,
instruction,
tuple(criteria.items()),
)
)
modes = ["separate_requests", "shared_video_serial_fields", "shared_video_batched_fields"]
report = {
"status": "two_field_runtime_probe_not_animal_recognition_validation",
"model": args.model,
"compute_dtype": "float32",
"state_layers": 35,
"timing_scope": "resident_model_fresh_inputs_including_media_processing_and_probabilities",
"input_note": "Four-second 224x224 synthetic color clips, with or without speech; no animals are visually present.",
"question_note": "Two independent grades of the proportion of sampled frames containing dogs or cats.",
"results": [],
}
for name, filename in [
("silent_video", "red-blue.mp4"),
("video_with_speech", "red-blue-speech.mp4"),
]:
state = State(videos=(str((args.media / filename).resolve()),))
def run(mode):
backend.branch_batch_size = 2 if mode == "shared_video_batched_fields" else 1
started = time.perf_counter()
if mode == "separate_requests":
scores, executions = [], []
for request in requests:
scores.extend(backend.score_batch(state, [request]))
executions.append(dict(backend.last_stats))
else:
scores = backend.score_batch(state, requests)
executions = [dict(backend.last_stats)]
probabilities = [softmax(score.logits) for score in scores]
return scores, {
"seconds": time.perf_counter() - started,
"execution": executions,
"probabilities": probabilities,
}
for mode in modes:
run(mode)
samples = {mode: [] for mode in modes}
comparisons = []
for repeat in range(args.repeats):
order = modes[repeat % 3 :] + modes[: repeat % 3]
if repeat % 2:
order = list(reversed(order))
scores = {}
for mode in order:
scores[mode], sample = run(mode)
samples[mode].append(sample)
comparisons.append(
{mode: compare(scores["separate_requests"], scores[mode]) for mode in modes[1:]}
)
medians = {
mode: statistics.median(sample["seconds"] for sample in records)
for mode, records in samples.items()
}
row = {
"input": name,
"fields": 2,
"median_seconds": medians,
"total_speedup": medians[modes[0]] / medians[modes[2]],
"batching_only_speedup": medians[modes[1]] / medians[modes[2]],
"samples": samples,
"comparisons": comparisons,
}
report["results"].append(row)
args.report.write_text(json.dumps(report, indent=2) + "\n")
print(
json.dumps(
{key: value for key, value in row.items() if key not in {"samples", "comparisons"}}
),
flush=True,
)
if __name__ == "__main__":
main()
|