training_components.py

1"""Warmstart and grouped validation primitives for the sealed-data integration."""
2import sys
3from pathlib import Path
4sys.path.insert(0,str(Path(__file__).resolve().parent.parent/'state-v2'))
5import torch
6from model import StateModel,loss_sum,optimizer
7from records import sha,digest
8PARENT_SHA='7b63bab559741893b632f456e5f9a8e48b15e48fc9c47c302734da9d613dfe85'
9def initialize(path):
10    if sha(path)!=PARENT_SHA:raise ValueError('wrong CM-6 parent')
11    saved=torch.load(path,map_location='cpu',weights_only=True)
12    if saved['format']!='chess-state-train-v2' or saved['identity_sha256']!=digest(saved['identity']):raise ValueError('invalid parent identity')
13    cfg=saved['model_config']
14    if any(cfg[k]!=v for k,v in dict(vocab_size=512,n_positions=1024,n_layer=12,n_head=12,n_embd=768).items()):raise ValueError('changed backbone')
15    net=StateModel(cfg);net.load_state_dict(saved['model_state_dict'],strict=True)
16    if sum(p.numel() for p in net.parameters())!=86235648:raise ValueError('parameter capacity changed')
17    return net,{'parent_sha256':PARENT_SHA,'parent_update':saved['cursor']['updates'],'parameters':86235648,'initialization':'all CM-6 weights copied; fresh AdamW and scheduler','parent_optimizer_preserved_but_not_loaded':True}
18def grouped_validation(net,batches,device):
19    """batches yields x,y,mask,source,phase,kind; metadata never enters the model."""
20    net.eval();totals={}
21    with torch.inference_mode():
22        for x,y,mask,source,phase,kind in batches:
23            x,y,mask=x.to(device),y.to(device),mask.to(device)
24            with torch.autocast('cuda',dtype=torch.bfloat16,enabled=str(device).startswith('cuda')):logits=net(x,mask)
25            losses=torch.nn.functional.cross_entropy(logits.float().reshape(-1,512),y.reshape(-1),reduction='none').reshape(-1,3).sum(-1)
26            hits=(logits.argmax(-1)==y).to(torch.float64)
27            tags=torch.stack((source,phase,kind),dim=1).cpu()
28            values=torch.cat((losses[:,None],hits),dim=1).double().cpu()
29            for tag in torch.unique(tags,dim=0):
30                selected=(tags==tag).all(1);key='/'.join(map(str,tag.tolist()))
31                acc=totals.setdefault(key,{'decisions':0,'loss_sum':0.,'correct':[0.,0.,0.]});v=values[selected].sum(0).tolist()
32                acc['decisions']+=int(selected.sum());acc['loss_sum']+=v[0];acc['correct']=[a+b for a,b in zip(acc['correct'],v[1:])]
33    net.train()
34    for acc in totals.values():
35        acc['prediction_targets']=3*acc['decisions'];acc['nll']=acc['loss_sum']/acc['prediction_targets'];acc['phase_accuracy']=[x/acc['decisions'] for x in acc['correct']]
36    return totals
37