| """Train one fifth-epoch, save an exact-resume checkpoint, then release the GPU.""" |
| from pathlib import Path |
| import sys |
| sys.path.insert(0,str(Path(__file__).resolve().parent)) |
| from scratch_model import * |
| import random,math,gc,fcntl,shutil,traceback,argparse,importlib.metadata,platform |
| from huggingface_hub import HfApi,CommitOperationAdd |
|
|
|
|
| def status(phase,**extra): |
| write(OUT/'status.json',dict(phase=phase,pid=os.getpid(),updated_unix=time.time(),**extra)) |
| print('SCRATCH_TRAIN',phase,json.dumps(extra),flush=True) |
|
|
|
|
| def rng_state(): |
| return dict(torch_rng=torch.get_rng_state(),cuda_rng=torch.cuda.get_rng_state_all(), |
| numpy_rng=np.random.get_state(),python_rng=random.getstate()) |
|
|
|
|
| def restore_rng(state): |
| torch.set_rng_state(state['torch_rng']);torch.cuda.set_rng_state_all(state['cuda_rng']) |
| np.random.set_state(state['numpy_rng']);random.setstate(state['python_rng']) |
|
|
|
|
| def seed(): |
| random.seed(CONFIG['seed']);np.random.seed(CONFIG['seed']);torch.manual_seed(CONFIG['seed']);torch.cuda.manual_seed_all(CONFIG['seed']) |
| torch.use_deterministic_algorithms(True) |
|
|
|
|
| class TrainingData: |
| def __init__(self): |
| label_columns=[f'p_{i:03d}' for i in range(153)]+['area_A2','volume_A3'] |
| self.pcqm=pd.read_parquet(DATA/'pcqm.parquet',columns=['split']) |
| self.cosmo=pd.read_parquet(DATA/'cosmo.parquet',columns=['split',*label_columns]) |
| self.p_cache=RawCache('pcqm');self.c_cache=RawCache('cosmo') |
| self.pt=np.flatnonzero(self.pcqm.split.to_numpy()=='train');self.pv=np.flatnonzero(self.pcqm.split.to_numpy()=='validation') |
| self.ct=np.flatnonzero(self.cosmo.split.to_numpy()=='train');self.cv=np.flatnonzero(self.cosmo.split.to_numpy()=='validation') |
| transform=json.loads((DATA/'cosmo_label_transform.json').read_text()) |
| y=self.cosmo[[f'p_{i:03d}' for i in range(153)]+['area_A2','volume_A3']].to_numpy(dtype=np.float64) |
| y=np.c_[np.log1p(y[:,:153]),np.log(y[:,153:])] |
| self.labels=torch.tensor((y-np.asarray(transform['mean']))/np.asarray(transform['scale']),dtype=torch.float32,device='cuda') |
| assert torch.isfinite(self.labels).all() |
| self.epoch_plans={e:self.plan(e) for e in range(1,CONFIG['epochs']+1)} |
| self.total_steps=sum(v[2] for v in self.epoch_plans.values()) |
| schedule=[];offset=0 |
| for epoch,(pp,cp,n) in self.epoch_plans.items(): |
| for fifth in range(1,6): |
| end=(fifth*n+4)//5;start=((fifth-1)*n+4)//5 |
| schedule.append(dict(point=len(schedule)+1,epoch=epoch,fifth=fifth,batch_start=start,batch_end=end, |
| batches=n,epoch_position=(epoch-1)+end/n,nominal_epoch=(epoch-1)+fifth/5,global_step=offset+end)) |
| offset+=n |
| assert len(schedule)==50 and schedule[-1]['global_step']==self.total_steps |
| write(OUT/'probe_schedule.json',dict(total_steps=self.total_steps,probes=50,points=schedule)) |
| self.schedule=schedule |
|
|
| def plan(self,epoch): |
| rng=np.random.default_rng(CONFIG['seed']+epoch) |
| pp=list(batches(self.pt,self.p_cache.lengths,rng,max_size=CONFIG['pcqm_batch_size'],atom_cap=CONFIG['pcqm_atom_cap'])) |
| cp=list(batches(self.ct,self.c_cache.lengths,rng,max_size=CONFIG['cosmo_batch_size'],atom_cap=CONFIG['cosmo_atom_cap'])) |
| assert np.array_equal(np.sort(np.concatenate(pp)),self.pt) |
| assert np.array_equal(np.sort(np.concatenate(cp)),self.ct) |
| return pp,cp,max(len(pp),len(cp)) |
|
|
|
|
| def update(model,optimizer,data,pi,ci): |
| optimizer.zero_grad(set_to_none=True) |
| raw,z,target=model.denoising(data.p_cache.sample(pi)) |
| first=uncertainty(raw[None],model.log_variance[:1]);first.backward() |
| log=dict(denoising_raw=float(raw.detach()),denoising_per_channel=float(raw.detach())/1088, |
| descriptor_std_pcqm=float(z.detach().std(dim=0,unbiased=False).mean()),weighted_denoising=float(first.detach())) |
| del raw,z,target,first |
| loss,z,pred,parts=model.cosmo(data.c_cache.sample(ci),data.labels[ci]) |
| second=uncertainty(loss[None],model.log_variance[1:]);second.backward() |
| log.update(cosmo=float(loss.detach()),descriptor_std_cosmo=float(z.detach().std(dim=0,unbiased=False).mean()), |
| weighted_cosmo=float(second.detach()),**{'cosmo_'+k:v for k,v in parts.items()}) |
| norm=nn.utils.clip_grad_norm_(model.parameters(),CONFIG['gradient_clip'],error_if_nonfinite=True) |
| assert all(math.isfinite(v) for v in log.values()),log |
| optimizer.step() |
| log.update(gradient_norm=float(norm),total=log['weighted_denoising']+log['weighted_cosmo'], |
| log_variance_denoising=float(model.log_variance[0].detach()),log_variance_cosmo=float(model.log_variance[1].detach())) |
| return log |
|
|
|
|
| @torch.no_grad() |
| def validate(model,data): |
| old=rng_state();torch.manual_seed(83026);torch.cuda.manual_seed_all(83026);model.eval() |
| result={} |
| try: |
| for name,ids,cache in [('pcqm',data.pv,data.p_cache),('cosmo',data.cv,data.c_cache)]: |
| total=0.;count=0;zs=np.zeros(640);zss=zs.copy();std=0.;atoms=0 |
| for group in batches(ids,cache.lengths,max_size=64,atom_cap=3600): |
| if name=='pcqm': |
| loss,z,target=model.denoising(cache.sample(group)) |
| valid=target[~(torch.arange(target.shape[1],device='cuda')[None]>=torch.tensor(cache.lengths[group],device='cuda')[:,None])] |
| total+=float(loss)*len(valid);atoms+=len(valid) |
| std+=float(valid.std(dim=0).mean())*len(valid) |
| else: |
| loss,z,_,_=model.cosmo(cache.sample(group),data.labels[group]);total+=float(loss)*len(group) |
| values=z.double().cpu().numpy();zs+=values.sum(0);zss+=np.square(values).sum(0);count+=len(group) |
| assert count==len(ids) |
| result[name+'_rows']=count |
| result[name+'_loss']=total/(atoms if name=='pcqm' else count) |
| result[name+'_descriptor_channel_std']=float(np.sqrt(np.maximum(zss/count-(zs/count)**2,0)).mean()) |
| if name=='pcqm':result['pcqm_target_batch_std']=std/atoms;result['pcqm_loss_per_channel']=result['pcqm_loss']/1088 |
| assert all(math.isfinite(v) for v in result.values()),result |
| return result |
| finally:restore_rng(old) |
|
|
|
|
| def checkpoint(point,model,optimizer,scheduler,global_step,metrics): |
| path=OUT/'checkpoints'/f'probe_{point:03d}';pending=path.with_name(path.name+'.partial') |
| if path.exists():raise RuntimeError(f'Checkpoint already exists: {path}') |
| pending.mkdir(parents=True,exist_ok=True) |
| for name,module in model.parts.items():torch.save(module.state_dict(),pending/name) |
| torch.save(dict(log_variance=model.log_variance.detach().cpu(),task_order=CONFIG['task_order']),pending/'task_weighting.pth') |
| torch.save(dict(point=point,global_step=global_step,optimizer=optimizer.state_dict(),scheduler=scheduler.state_dict(), |
| config=CONFIG,**rng_state()),pending/'training_state.pt') |
| write(pending/'metrics.json',metrics);write(pending/'config.json',CONFIG) |
| for source,name in [(OUT/'architecture.yaml','post_training_architecture_config.yaml'), |
| (DATA/'normalization.pt','normalization.pt'),(DATA/'cosmo_label_transform.json','cosmo_label_transform.json'), |
| (OUT/'preparation_complete.json','data_manifest.json')]:shutil.copyfile(source,pending/name) |
| write(pending/'manifest.json',dict(point=point,global_step=global_step, |
| files={p.name:dict(bytes=p.stat().st_size,sha256=sha(p)) for p in pending.iterdir() if p.is_file() and p.name!='manifest.json'})) |
| pending.rename(path);return path |
|
|
|
|
| def upload(path): |
| if (path/'hub_verified.json').exists():return json.loads((path/'hub_verified.json').read_text()) |
| api=HfApi();api.create_repo(MODEL_REPO,exist_ok=True) |
| files=[p for p in path.iterdir() if p.is_file() and p.name!='hub_verified.json'] |
| for attempt in range(5): |
| try: |
| commit=api.create_commit(repo_id=MODEL_REPO,operations=[CommitOperationAdd(path_in_repo=path.name+'/'+p.name,path_or_fileobj=str(p)) for p in files], |
| commit_message=f'Save exact-resume scratch checkpoint {path.name} before its BigSolDB probe') |
| entries={p.rfilename:p for p in api.model_info(MODEL_REPO,revision=commit.oid,files_metadata=True).siblings} |
| for p in files: |
| if entries[path.name+'/'+p.name].lfs:assert entries[path.name+'/'+p.name].lfs.sha256==sha(p) |
| result=dict(verified=True,repository=MODEL_REPO,revision=commit.oid,point=int(path.name.split('_')[-1])) |
| write(path/'hub_verified.json',result);return result |
| except Exception: |
| if attempt==4:raise |
| time.sleep(min(10*2**attempt,50)) |
|
|
|
|
| def scheduler_for(optimizer,data): |
| return torch.optim.lr_scheduler.OneCycleLR(optimizer,max_lr=[CONFIG['learning_rate_max'],CONFIG['learning_rate_uncertainty']],total_steps=data.total_steps) |
|
|
|
|
| def restore(path,model,optimizer,scheduler): |
| manifest=json.loads((path/'manifest.json').read_text()) |
| assert sha(path/'training_state.pt')==manifest['files']['training_state.pt']['sha256'] |
| state=torch.load(path/'training_state.pt',map_location='cpu',weights_only=False) |
| assert state['config']==CONFIG |
| optimizer.load_state_dict(state['optimizer']);scheduler.load_state_dict(state['scheduler']);restore_rng(state) |
| return state |
|
|
|
|
| def smoke(data): |
| import copy |
| seed();loads=[];original_load=torch.load |
| def record_load(path,*args,**kwargs):loads.append(str(path));return original_load(path,*args,**kwargs) |
| torch.load=record_load |
| try:model=ScratchJoint() |
| finally:torch.load=original_load |
| assert loads==[str(DATA/'normalization.pt')],loads |
| opt=optimizer_for(model);scheduler=scheduler_for(opt,data) |
| state=copy.deepcopy(model.state_dict());seeded_rng=rng_state() |
| pp,cp,_=data.epoch_plans[1];model.train() |
| update(model,opt,data,pp[0],cp[0]);scheduler.step() |
| first_gradients={name:any(p.grad is not None and p.grad.abs().sum()>0 for p in module.parameters()) for name,module in model.parts.items()} |
| |
| |
| gates=[v for k,v in state.items() if k.endswith('pair2bias.weight')] |
| assert gates and all(torch.count_nonzero(v)==0 for v in gates) |
| assert all(value for name,value in first_gradients.items() if name!='geometric_preprocessor.pth') |
| assert not torch.equal(model.encoder.state_dict()[next(iter(model.encoder.state_dict()))],state['encoder.'+next(iter(model.encoder.state_dict()))]) |
| |
| path=OUT/'checkpoints/probe_999';path.parent.mkdir(exist_ok=True) |
| if path.exists():shutil.rmtree(path) |
| saved=checkpoint(999,model,opt,scheduler,1,dict(smoke=True)) |
| expected=update(model,opt,data,pp[1],cp[1]);scheduler.step() |
| gradients={name:any(p.grad is not None and p.grad.abs().sum()>0 for p in module.parameters()) for name,module in model.parts.items()} |
| assert all(gradients.values()),gradients |
| after={k:v.detach().cpu().clone() for k,v in model.state_dict().items()} |
| del model,opt,scheduler;gc.collect();torch.cuda.empty_cache() |
| replay=ScratchJoint(saved);opt=optimizer_for(replay);scheduler=scheduler_for(opt,data);restore(saved,replay,opt,scheduler) |
| replay.train();actual=update(replay,opt,data,pp[1],cp[1]);scheduler.step() |
| assert expected==actual and all(torch.equal(v.cpu(),after[k]) for k,v in replay.state_dict().items()) |
| del replay,opt,scheduler,state,after;gc.collect();torch.cuda.empty_cache();shutil.rmtree(saved) |
| seed();model=ScratchJoint();opt=optimizer_for(model);scheduler=scheduler_for(opt,data) |
| initial=validate(model,data) |
| path=OUT/'checkpoints/probe_000' |
| if path.exists(): |
| existing=ScratchJoint(path) |
| assert all(torch.equal(v,existing.state_dict()[k]) for k,v in model.state_dict().items()) |
| del existing;gc.collect();torch.cuda.empty_cache() |
| else:path=checkpoint(0,model,opt,scheduler,0,dict(point=0,epoch_position=0,validation=initial)) |
| upload(path) |
| write(OUT/'smoke_validation.json',dict(passed=True,all_trainable_head_gradients=gradients, |
| first_step_gradients=first_gradients,author_zero_initialized_pair_bias_preserved=True, |
| initialization_loaded_only_training_statistics=loads,no_pretrained_rem3di_weights_loaded=True, |
| exact_optimizer_rng_next_step_replay=True,initial_validation=initial)) |
| source_names=['scratch_prepare.py','scratch_data.py','scratch_model.py','scratch_train.py','scratch_probe.py','scratch_driver.py','scratch_cache_initial.py','scratch_fast_cache.py', |
| 'eval_common.py','train_cosmo.py','multitask_data.py','bigsoldb_evaluate.py'] |
| write(OUT/'provenance.json',dict(initial_rem3di_checkpoint=None,initialization='Constructor initialization of all learned components', |
| architecture_sha256=sha(OUT/'architecture.yaml'), |
| foundation_sha256=sha(Path('/marimo/rem3di_training/MACE-POLAR-1-M.model')), |
| python=platform.python_version(),cuda=torch.version.cuda,gpu=torch.cuda.get_device_name(0), |
| versions={n:importlib.metadata.version(n) for n in ['torch','numpy','pandas','scikit-learn','rdkit','structured-data-models','mace-torch','e3nn']}, |
| source_sha256={n:sha(Path(__file__).with_name(n)) for n in source_names}, |
| legacy_source_sha256={str(p.relative_to(AUDIT/'author_legacy_source')):sha(p) |
| for p in sorted((AUDIT/'author_legacy_source/threedscriptors').rglob('*.py'))})) |
| ops=[CommitOperationAdd(path_in_repo='source/'+n,path_or_fileobj=str(Path(__file__).with_name(n))) for n in source_names] |
| ops += [CommitOperationAdd(path_in_repo=n,path_or_fileobj=str(OUT/n)) for n in ['smoke_validation.json','config.json','preparation_complete.json','probe_schedule.json','provenance.json']] |
| HfApi().create_commit(repo_id=MODEL_REPO,operations=ops,commit_message='Verify random initialization, real-data gradients and exact next-step resume before scratch training') |
|
|
|
|
| def train_point(point,data): |
| import wandb |
| path=OUT/'checkpoints'/f'probe_{point:03d}' |
| if path.exists():upload(path);return |
| previous=OUT/'checkpoints'/f'probe_{point-1:03d}' |
| seed();model=ScratchJoint(previous);opt=optimizer_for(model);scheduler=scheduler_for(opt,data) |
| state=restore(previous,model,opt,scheduler);global_step=state['global_step'] |
| info=data.schedule[point-1];pp,cp,n=data.epoch_plans[info['epoch']] |
| assert state['point']==point-1 |
| assert global_step==info['global_step']-(info['batch_end']-info['batch_start']) |
| wb_path=OUT/'wandb.json';wb=json.loads(wb_path.read_text()) if wb_path.exists() else dict(run_id=wandb.util.generate_id()) |
| write(wb_path,wb) |
| run=wandb.init(entity='queimo',project='rem3di-reproduction',id=wb['run_id'],resume='allow',name='scratch-pcqm-cosmo-ten-epochs', |
| group='rem3di-scratch-pcqm-cosmo',job_type='scratch-pretraining',config=CONFIG,dir=str(OUT),settings=wandb.Settings(init_timeout=90)) |
| write(wb_path,dict(run_id=run.id,url=run.url));restore_rng(state) |
| model.train();logs=[];start=time.perf_counter() |
| try: |
| for step in range(info['batch_start'],info['batch_end']): |
| log=update(model,opt,data,pp[step%len(pp)],cp[step%len(cp)]);scheduler.step();global_step+=1;logs.append(log) |
| if (step+1)%20==0 or step+1==info['batch_end']: |
| status('training',point=point,epoch=info['epoch'],batch=step+1,batches=n,wandb_url=run.url) |
| run.log(dict(optimizer_step=global_step,epoch_position=info['epoch']-1+(step+1)/n,**{'train/'+k:v for k,v in log.items()},lr=scheduler.get_last_lr()[0])) |
| assert global_step==info['global_step'] |
| validation=validate(model,data) |
| metrics={**info,'global_step':global_step,'train':{k:float(np.mean([r[k] for r in logs])) for k in logs[0]}, |
| 'validation':validation,'seconds':time.perf_counter()-start, |
| 'epoch_coverage':dict(pcqm_unique=len(data.pt),cosmo_unique=len(data.ct)) if info['fifth']==5 else None} |
| saved=checkpoint(point,model,opt,scheduler,global_step,metrics) |
| status('uploading_checkpoint',point=point,epoch=info['epoch']) |
| hub=upload(saved);write(OUT/'latest_point.json',dict(**metrics,hub=hub)) |
| run.log(dict(point=point,epoch_position=info['epoch_position'],**{'validation/'+k:v for k,v in validation.items()},hub_revision=hub['revision'])) |
| run.finish() |
| except Exception:run.finish(exit_code=1);raise |
|
|
|
|
| def main(): |
| parser=argparse.ArgumentParser();parser.add_argument('--point',type=int);parser.add_argument('--smoke',action='store_true');args=parser.parse_args() |
| lock=(OUT/'training.lock').open('a');fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) |
| assert json.loads((OUT/'preparation_complete.json').read_text())['verified'] |
| seed();data=TrainingData();write(OUT/'config.json',CONFIG) |
| if args.smoke: |
| if not (OUT/'smoke_validation.json').exists():smoke(data) |
| else: |
| assert json.loads((OUT/'smoke_validation.json').read_text())['passed'] |
| assert 1<=args.point<=50;train_point(args.point,data) |
|
|
|
|
| if __name__=='__main__': |
| try:main() |
| except Exception:status('failed_training',error=traceback.format_exc());raise |
|
|
