| import argparse |
| import hashlib |
| import json |
| import platform |
| import time |
| from importlib.metadata import version |
| from pathlib import Path |
| from common import ROOT, inside, rows, save_json, append, stable_seed, permit_compute |
| from grading import extract, correct, repetition, grade_ifeval |
| from resources import snapshot |
| from inference_queue import completed_outputs |
|
|
|
|
| def recover_generations(path): |
| if not path.exists(): |
| return [] |
| result, valid_bytes = [], 0 |
| with path.open('rb') as f: |
| size = path.stat().st_size |
| for line in f: |
| try: |
| result.append(json.loads(line)) |
| except (ValueError, UnicodeDecodeError): |
| if f.tell() != size or line.endswith(b'\n'): |
| raise |
| backup = path.with_name(path.name + '.interrupted_tail') |
| backup.write_bytes(line) |
| break |
| valid_bytes = f.tell() |
| if valid_bytes < path.stat().st_size: |
| with path.open('r+b') as f: |
| f.truncate(valid_bytes) |
| return result |
|
|
|
|
| def model_fingerprint(model): |
| files = sorted(list(model.glob('*.safetensors')) + list(model.glob('*.json')) + list(model.glob('*.jinja'))) |
| entries = [] |
| for file in files: |
| with file.open('rb') as f: |
| entries.append([file.name, hashlib.file_digest(f, 'sha256').hexdigest()]) |
| return entries |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--model', default=str(ROOT / 'models/original')) |
| parser.add_argument('--data', required=True) |
| parser.add_argument('--out', required=True) |
| parser.add_argument('--set', default='practice', choices=['practice', 'math500', 'gsm8k', 'ifeval']) |
| parser.add_argument('--samples', type=int) |
| parser.add_argument('--limit', type=int) |
| parser.add_argument('--max-tokens', type=int) |
| parser.add_argument('--config', default=str(ROOT / 'configs/protocol.json')) |
| parser.add_argument('--queue-limit', type=int) |
| args = parser.parse_args() |
| permit_compute() |
| from transformers import AutoTokenizer |
| from vllm import LLM, SamplingParams |
| model, data, out = inside(args.model), inside(args.data), inside(args.out) |
| protocol = json.loads(inside(args.config).read_text()) |
| if args.set != 'practice' and (args.limit or args.max_tokens): |
| raise ValueError('Сокращённый запуск финального бенчмарка запрещён; используй practice') |
| records = rows(data) |
| if args.limit: |
| records = records[:args.limit] |
| if len({str(r.get('id', r.get('key'))) for r in records}) != len(records): |
| raise ValueError('Повторяющиеся ID') |
| samples = args.samples or protocol['samples'].get(args.set, 1) |
| cap = args.max_tokens or protocol['max_tokens'] |
| if not 1 <= cap <= 32768 or samples < 1: |
| raise ValueError('Некорректный протокол') |
| out.mkdir(parents=True, exist_ok=True) |
| job = {'model': str(model), 'fingerprint': model_fingerprint(model), 'data': str(data), |
| 'data_sha256': hashlib.sha256(data.read_bytes()).hexdigest(), 'protocol': protocol, |
| 'samples': samples, 'max_tokens': cap, 'set': args.set, 'limit': args.limit, |
| 'runtime_versions': {x: version(x) for x in ('vllm', 'torch', 'transformers', 'math-verify')}, |
| 'platform': platform.platform(), 'seed_scheme': 'sha256(base_seed:item_id:sample)31'} |
| path = out / 'job.json' |
| if path.exists() and json.loads(path.read_text()) != job: |
| raise ValueError('В каталоге сохранён другой запуск') |
| save_json(path, job) |
| completed = recover_generations(out / 'generations.jsonl') |
| done = {(str(r['id']), r['sample']) for r in completed} |
| if len(done) != len(completed): |
| raise ValueError('Повторные записи генераций') |
| sample_order = protocol.get('sample_order', 'problem_major') |
| if sample_order not in ('problem_major', 'sample_major'): |
| raise ValueError('Неизвестный порядок образцов') |
| ordered = ([(r, s) for s in range(samples) for r in records] if sample_order == 'sample_major' |
| else [(r, s) for r in records for s in range(samples)]) |
| size = protocol.get('request_batch_size', 1) |
| batches = [ordered[i:i+size] for i in range(0, len(ordered), size)] |
| pending = [batch for batch in batches if any((str(r.get('id', r.get('key'))), s) not in done for r, s in batch)] |
| if args.queue_limit is not None: |
| if args.queue_limit < 1: |
| raise ValueError('queue-limit must be positive') |
| pending = pending[:args.queue_limit] |
| if pending: |
| tokenizer = AutoTokenizer.from_pretrained(str(model), local_files_only=True) |
| memory_options = {} |
| if protocol.get('kv_cache_memory_bytes') is not None: |
| memory_options['kv_cache_memory_bytes'] = protocol['kv_cache_memory_bytes'] |
| engine = LLM(model=str(model), dtype=protocol['dtype'], seed=protocol['seed'], |
| max_model_len=protocol['max_model_len'], gpu_memory_utilization=protocol['gpu_memory_utilization'], |
| max_num_seqs=protocol['max_num_seqs'], max_num_batched_tokens=protocol['max_num_batched_tokens'], |
| enable_prefix_caching=protocol['enable_prefix_caching'], enforce_eager=protocol['enforce_eager'], |
| limit_mm_per_prompt={'image': 0, 'video': 0}, **memory_options) |
| append(out/'resources.jsonl',snapshot('engine_initialized')) |
| for batch in pending: |
| prompts, samplings = [], [] |
| for row, sample in batch: |
| identity = str(row.get('id', row.get('key'))) |
| seed = stable_seed(identity, sample, protocol['seed']) |
| user = row['prompt'] if args.set == 'ifeval' else row['problem'] + '\n' + protocol['instruction'] |
| prompts.append(tokenizer.apply_chat_template([{'role': 'user', 'content': user}], tokenize=False, |
| add_generation_prompt=True, enable_thinking=True)) |
| samplings.append(SamplingParams(n=1, temperature=protocol['temperature'], top_p=protocol['top_p'], |
| top_k=protocol['top_k'], presence_penalty=protocol['presence_penalty'], |
| max_tokens=cap, seed=seed)) |
| start = time.time() |
| batch_tokens = 0 |
| for index,output in completed_outputs(engine,prompts,samplings,protocol.get('stream_completed',False)): |
| elapsed = time.time()-start |
| row,sample = batch[index] |
| prompt = prompts[index] |
| identity = str(row.get('id', row.get('key'))) |
| generation = output.outputs[0] |
| batch_tokens += len(generation.token_ids) |
| if generation.finish_reason not in ('stop','length'): |
| raise RuntimeError('Незавершённая генерация не засчитывается') |
| if (identity, sample) in done: |
| continue |
| generation = output.outputs[0] |
| prediction = extract(generation.text) |
| good = grade_ifeval(row, generation.text) if args.set == 'ifeval' else correct(prediction, row['answer']) |
| result = {'id': identity, 'sample': sample, 'seed': stable_seed(identity, sample, protocol['seed']), 'correct': bool(good), |
| 'prediction': prediction, 'prompt': prompt, 'text': generation.text, 'tokens': len(generation.token_ids), |
| 'prompt_token_ids': tokenizer.encode(prompt, add_special_tokens=False), |
| 'token_ids': list(generation.token_ids), |
| 'finish_reason': generation.finish_reason, 'batch_seconds': elapsed, |
| **repetition(generation.text)} |
| append(out / 'generations.jsonl', result) |
| completed.append(result) |
| print(json.dumps({k: result[k] for k in ('id', 'sample', 'correct', 'tokens', 'finish_reason')}), flush=True) |
| append(out/'resources.jsonl',snapshot('batch_completed')) |
| append(out/'timings.jsonl',{'requests':len(batch),'seconds':time.time()-start,'tokens':batch_tokens}) |
| n = len(completed) |
| if not n: |
| raise ValueError('Нет результатов') |
| save_json(out / 'summary.json', {'generations': n,'planned_generations':len(records)*samples, |
| 'complete':n==len(records)*samples,'problems':len({r['id'] for r in completed}), |
| 'planned_problems':len(records),'accuracy': sum(r['correct'] for r in completed)/n, |
| 'cut_share': sum(r['finish_reason'] == 'length' for r in completed)/n, |
| 'mean_tokens': sum(r['tokens'] for r in completed)/n, |
| 'loop_share': sum(r['repetitive_end'] for r in completed)/n, |
| 'duplicate_8gram_tail_mean': sum(r['duplicate_8gram_tail'] for r in completed)/n, |
| 'whitespace_tail_share': sum(r['whitespace_tail'] for r in completed)/n, |
| 'thinking_closed_share': sum(r['thinking_closed'] for r in completed)/n, |
| 'missing_box_share': sum(r['prediction'] is None for r in completed)/n if args.set != 'ifeval' else None, |
| 'note': 'Фиксированные группы запросов и стабильный seed на задачу и образец; численное совпадение с исходным vLLM не гарантируется.'}) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
