train.py

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