File size: 7,655 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 | import argparse
import json
import time
import numpy as np
import torch
from common import ROOT, inside, save_json, permit_compute
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--model', default=str(ROOT / 'models/original'))
parser.add_argument('--request', required=True)
parser.add_argument('--out', required=True)
args = parser.parse_args()
permit_compute()
from transformers import AutoModelForImageTextToText, AutoTokenizer
request = json.loads(inside(args.request).read_text())
out = inside(args.out)
if out.exists():
raise ValueError('Используй новый каталог активаций')
out.mkdir(parents=True)
save_json(out / 'request.json', request)
wanted = request.get('return', {})
attention_layers = wanted.get('attn_layers', [])
if any(x not in (3, 7, 11, 15, 19, 23) for x in attention_layers):
raise ValueError('Attention доступен только для full-attention слоёв')
tokenizer = AutoTokenizer.from_pretrained(args.model, local_files_only=True)
if 'token_ids' in request:
inputs = request['token_ids']
else:
texts = request.get('texts')
if texts is None:
texts = [tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=False,
enable_thinking=request.get('enable_thinking', True)) for chat in request['chats']]
inputs = [tokenizer.encode(text, add_special_tokens=False) for text in texts]
if not inputs or max(map(len, inputs)) > 36864 or sum(map(len, inputs)) > 524288:
raise ValueError('Превышен размер teacher-forced запроса')
windows = request.get('store_windows')
if windows is not None:
if len(windows) != len(inputs):
raise ValueError('Нужны окна для каждого текста')
for tokens, selections in zip(inputs, windows):
if not selections or any(not 0 <= left < right <= len(tokens) for left, right in selections):
raise ValueError('Некорректные окна сохранения')
if windows is not None and attention_layers:
raise ValueError('Оконное сохранение attention не поддерживается')
if attention_layers and (len(inputs) > 2 or max(map(len, inputs)) > 2048):
raise ValueError('Attention ограничен двумя текстами по 2048 токенов')
model = AutoModelForImageTextToText.from_pretrained(args.model, local_files_only=True, dtype=torch.bfloat16,
attn_implementation='eager' if attention_layers else 'sdpa').eval().to('cuda')
model.requires_grad_(False)
captures = {}
capture_chunk = True
handles = []
def capture(name, pre=False, attention=False):
def hook(module, arguments, output=None):
if not capture_chunk:
return
if pre:
value = arguments[0]
elif attention:
value = output[1] if isinstance(output, tuple) and len(output) > 1 else None
if value is None:
raise RuntimeError('Рантайм не вернул attention: нельзя подменять его пустым результатом')
else:
value = output[0] if isinstance(output, tuple) else output
captures[name] = value.detach().to('cpu', dtype=torch.float32).numpy()
return hook
language = model.model.language_model
for layer in wanted.get('hidden_states', []):
if not 0 <= layer <= 24:
raise ValueError('Некорректный номер residual stream')
module = language.embed_tokens if layer == 0 else language.layers[layer - 1]
handles.append(module.register_forward_hook(capture(f'hidden_{layer}')))
for layer in wanted.get('mlp_act', []):
handles.append(language.layers[layer].mlp.down_proj.register_forward_pre_hook(capture(f'mlp_act_{layer}', pre=True)))
for layer in attention_layers:
handles.append(language.layers[layer].self_attn.register_forward_hook(capture(f'attention_{layer}', attention=True)))
metadata = {'model': str(inside(args.model)), 'lengths': list(map(len, inputs)), 'files': [],
'hidden_definition': '0=embedding, l=raw residual after block l-1 (before final norm)',
'attention_note': 'Only full-attention layers; complete context, no attention chunks',
'logprob_definition': 'entry p scores token at position p+1', 'runtime': 'transformers'}
if windows is not None:
metadata['store_windows'] = windows
metadata['window_note'] = 'Full context is forwarded and cached; only overlapping chunks are saved, with one preceding token for requested log-probabilities.'
with torch.inference_mode():
for text_index, tokens in enumerate(inputs):
ids = torch.tensor([tokens], dtype=torch.long, device='cuda')
np.save(out / f'input_ids_{text_index}.npy', np.asarray(tokens, dtype=np.int64))
cache = None
step = len(tokens) if attention_layers else request.get('chunk_size', 256)
if not 1 <= step <= 2048:
raise ValueError('Некорректный chunk_size')
for start in range(0, len(tokens), step):
end = min(start + step, len(tokens))
captures.clear()
margin = int(wanted.get('token_logprobs', True))
capture_chunk = windows is None or any(start < right and end > max(0, left-margin)
for left, right in windows[text_index])
result = model(input_ids=ids[:, start:end], past_key_values=cache, use_cache=True,
output_attentions=bool(attention_layers), logits_to_keep=0 if capture_chunk else 1)
cache = result.past_key_values
n = min(end, len(tokens) - 1) - start
if capture_chunk and wanted.get('token_logprobs', True) and n > 0:
logits = result.logits[0, :n].float()
labels = ids[0, start + 1:start + 1 + n]
logp = logits.gather(1, labels[:, None]).squeeze(1) - torch.logsumexp(logits, dim=-1)
captures['token_logprobs'] = logp.cpu().numpy()
k = int(wanted.get('topk', 0))
if not 0 <= k <= 50:
raise ValueError('topk должен быть от 0 до 50')
if k:
values, indices = logits.topk(k, dim=-1)
captures['topk_logprobs'] = (values - torch.logsumexp(logits, dim=-1, keepdim=True)).cpu().numpy()
captures['topk_ids'] = indices.cpu().numpy()
del logits
for name, value in captures.items():
filename = f'text{text_index}_chunk{start}_{name}.npy'
np.save(out / filename, value)
metadata['files'].append({'name': filename, 'signal': name, 'text': text_index,
'token_start': start, 'token_end': end, 'shape': list(value.shape)})
del result
save_json(out / 'meta.json', metadata)
for handle in handles:
handle.remove()
metadata['complete'] = True
save_json(out / 'meta.json', metadata)
if __name__ == '__main__':
main()
|