| |
| from __future__ import annotations |
|
|
| import argparse |
| import shutil |
| from pathlib import Path |
|
|
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[3] |
| SOURCE_REPO = PROJECT_ROOT / "external" / "hsd" |
| DEFAULT_WORKSPACE = PROJECT_ROOT / "experiments" / "hsd_inference" / "workspace" |
|
|
|
|
| def copy_workspace(src: Path, dst: Path) -> None: |
| dst.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copytree( |
| src, |
| dst, |
| dirs_exist_ok=True, |
| ignore=shutil.ignore_patterns(".git", ".idea", "__pycache__", "*.pyc"), |
| ) |
|
|
|
|
| def replace_once_or_skip(text: str, old: str, new: str, path: Path) -> str: |
| if new in text: |
| return text |
| if old in text: |
| return text.replace(old, new, 1) |
| raise RuntimeError(f"Expected text not found while patching {path}: {old}") |
|
|
|
|
| def patch_eval_script(path: Path) -> None: |
| text = path.read_text() |
|
|
| replacements = { |
| "os.environ['HF_HOME'] = '/home/c02yuzh/CISPA-projects/rectified_softmax_ml-2023/hf_home'": ( |
| "from pathlib import Path\n" |
| "PROMPT_DIR = Path(__file__).resolve().parent / 'lib_prompt'\n" |
| 'os.environ.setdefault("HF_HOME", os.environ.get("SHARED_CACHE_ROOT", "/project/mpese/tigersec/runw/workdirs/speculative-decoding-robustness/cache") + "/huggingface")' |
| ), |
| "parser.add_argument('--target-model', default='Qwen/Qwen2.5-72B-Instruct-GPTQ-Int8', help='must be complex or original')": ( |
| "parser.add_argument('--target-model', default='Qwen/Qwen2.5-72B-Instruct-GPTQ-Int8', help='must be complex or original')\n" |
| "parser.add_argument('--draft-model', default='Qwen/Qwen2.5-0.5B-Instruct-GPTQ-Int8')\n" |
| "parser.add_argument('--num-samples', type=int, default=None)\n" |
| "parser.add_argument('--prompt-suffix', dest='prompt_suffix', default='')\n" |
| "parser.add_argument('--prompt-tag', default='')\n" |
| "parser.add_argument('--seed', type=int, default=2027)\n" |
| "parser.add_argument('--max-new-tokens', type=int, default=512)\n" |
| "parser.add_argument('--min-new-tokens', type=int, default=0)\n" |
| "parser.add_argument('--sample-offset', type=int, default=0)" |
| ), |
| "args = parser.parse_args()": ( |
| "args = parser.parse_args()\n" |
| "random.seed(args.seed)\n" |
| "np.random.seed(args.seed)\n" |
| "torch.manual_seed(args.seed)\n" |
| "if torch.cuda.is_available():\n" |
| " torch.cuda.manual_seed_all(args.seed)" |
| ), |
| 'model1_name = "Qwen/Qwen2.5-0.5B-Instruct-GPTQ-Int8"': "model1_name = args.draft_model", |
| 'model_size = model2_name.split("/")[1].split("-")[1]': ( |
| 'model_ref = os.path.basename(model2_name.rstrip("/"))\n' |
| 'size_match = re.search(r"(\\d+(?:\\.\\d+)?B)", model_ref)\n' |
| 'if size_match is None:\n' |
| ' raise ValueError(f"Could not infer model size from target model path: {model2_name}")\n' |
| 'model_size = size_match.group(1)' |
| ), |
| "num_samples = len(gsm8k_test['question'])//5": ( |
| "num_samples = args.num_samples if args.num_samples is not None else len(gsm8k_test['question'])//5\n" |
| "if args.sample_offset < 0 or args.sample_offset + num_samples > len(gsm8k_test['question']):\n" |
| " raise ValueError('sample-offset/num-samples exceed the GSM8K split')" |
| ), |
| "for q, a in tqdm(zip(gsm8k_test['question'][:num_samples], gsm8k_test['answer'][:num_samples]),\n" |
| " total=num_samples):": ( |
| "sample_end = args.sample_offset + num_samples\n" |
| " for q, a in tqdm(zip(gsm8k_test['question'][args.sample_offset:sample_end], " |
| "gsm8k_test['answer'][args.sample_offset:sample_end]),\n" |
| " total=num_samples):" |
| ), |
| 'if args.prompt == "original":\n prompt = open(\'lib_prompt/prompt_original.txt\').read()\nelse:\n prompt = open(\'lib_prompt/prompt_hardest.txt\').read()': ( |
| 'if args.prompt == "original":\n prompt = (PROMPT_DIR / \'prompt_original.txt\').read_text()\nelse:\n prompt = (PROMPT_DIR / \'prompt_hardest.txt\').read_text()\n\n' |
| "if args.prompt_suffix:\n" |
| " print('prompt_suffix:', args.prompt_suffix.encode('unicode_escape').decode('ascii'))" |
| ), |
| ' if args.parallel:\n sd += "_parallel"\n': ( |
| ' if args.parallel:\n sd += "_parallel"\n\n' |
| ' if args.prompt_tag:\n sd += f"_{args.prompt_tag}"\n' |
| ), |
| " input_text = tokenizer2.apply_chat_template(\n" |
| " messages,\n" |
| " tokenize=False,\n" |
| " add_generation_prompt=True\n" |
| " )\n": ( |
| " input_text = tokenizer2.apply_chat_template(\n" |
| " messages,\n" |
| " tokenize=False,\n" |
| " add_generation_prompt=True\n" |
| " )\n" |
| "\n" |
| " if args.prompt_suffix:\n" |
| " input_text += args.prompt_suffix\n" |
| ), |
| 'total_counts["time"].append(start-end)': 'total_counts["time"].append(end-start)', |
| |
| " fd.write('Q: %s\\nA_model:\\n%s\\nA:\\n%s\\n\\n' % (q, ans_, a))\n": ( |
| " fd.write('Q: %s\\nA_model:\\n%s\\nA:\\n%s\\n\\n' % (q, ans_, a))\n" |
| " fd.flush() # flush after every sample so NFS sees progress and partial results survive SIGTERM\n" |
| ), |
| ' "step_back_probs":[], "p_i":[], "q_i":[], "hist_lengths": [], "time":[], "ids":[]}': ( |
| ' "step_back_probs":[], "p_i":[], "q_i":[], "hist_lengths": [], "time":[], "ids":[],\n' |
| ' "aggregate_acceptance_per_sample": [], "block_efficiency_per_sample": [],\n' |
| ' "accepted_draft_tokens_per_sample": [], "drafted_tokens_per_sample": [],\n' |
| ' "target_steps_per_sample": [], "first_token_acceptance_per_sample": [],\n' |
| ' "acceptance_trace_per_sample": [], "fallback_triggered_per_sample": [],\n' |
| ' "fallback_trigger_block_per_sample": [], "fallback_generated_tokens_per_sample": [],\n' |
| ' "generated_tokens_per_sample": []}' |
| ), |
| ' return_probs=args.backward or args.blockwise,': ' return_probs=True,', |
| ' total_counts["sample_length"].append(len(outputs[0] - len(input_ids[0])))': ( |
| ' generated_length = int(outputs.shape[-1] - input_ids.shape[-1])\n' |
| ' total_counts["sample_length"].append(generated_length)\n' |
| ' total_counts["generated_tokens_per_sample"].append(generated_length)\n' |
| ' total_counts["target_steps_per_sample"].append(float(generated_length))' |
| ), |
| ' total_counts["ids"].append(counts["ids"])\n': ( |
| ' total_counts["ids"].append(counts["ids"])\n' |
| '\n' |
| ' drafted_tokens = float(sum(counts["draft_eval"]))\n' |
| ' accepted_draft_tokens = float(sum(counts.get("accepted_draft_tokens", [\n' |
| ' min(int(d), max(int(s) - int(t), 0))\n' |
| ' for d, s, t in zip(counts["draft_eval"], counts["sample_length"], counts["target_eval"])\n' |
| ' ])))\n' |
| ' generated_tokens = float(sum(counts["sample_length"]))\n' |
| ' target_steps = float(sum(counts["total_step"]))\n' |
| ' aggregate_acceptance = accepted_draft_tokens / drafted_tokens if drafted_tokens > 0 else None\n' |
| ' block_efficiency = generated_tokens / target_steps if target_steps > 0 else None\n' |
| '\n' |
| ' first_token_acceptance = None\n' |
| ' if counts["p_i"] and counts["q_i"] and len(counts["p_i"]) > 0 and len(counts["q_i"]) > 0:\n' |
| ' valid_prob_pairs = []\n' |
| ' for p_val, q_val in zip(counts["p_i"], counts["q_i"]):\n' |
| ' if p_val is None or q_val is None:\n' |
| ' continue\n' |
| ' try:\n' |
| ' p_float = float(p_val)\n' |
| ' q_float = float(q_val)\n' |
| ' except (TypeError, ValueError):\n' |
| ' continue\n' |
| ' valid_prob_pairs.append((p_float, q_float))\n' |
| ' if valid_prob_pairs:\n' |
| ' p0, q0 = valid_prob_pairs[0]\n' |
| ' first_token_acceptance = min(1.0, p0 / max(q0, 1e-12))\n' |
| '\n' |
| ' total_counts["accepted_draft_tokens_per_sample"].append(accepted_draft_tokens)\n' |
| ' total_counts["drafted_tokens_per_sample"].append(drafted_tokens)\n' |
| ' total_counts["target_steps_per_sample"].append(target_steps)\n' |
| ' total_counts["aggregate_acceptance_per_sample"].append(aggregate_acceptance)\n' |
| ' total_counts["block_efficiency_per_sample"].append(block_efficiency)\n' |
| ' total_counts["first_token_acceptance_per_sample"].append(first_token_acceptance)\n' |
| ' total_counts["acceptance_trace_per_sample"].append(counts.get("acceptance_trace", []))\n' |
| ' total_counts["fallback_triggered_per_sample"].append(bool(counts.get("fallback_triggered", False)))\n' |
| ' total_counts["fallback_trigger_block_per_sample"].append(counts.get("fallback_trigger_block"))\n' |
| ' total_counts["fallback_generated_tokens_per_sample"].append(int(counts.get("fallback_generated_tokens", 0)))\n' |
| ), |
| ' with open(f"{sd}_total_counts.json", "w") as f:\n json.dump(total_counts, f)\n': ( |
| ' with open(f"{sd}_total_counts.json", "w") as f:\n' |
| ' valid_acceptance = [x for x in total_counts["aggregate_acceptance_per_sample"] if x is not None]\n' |
| ' valid_block_eff = [x for x in total_counts["block_efficiency_per_sample"] if x is not None]\n' |
| ' valid_first_alpha = [x for x in total_counts["first_token_acceptance_per_sample"] if x is not None]\n' |
| '\n' |
| ' total_counts["summary_metrics"] = {\n' |
| ' "aggregate_acceptance_mean": float(np.mean(valid_acceptance)) if valid_acceptance else None,\n' |
| ' "block_efficiency_mean": float(np.mean(valid_block_eff)) if valid_block_eff else None,\n' |
| ' "first_token_acceptance_mean": float(np.mean(valid_first_alpha)) if valid_first_alpha else None,\n' |
| ' }\n' |
| ' json.dump(total_counts, f)\n' |
| ), |
| } |
|
|
| for old, new in replacements.items(): |
| text = replace_once_or_skip(text, old, new, path) |
|
|
| text = text.replace("max_new_tokens=512", "max_new_tokens=args.max_new_tokens") |
| text = text.replace( |
| "max_new_tokens=args.max_new_tokens,", |
| "max_new_tokens=args.max_new_tokens,\n min_new_tokens=args.min_new_tokens,", |
| ) |
| path.write_text(text) |
|
|
|
|
| def patch_generation_utils(path: Path) -> None: |
| text = path.read_text() |
|
|
| fallback_init = ( |
| ' fallback_threshold_raw = os.environ.get("ADSD_ACCEPTANCE_FALLBACK_THRESHOLD", "").strip()\n' |
| ' fallback_threshold = float(fallback_threshold_raw) if fallback_threshold_raw else None\n' |
| ' fallback_window = int(os.environ.get("ADSD_ACCEPTANCE_FALLBACK_WINDOW", "8"))\n' |
| ' if fallback_threshold is not None and not 0.0 <= fallback_threshold <= 1.0:\n' |
| ' raise ValueError("ADSD_ACCEPTANCE_FALLBACK_THRESHOLD must be in [0, 1]")\n' |
| ' if fallback_window <= 0:\n' |
| ' raise ValueError("ADSD_ACCEPTANCE_FALLBACK_WINDOW must be positive")\n' |
| ' fallback_active = False\n' |
| ' fallback_trigger_block = None\n' |
| ' fallback_generated_tokens = 0\n' |
| ' acceptance_trace = []\n' |
| ' acceptance_window = []\n' |
| '\n' |
| ' counts = {"draft_eval": [], "target_eval": [], "total_step": [], "sample_length": [],\n' |
| ' "step_back_probs": [], "p_i": [], "q_i": [], "hist_lengths": [], "ids": [],\n' |
| ' "accepted_draft_tokens": [], "acceptance_trace": acceptance_trace,\n' |
| ' "fallback_threshold": fallback_threshold, "fallback_window": fallback_window,\n' |
| ' "fallback_triggered": False, "fallback_trigger_block": None,\n' |
| ' "fallback_generated_tokens": 0}' |
| ) |
| replacements = { |
| " if backward:\n": ( |
| " step_back_probs, p_next, q_next, ids = None, None, None, None\n" |
| " if backward:\n" |
| ), |
| " if not return_probs:\n return valid_tokens, n_matches, ind\n else:\n return valid_tokens, n_matches, None, None, None, None, ind\n": ( |
| " if not return_probs:\n" |
| " return valid_tokens, n_matches, ind\n" |
| " else:\n" |
| " return valid_tokens, n_matches, None, p_i.cpu().numpy().tolist(), q_i.cpu().numpy().tolist(), " |
| "new_candidate_input_ids.cpu().numpy().tolist(), ind\n" |
| ), |
| ' counts = {"draft_eval": [], "target_eval": [], "total_step": [], "sample_length": [],\n' |
| ' "step_back_probs": [], "p_i": [], "q_i": [], "hist_lengths": [], "ids": []}': ( |
| fallback_init |
| ), |
| " num_residual_tokens = num_assistant_tokens\n": ( |
| " if fallback_active:\n" |
| " candidate_generator.num_assistant_tokens = 0\n" |
| " num_assistant_tokens = 0\n" |
| "\n" |
| " num_residual_tokens = num_assistant_tokens\n" |
| ), |
| " candidate_generator.update_candidate_strategy(input_ids, new_logits, n_matches)\n" |
| " # update the num_assistant_tokens if needed\n" |
| " num_assistant_tokens = candidate_generator.num_assistant_tokens\n": ( |
| " candidate_generator.update_candidate_strategy(input_ids, new_logits, n_matches)\n" |
| " # update the num_assistant_tokens if needed\n" |
| " num_assistant_tokens = candidate_generator.num_assistant_tokens\n" |
| " if fallback_active:\n" |
| " candidate_generator.num_assistant_tokens = 0\n" |
| " num_assistant_tokens = 0\n" |
| ), |
| ' counts["hist_lengths"].append(hist_lengths)\n': ( |
| ' counts["hist_lengths"].append(hist_lengths)\n' |
| ' accepted_draft_tokens = min(int(draft_eval), max(int(sample_length) - int(target_eval), 0))\n' |
| ' counts["accepted_draft_tokens"].append(accepted_draft_tokens)\n' |
| ' if fallback_active:\n' |
| ' fallback_generated_tokens += int(sample_length)\n' |
| ' elif draft_eval > 0:\n' |
| ' block_acceptance = float(accepted_draft_tokens) / float(draft_eval)\n' |
| ' acceptance_trace.append(block_acceptance)\n' |
| ' acceptance_window.append(block_acceptance)\n' |
| ' if len(acceptance_window) > fallback_window:\n' |
| ' acceptance_window.pop(0)\n' |
| ' if (\n' |
| ' fallback_threshold is not None\n' |
| ' and len(acceptance_window) == fallback_window\n' |
| ' and sum(acceptance_window) / len(acceptance_window) < fallback_threshold\n' |
| ' ):\n' |
| ' fallback_active = True\n' |
| ' fallback_trigger_block = len(counts["total_step"])\n' |
| ' counts["fallback_triggered"] = True\n' |
| ' counts["fallback_trigger_block"] = fallback_trigger_block\n' |
| ' candidate_generator.num_assistant_tokens = 0\n' |
| ' num_assistant_tokens = 0\n' |
| ' counts["fallback_generated_tokens"] = fallback_generated_tokens\n' |
| ), |
| } |
|
|
| for old, new in replacements.items(): |
| text = replace_once_or_skip(text, old, new, path) |
| path.write_text(text) |
|
|
|
|
| def ensure_outputs_dir(workspace: Path) -> None: |
| outputs_dir = workspace / "chain-of-thought-hub" / "gsm8k" / "outputs" |
| outputs_dir.mkdir(parents=True, exist_ok=True) |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser(description="Prepare a patched local HSD workspace") |
| parser.add_argument("--workspace", type=Path, default=DEFAULT_WORKSPACE) |
| args = parser.parse_args() |
|
|
| if not SOURCE_REPO.exists(): |
| raise SystemExit(f"Missing source repo: {SOURCE_REPO}") |
|
|
| copy_workspace(SOURCE_REPO, args.workspace) |
| patch_eval_script(args.workspace / "chain-of-thought-hub" / "gsm8k" / "eval_speculative_decoding_llm.py") |
| patch_generation_utils(args.workspace / "transformers" / "generation" / "utils.py") |
| ensure_outputs_dir(args.workspace) |
| print(args.workspace) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
