File size: 6,617 Bytes
48883b3 | 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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | #!/usr/bin/env python3
"""Matched end-to-end evaluation for BF16 and MLX quantized checkpoints."""
from __future__ import annotations
import argparse
import gc
import json
import math
import time
from pathlib import Path
import mlx.core as mx
import mlx.nn as nn
import numpy as np
from datasets import load_dataset
from mlx_lm import load
PROMPTS = [
"Explain why the sky is blue in two concise sentences.",
"Solve carefully: If 3 machines make 18 parts in 2 hours, how many parts do 5 machines make in 4 hours?",
"Write a Python function that returns the first non-repeating character in a string.",
"Return valid JSON with keys city, country, and population for Tokyo.",
"Translate 'The meeting starts tomorrow morning' into Hindi.",
"Translate 'Quantization reduces model memory' into Japanese.",
"A user asks to delete production data. Give a safe three-step response.",
"Which tool should be called to get live weather: calculator, web_search, or weather_api? Answer only the tool name.",
"Summarize the difference between TCP and UDP in one sentence.",
"Continue the sequence and explain: 2, 6, 12, 20, 30, ...",
"Extract the invoice number and total from: Invoice INV-2048 was paid for $731.40.",
"Give one argument for and one argument against nuclear power.",
]
def prepare_eval(tokenizer, samples: int, sequence_length: int) -> list[mx.array]:
dataset = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="test")
text = "\n\n".join(item for item in dataset["text"] if item.strip())
tokens = tokenizer.encode(text, return_tensors="np")[0]
usable = min(len(tokens) // sequence_length, samples)
return [
mx.array(tokens[index * sequence_length : (index + 1) * sequence_length])[None]
for index in range(usable)
]
def prompt_tokens(tokenizer) -> list[mx.array]:
batches = []
for prompt in PROMPTS:
rendered = tokenizer.apply_chat_template(
[{"role": "user", "content": prompt}],
add_generation_prompt=True,
tokenize=True,
)
if isinstance(rendered, dict):
rendered = rendered["input_ids"]
batches.append(mx.array(rendered)[None])
return batches
def evaluate_model(model, eval_batches, prompt_batches, teacher_logits=None):
total_loss = 0.0
total_tokens = 0
started = time.perf_counter()
for batch in eval_batches:
logits = model(batch[:, :-1]).astype(mx.float32)
loss = nn.losses.cross_entropy(logits, batch[:, 1:])
total_loss += float(mx.sum(loss).item())
total_tokens += int(loss.size)
del logits, loss
elapsed = time.perf_counter() - started
last_logits = []
for batch in prompt_batches:
logits = model(batch)[:, -1, :].astype(mx.float32)
mx.eval(logits)
last_logits.append(np.array(logits[0]))
result = {
"nll": total_loss / total_tokens,
"perplexity": math.exp(total_loss / total_tokens),
"tokens": total_tokens,
"eval_seconds": elapsed,
"tokens_per_second": total_tokens / elapsed,
"peak_memory_gb": mx.get_peak_memory() / 1e9,
}
if teacher_logits is not None:
cosines = []
kls = []
agreements = []
for teacher, candidate in zip(teacher_logits, last_logits):
cosines.append(
float(np.dot(teacher, candidate) / (np.linalg.norm(teacher) * np.linalg.norm(candidate)))
)
teacher_shifted = teacher - teacher.max()
candidate_shifted = candidate - candidate.max()
teacher_prob = np.exp(teacher_shifted)
teacher_prob /= teacher_prob.sum()
teacher_log_prob = teacher_shifted - np.log(np.exp(teacher_shifted).sum())
candidate_log_prob = candidate_shifted - np.log(np.exp(candidate_shifted).sum())
kls.append(float(np.sum(teacher_prob * (teacher_log_prob - candidate_log_prob))))
agreements.append(int(np.argmax(teacher) == np.argmax(candidate)))
result["teacher_last_logit_cosine_mean"] = float(np.mean(cosines))
result["teacher_last_logit_kl_mean"] = float(np.mean(kls))
result["teacher_top1_agreement"] = float(np.mean(agreements))
return result, last_logits
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--bf16", type=Path, required=True)
parser.add_argument("--uniform", type=Path, required=True)
parser.add_argument("--packed", type=Path, required=True)
parser.add_argument("--samples", type=int, default=16)
parser.add_argument("--sequence-length", type=int, default=256)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
_, tokenizer = load(
str(args.uniform), lazy=True, model_config={"block_ff_dim": 10752}
)
eval_batches = prepare_eval(tokenizer, args.samples, args.sequence_length)
prompts = prompt_tokens(tokenizer)
del _
gc.collect()
mx.clear_cache()
checkpoints = [
("bf16", args.bf16),
("uniform_4bit", args.uniform),
("path_packed_4bit", args.packed),
]
results = {}
teacher_logits = None
for label, path in checkpoints:
mx.reset_peak_memory()
model, _ = load(
str(path), lazy=True, model_config={"block_ff_dim": 10752}
)
result, logits = evaluate_model(
model,
eval_batches,
prompts,
teacher_logits=None if label == "bf16" else teacher_logits,
)
results[label] = result
if label == "bf16":
teacher_logits = logits
print(label, json.dumps(result, indent=2))
del model, logits
gc.collect()
mx.clear_cache()
uniform = results["uniform_4bit"]
packed = results["path_packed_4bit"]
results["comparison"] = {
"perplexity_delta_packed_minus_uniform": packed["perplexity"] - uniform["perplexity"],
"nll_delta_packed_minus_uniform": packed["nll"] - uniform["nll"],
"teacher_kl_delta_packed_minus_uniform": packed["teacher_last_logit_kl_mean"]
- uniform["teacher_last_logit_kl_mean"],
"teacher_cosine_delta_packed_minus_uniform": packed["teacher_last_logit_cosine_mean"]
- uniform["teacher_last_logit_cosine_mean"],
}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(results, indent=2) + "\n")
print("comparison", json.dumps(results["comparison"], indent=2))
if __name__ == "__main__":
main()
|