File size: 8,633 Bytes
093306a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """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))
# ---- restore (strict) ----
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
# stop decision (rank0 time, broadcast)
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
# freeze N after calibration (warmup only so far, LR independent of N)
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()
|