#!/usr/bin/env python3 """ADSD inference with the HSD speculative decoding engine.""" import argparse import hashlib import importlib.util import json import os from pathlib import Path import shutil import subprocess import sys ROOT = Path(__file__).resolve().parents[2] HERE = Path(__file__).resolve().parent sys.path.insert(0, str(HERE)) from native_hooks import atomic def require(condition, message): if not condition: raise ValueError(message) def digest(path): h = hashlib.sha256() with Path(path).open('rb') as f: for block in iter(lambda: f.read(8 * 1024 * 1024), b''): h.update(block) return h.hexdigest() def module(path): spec = importlib.util.spec_from_file_location('workspace_tools', path) result = importlib.util.module_from_spec(spec) spec.loader.exec_module(result) return result def prepare(args): """Copy the installed package; never patch a shared Python installation.""" import importlib.metadata installed = Path(importlib.metadata.distribution('transformers').locate_file('transformers')) runtime = args.runtime.resolve() runtime.mkdir(parents=True, exist_ok=False) overlay = runtime / 'overlay' / 'transformers' shutil.copytree(installed, overlay, ignore=shutil.ignore_patterns('__pycache__', '*.pyc')) helper = module(ROOT / 'scripts/eval/internal/prepare_workspace.py') source = ROOT / 'external/hsd' workspace = runtime / 'workspace' helper.copy_workspace(source, workspace) evaluator = workspace / 'chain-of-thought-hub/gsm8k/eval_speculative_decoding_llm.py' helper.patch_eval_script(evaluator) helper.patch_generation_utils(workspace / 'transformers/generation/utils.py') helper.ensure_outputs_dir(workspace) for src in (workspace / 'transformers').rglob('*.py'): dest = overlay / src.relative_to(workspace / 'transformers') dest.parent.mkdir(parents=True, exist_ok=True) dest.write_text(src.read_text().expandtabs(4)) text = evaluator.read_text() # A snapshot path's basename is a revision hash, not its model name. text = text.replace('model_ref = os.path.basename(model2_name.rstrip("/"))', 'model_ref = model2_name') text = text.replace("gsm8k = load_dataset('gsm8k', 'main')", '''from datasets import Dataset, DatasetDict gsm8k = DatasetDict({split: Dataset.from_list([ json.loads(line) for line in open(os.environ['ADSD_' + split.upper() + '_JSONL']) if line.strip()]) for split in ('train', 'test')})''') marker = ' end = time.time()\n' require(text.count(marker) == 1, 'Expected one evaluator insertion point') text = text.replace(marker, marker + ' from native_hooks import capture\n capture(globals())\n') marker = ' print("start training")' require(text.count(marker) == 1, 'Expected one evaluator insertion point') text = text.replace(marker, ' from native_hooks import policy\n policy(globals())\n' + marker) evaluator.write_text(text) atomic(runtime / 'manifest.json', { 'transformers_version': importlib.metadata.version('transformers'), 'base_package': str(installed), 'evaluator_sha256': digest(evaluator), 'overlay_hashes': {str(p.relative_to(overlay)): digest(p) for p in overlay.rglob('*.py')}, 'note': 'Transformers runtime with HSD inference components.'}) def environment(args): env = os.environ.copy() cache_base = getattr(args, 'cache_root', None) or os.environ.get('ADSD_CACHE_ROOT') cache = Path(cache_base).expanduser().resolve() / args.output.name if cache_base else args.output.resolve() / 'cache' for key in ('HF_HOME', 'HUGGINGFACE_HUB_CACHE', 'TRANSFORMERS_CACHE', 'HF_DATASETS_CACHE', 'TORCH_HOME', 'XDG_CACHE_HOME', 'TRITON_CACHE_DIR', 'TORCH_EXTENSIONS_DIR', 'SHARED_CACHE_ROOT'): env[key] = str(cache / key.lower()) for key in ('HF_HUB_OFFLINE', 'TRANSFORMERS_OFFLINE', 'HF_DATASETS_OFFLINE', 'PYTHONNOUSERSITE', 'PYTHONDONTWRITEBYTECODE', 'PYTHONUNBUFFERED'): env[key] = '1' env['TOKENIZERS_PARALLELISM'] = 'false' env.pop('ADSD_ACCEPTANCE_FALLBACK_THRESHOLD', None) env['PYTHONPATH'] = os.pathsep.join([str(args.runtime.resolve() / 'overlay'), str(HERE), str(ROOT)]) env['ADSD_TRAIN_JSONL'] = str(args.train.resolve()) env['ADSD_TEST_JSONL'] = str(args.test.resolve()) env['ADSD_RUN_OUTPUT'] = str(args.output.resolve()) return env def provenance(args): args.output = args.output.resolve() args.output.mkdir(parents=True, exist_ok=False) if not (args.runtime / 'manifest.json').is_file(): raise ValueError('Run prepare first with this Python environment') manifest = {'command': sys.argv, 'args': {k: str(v) if isinstance(v, Path) else v for k, v in vars(args).items()}, 'runtime_manifest_sha256': digest(args.runtime / 'manifest.json'), 'dataset_sha256': {split: digest(getattr(args, split)) for split in ('train', 'test')}, 'model_files': {}} for role in ('target', 'draft'): path = getattr(args, role).resolve() if not (path / 'config.json').is_file(): raise ValueError(f'cache preload incomplete: {path}') required = list(path.glob('*.safetensors')) if not required: raise ValueError(f'cache preload incomplete: no safetensors in {path}') manifest['model_files'][role] = { p.name: {'bytes': p.stat().st_size, 'sha256': digest(p)} for p in sorted(path.iterdir()) if p.is_file() and not p.name.startswith('.')} source_paths = [HERE/'run.py', HERE/'native_hooks.py', ROOT/'scripts/eval/internal/prepare_workspace.py'] manifest['source_file_sha256'] = {str(p.relative_to(ROOT)):digest(p) for p in source_paths if p.exists()} manifest['source_commit_hint'] = os.environ.get('ADSD_SOURCE_COMMIT') if args.attack: manifest['attack_sha256'] = digest(args.attack) manifest['gpu'] = subprocess.check_output(['nvidia-smi', '--query-gpu=name,uuid,driver_version', '--format=csv,noheader'], text=True).strip() atomic(args.output / 'manifest.json', manifest) def evaluate(args): provenance(args) base = args.runtime.resolve() / 'workspace/chain-of-thought-hub/gsm8k' work = args.output / 'native' work.mkdir() shutil.copytree(base / 'lib_prompt', work / 'lib_prompt') shutil.copy2(base / 'eval_speculative_decoding_llm.py', work) (work / 'outputs').mkdir() suffix = '' if args.attack: best = json.loads(args.attack.read_text()) suffix = best.get('suffix_text') if not isinstance(suffix, str) or not suffix: raise ValueError('Attack artifact lacks nonempty suffix/prefix text') command = [sys.executable, str(work / 'eval_speculative_decoding_llm.py'), '--target-model', str(args.target.resolve()), '--draft-model', str(args.draft.resolve()), '--num-samples', str(args.samples), '--max-new-tokens', str(args.cap), '--sample-offset', str(args.sample_offset), '--seed', str(args.seed), '--gamma', '10', '--temperature', '1.0', '--prompt-tag', 'inference'] if args.mode == 'tokenwise': command.append('--speculative') if suffix: command += ['--prompt-suffix', suffix] atomic(args.output / 'command.json', command) with (args.output / 'stdout.log').open('w') as log: completed = subprocess.run(command, cwd=work, env=environment(args), stdout=log, stderr=subprocess.STDOUT) rows = sorted(args.output.glob('row_*.json')) ok = completed.returncode == 0 and len(rows) == args.samples if not ok: atomic(args.output / 'status.json', {'status': 'failed', 'exit_code': completed.returncode, 'rows': len(rows)}) raise SystemExit(completed.returncode or 1) try: summarize(args.output) except Exception as exc: atomic(args.output / 'status.json', {'status': 'failed', 'stage': 'summarization', 'error': str(exc), 'exit_code': 1, 'rows': len(rows)}) raise atomic(args.output / 'progress.json', {'status': 'complete', 'completed': len(rows), 'expected': args.samples}) atomic(args.output / 'status.json', {'status': 'complete', 'exit_code': 0, 'rows': len(rows)}) def summarize(path): rows = [json.loads(p.read_text()) for p in sorted(path.glob('row_*.json'))] def total(row, field): return sum(row['counts'].get(field, [])) for row in rows: require(row['elapsed_seconds'] > 0, 'Nonpositive generation time') if row['counts']: require(total(row, 'sample_length') == len(row['generated_ids']), 'Native token accounting mismatch') require(0 <= total(row, 'accepted_draft_tokens') <= total(row, 'draft_eval'), 'Invalid accepted-token accounting') seconds = sum(r['elapsed_seconds'] for r in rows) tokens = sum(len(r['generated_ids']) for r in rows) drafted = sum(total(r, 'draft_eval') for r in rows) accepted = sum(total(r, 'accepted_draft_tokens') for r in rows) calls = sum(total(r, 'total_step') for r in rows) block = [len(r['generated_ids']) / total(r, 'total_step') for r in rows if total(r, 'total_step')] acceptance = [total(r, 'accepted_draft_tokens') / total(r, 'draft_eval') for r in rows if total(r, 'draft_eval')] summary = dict(samples=len(rows), correct=sum(r['correct_last_number'] for r in rows), accuracy=sum(r['correct_last_number'] for r in rows)/len(rows), latency_mean=seconds/len(rows), generated_tokens_mean=tokens/len(rows), throughput_global=tokens/seconds, capped=sum(r['reached_cap'] for r in rows), aggregate_acceptance_global=accepted/drafted if drafted else None, aggregate_acceptance_mean=sum(acceptance)/len(acceptance) if acceptance else None, block_efficiency_mean=sum(block)/len(block) if block else None, block_efficiency_global=tokens/calls if calls else None, scoring='last number of generated continuation; native full-text parser retained in stdout.log') import re native_score = re.findall(r'num_q (\d+) correct (\d+) ratio ([\d.]+)', (path / 'stdout.log').read_text(errors='replace')) if native_score: n, correct, rounded = native_score[-1] summary['native_parser'] = dict(samples=int(n), correct=int(correct), accuracy=int(correct)/int(n), reported_rounded=float(rounded)) atomic(path / 'summary.json', summary) return summary def main(): parser = argparse.ArgumentParser(description=__doc__) sub = parser.add_subparsers(dest='action', required=True) p = sub.add_parser('prepare') p.add_argument('--runtime', type=Path, required=True) p = sub.add_parser('evaluate') p.add_argument('--cache-root', type=Path, help='Cache directory; defaults to ADSD_CACHE_ROOT') for name in ('runtime', 'output', 'target', 'draft', 'train', 'test'): p.add_argument('--' + name, type=Path, required=True) p.add_argument('--attack', type=Path, help='Fixed suffix JSON; omit for benign inference') p.add_argument('--seed', type=int, default=2027) p.add_argument('--samples', type=int, default=263) p.add_argument('--sample-offset', type=int, default=0) p.add_argument('--cap', type=int, default=512) p.add_argument('--mode', choices=['tokenwise', 'target_only'], default='tokenwise') args = parser.parse_args() {'prepare': prepare, 'evaluate': evaluate}[args.action](args) if __name__ == '__main__': main()