train.py

1"""One sealed curriculum pass, two ranks, original eight-hour deadline on resume."""
2import argparse,json,math,os,sys,time,random,shutil,subprocess
3from pathlib import Path
4from contextlib import nullcontext
5from datetime import timedelta
6import numpy as np
7import torch
8import torch.distributed as dist
9from torch.nn.parallel import DistributedDataParallel as DDP
10from composite import open_curriculum,Plan,start_prefetch,ExposureLedger,phase,tensors
11from checkpoint import capture_rng,restore_rng
12from training_components import initialize,optimizer
13from deadline import Deadline,read_deadline
14from records import sha,digest,atomic_json,Records as SmallRecords
15from runtime_records import Records
16sys.path.insert(0,str(Path(__file__).resolve().parent.parent/'elite-continuation-v1'))
17from training_core import atomic_save,launch_owned_workers
18from curriculum_loss import weighted_loss
19from exact_metrics import summarize
20
21def verify(cfg,args):
22    for path,expected in cfg['runtime_sha256'].items():
23        if sha(path)!=expected:raise ValueError('runtime changed: '+path)
24    if sha(Path(cfg['dataset'])/'manifest.json')!=cfg['sealed_dataset_manifest_sha256'] or sha(Path(cfg['dataset'])/'schedule.bin')!=cfg['schedule_sha256']:raise ValueError('sealed curriculum changed')
25    if sha(cfg['parent'])!=cfg['parent_sha256']:raise ValueError('parent changed')
26    for path,expected in cfg['validation_sha256'].items():
27        if sha(path)!=expected:raise ValueError('validation changed')
28    if args.resume and (not args.resume_sha256 or sha(args.resume)!=args.resume_sha256):raise ValueError('resume hash')
29
30def capacity(cfg,extra=1_200_000_000):
31    out=Path(cfg['output']);used=sum(p.stat().st_size for p in out.iterdir() if p.is_file())
32    if used+extra>cfg['maximum_output_bytes']:raise RuntimeError('checkpoint output storage cap')
33    if shutil.disk_usage(out).free<cfg['native_free_reserve_bytes']+extra:raise RuntimeError('native storage reserve')
34
35def resource_check(cfg):
36    capacity(cfg)
37    if shutil.disk_usage(cfg['output']).free<cfg['native_free_reserve_bytes']+cfg['maximum_output_bytes']:raise RuntimeError('full run storage reservation')
38    health=json.loads(Path('/work/littleply-serving/state/health.json').read_text())
39    if not health.get('ready') or time.time()-health['updated']>180:raise RuntimeError('live bot not freshly ready')
40    gpu=subprocess.run(['nvidia-smi','--query-gpu=index,memory.free','--format=csv,noheader,nounits'],capture_output=True,text=True,check=True,timeout=15).stdout
41    free={int(line.split(',')[0]):int(line.split(',')[1]) for line in gpu.splitlines()}
42    if set(free)!={0,1} or min(free.values())<16384:raise RuntimeError('both GPUs require16GiB before mandatory probe')
43    atomic_json(Path(cfg['output'])/'resource-preflight.json',{'unix':time.time(),'free_mib':free,'bot_worker':health.get('worker_pid'),'bot_checkpoint':health.get('checkpoint_sha256'),'bot_ready':True,'native_free_bytes':shutil.disk_usage(cfg['output']).free})
44
45def probe(net,data,device,rank,clock):
46    reserve=[torch.empty_like(p) for p in net.parameters() for _ in range(2)];out=[]
47    for micro in (16,32,64):
48        if clock.remaining()<600:raise TimeoutError('insufficient original probe budget')
49        rows=data.take(np.arange(micro));x,y,m=[v.to(device) for v in tensors(rows)]
50        torch.cuda.reset_peak_memory_stats(device);began=time.monotonic()
51        for _ in range(2):
52            with torch.autocast('cuda',dtype=torch.bfloat16):loss=torch.nn.functional.cross_entropy(net(x,m).float().reshape(-1,512),y.reshape(-1))
53            loss.backward();net.zero_grad(set_to_none=True)
54        torch.cuda.synchronize();free,_=torch.cuda.mem_get_info(device)
55        good=torch.tensor(int(free>=4*1024**3),device=device);dist.all_reduce(good,op=dist.ReduceOp.MIN)
56        if not good.item():raise RuntimeError('probe headroom below4GiB; no fallback')
57        out.append(dict(rank=rank,micro=micro,seconds=time.monotonic()-began,free_bytes=free,peak_reserved_bytes=torch.cuda.max_memory_reserved(device)))
58    del reserve
59    return out
60
61def evaluate(net,cfg,device):
62    net.eval();groups={};total_loss=0.;total_n=0;exact_outcomes=[]
63    with torch.inference_mode():
64        for spec in cfg['validation_groups']:
65            cls=Records if spec['kind']==0 else SmallRecords
66            source=cls(spec['root'],spec['source'],spec['split'],verify=True)
67            metadata=None
68            if spec['kind']==1:
69                metadata=[json.loads(x) for x in Path(spec['tie_metadata']).read_text().splitlines()]
70            for start in range(0,len(source),64):
71                rows=source.take(range(start,min(start+64,len(source))));x,y,m=[v.to(device) for v in tensors(rows)]
72                with torch.autocast('cuda',dtype=torch.bfloat16,enabled=device.type=='cuda'):logits=net(x,m)
73                losses=torch.nn.functional.cross_entropy(logits.float().reshape(-1,512),y.reshape(-1),reduction='none').reshape(-1,3).sum(-1).cpu()
74                correct=(logits.argmax(-1)==y).sum(0).cpu().tolist()
75                a=groups.setdefault(spec['name'],dict(decisions=0,loss_sum=0.,correct=[0,0,0],kind=spec['kind'],source=spec['source'],phase_counts=[0,0,0],phase_loss_sums=[0.,0.,0.],raw_tie_preserving=0))
76                a['decisions']+=len(rows);a['loss_sum']+=float(losses.sum());a['correct']=[v+w for v,w in zip(a['correct'],correct)]
77                a['phase_counts']=[v+int(np.sum(phase(rows)==i)) for i,v in enumerate(a['phase_counts'])]
78                a['phase_loss_sums']=[v+float(losses[torch.from_numpy(phase(rows)==i)].sum()) for i,v in enumerate(a['phase_loss_sums'])]
79                total_loss+=float(losses.sum());total_n+=len(rows)
80                if metadata is not None:
81                    prefix=x[:,:120];mask=m[:,:120];generated=[]
82                    for _ in range(3):
83                        with torch.autocast('cuda',dtype=torch.bfloat16,enabled=device.type=='cuda'):token=net.next_logits(prefix,mask).argmax(-1)
84                        generated.append(token.cpu().tolist());prefix=torch.cat((prefix,token[:,None]),1);mask=torch.cat((mask,torch.ones_like(token[:,None],dtype=mask.dtype)),1)
85                    for i,triple in enumerate(zip(*generated)):
86                        item=metadata[int(rows[i]['game'])];hit=int(list(triple) in item['tied_target_triples']);a['raw_tie_preserving']+=hit
87                        exact_outcomes.append({'key':item['position_key'],'game':item['source_sequence_sha256'],'correct':hit})
88            if hasattr(source,'close'):source.close()
89    if not total_n:raise ValueError('empty frozen validation')
90    for a in groups.values():
91        a['nll']=a['loss_sum']/(3*a['decisions']);a['phase_accuracy']=[v/a['decisions'] for v in a['correct']]
92        a['raw_tie_preserving_fraction']=a['raw_tie_preserving']/a['decisions'] if a['kind']==1 else None
93        a['phase_nll']=[s/(3*n) if n else None for s,n in zip(a['phase_loss_sums'],a['phase_counts'])]
94    net.train();return dict(nll=total_loss/(3*total_n),decisions=total_n,groups=groups,exact_outcomes=exact_outcomes,exact_cluster_accuracy=summarize(exact_outcomes),nll_definition='canonical-label CE; exact theory additionally reports autoregressive raw tie-set preservation')
95
96def worker(cfg,args):
97    rank=int(os.environ['RANK']);world=int(os.environ['WORLD_SIZE'])
98    if world!=2:raise ValueError('two ranks required')
99    device=torch.device('cuda',rank);torch.cuda.set_device(device);torch.set_num_threads(4)
100    random.seed(cfg['seed']+rank);np.random.seed(cfg['seed']+rank);torch.manual_seed(cfg['seed']+rank);torch.cuda.manual_seed(cfg['seed']+rank)
101    deadline_value=read_deadline(os.environ['CHESS_JOB_DEADLINE']);clock=Deadline(deadline_value)
102    dist.init_process_group('nccl',init_method='file://'+os.environ['CHESS_DDP_INIT_FILE'],rank=rank,world_size=world,timeout=timedelta(seconds=180))
103    data=None;prefetch=None
104    try:
105        data=open_curriculum(cfg['dataset'],cfg['adapter_directory']);ledger=ExposureLedger(data);plan=Plan(data,rank,micro=64)
106        raw,initialization=initialize(cfg['parent']);raw.to(device);net=DDP(raw,device_ids=[rank],gradient_as_bucket_view=True)
107        proofs=probe(net,data,device,rank,clock);both=[None,None];dist.all_gather_object(both,proofs)
108        out=Path(cfg['output']);opt=optimizer(raw,cfg['lr'],fused=True);accum=4;total_updates=math.ceil(plan.steps/accum)
109        def scale(step):
110            if step<cfg['warmup_updates']:return (step+1)/cfg['warmup_updates']
111            f=min(1.,(step-cfg['warmup_updates'])/max(1,total_updates-cfg['warmup_updates']))
112            return .1+.9*.5*(1+math.cos(math.pi*f))
113        sched=torch.optim.lr_scheduler.LambdaLR(opt,scale);cursor=updates=0;best=float('inf');best_ref=None;best_validation=None;initial_validation=None;latest=None;group_totals=torch.zeros((12,3),dtype=torch.float64,device=device)
114        attempt='initial' if not args.resume else args.resume_sha256[:16]
115        if args.resume:
116            saved=torch.load(args.resume,map_location='cpu',weights_only=True)
117            if saved['format']!='cm6-curriculum-checkpoint-v1' or saved['identity']!=cfg or saved['deadline']!=deadline_value or saved['terminal']:raise ValueError('incompatible/terminal resume')
118            if len(saved['rank_rng'])!=2:raise ValueError('missing rank RNG')
119            consumed=ledger.verify(saved['exposures']);cursor=saved['microstep'];updates=saved['updates']
120            if plan.cursor(cursor)!=consumed or updates!=math.ceil(cursor/accum):raise ValueError('cursor/update mismatch')
121            raw.load_state_dict(saved['model_state_dict']);opt.load_state_dict(saved['optimizer_state_dict']);sched.load_state_dict(saved['scheduler_state_dict'])
122            best=saved['best_nll'];latest=saved['latest_validation'];group_totals.copy_(saved['rank_group_totals'][rank].to(device));restore_rng(saved['rank_rng'][rank],device)
123            best_validation=saved['best_validation']
124            initial_validation=saved['initial_validation']
125            best_ref=({'path':str(Path(args.resume).resolve()),'sha256':args.resume_sha256} if saved['is_best_snapshot'] else saved['best_checkpoint'])
126            if best_ref is None or sha(best_ref['path'])!=best_ref['sha256']:raise ValueError('missing or changed immutable best checkpoint')
127            del saved
128        if rank==0:atomic_json(out/('probe-resume.json' if args.resume else 'probe.json'),{'ranks':both,'initialization':initialization})
129        dummy,prefetch=start_prefetch(data,plan,cursor,True)
130        log=(out/'metrics.jsonl').open('a',buffering=1) if rank==0 else None
131        if rank==0:log.write(json.dumps({'event':'attempt_start','attempt':attempt,'unix':time.time(),'committed_schedule_cursor':plan.cursor(cursor),'resume_sha256':args.resume_sha256,'accounting':'cursor is restored committed exposure count; superseded attempt logs retain possibly replayed uncommitted work'})+'\n')
132        def checkpoint(name,terminal=False,is_best=False):
133            rngs=[None,None];dist.all_gather_object(rngs,capture_rng(device));rank_totals=[None,None];dist.all_gather_object(rank_totals,group_totals.cpu())
134            if rank==0:
135                capacity(cfg)
136                atomic_save(out/name,{'format':'cm6-curriculum-checkpoint-v1','identity':cfg,'identity_sha256':digest(cfg),'deadline':deadline_value,'terminal':terminal,
137                    'model_config':raw.net.config.to_dict(),'model_state_dict':raw.state_dict(),'optimizer_state_dict':opt.state_dict(),'scheduler_state_dict':sched.state_dict(),
138                    'rank_rng':rngs,'rank_group_totals':rank_totals,'exposures':ledger.snapshot(plan.cursor(cursor)),'microstep':cursor,'updates':updates,'best_nll':best,'latest_validation':latest,
139                    'best_checkpoint':best_ref,'is_best_snapshot':is_best,'attempt':attempt,'best_validation':best_validation,'initial_validation':initial_validation})
140            dist.barrier()
141        def validation():
142            nonlocal latest,best,best_ref,best_validation,initial_validation
143            improved=False
144            if rank==0:
145                began=time.time();latest=evaluate(raw,cfg,device);latest.update(unix=time.time(),updates=updates,seconds=time.time()-began)
146                improved=latest['nll']<best;best=min(best,latest['nll']);log.write(json.dumps(dict(event='validation',**latest))+'\n')
147                if updates==0:atomic_json(out/'initial-validation.json',latest)
148            flag=torch.tensor(int(improved),device=device);dist.broadcast(flag,0)
149            state=[latest,best];dist.broadcast_object_list(state,0);latest,best=state
150            if updates==0:initial_validation=latest
151            if flag.item():
152                best_validation=latest
153                name=f'best-{updates}-{attempt}.pt'
154                checkpoint(name,is_best=True)
155                ref=[{'path':str(out/name),'sha256':sha(out/name)} if rank==0 else None];dist.broadcast_object_list(ref,0);best_ref=ref[0]
156                if rank==0:atomic_json(out/'best.json',best_ref)
157        if not args.resume:validation()
158        last_time=time.monotonic();last_cursor=plan.cursor(cursor)
159        while cursor<plan.steps:
160            stop=torch.tensor(int(clock.soft_stop()),device=device);dist.all_reduce(stop,op=dist.ReduceOp.MAX)
161            if stop.item():break
162            end=min(cursor+accum,plan.steps);denom=3*sum(plan.global_decisions(i) for i in range(cursor,end));opt.zero_grad(set_to_none=True)
163            wait_seconds=load_seconds=0.;began=time.monotonic()
164            for step in range(cursor,end):
165                t=time.monotonic();batch,tags,load=prefetch.take(step);wait_seconds+=time.monotonic()-t;load_seconds+=load;empty=len(batch[0])==0
166                x,y,mask=[v.to(device,non_blocking=True) for v in (dummy if empty else batch)]
167                with net.no_sync() if step<end-1 else nullcontext():
168                    with torch.autocast('cuda',dtype=torch.bfloat16):logits=net(x,mask)
169                    loss=torch.nn.functional.cross_entropy(logits.float().reshape(-1,512),y.reshape(-1),reduction='none').reshape(-1,3)
170                    weighted_loss(loss,world,denom,empty).backward()
171                if not empty:
172                    tag=tags.to(device);group_totals[:,0].scatter_add_(0,tag,torch.ones(len(tag),dtype=torch.float64,device=device));group_totals[:,1].scatter_add_(0,tag,loss.detach().sum(-1).double());group_totals[:,2].scatter_add_(0,tag,(logits.argmax(-1)==y).sum(-1).double())
173            torch.nn.utils.clip_grad_norm_(raw.parameters(),1.,error_if_nonfinite=True);opt.step();sched.step();cursor=end;updates+=1
174            if updates%50==0:
175                sums=group_totals.clone();dist.all_reduce(sums);torch.cuda.synchronize();now=time.monotonic();free,_=torch.cuda.mem_get_info(device)
176                mem=torch.tensor([free,torch.cuda.max_memory_reserved(device)],dtype=torch.int64,device=device);mems=[torch.empty_like(mem) for _ in range(2)];dist.all_gather(mems,mem)
177                if rank==0:log.write(json.dumps({'event':'training','attempt':attempt,'unix':time.time(),'updates':updates,'schedule_cursor':plan.cursor(cursor),'prediction_targets':3*plan.cursor(cursor),'window_decisions_per_second':(plan.cursor(cursor)-last_cursor)/(now-last_time),'last_update_wall_seconds':now-began,'cpu_wait_seconds':wait_seconds,'cpu_load_seconds':load_seconds,'lr':sched.get_last_lr()[0],'remaining_seconds':clock.remaining(),'source_phase_kind_totals':sums.cpu().tolist(),'rank_free_peak_bytes':[v.cpu().tolist() for v in mems]})+'\n')
178                last_time=now;last_cursor=plan.cursor(cursor)
179            if updates%cfg['validation_every_updates']==0 and clock.remaining()>600:validation()
180            if updates%cfg['checkpoint_every_updates']==0 and clock.remaining()>300:checkpoint('progress.pt')
181        if clock.remaining()>360:validation()
182        checkpoint('last.pt',True)
183        if rank==0:
184            initial=initial_validation
185            atomic_json(out/'exact-validation-comparison.json',{'initial':initial['exact_cluster_accuracy'],'selected':best_validation['exact_cluster_accuracy'],
186                'paired_game_cluster_delta':summarize(best_validation['exact_outcomes'],initial['exact_outcomes']),'selected_checkpoint':best_ref,'selection':'pooled frozen canonical-label validation NLL, including step zero'})
187            atomic_json(out/'terminal.json',{'status':'complete','reason':'single_schedule_complete' if cursor==plan.steps else 'original_deadline_soft_stop','updates':updates,'schedule_cursor':plan.cursor(cursor),'remaining_schedule_entries':len(data.schedule)-plan.cursor(cursor),'prediction_targets':3*plan.cursor(cursor),'best_nll':best,'latest_validation':latest,'best_checkpoint':best_ref,'last_sha256':sha(out/'last.pt'),'finished_unix':time.time(),'no_promotion':True});log.close()
188    finally:
189        if prefetch:prefetch.close()
190        if data:
191            for source in data.sources.values():
192                if hasattr(source,'close'):source.close()
193        dist.destroy_process_group()
194if __name__=='__main__':
195    p=argparse.ArgumentParser();p.add_argument('config');p.add_argument('--resume');p.add_argument('--resume-sha256');args=p.parse_args();cfg=json.loads(Path(args.config).read_text())
196    if 'CHESS_JOB_DEADLINE' not in os.environ:raise ValueError('reviewed owner required')
197    if 'RANK' in os.environ:worker(cfg,args)
198    else:
199        verify(cfg,args);resource_check(cfg);extra=[] if not args.resume else ['--resume',args.resume,'--resume-sha256',args.resume_sha256]
200        raise SystemExit(launch_owned_workers(Path(__file__),[args.config,*extra]))
201