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"], }) @torch.no_grad() 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, :] @torch.no_grad() 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)