| """Stasis 1 sixth tranche: 2-GPU DDP continuation from fifth recovery step 4,126. |
| Launch: torchrun --standalone --nproc_per_node=2 train_sixth.py --config RUN_CONFIG.json |
| Model code: frozen_stasis (unchanged sha pins). Never AutoModelForCausalLM.""" |
| import argparse, json, math, os, sys, time |
| import numpy as np, torch, torch.distributed as dist |
| from torch.nn.parallel import DistributedDataParallel as DDP |
| sys.path.insert(0,os.path.dirname(os.path.abspath(__file__))) |
| from frozen_stasis.config import ModelConfig |
| from frozen_stasis.model import StasisDecoder |
| import sixth_ckpt as C |
|
|
| W=2049; TARGETS=2048 |
|
|
| def lr_at(u,c,N): |
| if u<c['warmup_updates']: return c['peak_lr']*(u+1)/c['warmup_updates'] |
| if N is None: return c['peak_lr'] |
| p=min(1.0,(u-c['warmup_updates'])/max(1,N-c['warmup_updates'])) |
| return c['end_lr']+0.5*(c['peak_lr']-c['end_lr'])*(1+math.cos(math.pi*p)) |
|
|
| def main(): |
| ap=argparse.ArgumentParser(); ap.add_argument('--config',required=True); a=ap.parse_args() |
| c=json.load(open(a.config)) |
| backend=c.get('backend','nccl') |
| dist.init_process_group(backend) |
| rank,world=dist.get_rank(),dist.get_world_size() |
| cuda=backend=='nccl' |
| dev=torch.device(f'cuda:{rank}') if cuda else torch.device('cpu') |
| if cuda: torch.cuda.set_device(dev) |
| windows_per_update=c['windows_per_update']; micro=c['microbatch'] |
| C.need(windows_per_update%(world*micro)==0,'windows/update must split over ranks*microbatch') |
| accum=windows_per_update//(world*micro) |
| out=c['out_dir']; os.makedirs(out,exist_ok=True) |
| log=open(os.path.join(out,'train_log.jsonl'),'a') if rank==0 else None |
| def emit(**k): |
| if rank==0: log.write(json.dumps(dict(t=round(time.time(),3),**k))+'\n'); log.flush() |
| torch.backends.cuda.matmul.allow_tf32=False; torch.backends.cudnn.allow_tf32=False |
| model=StasisDecoder(ModelConfig(**c.get('model',{}))).to(dev) |
| C.need(sum(p.numel() for p in model.parameters())==109529856 or c.get('allow_test_model'),'parameter count') |
| opt,names=C.make_optimizer(model,c['peak_lr'],fused=cuda and c.get('fused_adamw',True)) |
| scaler=torch.amp.GradScaler('cuda' if cuda else 'cpu',enabled=c.get('scaler',True)) |
| |
| resume=c.get('resume_path') |
| src=resume if resume else c['parent_path'] |
| payload=C.load_payload(src,expected_sha=None if resume else c.get('parent_sha256'),expected_bytes=None if resume else c.get('parent_bytes')) |
| step,cursor=C.restore(payload,model,opt,names,scaler,rank=rank) |
| if resume: |
| C.need(payload['binding']['run_id']==c['run_id'],'resume run_id'); start=cursor['tranche_update']; N=payload['binding'].get('frozen_updates') |
| base=payload['binding']['parent_cursor'] |
| else: |
| C.need(step==c['parent_optimizer_step'] and cursor['optimizer_step']==step,'parent optimizer step'); start=0; N=None; base=cursor |
| del payload |
| for g in opt.param_groups: g['lr']=lr_at(start,c,N) |
| ddp=DDP(model,device_ids=[rank] if cuda else None,gradient_as_bucket_view=True,broadcast_buffers=False) |
| if c.get('compile'): ddp=torch.compile(ddp) |
| data=np.memmap(c['train_path'],dtype='<u2',mode='r'); C.need(len(data)%W==0,'train.bin window multiple') |
| data=data.reshape(-1,W); max_updates=min(len(data)//windows_per_update,c.get('max_updates',10**9)) |
| if N is not None: C.need(N<=max_updates,'frozen N exceeds data') |
| stop_at=c['train_stop_epoch']; save_every=c['save_every_seconds']; last_save=time.time() |
| binding=dict(schema='stasis.sixth_tranche.binding.v1',run_id=c['run_id'],parent_sha256=c.get('parent_sha256'),train_sha256=c.get('train_sha256'), |
| gates_sha256=c.get('gates_sha256'),schedule=dict(kind='sixth_rewarm_cosine',warmup_updates=c['warmup_updates'],peak_lr=c['peak_lr'],end_lr=c['end_lr']), |
| frozen_updates=N,parent_cursor=base,windows_per_update=windows_per_update,world_size=world,microbatch=micro) |
| def cur(u): |
| t=u*windows_per_update*TARGETS |
| return dict(tranche_update=u,optimizer_step=base['optimizer_step']+u,cumulative_recovery_updates=base['cumulative_recovery_updates']+u, |
| tranche_targets=t,cumulative_recovery_targets=base['cumulative_recovery_targets']+t,lineage_targets=base['lineage_targets']+t,next_tranche_window=u*windows_per_update) |
| def save(u,tag): |
| if rank==0: |
| binding['frozen_updates']=N |
| p=C.capture(ddp.module if hasattr(ddp,'module') else model,opt,names,scaler,binding=binding,cursor=cur(u)) |
| path=os.path.join(out,f'recovery-{tag}.pt'); t0=time.time(); h=C.atomic_save(p,path) |
| json.dump(dict(file=os.path.basename(path),sha256=h,bytes=os.path.getsize(path),cursor=cur(u)),open(path+'.manifest.json','w')) |
| emit(event='save',tag=tag,update=u,sha256=h,seconds=round(time.time()-t0,1)) |
| dist.barrier() |
| emit(event='start',rank_world=world,start_update=start,max_updates=max_updates,frozen_updates=N,accum=accum,micro=micro,parent_step=step) |
| u=start; times=[]; consecutive_skips=0; t_update=time.time() |
| raw=ddp.module if hasattr(ddp,'module') else model |
| while True: |
| if N is not None and u>=N: break |
| if u>=max_updates: break |
| |
| flag=torch.tensor([1 if time.time()>=stop_at else 0],device=dev); dist.broadcast(flag,0) |
| if flag.item(): emit(event='time_stop',update=u); break |
| lr=lr_at(u,c,N) |
| for g in opt.param_groups: g['lr']=lr |
| base_w=u*windows_per_update+rank*accum*micro |
| tot=torch.zeros((),device=dev,dtype=torch.float64) |
| for m in range(accum): |
| x=torch.from_numpy(data[base_w+m*micro:base_w+(m+1)*micro].astype(np.int64)).to(dev,non_blocking=True) |
| ctx=ddp.no_sync() if m<accum-1 else torch.enable_grad() |
| with ctx: |
| with torch.autocast('cuda' if cuda else 'cpu',dtype=torch.float16 if cuda else torch.bfloat16,enabled=c.get('amp',cuda)): |
| _,loss=ddp(x[:,:-1],x[:,1:]) |
| scaler.scale(loss/accum).backward() |
| tot+=loss.detach().double() |
| scaler.unscale_(opt) |
| norm=torch.nn.utils.clip_grad_norm_(raw.parameters(),1.0) |
| before=scaler.get_scale(); scaler.step(opt); scaler.update(); after=scaler.get_scale() |
| opt.zero_grad(set_to_none=True) |
| skipped=bool(after<before) if scaler.is_enabled() else False |
| dist.all_reduce(tot); mean_loss=float(tot)/(accum*world) |
| consecutive_skips=consecutive_skips+1 if skipped else 0 |
| u+=1; now=time.time(); dt=now-t_update; t_update=now; times.append(dt) |
| emit(event='update',update=u,loss=round(mean_loss,6),lr=lr,grad_norm=float(norm),scale=after,skipped=skipped,seconds=round(dt,3)) |
| C.need(math.isfinite(mean_loss) or skipped,'nonfinite loss without skip') |
| if consecutive_skips>=c.get('max_consecutive_skips',25): |
| emit(event='abort_skips',update=u); save(u,'abort'); break |
| |
| if N is None and u==start+c['calibration_updates']: |
| spu=torch.tensor([float(np.median(times[3:]))],device=dev); dist.broadcast(spu,0); spu=float(spu.item()) |
| remaining=stop_at-time.time()-c['reserve_seconds_for_saves'] |
| N=int(min(max_updates,u+math.floor(remaining/(spu*c['throughput_safety'])))) |
| C.need(N>c['warmup_updates'],'not enough time for schedule') |
| binding['frozen_updates']=N |
| if rank==0: json.dump(dict(frozen_updates=N,seconds_per_update_median=spu,targets=N*windows_per_update*TARGETS,frozen_at_update=u, |
| remaining_seconds=remaining,safety=c['throughput_safety'],utc=time.time()),open(os.path.join(out,'SCHEDULE_FROZEN.json'),'w'),indent=1) |
| emit(event='schedule_frozen',N=N,spu=spu) |
| sflag=torch.tensor([1 if time.time()-last_save>=save_every else 0],device=dev); dist.broadcast(sflag,0) |
| if sflag.item(): save(u,'latest'); last_save=time.time() |
| save(u,'final') |
| if rank==0: |
| mo=dict(schema='stasis.sixth_tranche.model_only.v1',model={k:v.detach().cpu().clone() for k,v in raw.state_dict().items()},cursor=cur(u),binding=binding) |
| h=C.atomic_save(mo,os.path.join(out,'endpoint.model_only.pt')) |
| json.dump(dict(completed_updates=u,frozen_updates=N,schedule_complete=(N is not None and u>=N),endpoint_sha256=h,cursor=cur(u)),open(os.path.join(out,'TRAIN_STATUS.json'),'w'),indent=1) |
| emit(event='done',update=u) |
| dist.barrier(); dist.destroy_process_group() |
|
|
| if __name__=='__main__': main() |
|
|
