| 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() |
|
|
