Upload Qwen3-0.6B-Math-Surgery-v2 (76.67% (23/30, 0 regressions) - Safetensors + GGUF Q8_0 & F16 + JSONs)
f848473 verified | import os | |
| import re | |
| import gc | |
| import time | |
| import json | |
| import ctypes | |
| import subprocess | |
| import sys | |
| import torch | |
| import torch.nn.functional as F | |
| from safetensors.torch import load_file | |
| from benchmark_problems import BENCHMARK_PROBLEMS | |
| torch.set_num_threads(2) | |
| MODEL_DIR = "/home/user/.cache/qwen3-0.6b" | |
| TOK_CACHE_PATH = "/home/user/.cache/benchmark_tokens.pt" | |
| def ensure_tokenized_cache(): | |
| if os.path.exists(TOK_CACHE_PATH): | |
| return | |
| # Run tokenization in a separate short-lived subprocess so Rust tokenizer heap never touches main process | |
| cmd = [ | |
| sys.executable, | |
| "-c", | |
| ( | |
| "import torch; from transformers import AutoTokenizer; " | |
| "from benchmark_problems import BENCHMARK_PROBLEMS; " | |
| f"tok = AutoTokenizer.from_pretrained('{MODEL_DIR}'); " | |
| "pids = [tok.encode(tok.apply_chat_template([{'role':'user','content':p['question']+'\\nWrite only the brief calculation (no markdown headers) and put the final number in \\\\boxed{}.'}], tokenize=False, add_generation_prompt=True, enable_thinking=False), add_special_tokens=False) for p in BENCHMARK_PROBLEMS]; " | |
| f"torch.save({{'prompt_ids_list': pids, 'eos_id': tok.eos_token_id, 'pad_id': tok.pad_token_id}}, '{TOK_CACHE_PATH}')" | |
| ), | |
| ] | |
| subprocess.run(cmd, check=True) | |
| def decode_in_subprocess(all_gen_ids): | |
| tmp_in = "/home/user/.cache/tmp_gen_ids.pt" | |
| tmp_out = "/home/user/.cache/tmp_decoded.json" | |
| torch.save(all_gen_ids, tmp_in) | |
| cmd = [ | |
| sys.executable, | |
| "-c", | |
| ( | |
| "import json, torch; from transformers import AutoTokenizer; " | |
| f"tok = AutoTokenizer.from_pretrained('{MODEL_DIR}'); " | |
| f"gids = torch.load('{tmp_in}', weights_only=True); " | |
| "txts = [tok.decode(g, skip_special_tokens=True).strip() for g in gids]; " | |
| f"open('{tmp_out}', 'w').write(json.dumps(txts))" | |
| ), | |
| ] | |
| subprocess.run(cmd, check=True) | |
| with open(tmp_out, "r") as f: | |
| return json.load(f) | |
| def extract_answer(text): | |
| boxed = re.findall(r"\\boxed\{([^}]*)\}", text) | |
| if boxed: | |
| cand = boxed[-1].strip().replace(",", "").replace("$", "").replace(" ", "") | |
| m = re.findall(r"-?\d+(?:\.\d+)?", cand) | |
| if m: | |
| return m[-1] | |
| boxed_open = re.findall(r"\\boxed\{(-?\d+(?:\.\d+)?)", text) | |
| if boxed_open: | |
| return boxed_open[-1] | |
| nums = re.findall(r"-?\d+(?:,\d{3})*(?:\.\d+)?", text) | |
| if nums: | |
| return nums[-1].replace(",", "") | |
| return None | |
| def rms_norm(x, w, eps=1e-6): | |
| dtype = x.dtype | |
| x_fp32 = x.to(torch.float32) | |
| var = x_fp32.pow(2).mean(-1, keepdim=True) | |
| return w * (x_fp32 * torch.rsqrt(var + eps)).to(dtype) | |
| def rotate_half(x): | |
| x1 = x[..., : x.shape[-1] // 2] | |
| x2 = x[..., x.shape[-1] // 2 :] | |
| return torch.cat((-x2, x1), dim=-1) | |
| class Qwen3FastRunner: | |
| """ | |
| Pure PyTorch implementation of Qwen3ForCausalLM (exact bit-for-bit parity with HuggingFace transformers) | |
| that avoids loading the 160MB transformers library in the inference process. | |
| """ | |
| def __init__(self, sd): | |
| self.num_layers = 28 | |
| self.num_heads = 16 | |
| self.num_kv_heads = 8 | |
| self.head_dim = 128 | |
| self.inv_freq = 1.0 / (1000000.0 ** (torch.arange(0, 128, 2, dtype=torch.float32) / 128.0)) | |
| self.embed = sd["model.embed_tokens.weight"] | |
| if sd.get("_untied_lm_head", False): | |
| self.lm_head = sd["lm_head.weight"] | |
| else: | |
| self.lm_head = self.embed | |
| self.norm = sd["model.norm.weight"] | |
| self.layers = [] | |
| for i in range(self.num_layers): | |
| p = f"model.layers.{i}." | |
| self.layers.append({ | |
| "in_norm": sd[p + "input_layernorm.weight"], | |
| "q_proj": sd[p + "self_attn.q_proj.weight"], | |
| "k_proj": sd[p + "self_attn.k_proj.weight"], | |
| "v_proj": sd[p + "self_attn.v_proj.weight"], | |
| "o_proj": sd[p + "self_attn.o_proj.weight"], | |
| "q_norm": sd[p + "self_attn.q_norm.weight"], | |
| "k_norm": sd[p + "self_attn.k_norm.weight"], | |
| "post_norm": sd[p + "post_attention_layernorm.weight"], | |
| "gate_proj": sd[p + "mlp.gate_proj.weight"], | |
| "up_proj": sd[p + "mlp.up_proj.weight"], | |
| "down_proj": sd[p + "mlp.down_proj.weight"], | |
| }) | |
| def forward_step(self, input_ids, position_ids, kv_cache, attn_mask): | |
| B, S = input_ids.shape | |
| h = F.embedding(input_ids, self.embed) | |
| freqs = position_ids[..., None].float() * self.inv_freq | |
| emb = torch.cat((freqs, freqs), dim=-1) | |
| cos = emb.cos().to(h.dtype).unsqueeze(1) | |
| sin = emb.sin().to(h.dtype).unsqueeze(1) | |
| for i, lyr in enumerate(self.layers): | |
| residual = h | |
| x = rms_norm(h, lyr["in_norm"]) | |
| q = F.linear(x, lyr["q_proj"]).view(B, S, self.num_heads, self.head_dim) | |
| k = F.linear(x, lyr["k_proj"]).view(B, S, self.num_kv_heads, self.head_dim) | |
| v = F.linear(x, lyr["v_proj"]).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2) | |
| q = rms_norm(q, lyr["q_norm"]).transpose(1, 2) | |
| k = rms_norm(k, lyr["k_norm"]).transpose(1, 2) | |
| q = (q * cos) + (rotate_half(q) * sin) | |
| k = (k * cos) + (rotate_half(k) * sin) | |
| if kv_cache[i] is not None: | |
| pk, pv = kv_cache[i] | |
| k = torch.cat([pk, k], dim=2) | |
| v = torch.cat([pv, v], dim=2) | |
| kv_cache[i] = (k, v) | |
| attn_out = F.scaled_dot_product_attention( | |
| q, k, v, attn_mask=attn_mask, enable_gqa=True | |
| ) | |
| attn_out = attn_out.transpose(1, 2).contiguous().view(B, S, -1) | |
| h = residual + F.linear(attn_out, lyr["o_proj"]) | |
| residual = h | |
| x = rms_norm(h, lyr["post_norm"]) | |
| mlp_out = F.silu(F.linear(x, lyr["gate_proj"])) * F.linear(x, lyr["up_proj"]) | |
| h = residual + F.linear(mlp_out, lyr["down_proj"]) | |
| h_last = rms_norm(h[:, -1:, :], self.norm) | |
| logits = F.linear(h_last, self.lm_head) | |
| return logits[:, 0, :] | |
| def generate_batch(self, batch_prompt_ids, eos_id, pad_id, max_new_tokens=95): | |
| B = len(batch_prompt_ids) | |
| max_len = max(len(x) for x in batch_prompt_ids) | |
| input_ids = torch.full((B, max_len), pad_id, dtype=torch.long) | |
| valid_mask = torch.zeros((B, max_len), dtype=torch.bool) | |
| for i, ids in enumerate(batch_prompt_ids): | |
| L = len(ids) | |
| input_ids[i, max_len - L :] = torch.tensor(ids, dtype=torch.long) | |
| valid_mask[i, max_len - L :] = True | |
| pos_ids = (valid_mask.long().cumsum(dim=-1) - 1).clamp(min=0) | |
| causal = torch.tril(torch.ones((max_len, max_len), dtype=torch.bool)) | |
| attn_mask = causal.unsqueeze(0).unsqueeze(0) & valid_mask[:, None, None, :] | |
| kv_cache = [None] * self.num_layers | |
| active_indices = list(range(B)) | |
| gen_tokens = [[] for _ in range(B)] | |
| cur_ids = input_ids | |
| cur_pos = pos_ids | |
| cur_attn_mask = attn_mask | |
| for step in range(max_new_tokens): | |
| if not active_indices: | |
| break | |
| logits = self.forward_step(cur_ids, cur_pos, kv_cache, cur_attn_mask) | |
| next_tokens = logits.argmax(dim=-1) | |
| still_active_local = [] | |
| new_active_indices = [] | |
| for local_i, orig_i in enumerate(active_indices): | |
| tok = int(next_tokens[local_i].item()) | |
| if tok == eos_id: | |
| continue | |
| gen_tokens[orig_i].append(tok) | |
| # Early stop once "}" after "\boxed{" is generated and newline/period follows | |
| # Wait: NO early stop heuristics! Let the model naturally emit eos_id so decoding is 100% standard greedy! | |
| still_active_local.append(local_i) | |
| new_active_indices.append(orig_i) | |
| if not still_active_local: | |
| break | |
| if len(still_active_local) < len(active_indices): | |
| idx = torch.tensor(still_active_local, dtype=torch.long) | |
| next_tokens = next_tokens.index_select(0, idx) | |
| valid_mask = valid_mask.index_select(0, idx) | |
| cur_pos = cur_pos.index_select(0, idx) | |
| for l_idx in range(self.num_layers): | |
| pk, pv = kv_cache[l_idx] | |
| kv_cache[l_idx] = (pk.index_select(0, idx), pv.index_select(0, idx)) | |
| active_indices = new_active_indices | |
| cur_ids = next_tokens.unsqueeze(1) | |
| cur_pos = cur_pos[:, -1:] + 1 | |
| valid_mask = torch.cat([valid_mask, torch.ones((valid_mask.shape[0], 1), dtype=torch.bool)], dim=1) | |
| cur_attn_mask = valid_mask[:, None, None, :] | |
| return gen_tokens | |
| def evaluate_state_dict(sd, problems=BENCHMARK_PROBLEMS, max_new_tokens=95, batch_size=6, verbose=True): | |
| ensure_tokenized_cache() | |
| tok_data = torch.load(TOK_CACHE_PATH, weights_only=True) | |
| prompt_ids_list = tok_data["prompt_ids_list"] | |
| eos_id = tok_data["eos_id"] | |
| pad_id = tok_data["pad_id"] | |
| # Sort indices by prompt length to minimize padding waste | |
| order = sorted(range(len(problems)), key=lambda i: len(prompt_ids_list[i])) | |
| runner = Qwen3FastRunner(sd) | |
| t0 = time.time() | |
| ordered_gen_ids = [] | |
| num_batches = (len(problems) + batch_size - 1) // batch_size | |
| for b_idx, start in enumerate(range(0, len(problems), batch_size)): | |
| end = min(start + batch_size, len(problems)) | |
| batch_indices = order[start:end] | |
| batch_prompts = [prompt_ids_list[i] for i in batch_indices] | |
| tb0 = time.time() | |
| gen_ids = runner.generate_batch(batch_prompts, eos_id, pad_id, max_new_tokens=max_new_tokens) | |
| ordered_gen_ids.extend(gen_ids) | |
| if verbose: | |
| avg_t = sum(len(g) for g in gen_ids) / len(gen_ids) | |
| print(f" Batch {b_idx+1}/{num_batches} ({end-start} probs, avg {avg_t:.1f} toks) in {time.time()-tb0:.1f}s", flush=True) | |
| elapsed = time.time() - t0 | |
| # Restore original problem order | |
| all_gen_ids = [None] * len(problems) | |
| for orig_idx, g_ids in zip(order, ordered_gen_ids): | |
| all_gen_ids[orig_idx] = g_ids | |
| del runner | |
| gc.collect() | |
| responses = decode_in_subprocess(all_gen_ids) | |
| results = [] | |
| correct = 0 | |
| for prob, resp, g_ids in zip(problems, responses, all_gen_ids): | |
| pred = extract_answer(resp) | |
| is_ok = (pred == str(prob["answer"])) | |
| if is_ok: | |
| correct += 1 | |
| results.append({ | |
| "id": prob["id"], | |
| "category": prob["category"], | |
| "question": prob["question"], | |
| "expected": str(prob["answer"]), | |
| "predicted": pred, | |
| "correct": is_ok, | |
| "num_tokens": len(g_ids), | |
| "response": resp, | |
| }) | |
| summary = { | |
| "total": len(problems), | |
| "correct": correct, | |
| "accuracy": round(100.0 * correct / len(problems), 2), | |
| "elapsed_sec": round(elapsed, 2), | |
| "results": results, | |
| } | |
| return summary | |
| if __name__ == "__main__": | |
| print("Evaluating Baseline Qwen3-0.6B on 30-problem Math Benchmark...", flush=True) | |
| sd = load_file(os.path.join(MODEL_DIR, "model.safetensors"), device="cpu") | |
| summary = evaluate_state_dict(sd, BENCHMARK_PROBLEMS, max_new_tokens=95, batch_size=6, verbose=True) | |
| with open("/home/user/baseline_results.json", "w") as f: | |
| json.dump(summary, f, indent=2) | |
| print(f"\n==========================================", flush=True) | |
| print(f"BASELINE QWEN3-0.6B: {summary['correct']}/{summary['total']} ({summary['accuracy']}%) in {summary['elapsed_sec']}s", flush=True) | |
| print(f"==========================================", flush=True) | |
| for r in summary["results"]: | |
| status = "OK" if r["correct"] else "FAIL" | |
| print(f"[{status:4s}] {r['id']:8s} ({r['category']:24s}) | Exp: {r['expected']:>6s} | Pred: {str(r['predicted']):>6s} | Toks: {r['num_tokens']:3d}", flush=True) | |