File size: 9,928 Bytes
fc32096 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | 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()
|