train_chess.py

1"""Bounded two-GPU pilot training. No network, downloads, or test-set tuning."""
2from collections import Counter
3from datetime import timedelta
4import hashlib
5import json
6import math
7import os
8from pathlib import Path
9import random
10import subprocess
11import threading
12import time
13import argparse
14import sys
15import torch
16import torch.distributed as dist
17from torch.nn.parallel import DistributedDataParallel as DDP
18import torch.nn.functional as F
19import training_core as c
20from training_deadline import Deadline,read_deadline
21from chessdata import State,BOS,EOS,VERSION,decode_move
22from chessdata.reader import games
23from chessdata.convert import file_hash
24
25
26def log(root,event):
27    event={'time_unix':time.time(),**event}
28    with (root/'events.jsonl').open('a') as f:f.write(json.dumps(event)+'\n')
29    print(json.dumps(event),flush=True)
30
31
32def telemetry(root,stop):
33    with (root/'gpu-telemetry.jsonl').open('a') as f:
34        while not stop.is_set():
35            r=subprocess.run(['nvidia-smi','--query-gpu=index,utilization.gpu,memory.used,power.draw','--format=csv,noheader,nounits'],capture_output=True,text=True)
36            f.write(json.dumps({'time_unix':time.time(),'gpus':r.stdout.strip()})+'\n');f.flush();stop.wait(1)
37
38
39@torch.no_grad()
40def evaluate(net,data,device,micro=8,deadline=None,legality_games=None):
41    """Raw teacher-forced metrics. No logits are masked/renormalized.
42
43    Triple argmax is teacher-forced: destination sees the GOLD source, promotion
44    sees GOLD source+destination. This is not a free-running legal-move rate.
45    """
46    net.eval();total=Counter();phase_n=Counter();phase_correct=Counter();start=time.perf_counter()
47    if not data:raise ValueError('validation data must be nonempty')
48    sample_size=min(len(data),legality_games) if legality_games is not None else len(data)
49    if sample_size<1:raise ValueError('legality sample must be positive')
50    legality_indices={i*len(data)//sample_size for i in range(sample_size)}
51    for offset in range(0,len(data),micro):
52        if deadline is not None and deadline.soft_stop():
53            net.train()
54            return None  # Partial validation never changes best/patience.
55        samples=data[offset:offset+micro];x,y,mask=c.collate(samples,device)
56        with torch.autocast('cuda',dtype=torch.bfloat16):logits=net(x,attention_mask=mask,use_cache=False).logits
57        total['nll_sum']+=F.cross_entropy(logits.float().reshape(-1,128),y.reshape(-1),ignore_index=-100,reduction='sum').item()
58        probs=logits.float().softmax(-1).cpu();pred=probs.argmax(-1);gold=y.cpu();valid=gold.ne(-100)
59        total['targets']+=int(valid.sum())
60        total['reserved_probability_sum']+=float((probs[:,:,72:].sum(-1)*valid).sum())
61        positions=torch.arange(gold.shape[1]).unsqueeze(0)
62        for phase,index in (('source',0),('destination',1),('promotion',2),('eos',None)):
63            selected=valid & (gold.eq(EOS) if index is None else (positions.remainder(3).eq(index)&gold.ne(EOS)))
64            phase_n[phase]+=int(selected.sum());phase_correct[phase]+=int((selected&pred.eq(gold)).sum())
65        for row,tokens in enumerate(samples):
66            if offset+row not in legality_indices:continue
67            total['legality_games']+=1
68            state=State();state.feed(BOS)
69            for pos,target in enumerate(tokens[1:]):
70                total['legality_targets']+=1
71                legal=torch.tensor(state.mask(legal=True),dtype=torch.bool)
72                total['illegal_probability_sum']+=probs[row,pos,~legal].sum().item()
73                if pos%3==0 and target!=EOS:
74                    total['moves']+=1
75                    try:
76                        move=decode_move(pred[row,pos:pos+3].tolist())
77                        total['legal_triples']+=int(move in state.board.legal_moves)
78                    except ValueError:pass
79                state.feed(target)
80    net.train()
81    return {'nll':total['nll_sum']/total['targets'],'valid_targets':total['targets'],'games':len(data),
82            'phase_accuracy':{p:phase_correct[p]/n if n else None for p,n in phase_n.items()},
83            'phase_counts':dict(phase_n),'teacher_forced_triple_legality':total['legal_triples']/max(1,total['moves']),
84            'raw_reserved_probability':total['reserved_probability_sum']/total['targets'],
85            'raw_illegal_token_probability':total['illegal_probability_sum']/max(1,total['legality_targets']),
86            'legality_games':total['legality_games'],'legality_targets':total['legality_targets'],
87            'legality_subset':'N evenly spaced game indices across immutable validation shard order; N=None means all',
88            'protocol':'raw128 teacher-forced next-token logits; triple destination/promotion conditioned on gold prefix; not free-running',
89            'seconds':time.perf_counter()-start}
90
91
92def terminal_resume_status(saved, fallback_epochs):
93    """Patience/finished schedule are terminal; a time-budget pause is resumable."""
94    if saved['metrics']['bad_checks']>=5:
95        return 'early_stop_5_checks'
96    maximum=saved.get('identity',{}).get('max_updates')
97    if maximum is not None and saved['cursor']['step']>=maximum:
98        return 'max_updates'
99    scheduled=saved.get('identity',{}).get('epochs_schedule',fallback_epochs)
100    if saved['cursor']['epoch']>=scheduled:
101        return 'max_epochs'
102    return None
103
104
105def main(args):
106    # Read and verify the checkpoint on CPU before touching CUDA, DDP, caches or
107    # reports. Completed runs are idempotent no-ops, even through torchrun. No
108    # implicit patience reset or fresh training allowance is granted by --resume.
109    resume_saved=c.load_checkpoint(args.resume) if args.resume else None
110    if resume_saved is not None:
111        terminal=terminal_resume_status(resume_saved,args.epochs)
112        if terminal:
113            receipt={'status':terminal,'resume_action':'no_op_terminal',
114                     'cursor':resume_saved['cursor'],'best_step':resume_saved['metrics']['best_step'],
115                     'bad_checks':resume_saved['metrics']['bad_checks'],
116                     'optimizer_updates':0,'validation_checks':0,'checkpoint_written':False}
117            if int(os.environ.get('RANK','0'))==0:print(json.dumps(receipt),flush=True)
118            return receipt
119    deadline=Deadline(read_deadline(os.environ['CHESS_JOB_DEADLINE'])) if os.environ.get('CHESS_JOB_DEADLINE') else None
120    if deadline is not None and deadline.soft_stop():raise TimeoutError('job deadline reached before CUDA setup')
121    rank=int(os.environ.get('RANK','0'));local=int(os.environ.get('LOCAL_RANK','0'));world=int(os.environ.get('WORLD_SIZE','1'))
122    if world!=2:raise ValueError('baseline requires exactly two DDP workers')
123    torch.set_num_threads(4);torch.cuda.set_device(local);device=f'cuda:{local}'
124    rendezvous=os.environ.get('CHESS_DDP_INIT_FILE')
125    if rendezvous:
126        dist.init_process_group('nccl',init_method='file://'+str(Path(rendezvous).resolve()),rank=rank,world_size=world,timeout=timedelta(seconds=120))
127    else:dist.init_process_group('nccl',timeout=timedelta(seconds=120))
128    root=Path(args.output);root.mkdir(parents=True,exist_ok=True)
129    stop=threading.Event();monitor=None;status='started';started=time.perf_counter()
130    try:
131        # Rank 0 validates and caches train/validation only. Reader preflights hashes
132        # of all splits, but the test tokens are never forwarded to the model.
133        manifest=json.loads((Path(args.dataset)/'manifest.json').read_text())
134        if rank==0:
135            if getattr(args,'integrity_receipt',None):
136                from verified_training_data import load
137                data=load(args.dataset,args.integrity_receipt)
138            else:data={s:[tokens for tokens,meta in games(args.dataset,s)] for s in ('train','validation')}
139            if not data['train'] or not data['validation']:raise ValueError('nonempty train and validation required')
140            c.atomic_save(root/'data-cache.pt',data)
141            log(root,{'stage':'data_validated','train_games':len(data['train']),'validation_games':len(data['validation']),
142                      'train_targets':sum(len(s)-1 for s in data['train']),'test_evaluated':False})
143        dist.barrier()
144        data=torch.load(root/'data-cache.pt',weights_only=True);train=data['train'];val=data['validation']
145        lengths=[len(s)-1 for s in train]
146        torch.manual_seed(args.seed);random.seed(args.seed)
147        net=c.model(device=device)
148        count=sum(p.numel() for p in net.parameters())
149        if count!=85940736:raise ValueError(f'baseline parameter mismatch: {count}')
150        if any(p.dtype!=torch.float32 for p in net.parameters()):raise ValueError('expected FP32 weights')
151        wrapped=DDP(net,device_ids=[local],gradient_as_bucket_view=True)
152        opt=c.optimizer(net,args.lr,fused=True)
153        step_count=len(c.groups(c.batches(lengths,0,args.seed,args.micro,world),args.accum))
154        schedule_steps=getattr(args,'schedule_steps',None) or args.epochs*step_count
155        warmup_steps=getattr(args,'warmup_steps',None) or max(1,round(schedule_steps*.1))
156        if warmup_steps>=schedule_steps:raise ValueError('warmup must be shorter than scheduler horizon')
157        sched=c.scheduler(opt,schedule_steps,warmup_steps)
158        identity={'tokenizer':{'version':VERSION,'sha256':file_hash(Path(__file__).parent/'chessdata'/'tokenizer.py')},
159                  'model':c.BASELINE,'parameters':count,'dataset_fingerprint':manifest['fingerprint'],
160                  'training_code':{f:file_hash(Path(__file__).parent/f) for f in ('training_core.py','train_chess.py','training_deadline.py','verified_training_data.py')},
161                  'world_size':world,'seed':args.seed,'microbatch':args.micro,'accumulation':args.accum,
162                  'launcher':'owned-process-group-file-store' if rendezvous else 'external-env-rendezvous',
163                  'lr':args.lr,'epochs_schedule':args.epochs,'optimizer':'FP32 fused AdamW betas .9/.95 matrix wd .1 clip1',
164                  'precision':'FP32 weights/gradients/optimizer; BF16 autocast SDPA','torch':str(torch.__version__),
165                  'sampling':'seeded length buckets, disjoint full games; right pad; zero-target dummy rank on partial global batches'}
166        if deadline is not None:identity['total_job_deadline']=deadline.value
167        identity['validation_every_steps']=getattr(args,'validation_every_steps',None)
168        identity['legality_eval_games']=getattr(args,'legality_eval_games',None)
169        identity['schedule_steps']=schedule_steps;identity['warmup_steps']=warmup_steps
170        identity['max_updates']=getattr(args,'max_updates',None)
171        if getattr(args,'integrity_receipt',None):identity['integrity_receipt_sha256']=file_hash(args.integrity_receipt)
172        cursor={'epoch':0,'next_group':0,'step':0};history=[];best=float('inf');bad=0;best_step=None
173        if args.resume:
174            saved=resume_saved
175            if saved['identity']!=identity:raise ValueError('incompatible checkpoint identity')
176            net.load_state_dict(saved['model_state_dict'])
177            opt.load_state_dict(saved['optimizer_state_dict']);sched.load_state_dict(saved['scheduler_state_dict'])
178            cursor=saved['cursor'];history=saved['metrics']['history'];best=saved['metrics']['best_nll'];bad=saved['metrics']['bad_checks'];best_step=saved['metrics']['best_step']
179            c.restore_rng(saved['rng_by_rank'][rank],device)
180        dist.barrier();torch.cuda.synchronize();gpu_start=time.perf_counter();torch.cuda.reset_peak_memory_stats()
181        if rank==0:
182            monitor=threading.Thread(target=telemetry,args=(root,stop),daemon=True);monitor.start()
183            (root/'identity.json').write_text(json.dumps(identity,indent=2))
184        if args.mode=='probe':
185            longest=max(train,key=len);probe=[longest]*16;results=[]
186            for micro in (2,8):
187                torch.cuda.reset_peak_memory_stats();times=[]
188                for rep in range(3):
189                    opt.zero_grad(set_to_none=True);dist.barrier();torch.cuda.synchronize();t=time.perf_counter()
190                    values=c.backward_group(wrapped,probe,[list(range(micro*world))],micro,rank,world,device,True)
191                    torch.nn.utils.clip_grad_norm_(net.parameters(),1.0);opt.step();torch.cuda.synchronize();elapsed=time.perf_counter()-t
192                    if rep:times.append(elapsed)
193                memory=[None]*world;dist.all_gather_object(memory,{'rank':rank,'peak_allocated':torch.cuda.max_memory_allocated(),'peak_reserved':torch.cuda.max_memory_reserved()})
194                results.append({'microbatch_per_gpu':micro,'sequence_length':len(longest)-1,'timed_steps':2,
195                                'step_seconds':times,'useful_targets_per_second':values['valid_targets']*2/sum(times),
196                                'padded_tokens_per_second':values['padded_tokens']*2/sum(times),'memory':memory})
197            if rank==0:
198                result={'status':'passed','parameters':count,'results':results,'gpu_work_wall_seconds':time.perf_counter()-gpu_start,
199                        'scope':'six optimizer steps total, one warmup per shape; repeated longest pilot game; discarded weights; not production throughput'}
200                (root/'probe-results.json').write_text(json.dumps(result,indent=2));log(root,{'stage':'probe_complete',**result})
201            return
202        initial=[evaluate(net,val,device,deadline=deadline,legality_games=getattr(args,'legality_eval_games',None)) if rank==0 and not args.resume else None]
203        dist.broadcast_object_list(initial,src=0)
204        if rank==0 and initial[0]:log(root,{'stage':'initial_validation',**initial[0]})
205        totals=Counter();stop_training=False
206        def save_state(improved):
207            rngs=[None]*world;dist.all_gather_object(rngs,c.rng_state(device))
208            if rank==0:
209                metrics={'history':history,'best_nll':best,'bad_checks':bad,'best_step':best_step,'test_evaluated':False}
210                c.atomic_save(root/'last.pt',c.checkpoint_payload(wrapped,opt,sched,identity,cursor,rngs,metrics))
211                if improved:c.link_best(root/'last.pt',root/'best.pt')
212            dist.barrier()
213        while cursor['epoch']<args.epochs and not stop_training:
214            epoch=cursor['epoch'];plan=c.groups(c.batches(lengths,epoch,args.seed,args.micro,world),args.accum)
215            epoch_totals=Counter();epoch_start=time.perf_counter();completed=False
216            for group_index in range(cursor['next_group'],len(plan)):
217                maximum=getattr(args,'max_updates',None)
218                reason='max_updates' if maximum is not None and cursor['step']>=maximum else None
219                if reason is None and (deadline.soft_stop() if deadline is not None else time.perf_counter()-gpu_start>=args.max_seconds):reason='time_budget'
220                should_stop=[reason if rank==0 else None]
221                dist.broadcast_object_list(should_stop,src=0)
222                if should_stop[0]:status=should_stop[0];stop_training=True;break
223                opt.zero_grad(set_to_none=True);torch.cuda.synchronize();t=time.perf_counter()
224                values=c.backward_group(wrapped,train,plan[group_index],args.micro,rank,world,device,True)
225                norm=torch.nn.utils.clip_grad_norm_(net.parameters(),1.,error_if_nonfinite=True)
226                opt.step();sched.step();torch.cuda.synchronize();elapsed=time.perf_counter()-t
227                values['train_seconds']=elapsed;totals.update(values);epoch_totals.update(values)
228                cursor={'epoch':epoch,'next_group':group_index+1,'step':cursor['step']+1}
229                if rank==0 and cursor['step']%10==0:
230                    log(root,{'stage':'train','epoch':epoch+1,'step':cursor['step'],'nll':values['loss_sum']/values['valid_targets'],
231                              'lr':sched.get_last_lr()[0],'gradient_norm_before_clip':float(norm),
232                              'useful_targets_per_second':values['valid_targets']/elapsed,'padded_tokens_per_second':values['padded_tokens']/elapsed,
233                              'padding_fraction':1-values['valid_targets']/values['padded_tokens'],'gpu_work_wall_seconds':time.perf_counter()-gpu_start})
234                cadence=getattr(args,'validation_every_steps',None)
235                if maximum is not None and cursor['step']>=maximum:
236                    status='max_updates';stop_training=True;break
237                if cadence and cursor['step']%cadence==0:break
238            if cursor['next_group']==len(plan):
239                completed=True;cursor={'epoch':epoch+1,'next_group':0,'step':cursor['step']}
240            dist.barrier()
241            check=[evaluate(net,val,device,deadline=deadline,legality_games=getattr(args,'legality_eval_games',None)) if rank==0 else None];dist.broadcast_object_list(check,src=0)
242            if check[0] is None:
243                status='total_deadline_soft_stop';stop_training=True
244                save_state(False)
245                break
246            metrics={'epoch':epoch+1,'complete_epoch':completed,'step':cursor['step'],'validation':check[0],
247                     'train_nll':epoch_totals['loss_sum']/max(1,epoch_totals['valid_targets']),
248                     'useful_targets_per_second':epoch_totals['valid_targets']/max(1e-9,epoch_totals['train_seconds']),
249                     'padded_tokens_per_second':epoch_totals['padded_tokens']/max(1e-9,epoch_totals['train_seconds']),
250                     'padding_fraction':1-epoch_totals['valid_targets']/max(1,epoch_totals['padded_tokens'])}
251            history.append(metrics);improved=check[0]['nll']<best-1e-4
252            if improved:best=check[0]['nll'];bad=0;best_step=cursor['step']
253            else:bad+=1
254            if rank==0:log(root,{'stage':'validation','best_nll':best,'bad_checks':bad,**metrics})
255            save_state(improved)
256            if bad>=5:status='early_stop_5_checks';stop_training=True
257        if status=='started':status='max_epochs'
258        memory=[None]*world;dist.all_gather_object(memory,{'rank':rank,'peak_allocated':torch.cuda.max_memory_allocated(),'peak_reserved':torch.cuda.max_memory_reserved()})
259        if rank==0:
260            result={'status':status,'cursor':cursor,'best_step':best_step,'best_validation_nll':best,'history':history,
261                    'initial_validation':initial[0],'totals':dict(totals),'memory':memory,'gpu_work_wall_seconds':time.perf_counter()-gpu_start,
262                    'whole_process_seconds':time.perf_counter()-started,'test_evaluated':False,
263                    'best_checkpoint':'best.pt' if (root/'best.pt').exists() else None,'best_checkpoint_sha256':file_hash(root/'best.pt') if (root/'best.pt').exists() else None,
264                    'last_checkpoint_sha256':file_hash(root/'last.pt'),'identity_sha256':c.digest_json(identity)}
265            (root/'results.json').write_text(json.dumps(result,indent=2));log(root,{'stage':'finished','status':status,'best_step':best_step,'best_nll':best})
266    except BaseException as error:
267        (root/f'failure-rank{rank}.json').write_text(json.dumps({'error':repr(error),'time_unix':time.time()}))
268        raise
269    finally:
270        stop.set()
271        if monitor:monitor.join(timeout=3)
272        dist.destroy_process_group()
273
274
275if __name__=='__main__':
276    p=argparse.ArgumentParser();p.add_argument('--dataset',default='lichess-pilot-202609-20261028/training-context1024')
277    p.add_argument('--output',required=True);p.add_argument('--mode',choices=['probe','train'],default='train')
278    p.add_argument('--micro',type=int,default=8);p.add_argument('--accum',type=int,default=2);p.add_argument('--lr',type=float,default=3e-4)
279    p.add_argument('--epochs',type=int,default=20);p.add_argument('--seed',type=int,default=2026);p.add_argument('--max-seconds',type=int,default=420);p.add_argument('--resume')
280    p.add_argument('--validation-every-steps',type=int);p.add_argument('--legality-eval-games',type=int)
281    p.add_argument('--integrity-receipt')
282    p.add_argument('--schedule-steps',type=int);p.add_argument('--warmup-steps',type=int)
283    p.add_argument('--max-updates',type=int)
284    p.add_argument('--launch-ddp',action='store_true')
285    args=p.parse_args()
286    if not 1<=args.epochs<=20 or not 1<=args.max_seconds<=420 or args.micro<1 or args.accum<1:p.error('bounded calibration limits violated')
287    if args.validation_every_steps is not None and args.validation_every_steps<1:p.error('invalid validation cadence')
288    if args.legality_eval_games is not None and args.legality_eval_games<1:p.error('legality sample must be positive')
289    if any(v is not None and v<1 for v in (args.schedule_steps,args.warmup_steps)):p.error('scheduler steps must be positive')
290    if args.max_updates is not None and args.max_updates<1:p.error('max updates must be positive')
291    if args.launch_ddp:
292        raise SystemExit(c.launch_owned_workers(Path(__file__),[arg for arg in sys.argv[1:] if arg!='--launch-ddp']))
293    main(args)
294