train.py

1"""Reviewed finite-schedule completion with a separately authorized deadline and loss selection."""
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 stage_data import open_curriculum,Plan,start_prefetch,ExposureLedger,phase,tensors
11from checkpoint import capture_rng,restore_rng
12from model import StateModel,optimizer
13from resume_state import validate_source,restore_optimizer_scheduler,select_loss
14from optimizer_policy import lr_factor
15from evaluation import evaluate
16from deadline import Deadline,read_deadline
17from records import sha,digest,atomic_json,Records as SmallRecords
18from runtime_records import Records
19sys.path.append(str(Path(__file__).resolve().parent.parent/'elite-continuation-v1'))
20from training_core import atomic_save,launch_owned_workers
21from curriculum_loss import weighted_loss
22from exact_metrics import summarize
23
24def verify(cfg,args):
25    if cfg.get('draft_only'):raise ValueError('draft configuration is not launchable')
26    for path,expected in cfg['runtime_sha256'].items():
27        if sha(path)!=expected:raise ValueError('runtime changed: '+path)
28    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')
29    if sha(cfg['source_checkpoint'])!=cfg['source_checkpoint_sha256']:raise ValueError('resume source changed')
30    for path,expected in cfg['validation_sha256'].items():
31        if sha(path)!=expected:raise ValueError('validation changed')
32    if args.resume and (not args.resume_sha256 or sha(args.resume)!=args.resume_sha256):raise ValueError('resume hash')
33
34def capacity(cfg,extra=1_200_000_000):
35    out=Path(cfg['output']);used=sum(p.stat().st_size for p in Path(cfg['original_root']).rglob('*') if p.is_file())
36    if used+extra>cfg['maximum_output_bytes']:raise RuntimeError('checkpoint output storage cap')
37    if shutil.disk_usage(out).free<cfg['native_free_reserve_bytes']+extra:raise RuntimeError('native storage reserve')
38
39def resource_check(cfg):
40    capacity(cfg)
41    if shutil.disk_usage(cfg['output']).free<cfg['native_free_reserve_bytes']+cfg['maximum_output_bytes']:raise RuntimeError('full run storage reservation')
42    health=json.loads(Path('/work/littleply-serving/state/health.json').read_text())
43    if not health.get('ready') or time.time()-health['updated']>180:raise RuntimeError('live bot not freshly ready')
44    gpu=subprocess.run(['nvidia-smi','--query-gpu=index,memory.free','--format=csv,noheader,nounits'],capture_output=True,text=True,check=True,timeout=15).stdout
45    free={int(line.split(',')[0]):int(line.split(',')[1]) for line in gpu.splitlines()}
46    if set(free)!={0,1} or min(free.values())<16384:raise RuntimeError('both GPUs require16GiB before mandatory probe')
47    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})
48
49def probe(net,data,device,rank,clock):
50    reserve=[torch.empty_like(p) for p in net.parameters() for _ in range(2)];out=[]
51    for micro in (16,32,64):
52        if clock.remaining()<600:raise TimeoutError('insufficient completion-stage probe budget')
53        rows=data.take(np.arange(micro));x,y,m=[v.to(device) for v in tensors(rows)]
54        torch.cuda.reset_peak_memory_stats(device);began=time.monotonic()
55        for _ in range(2):
56            with torch.autocast('cuda',dtype=torch.bfloat16):loss=torch.nn.functional.cross_entropy(net(x,m).float().reshape(-1,512),y.reshape(-1))
57            loss.backward();net.zero_grad(set_to_none=True)
58        torch.cuda.synchronize();free,_=torch.cuda.mem_get_info(device)
59        good=torch.tensor(int(free>=4*1024**3),device=device);dist.all_reduce(good,op=dist.ReduceOp.MIN)
60        if not good.item():raise RuntimeError('probe headroom below4GiB; no fallback')
61        out.append(dict(rank=rank,micro=micro,seconds=time.monotonic()-began,free_bytes=free,peak_reserved_bytes=torch.cuda.max_memory_reserved(device)))
62    del reserve
63    return out
64
65def worker(cfg,args):
66    rank=int(os.environ['RANK']);world=int(os.environ['WORLD_SIZE'])
67    if world!=2:raise ValueError('two ranks required')
68    device=torch.device('cuda',rank);torch.cuda.set_device(device);torch.set_num_threads(4)
69    random.seed(cfg['seed']+rank);np.random.seed(cfg['seed']+rank);torch.manual_seed(cfg['seed']+rank);torch.cuda.manual_seed(cfg['seed']+rank)
70    deadline_value=read_deadline(os.environ['CHESS_JOB_DEADLINE']);clock=Deadline(deadline_value)
71    dist.init_process_group('nccl',init_method='file://'+os.environ['CHESS_DDP_INIT_FILE'],rank=rank,world_size=world,timeout=timedelta(seconds=180))
72    data=None;prefetch=None;baseline_ref=None
73    try:
74        data=open_curriculum(cfg['dataset'],cfg['adapter_directory']);ledger=ExposureLedger(data);plan=Plan(data,rank,micro=64)
75        saved=torch.load(cfg['source_checkpoint'],map_location='cpu',weights_only=True)
76        validate_source(saved,cfg,deadline_value,ledger,plan)
77        raw=StateModel(saved['model_config']);raw.load_state_dict(saved['model_state_dict'],strict=True);raw.to(device);net=DDP(raw,device_ids=[rank],gradient_as_bucket_view=True)
78        initialization={'source_checkpoint_sha256':cfg['source_checkpoint_sha256'],'source_updates':saved['updates'],'source_cursor':saved['exposures']['cursor'],'resume':'model/Adam/scheduler/sampler/both-rank RNG; separately authorized completion deadline'}
79        proofs=probe(net,data,device,rank,clock);both=[None,None];dist.all_gather_object(both,proofs)
80        out=Path(cfg['output']);opt=optimizer(raw,cfg['lr'],fused=True);accum=4;total_updates=math.ceil(plan.steps/accum)
81        def scale(step):return lr_factor(step,total_updates,cfg['warmup_updates'])
82        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((27,3),dtype=torch.float64,device=device)
83        attempt='completion'
84        opt.retention_proof=restore_optimizer_scheduler(opt,sched,saved)
85        cursor=saved['microstep'];updates=saved['updates'];latest=saved['latest_validation'];best=saved['best_nll'];best_validation=saved['best_validation'];initial_validation=saved['initial_validation']
86        group_totals.copy_(saved['rank_group_totals'][rank].to(device));restore_rng(saved['rank_rng'][rank],device)
87        del saved
88        if args.resume:raise ValueError('single-attempt stage; no resume')
89        guard_failed=[];baseline_ref=None;baseline_validation=None
90        if rank==0:atomic_json(out/('probe-resume.json' if args.resume else 'probe.json'),{'ranks':both,'initialization':initialization})
91        dummy,prefetch=start_prefetch(data,plan,cursor,True)
92        log=(out/'metrics.jsonl').open('a',buffering=1) if rank==0 else None
93        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')
94        def checkpoint(name,terminal=False,is_best=False):
95            rngs=[None,None];dist.all_gather_object(rngs,capture_rng(device));rank_totals=[None,None];dist.all_gather_object(rank_totals,group_totals.cpu())
96            if rank==0:
97                capacity(cfg)
98                atomic_save(out/name,{'format':'cm7-theory-completion-v1','identity':cfg,'identity_sha256':digest(cfg),'deadline':deadline_value,'terminal':terminal,
99                    '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(),
100                    'rank_rng':rngs,'rank_group_totals':rank_totals,'exposures':ledger.snapshot(plan.cursor(cursor)),'microstep':cursor,'updates':updates,'best_nll':best,'latest_validation':latest,
101                    'best_checkpoint':best_ref,'is_best_snapshot':is_best,'attempt':attempt,'best_validation':best_validation,'initial_validation':initial_validation,'retention_failures':guard_failed,'baseline_checkpoint':baseline_ref,'optimizer_retention':opt.retention_proof})
102            dist.barrier()
103        def validation():
104            nonlocal latest,best,best_ref,best_validation
105            improved=False
106            if rank==0:
107                began=time.time();latest=evaluate(raw,cfg,device);latest.update(unix=time.time(),updates=updates,seconds=time.time()-began)
108                improved=select_loss(latest,best);log.write(json.dumps(dict(event='validation',retention_failures=[],accuracy_diagnostic_only=True,**latest))+'\n')
109            state=[latest,improved];dist.broadcast_object_list(state,0);latest,improved=state
110            if improved:
111                best_validation=latest;best=latest['nll'];name=f'best-{updates}-{attempt}.pt';checkpoint(name,is_best=True)
112                ref=[{'path':str(out/name),'sha256':sha(out/name)} if rank==0 else None];dist.broadcast_object_list(ref,0);best_ref=ref[0]
113                if rank==0:atomic_json(out/'best.json',best_ref)
114        # Carry the previous loss-selected best independently of last-state resume.
115        if rank==0:
116            prior=torch.load(cfg['prior_selected']['path'],map_location='cpu',weights_only=True)
117            if sha(cfg['prior_selected']['path'])!=cfg['prior_selected']['sha256'] or prior['latest_validation']['nll']!=best:raise ValueError('prior selected loss/hash')
118            prior.update(format='cm7-theory-completion-v1',identity=cfg,identity_sha256=digest(cfg),deadline=deadline_value,is_best_snapshot=True,terminal=False,attempt='inherited-loss-best')
119            capacity(cfg);atomic_save(out/'best-inherited.pt',prior);del prior
120        dist.barrier()
121        ref=[{'path':str(out/'best-inherited.pt'),'sha256':sha(out/'best-inherited.pt')} if rank==0 else None];dist.broadcast_object_list(ref,0);best_ref=baseline_ref=ref[0]
122        if rank==0:
123            from finalization import atomic
124            atomic(out/'baseline.json',{'checkpoint':baseline_ref,'config_sha256':digest(cfg)})
125            atomic_json(out/'best.json',best_ref);atomic_json(out/'initial-validation.json',initial_validation)
126            atomic_json(out/'resume-receipt.json',{'updates':updates,'microstep':cursor,'cursor':plan.cursor(cursor),'optimizer':opt.retention_proof,'source_checkpoint_sha256':cfg['source_checkpoint_sha256'],'authorized_completion_deadline':deadline_value,'prior_selected':cfg['prior_selected']})
127        dist.barrier()
128        last_time=time.monotonic();last_cursor=plan.cursor(cursor)
129        while cursor<plan.steps and not guard_failed:
130            stop=torch.tensor(int(clock.soft_stop()),device=device);dist.all_reduce(stop,op=dist.ReduceOp.MAX)
131            if stop.item():break
132            start_cursor=cursor;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)
133            wait_seconds=load_seconds=0.;began=time.monotonic()
134            for step in range(cursor,end):
135                t=time.monotonic();batch,tags,load=prefetch.take(step);wait_seconds+=time.monotonic()-t;load_seconds+=load;empty=len(batch[0])==0
136                x,y,mask=[v.to(device,non_blocking=True) for v in (dummy if empty else batch)]
137                with net.no_sync() if step<end-1 else nullcontext():
138                    with torch.autocast('cuda',dtype=torch.bfloat16):logits=net(x,mask)
139                    loss=torch.nn.functional.cross_entropy(logits.float().reshape(-1,512),y.reshape(-1),reduction='none').reshape(-1,3)
140                    weighted_loss(loss,world,denom,empty).backward()
141                if not empty:
142                    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())
143            torch.nn.utils.clip_grad_norm_(raw.parameters(),.5,error_if_nonfinite=True);opt.step();sched.step();cursor=end;updates+=1
144            if updates%50==0:
145                sums=group_totals.clone();dist.all_reduce(sums);torch.cuda.synchronize();now=time.monotonic();free,_=torch.cuda.mem_get_info(device)
146                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)
147                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')
148                last_time=now;last_cursor=plan.cursor(cursor)
149            boundary=any(plan.cursor(start_cursor)<b<=plan.cursor(cursor) for b in cfg['pass_boundaries'])
150            if (updates in (25,50) or updates%cfg['validation_every_updates']==0 or boundary) and clock.remaining()>120:validation()
151            if updates%cfg['checkpoint_every_updates']==0 and clock.remaining()>300:
152                name='diagnostic-1500.pt' if updates==1500 else 'progress.pt'
153                checkpoint(name)
154                if updates==1500 and rank==0:
155                    atomic_json(out/'diagnostic-1500.json',{'updates':1500,'config_sha256':digest(cfg),'checkpoint':{'path':str(out/name),'sha256':sha(out/name)}})
156        if not guard_failed and clock.remaining()>120 and latest['updates']!=updates:validation()
157        checkpoint('last.pt',True)
158        if rank==0:
159            initial=initial_validation
160            atomic_json(out/'exact-validation-comparison.json',{'initial':initial['exact_cluster_accuracy'],'selected':best_validation['exact_cluster_accuracy'],
161                'paired_game_cluster_delta':summarize(best_validation['exact_outcomes'],initial['exact_outcomes']),'selected_checkpoint':best_ref,'selection':'minimum fixed pooled canonical-label NLL; all quality accuracy/loss thresholds removed; finite numerical safety remains'})
162            atomic_json(out/'terminal.json',{'status':'complete','reason':'retention_guard' if guard_failed else ('single_schedule_complete' if cursor==plan.steps else 'original_deadline_soft_stop'),'retention_failures':guard_failed,'baseline_checkpoint':baseline_ref,'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()
163    except BaseException as error:
164        # No collectives or new GPU work on an uncertain rank failure. The
165        # immutable prepass remains the selected fallback; guardian owns cleanup.
166        if rank==0 and baseline_ref is not None:
167            atomic_json(Path(cfg['output'])/'best.json',baseline_ref)
168            atomic_json(Path(cfg['output'])/'failure.json',{'error_type':type(error).__name__,'error':str(error),'selected_fallback':baseline_ref,'unix':time.time(),'uncommitted_weights_not_certified':True,'no_promotion':True})
169        raise
170    finally:
171        if prefetch:prefetch.close()
172        if data:
173            for source in data.sources.values():
174                if hasattr(source,'close'):source.close()
175        dist.destroy_process_group()
176if __name__=='__main__':
177    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())
178    if 'CHESS_JOB_DEADLINE' not in os.environ:raise ValueError('reviewed owner required')
179    if 'RANK' in os.environ:worker(cfg,args)
180    else:
181        verify(cfg,args);resource_check(cfg);extra=[] if not args.resume else ['--resume',args.resume,'--resume-sha256',args.resume_sha256]
182        raise SystemExit(launch_owned_workers(Path(__file__),[args.config,*extra]))
183