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 collective_timeout(seconds):
106    """Allow rank-0 full-corpus I/O/evaluation; the job guardian stays independent."""
107    if type(seconds) is not int or not 60<=seconds<=1800:
108        raise ValueError('collective timeout must be an integer in [60, 1800] seconds')
109    return timedelta(seconds=seconds)
110
111
112def main(args):
113    # Read and verify the checkpoint on CPU before touching CUDA, DDP, caches or
114    # reports. Completed runs are idempotent no-ops, even through torchrun. No
115    # implicit patience reset or fresh training allowance is granted by --resume.
116    resume_saved=c.load_checkpoint(args.resume) if args.resume else None
117    if resume_saved is not None:
118        terminal=terminal_resume_status(resume_saved,args.epochs)
119        if terminal:
120            receipt={'status':terminal,'resume_action':'no_op_terminal',
121                     'cursor':resume_saved['cursor'],'best_step':resume_saved['metrics']['best_step'],
122                     'bad_checks':resume_saved['metrics']['bad_checks'],
123                     'optimizer_updates':0,'validation_checks':0,'checkpoint_written':False}
124            if int(os.environ.get('RANK','0'))==0:print(json.dumps(receipt),flush=True)
125            return receipt
126    deadline=Deadline(read_deadline(os.environ['CHESS_JOB_DEADLINE'])) if os.environ.get('CHESS_JOB_DEADLINE') else None
127    if deadline is not None and deadline.soft_stop():raise TimeoutError('job deadline reached before CUDA setup')
128    rank=int(os.environ.get('RANK','0'));local=int(os.environ.get('LOCAL_RANK','0'));world=int(os.environ.get('WORLD_SIZE','1'))
129    if world!=2:raise ValueError('baseline requires exactly two DDP workers')
130    torch.set_num_threads(4);torch.cuda.set_device(local);device=f'cuda:{local}'
131    rendezvous=os.environ.get('CHESS_DDP_INIT_FILE')
132    timeout_seconds=getattr(args,'collective_timeout_seconds',900)
133    timeout=collective_timeout(timeout_seconds)
134    if rendezvous:
135        dist.init_process_group('nccl',init_method='file://'+str(Path(rendezvous).resolve()),rank=rank,world_size=world,timeout=timeout)
136    else:dist.init_process_group('nccl',timeout=timeout)
137    root=Path(args.output);root.mkdir(parents=True,exist_ok=True)
138    stop=threading.Event();monitor=None;status='started';started=time.perf_counter()
139    try:
140        # Rank 0 validates and caches train/validation only. Reader preflights hashes
141        # of all splits, but the test tokens are never forwarded to the model.
142        manifest=json.loads((Path(args.dataset)/'manifest.json').read_text())
143        if rank==0:
144            if getattr(args,'integrity_receipt',None):
145                from verified_training_data import load
146                data=load(args.dataset,args.integrity_receipt)
147            else:data={s:[tokens for tokens,meta in games(args.dataset,s)] for s in ('train','validation')}
148            if not data['train'] or not data['validation']:raise ValueError('nonempty train and validation required')
149            c.atomic_save(root/'data-cache.pt',data)
150            log(root,{'stage':'data_validated','train_games':len(data['train']),'validation_games':len(data['validation']),
151                      'train_targets':sum(len(s)-1 for s in data['train']),'test_evaluated':False})
152        dist.barrier()
153        data=torch.load(root/'data-cache.pt',weights_only=True);train=data['train'];val=data['validation']
154        lengths=[len(s)-1 for s in train]
155        torch.manual_seed(args.seed);random.seed(args.seed)
156        net=c.model(device=device)
157        count=sum(p.numel() for p in net.parameters())
158        if count!=85940736:raise ValueError(f'baseline parameter mismatch: {count}')
159        if any(p.dtype!=torch.float32 for p in net.parameters()):raise ValueError('expected FP32 weights')
160        wrapped=DDP(net,device_ids=[local],gradient_as_bucket_view=True)
161        opt=c.optimizer(net,args.lr,fused=True)
162        bucket=getattr(args,'bucket',2048)
163        step_count=len(c.groups(c.batches(lengths,0,args.seed,args.micro,world,bucket),args.accum))
164        schedule_steps=getattr(args,'schedule_steps',None) or args.epochs*step_count
165        warmup_steps=getattr(args,'warmup_steps',None) or max(1,round(schedule_steps*.1))
166        if warmup_steps>=schedule_steps:raise ValueError('warmup must be shorter than scheduler horizon')
167        sched=c.scheduler(opt,schedule_steps,warmup_steps)
168        identity={'tokenizer':{'version':VERSION,'sha256':file_hash(Path(__file__).parent/'chessdata'/'tokenizer.py')},
169                  'model':c.BASELINE,'parameters':count,'dataset_fingerprint':manifest['fingerprint'],
170                  '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')},
171                  'world_size':world,'seed':args.seed,'microbatch':args.micro,'accumulation':args.accum,
172                  'launcher':'owned-process-group-file-store' if rendezvous else 'external-env-rendezvous',
173                  'lr':args.lr,'epochs_schedule':args.epochs,'optimizer':'FP32 fused AdamW betas .9/.95 matrix wd .1 clip1',
174                  'precision':'FP32 weights/gradients/optimizer; BF16 autocast SDPA','torch':str(torch.__version__),
175                  'sampling':'seeded length buckets, disjoint full games; right pad; zero-target dummy rank on partial global batches'}
176        if deadline is not None:identity['total_job_deadline']=deadline.value
177        identity['validation_every_steps']=getattr(args,'validation_every_steps',None)
178        identity['legality_eval_games']=getattr(args,'legality_eval_games',None)
179        identity['schedule_steps']=schedule_steps;identity['warmup_steps']=warmup_steps
180        identity['max_updates']=getattr(args,'max_updates',None)
181        identity['collective_timeout_seconds']=timeout_seconds
182        identity['length_bucket_games']=bucket
183        if getattr(args,'integrity_receipt',None):identity['integrity_receipt_sha256']=file_hash(args.integrity_receipt)
184        cursor={'epoch':0,'next_group':0,'step':0};history=[];best=float('inf');bad=0;best_step=None
185        if args.resume:
186            saved=resume_saved
187            if saved['identity']!=identity:raise ValueError('incompatible checkpoint identity')
188            net.load_state_dict(saved['model_state_dict'])
189            opt.load_state_dict(saved['optimizer_state_dict']);sched.load_state_dict(saved['scheduler_state_dict'])
190            cursor=saved['cursor'];history=saved['metrics']['history'];best=saved['metrics']['best_nll'];bad=saved['metrics']['bad_checks'];best_step=saved['metrics']['best_step']
191            c.restore_rng(saved['rng_by_rank'][rank],device)
192        dist.barrier();torch.cuda.synchronize();gpu_start=time.perf_counter();torch.cuda.reset_peak_memory_stats()
193        if rank==0:
194            monitor=threading.Thread(target=telemetry,args=(root,stop),daemon=True);monitor.start()
195            (root/'identity.json').write_text(json.dumps(identity,indent=2))
196        if args.mode=='probe':
197            longest=max(train,key=len);micro=args.micro
198            probe=[longest]*(micro*world*args.accum);results=[]
199            representative=c.groups(c.batches(lengths,0,args.seed,micro,world,bucket),args.accum)
200            for shape in ('representative','longest'):
201                torch.cuda.reset_peak_memory_stats();times=[];target_counts=[];padded_counts=[]
202                for rep in range(3):
203                    if deadline is not None and deadline.soft_stop():raise TimeoutError('total deadline reached during probe')
204                    samples=train if shape=='representative' else probe
205                    group=(representative[min(len(representative)-1,rep*len(representative)//3)] if shape=='representative'
206                           else [list(range(i*micro*world,(i+1)*micro*world)) for i in range(args.accum)])
207                    opt.zero_grad(set_to_none=True);dist.barrier();torch.cuda.synchronize();t=time.perf_counter()
208                    values=c.backward_group(wrapped,samples,group,micro,rank,world,device,True)
209                    torch.nn.utils.clip_grad_norm_(net.parameters(),1.0);opt.step();torch.cuda.synchronize();elapsed=time.perf_counter()-t
210                    if rep:times.append(elapsed);target_counts.append(values['valid_targets']);padded_counts.append(values['padded_tokens'])
211                memory=[None]*world;free,total=torch.cuda.mem_get_info()
212                dist.all_gather_object(memory,{'rank':rank,'peak_allocated':torch.cuda.max_memory_allocated(),'peak_reserved':torch.cuda.max_memory_reserved(),
213                    'device_free_bytes_after_shape':free,'device_total_bytes':total})
214                results.append({'shape':shape,'microbatch_per_gpu':micro,'accumulation':args.accum,'effective_games_per_update':micro*world*args.accum,
215                                'maximum_sequence_length':len(longest)-1,'timed_steps':2,'step_seconds':times,'valid_targets':target_counts,
216                                'useful_targets_per_second':sum(target_counts)/sum(times),'padded_tokens_per_second':sum(padded_counts)/sum(times),'memory':memory})
217            if rank==0:
218                result={'status':'passed','parameters':count,'results':results,'gpu_work_wall_seconds':time.perf_counter()-gpu_start,
219                        'scope':'six optimizer steps total; representative seeded groups and repeated longest full game; one warmup per shape; discarded weights; not full-run throughput'}
220                (root/'probe-results.json').write_text(json.dumps(result,indent=2));log(root,{'stage':'probe_complete',**result})
221            return
222        initial=[evaluate(net,val,device,deadline=deadline,legality_games=getattr(args,'legality_eval_games',None)) if rank==0 and not args.resume else None]
223        dist.broadcast_object_list(initial,src=0)
224        if rank==0 and initial[0]:log(root,{'stage':'initial_validation',**initial[0]})
225        totals=Counter();stop_training=False
226        def save_state(improved):
227            rngs=[None]*world;dist.all_gather_object(rngs,c.rng_state(device))
228            if rank==0:
229                metrics={'history':history,'best_nll':best,'bad_checks':bad,'best_step':best_step,'test_evaluated':False}
230                c.atomic_save(root/'last.pt',c.checkpoint_payload(wrapped,opt,sched,identity,cursor,rngs,metrics))
231                if improved:c.link_best(root/'last.pt',root/'best.pt')
232            dist.barrier()
233        while cursor['epoch']<args.epochs and not stop_training:
234            epoch=cursor['epoch'];plan=c.groups(c.batches(lengths,epoch,args.seed,args.micro,world,bucket),args.accum)
235            epoch_totals=Counter();epoch_start=time.perf_counter();completed=False
236            for group_index in range(cursor['next_group'],len(plan)):
237                maximum=getattr(args,'max_updates',None)
238                reason='max_updates' if maximum is not None and cursor['step']>=maximum else None
239                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'
240                should_stop=[reason if rank==0 else None]
241                dist.broadcast_object_list(should_stop,src=0)
242                if should_stop[0]:status=should_stop[0];stop_training=True;break
243                opt.zero_grad(set_to_none=True);torch.cuda.synchronize();t=time.perf_counter()
244                values=c.backward_group(wrapped,train,plan[group_index],args.micro,rank,world,device,True)
245                norm=torch.nn.utils.clip_grad_norm_(net.parameters(),1.,error_if_nonfinite=True)
246                opt.step();sched.step();torch.cuda.synchronize();elapsed=time.perf_counter()-t
247                values['train_seconds']=elapsed;totals.update(values);epoch_totals.update(values)
248                cursor={'epoch':epoch,'next_group':group_index+1,'step':cursor['step']+1}
249                if rank==0 and cursor['step']%10==0:
250                    log(root,{'stage':'train','epoch':epoch+1,'step':cursor['step'],'nll':values['loss_sum']/values['valid_targets'],
251                              'lr':sched.get_last_lr()[0],'gradient_norm_before_clip':float(norm),
252                              'useful_targets_per_second':values['valid_targets']/elapsed,'padded_tokens_per_second':values['padded_tokens']/elapsed,
253                              'padding_fraction':1-values['valid_targets']/values['padded_tokens'],'gpu_work_wall_seconds':time.perf_counter()-gpu_start})
254                cadence=getattr(args,'validation_every_steps',None)
255                if maximum is not None and cursor['step']>=maximum:
256                    status='max_updates';stop_training=True;break
257                if cadence and cursor['step']%cadence==0:break
258            if cursor['next_group']==len(plan):
259                completed=True;cursor={'epoch':epoch+1,'next_group':0,'step':cursor['step']}
260            dist.barrier()
261            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)
262            if check[0] is None:
263                status='total_deadline_soft_stop';stop_training=True
264                save_state(False)
265                break
266            metrics={'epoch':epoch+1,'complete_epoch':completed,'step':cursor['step'],'validation':check[0],
267                     'train_nll':epoch_totals['loss_sum']/max(1,epoch_totals['valid_targets']),
268                     'useful_targets_per_second':epoch_totals['valid_targets']/max(1e-9,epoch_totals['train_seconds']),
269                     'padded_tokens_per_second':epoch_totals['padded_tokens']/max(1e-9,epoch_totals['train_seconds']),
270                     'padding_fraction':1-epoch_totals['valid_targets']/max(1,epoch_totals['padded_tokens'])}
271            history.append(metrics);improved=check[0]['nll']<best-1e-4
272            if improved:best=check[0]['nll'];bad=0;best_step=cursor['step']
273            else:bad+=1
274            if rank==0:log(root,{'stage':'validation','best_nll':best,'bad_checks':bad,**metrics})
275            save_state(improved)
276            if bad>=5:status='early_stop_5_checks';stop_training=True
277        if status=='started':status='max_epochs'
278        memory=[None]*world;dist.all_gather_object(memory,{'rank':rank,'peak_allocated':torch.cuda.max_memory_allocated(),'peak_reserved':torch.cuda.max_memory_reserved()})
279        if rank==0:
280            result={'status':status,'cursor':cursor,'best_step':best_step,'best_validation_nll':best,'history':history,
281                    'initial_validation':initial[0],'totals':dict(totals),'memory':memory,'gpu_work_wall_seconds':time.perf_counter()-gpu_start,
282                    'whole_process_seconds':time.perf_counter()-started,'test_evaluated':False,
283                    '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,
284                    'last_checkpoint_sha256':file_hash(root/'last.pt'),'identity_sha256':c.digest_json(identity)}
285            (root/'results.json').write_text(json.dumps(result,indent=2));log(root,{'stage':'finished','status':status,'best_step':best_step,'best_nll':best})
286    except BaseException as error:
287        (root/f'failure-rank{rank}.json').write_text(json.dumps({'error':repr(error),'time_unix':time.time()}))
288        raise
289    finally:
290        stop.set()
291        if monitor:monitor.join(timeout=3)
292        dist.destroy_process_group()
293
294
295if __name__=='__main__':
296    p=argparse.ArgumentParser();p.add_argument('--dataset',default='lichess-pilot-202609-20261028/training-context1024')
297    p.add_argument('--output',required=True);p.add_argument('--mode',choices=['probe','train'],default='train')
298    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)
299    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')
300    p.add_argument('--validation-every-steps',type=int);p.add_argument('--legality-eval-games',type=int)
301    p.add_argument('--integrity-receipt')
302    p.add_argument('--schedule-steps',type=int);p.add_argument('--warmup-steps',type=int)
303    p.add_argument('--max-updates',type=int)
304    p.add_argument('--bucket',type=int,default=2048)
305    p.add_argument('--collective-timeout-seconds',type=int,default=900)
306    p.add_argument('--launch-ddp',action='store_true')
307    args=p.parse_args()
308    try:collective_timeout(args.collective_timeout_seconds)
309    except ValueError as error:p.error(str(error))
310    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')
311    if args.validation_every_steps is not None and args.validation_every_steps<1:p.error('invalid validation cadence')
312    if args.legality_eval_games is not None and args.legality_eval_games<1:p.error('legality sample must be positive')
313    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')
314    if args.max_updates is not None and args.max_updates<1:p.error('max updates must be positive')
315    if args.bucket<args.micro*2 or args.bucket>8192 or args.bucket%(args.micro*2):p.error('bucket must fit whole global microbatches and be <=8192')
316    if args.launch_ddp:
317        raise SystemExit(c.launch_owned_workers(Path(__file__),[arg for arg in sys.argv[1:] if arg!='--launch-ddp']))
318    main(args)
319