optimizer_policy.py

1"""Retain matching Adam state, but explicitly reset stage LR/scheduler bases."""
2import math,hashlib,json
3import torch
4def restore_moments(model,opt,saved,peak=2e-6):
5    if peak!=2e-6:raise ValueError('unreviewed peak LR')
6    groups=saved['param_groups'];names={id(p):n for n,p in model.named_parameters()};mapping=[];steps=[]
7    if len(groups)!=len(opt.param_groups):raise ValueError('optimizer groups changed')
8    if set(saved['state'])!={key for group in groups for key in group['params']}:raise ValueError('optimizer state coverage')
9    for old,new in zip(groups,opt.param_groups):
10        if len(old['params'])!=len(new['params']):raise ValueError('optimizer parameter count')
11        for key,param in zip(old['params'],new['params']):
12            state=saved['state'].get(key)
13            if not state or set(state)!={'step','exp_avg','exp_avg_sq'}:raise ValueError('matching populated Adam moments required')
14            if state['exp_avg'].shape!=param.shape or state['exp_avg_sq'].shape!=param.shape:raise ValueError('moment shape mismatch')
15            if not torch.isfinite(state['exp_avg']).all() or not torch.isfinite(state['exp_avg_sq']).all() or (state['exp_avg_sq']<0).any() or not math.isfinite(float(state['step'])) or float(state['step'])<1:raise ValueError('invalid Adam state')
16            mapping.append([names[id(param)],list(param.shape)]);steps.append(float(state['step']))
17    opt.load_state_dict(saved)
18    for group in opt.param_groups:
19        group['lr']=peak;group['initial_lr']=peak
20        if tuple(group['betas'])!=(.9,.95) or group['eps']!=1e-8:raise ValueError('optimizer hyperparameters changed')
21    return {'mapping_sha256':hashlib.sha256(json.dumps(mapping).encode()).hexdigest(),'parameters':len(mapping),'minimum_step':min(steps),'maximum_step':max(steps),'moments_retained':True,'peak_lr':peak}
22def lr_factor(step,total,warmup=50):
23    if total<=warmup:raise ValueError('schedule too short for frozen warmup')
24    if step<warmup:return .25+.75*step/warmup
25    progress=min(1.,(step-warmup)/(total-warmup));return .25+.75*.5*(1+math.cos(math.pi*progress))
26
27class Guard:
28    def __init__(self,baseline):self.baseline=baseline
29    def reasons(self,current):
30        b=self.baseline;reasons=[]
31        scalar=('ordinary_nll','exact_accuracy','legality','book_training_recall','book_validation_pair_nll')
32        values=[current[k] for k in scalar]+list(current['source_nll'].values())+[x for x in current['phase_nll'].values() if x is not None]
33        if not current.get('finite',True) or not all(math.isfinite(x) for x in values):reasons.append('nonfinite')
34        if set(current['source_nll'])!=set(b['source_nll']) or set(current['phase_nll'])!=set(b['phase_nll']):reasons.append('diagnostic_coverage_changed')
35        if current['ordinary_nll']>b['ordinary_nll']+.01:reasons.append('ordinary_nll')
36        for name,value in current['source_nll'].items():
37            if value>b['source_nll'][name]+.015:reasons.append('source_nll:'+name)
38        for name,value in current['phase_nll'].items():
39            if value is not None and b['phase_nll'].get(name) is not None and value>b['phase_nll'][name]+.02:reasons.append('phase_nll:'+name)
40        if current['exact_accuracy']<b['exact_accuracy']-.02:reasons.append('exact_preserving')
41        if current['legality']<b['legality']-.01:reasons.append('raw_legality')
42        return reasons
43def selection_key(metrics):return (-metrics['book_training_recall'],metrics['book_validation_pair_nll'])
44