model.py

1"""State-v2 GPT-2: unchanged backbone/position capacity, three supervised logits."""
2import hashlib,json
3from pathlib import Path
4import torch
5import torch.nn.functional as F
6from transformers import GPT2Config,GPT2LMHeadModel
7from codec import VOCAB,INPUT,TARGET_POSITIONS
8from records import sha,digest
9
10BASE=dict(vocab_size=VOCAB,n_positions=1024,n_ctx=1024,n_embd=768,n_layer=12,n_head=12,
11    resid_pdrop=0.,embd_pdrop=0.,attn_pdrop=0.,use_cache=False,bos_token_id=69,eos_token_id=70,pad_token_id=71)
12PARENT_SHA='c0a3f67ac4b364276028950b2010c5644e03efb3ca0f1228b32c36b8b5340238'
13
14class StateModel(torch.nn.Module):
15    def __init__(self,config=None):
16        super().__init__();cfg=GPT2Config(**(BASE if config is None else config));cfg._attn_implementation='sdpa'
17        self.net=GPT2LMHeadModel(cfg)
18        self.register_buffer('positions',torch.arange(INPUT),persistent=False)
19    def forward(self,x,attention):
20        if x.ndim!=2 or x.shape[1]!=INPUT or attention.shape!=x.shape:raise ValueError('state-v2 expects exactly122 fixed slots')
21        hidden=self.net.transformer(input_ids=x,attention_mask=attention,position_ids=self.positions.expand(x.shape[0],-1),use_cache=False).last_hidden_state
22        return self.net.lm_head(hidden[:,TARGET_POSITIONS,:])
23    def next_logits(self,x,attention):
24        # Inference uses SAME absolute slots; 120 prompt then 1/2 generatedtokens.
25        if x.ndim!=2 or not 120<=x.shape[1]<=INPUT or attention.shape!=x.shape:raise ValueError('inference phase')
26        h=self.net.transformer(input_ids=x,attention_mask=attention,position_ids=self.positions[:x.shape[1]].expand(x.shape[0],-1),use_cache=False).last_hidden_state
27        return self.net.lm_head(h[:,-1,:])
28
29def initialize(path,seed=20261029,expected=PARENT_SHA):
30    if sha(path)!=expected:raise ValueError('wrong parent checkpoint')
31    saved=torch.load(path,map_location='cpu',weights_only=True)
32    cfg=saved['model_config'];state=saved['model_state_dict']
33    for k,v in [('vocab_size',128),('n_positions',1024),('n_embd',768),('n_layer',12),('n_head',12)]:
34        if cfg[k]!=v:raise ValueError('parent architecture mismatch')
35    if saved['identity_sha256']!=digest(saved['identity']):raise ValueError('parent identity mismatch')
36    torch.manual_seed(seed);net=StateModel();new=net.net.state_dict()
37    if set(state)!=set(new):raise ValueError('parent keys mismatch')
38    for name,tensor in state.items():
39        if name in ('transformer.wte.weight','lm_head.weight'):
40            if tensor.shape!=(128,768):raise ValueError('parent embedding shape')
41            new[name][:128].copy_(tensor)
42        else:
43            if tensor.shape!=new[name].shape:raise ValueError('backbone shape mismatch')
44            new[name].copy_(tensor)
45    net.net.load_state_dict(new,strict=True)
46    if net.net.lm_head.weight.data_ptr()!=net.net.transformer.wte.weight.data_ptr():raise ValueError('tied embeddings lost')
47    for name,tensor in state.items():
48        actual=net.net.state_dict()[name]
49        if not torch.equal(actual[:128] if name in ('transformer.wte.weight','lm_head.weight') else actual,tensor):raise ValueError('parent values changed')
50    count=sum(p.numel() for p in net.parameters())
51    if count!=86_235_648:raise ValueError('parameter capacity mismatch')
52    return net,{'parent_sha256':expected,'parameters':count,'copied_embedding_rows':128,'new_embedding_rows':384,
53        'position_allocation':1024,'operating_positions':122,'unused_position_rows':902,'optimizer':'fresh AdamW; not exact optimizer resume'}
54
55def loss_sum(logits,labels):
56    if logits.shape[:-1]!=labels.shape or labels.shape[-1]!=3:raise ValueError('three manually shifted targets required')
57    return F.cross_entropy(logits.float().reshape(-1,VOCAB),labels.reshape(-1),reduction='sum')
58
59def optimizer(net,lr,fused=False):
60    groups=[{'params':[p for p in net.parameters() if p.ndim>=2],'weight_decay':.1},
61            {'params':[p for p in net.parameters() if p.ndim<2],'weight_decay':0.}]
62    return torch.optim.AdamW(groups,lr=lr,betas=(.9,.95),eps=1e-8,fused=fused)
63