| """Complete joint epochs, exact-resume checkpoints, W&B and verified Hub backups.""" |
| from pathlib import Path |
| import sys |
| sys.path.insert(0,str(Path(__file__).resolve().parent)) |
| from multitask_model import * |
| import argparse,ast,copy,math,random |
| os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG',':4096:8') |
|
|
|
|
| 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 check_monroe_losses(): |
| """Compare losses and gradients with the pinned author's actual implementation.""" |
| import openadmet_evaluate as runtime |
| runtime.restore_monroe() |
| from monroe.mtl.UW import UW |
| source=WORK/'external/Monroe/monroe/train/loss.py' |
| tree=ast.parse(source.read_text());function=next(n for n in tree.body if isinstance(n,ast.FunctionDef) and n.name=='off_diag_cov_loss') |
| namespace={'torch':torch};exec(compile(ast.Module(body=[function],type_ignores=[]),str(source),'exec'),namespace) |
| torch.manual_seed(42026) |
| z=torch.randn(19,37,requires_grad=True);z2=z.detach().clone().requires_grad_() |
| ref_cov=namespace['off_diag_cov_loss'](z);our_cov=off_diag_cov_loss(z2) |
| ref_cov.backward();our_cov.backward() |
| assert torch.equal(ref_cov,our_cov) and torch.equal(z.grad,z2.grad) |
| reference=UW(task_types={k:None for k in CONFIG['task_order']},encoder_class=None,decoders=nn.Identity(),device='cpu') |
| reference.init_param() |
| losses=torch.tensor([.2,.8,1.3],requires_grad=True) |
| ours=losses.detach().clone().requires_grad_();scale=nn.Parameter(reference.loss_scale.detach().clone()) |
| ref=reference.backward(losses,extra_losses=torch.tensor([.125])) |
| value=uncertainty_objective(ours,scale)+.125;value.backward() |
| assert torch.equal(ref,value) and torch.equal(reference.loss_scale.grad,scale.grad) and torch.equal(losses.grad,ours.grad) |
| result=dict(passed=True,monroe_revision=MONROE_REV,uw_value=float(value.detach()), |
| uw_value_and_gradients_exact=True,covariance_value_and_gradients_exact=True, |
| uw_source_sha256=sha(WORK/'external/Monroe/monroe/mtl/UW.py'),covariance_source_sha256=sha(source), |
| no_normalization_of_objective_weights=True,decorrelation_outside_uncertainty=True, |
| decorrelation_weight=CONFIG['decorrelation_weight']) |
| write(OUT/'loss_implementation_validation.json',result) |
| return result |
|
|
|
|
| class TrainingData: |
| def __init__(self): |
| self.cosmo=cosmo_frame();self.pairs=pd.read_parquet(DATA/'combisolv_pairs.parquet') |
| self.c_cache=RawCache('cosmo');self.q_cache=RawCache('combisolv') |
| self.tr=np.flatnonzero(self.cosmo.joint_split.to_numpy()=='train');self.va=np.flatnonzero(self.cosmo.joint_split.to_numpy()=='validation') |
| self.qtr=np.flatnonzero(self.pairs.split.to_numpy()=='train');self.qva=np.flatnonzero(self.pairs.split.to_numpy()=='validation') |
| transform=json.loads((DATA/'cosmo_label_transform.json').read_text()) |
| self.c_mean=np.asarray(transform['mean']);self.c_scale=np.asarray(transform['scale']) |
| labels=self.cosmo[[f'p_{i:03d}' for i in range(153)]+['area_A2','volume_A3']].to_numpy(dtype=np.float64) |
| loglabels=np.c_[np.log1p(labels[:,:153]),np.log(labels[:,153:])] |
| self.c_target=torch.tensor((loglabels-self.c_mean)/self.c_scale,dtype=torch.float32,device='cuda') |
| self.q_transform=json.loads((DATA/'combisolv_label_transform.json').read_text()) |
| self.q_target=torch.tensor((self.pairs.dGsolv_kcal_mol.to_numpy()-self.q_transform['mean'])/self.q_transform['scale'],dtype=torch.float32,device='cuda') |
| assert len(self.c_cache.lengths)==len(self.cosmo) and len(self.q_target)==1000000 |
| assert torch.isfinite(self.c_target).all() and torch.isfinite(self.q_target).all() |
|
|
| def plans(self,epoch): |
| rng=np.random.default_rng(CONFIG['seed']+epoch) |
| cp=list(batches(self.tr,self.c_cache.lengths,rng,max_size=CONFIG['cosmo_batch_size'],atom_cap=CONFIG['cosmo_atom_cap'])) |
| order=rng.permutation(self.qtr) |
| qp=[order[i:i+CONFIG['qm_pair_batch_size']] for i in range(0,len(order),CONFIG['qm_pair_batch_size'])] |
| assert sorted(i for batch in cp for i in batch)==self.tr.tolist() |
| assert np.array_equal(np.sort(np.concatenate(qp)),self.qtr) |
| return cp,qp,max(len(cp),len(qp)) |
|
|
|
|
| def update(model,optimizer,data,ci,qi): |
| optimizer.zero_grad(set_to_none=True) |
| losses,cov,z,_,parts=model.cosmo_denoising(data.c_cache.sample(ci),data.c_target[ci]) |
| first=uncertainty_objective(losses,model.log_variance[:2])+CONFIG['decorrelation_weight']*.5*cov |
| |
| |
| first.backward() |
| log=dict(denoising=float(losses[0].detach()),denoising_raw=float(losses[0].detach())*1088, |
| cosmo=float(losses[1].detach()),cov_cosmo=float(cov.detach()), |
| descriptor_std_cosmo=float(z.detach().float().std(dim=0).mean()),**{'cosmo_'+k:v for k,v in parts.items()}) |
| first_value=float(first.detach()) |
| del losses,cov,z,first |
| qm,cov,pred,z=model.qm(data.q_cache,qi,data.pairs,data.q_target) |
| second=uncertainty_objective(qm[None],model.log_variance[2:])+CONFIG['decorrelation_weight']*.5*cov |
| second.backward() |
| log.update(combisolv=float(qm.detach()),cov_combisolv=float(cov.detach()), |
| descriptor_std_combisolv=float(z.detach().float().std(dim=0).mean()), |
| total=first_value+float(second.detach())) |
| 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['gradient_norm']=float(norm) |
| log['combisolv_rmse_kcal_mol']=math.sqrt(log['combisolv'])*data.q_transform['scale'] |
| for i,task in enumerate(CONFIG['task_order']): |
| log['log_variance_'+task]=float(model.log_variance[i].detach()) |
| log['precision_'+task]=float((1/(2*model.log_variance[i].exp())).detach()) |
| return log |
|
|
|
|
| @torch.no_grad() |
| def validate(model,data): |
| old_rng=rng_state();torch.manual_seed(83026);torch.cuda.manual_seed_all(83026) |
| model.eval();sums=np.zeros(3);count=0 |
| for ids in batches(data.va,data.c_cache.lengths,max_size=64,atom_cap=3600): |
| losses,cov,_,_,_=model.cosmo_denoising(data.c_cache.sample(ids),data.c_target[ids]) |
| sums+=np.asarray([float(losses[0]),float(losses[1]),float(cov)])*len(ids);count+=len(ids) |
| assert count==len(data.va) |
| q_sse=0.;q_mae=0.;q_count=0 |
| for start in range(0,len(data.qva),CONFIG['qm_pair_batch_size']): |
| ids=data.qva[start:start+CONFIG['qm_pair_batch_size']] |
| loss,_,pred,_=model.qm(data.q_cache,ids,data.pairs,data.q_target) |
| diff=(pred-data.q_target[ids])*data.q_transform['scale'] |
| q_sse+=float(diff.double().square().sum());q_mae+=float(diff.double().abs().sum());q_count+=len(ids) |
| assert q_count==len(data.qva) |
| restore_rng(old_rng) |
| result=dict(denoising=float(sums[0]/count),denoising_raw=float(sums[0]/count*1088), |
| cosmo=float(sums[1]/count),cov_cosmo=float(sums[2]/count), |
| combisolv_rmse_kcal_mol=math.sqrt(q_sse/q_count),combisolv_mae_kcal_mol=q_mae/q_count, |
| cosmo_rows=count,combisolv_rows=q_count) |
| assert all(math.isfinite(v) for v in result.values()),result |
| return result |
|
|
|
|
| def save_checkpoint(path,model,optimizer,scheduler,epoch,global_step,metrics): |
| pending=path.with_name(path.name+'.partial');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(epoch=epoch,global_step=global_step,optimizer=optimizer.state_dict(),scheduler=scheduler.state_dict(), |
| config=CONFIG,dataset=json.loads((OUT/'prepared_data_hub.json').read_text()),**rng_state()),pending/'training_state.pt') |
| write(pending/'metrics.json',metrics);write(pending/'config.json',CONFIG) |
| for filename in ['cosmo_label_transform.json','combisolv_label_transform.json','manifest.json']: |
| shutil.copyfile(DATA/filename,pending/('data_manifest.json' if filename=='manifest.json' else filename)) |
| shutil.copyfile(OUT/'initialization.json',pending/'initialization.json') |
| (pending/'post_training_architecture_config.yaml').write_text((common.ORIGINAL/'post_training_architecture_config.yaml').read_text()) |
| write(pending/'manifest.json',dict(epoch=epoch,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'})) |
| if path.exists():shutil.rmtree(path) |
| pending.rename(path) |
|
|
|
|
| def upload_checkpoint(path): |
| marker=path/'hub_verified.json' |
| if marker.exists():return json.loads(marker.read_text()) |
| api=HfApi();api.create_repo(MODEL_REPO,repo_type='model',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'Preserve complete joint checkpoint {path.name}: all heads, learned task weights and exact resume state') |
| entries={s.rfilename:s for s in api.model_info(MODEL_REPO,revision=commit.oid,files_metadata=True).siblings} |
| for p in files: |
| if p.suffix not in ['.pth','.pt']:continue |
| remote=entries[path.name+'/'+p.name] |
| assert remote.lfs and remote.lfs.sha256==sha(p),p.name |
| result=dict(repository=MODEL_REPO,revision=commit.oid,epoch=int(path.name.split('_')[-1]),verified=True) |
| write(marker,result);return result |
| except Exception: |
| if attempt==4:raise |
| print('HUB_UPLOAD_RETRY',traceback.format_exc(),flush=True);time.sleep(min(10*2**attempt,50)) |
|
|
|
|
| def smoke(data): |
| verified_path=OUT/'loss_implementation_validation.json' |
| if verified_path.exists(): |
| verified=json.loads(verified_path.read_text()) |
| assert verified['passed'] |
| assert verified['uw_source_sha256']==sha(WORK/'external/Monroe/monroe/mtl/UW.py') |
| assert verified['covariance_source_sha256']==sha(WORK/'external/Monroe/monroe/train/loss.py') |
| else:check_monroe_losses() |
| model=MultiTask();optimizer=optimizer_for(model) |
| scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer,lambda e:1.) |
| model.eval();baseline=RawCache('bigsoldb') |
| with torch.no_grad(): |
| ref=np.load(WORK/'bigsoldb_evaluation/features/cosmo_epoch10.npz')['X'] |
| ours=np.empty_like(ref) |
| for ids in common.batch_indices(np.arange(len(ref)),baseline.lengths,max_size=24,atom_cap=2400): |
| _,_,z=model.encode(baseline.sample(ids));ours[ids]=z.cpu().numpy() |
| difference=np.abs(ours-ref) |
| assert np.allclose(ours,ref,atol=2e-4,rtol=2e-4),(float(difference.max()),float(difference.mean())) |
| initial_difference=dict(molecules=len(ref),max_absolute=float(difference.max()),mean_absolute=float(difference.mean())) |
| frozen={n:b.cpu().clone() for n,b in model.pre.atomic_preprocessor.invariant_normalization.named_buffers()} |
| cp,qp,_=data.plans(1);ci=cp[10];qi=qp[10] |
| initial=next(model.encoder.parameters()).detach().clone() |
| model.train();torch.cuda.reset_peak_memory_stats();t=time.perf_counter() |
| first=update(model,optimizer,data,ci,qi);duration=time.perf_counter()-t |
| assert not torch.equal(initial,next(model.encoder.parameters()).detach()) |
| gradients={name:any(p.grad is not None and bool((p.grad!=0).any()) for p in module.parameters()) for name,module in model.parts.items()} |
| assert all(gradients.values()),gradients |
| assert model.log_variance.grad is not None and bool((model.log_variance.grad!=0).all()) |
| temp=OUT/'smoke_checkpoint';save_checkpoint(temp,model,optimizer,scheduler,0,1,first) |
| expected=update(model,optimizer,data,ci,qi) |
| expected_state={n:v.detach().cpu().clone() for n,v in model.state_dict().items()} |
| model.load_parts(temp);state=torch.load(temp/'training_state.pt',map_location='cpu',weights_only=False) |
| optimizer.load_state_dict(state['optimizer']);scheduler.load_state_dict(state['scheduler']);restore_rng(state) |
| replay_started=time.perf_counter() |
| actual=update(model,optimizer,data,ci,qi) |
| steady_duration=time.perf_counter()-replay_started |
| assert all(np.isclose(actual[k],expected[k],atol=1e-7,rtol=1e-7) for k in expected),(expected,actual) |
| assert all(torch.equal(v.detach().cpu(),expected_state[n]) for n,v in model.state_dict().items()),'Resume changed next-step parameters' |
| for n,b in model.pre.atomic_preprocessor.invariant_normalization.named_buffers():assert torch.equal(b.cpu(),frozen[n]) |
| peak=torch.cuda.max_memory_allocated()/1e9 |
| result=dict(passed=True,all_head_gradients=gradients,encoder_updated=True,task_weights_updated=True, |
| optimizer_and_rng_resume_next_step_exact=True,normalization_buffers_unchanged=True, |
| initial_bigsoldb_embedding_match=initial_difference,real_batch_loss=first,seconds_per_joint_step=duration, |
| steady_seconds_per_joint_step=steady_duration, |
| smoke_cosmo_molecules=len(ci),smoke_qm_pairs=len(qi),peak_gpu_GB=peak,config=CONFIG) |
| write(OUT/'smoke_validation.json',result);shutil.rmtree(temp) |
| status('smoke_passed',seconds_per_joint_step=duration,peak_gpu_GB=peak,initial_embedding_difference=initial_difference) |
|
|
|
|
| def main(): |
| parser=argparse.ArgumentParser();parser.add_argument('--smoke',action='store_true');parser.add_argument('--check-losses',action='store_true');args=parser.parse_args() |
| OUT.mkdir(exist_ok=True) |
| lock=(OUT/'training.lock').open('a');fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) |
| if args.check_losses:print(check_monroe_losses());return |
| assert (OUT/'preparation_complete.json').exists(),'Data preparation has not finished' |
| |
| torch.use_deterministic_algorithms(True) |
| torch.manual_seed(CONFIG['seed']);torch.cuda.manual_seed_all(CONFIG['seed']);np.random.seed(CONFIG['seed']);random.seed(CONFIG['seed']) |
| data=TrainingData() |
| if args.smoke:smoke(data);return |
| smoke_result=json.loads((OUT/'smoke_validation.json').read_text()) |
| assert smoke_result['passed'] and smoke_result['config']==CONFIG |
| checkpoints=sorted((OUT/'checkpoints').glob('epoch_[0-9][0-9]')) if (OUT/'checkpoints').exists() else [] |
| model=MultiTask(checkpoints[-1] if checkpoints else None);optimizer=optimizer_for(model) |
| scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer,lambda e:.1+.9*.5*(1+math.cos(math.pi*e/CONFIG['epochs']))) |
| frozen={n:b.cpu().clone() for n,b in model.pre.atomic_preprocessor.invariant_normalization.named_buffers()} |
| start=1;global_step=0 |
| if checkpoints: |
| state=torch.load(checkpoints[-1]/'training_state.pt',map_location='cpu',weights_only=False) |
| assert state['config']==CONFIG |
| assert state['dataset']==json.loads((OUT/'prepared_data_hub.json').read_text()) |
| optimizer.load_state_dict(state['optimizer']);scheduler.load_state_dict(state['scheduler']);restore_rng(state) |
| start=state['epoch']+1;global_step=state['global_step'];upload_checkpoint(checkpoints[-1]) |
| write(OUT/'resume_state.json',dict(restored_epoch=start-1,optimizer_scheduler_rng_restored=True,global_step=global_step)) |
| import wandb |
| 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',mode='online', |
| name='pcqm2-cosmo10-combisolv-three-task-5epochs',group='rem3di-cosmo-combisolv',job_type='multitask-pretraining', |
| config={**CONFIG,'dataset':json.loads((DATA/'manifest.json').read_text()),'initialization':json.loads((OUT/'initialization.json').read_text())}, |
| dir=str(OUT),settings=wandb.Settings(init_timeout=90)) |
| write(wb_path,dict(run_id=run.id,url=run.url));write(OUT/'config.json',CONFIG) |
| run.define_metric('global_step');run.define_metric('train/*',step_metric='global_step') |
| run.define_metric('epoch');run.define_metric('validation/*',step_metric='epoch') |
| if checkpoints:restore_rng(state) |
| try: |
| for epoch in range(start,CONFIG['epochs']+1): |
| model.train();cp,qp,total=data.plans(epoch);started=time.time();history=[];seen_c=seen_q=0 |
| status('training',epoch=epoch,epochs=CONFIG['epochs'],batches=total,cosmo_batches=len(cp),qm_batches=len(qp),wandb_url=run.url) |
| for bi in range(total): |
| ci=cp[bi%len(cp)];qi=qp[bi%len(qp)] |
| log=update(model,optimizer,data,ci,qi);history.append(log);global_step+=1;seen_c+=len(ci);seen_q+=len(qi) |
| if bi==0 or (bi+1)%20==0 or bi+1==total: |
| run.log(dict(global_step=global_step,epoch=epoch,**{'train/'+k:v for k,v in log.items()})) |
| status('training',epoch=epoch,epochs=CONFIG['epochs'],batch=bi+1,batches=total,global_step=global_step, |
| cosmo_molecules_seen=seen_c,combisolv_pairs_seen=seen_q,seconds=time.time()-started, |
| losses={k:log[k] for k in ['denoising','cosmo','combisolv','total']},wandb_url=run.url) |
| assert seen_c>=len(data.tr) and seen_q>=len(data.qtr) |
| status('validating',epoch=epoch,wandb_url=run.url) |
| validation=validate(model,data) |
| for n,b in model.pre.atomic_preprocessor.invariant_normalization.named_buffers():assert torch.equal(b.cpu(),frozen[n]) |
| scheduler.step() |
| result=dict(epoch=epoch,global_step=global_step,seconds=time.time()-started, |
| train={k:float(np.mean([h[k] for h in history])) for k in history[0]},validation=validation, |
| coverage=dict(cosmo_unique_rows=len(data.tr),combisolv_unique_rows=len(data.qtr), |
| cosmo_rows_with_repeats=seen_c,combisolv_rows_with_repeats=seen_q,every_training_row_covered=True), |
| task_log_variances=model.log_variance.detach().cpu().tolist()) |
| target=OUT/'checkpoints'/f'epoch_{epoch:02d}' |
| save_checkpoint(target,model,optimizer,scheduler,epoch,global_step,result) |
| status('uploading_checkpoint',epoch=epoch,wandb_url=run.url) |
| hub=upload_checkpoint(target) |
| write(OUT/'latest_epoch.json',dict(**result,hub=hub,wandb_url=run.url)) |
| run.log(dict(epoch=epoch,**{'validation/'+k:v for k,v in validation.items()}, |
| **{'epoch/'+k:v for k,v in result['train'].items()},**{'backup/verified_epoch':epoch,'backup/revision':hub['revision']})) |
| |
| |
| write(OUT/'training_complete.json',dict(epochs=CONFIG['epochs'],wandb_url=run.url,completed_unix=time.time())) |
| run.summary.update(dict(training_complete=True,epochs=CONFIG['epochs']));run.finish() |
| status('training_complete',epochs=CONFIG['epochs']) |
| except Exception: |
| status('failed_training',error=traceback.format_exc());run.finish(exit_code=1);raise |
|
|
|
|
| if __name__=='__main__': |
| try:main() |
| except Exception:status('failed_training',error=traceback.format_exc());raise |
|
|
