| """Ten-epoch supervised COSMO continuation of the selected PCQM epoch 2.""" |
| import argparse |
| import fcntl |
| import gc |
| import hashlib |
| import json |
| import math |
| import os |
| from pathlib import Path |
| import random |
| import shutil |
| import sys |
| import time |
| import traceback |
|
|
| WORK = Path(__file__).resolve().parent.parent |
| AUDIT = Path('/marimo/rem3di_training/corrected_20261027/evaluation/paper_audit') |
| ORIGINAL = AUDIT/'author_checkpoint/exp138/polar' |
| sys.path[:0] = [str(AUDIT),str(AUDIT/'author_legacy_source'),'/marimo/Rem3Di'] |
| os.environ.update(CUDA_VISIBLE_DEVICES='0',OMP_NUM_THREADS='8',MKL_NUM_THREADS='8',WANDB_MODE='online') |
| from dotenv import load_dotenv |
| load_dotenv('/marimo/.env',override=False) |
| import numpy as np |
| import pandas as pd |
| import torch |
| from torch import nn |
| import yaml |
| from huggingface_hub import HfApi, hf_hub_download, snapshot_download, CommitOperationAdd |
| from threedscriptors.configuration.architecture_config import EncoderOnlyArchitectureConfig |
| from threedscriptors.data_handling.sample import Sample as LegacySample |
| from remedi.data_handling.sample import Sample, yield_molecules_collate_fn |
| from build_author_mace_runtime import build_author_mace_runtime |
|
|
| torch.serialization.add_safe_globals([slice]) |
| torch.set_num_threads(8) |
|
|
|
|
| def write(path,value): |
| path.parent.mkdir(parents=True,exist_ok=True) |
| tmp=path.with_suffix(path.suffix+'.partial') |
| tmp.write_text(json.dumps(value,indent=2,allow_nan=False));tmp.replace(path) |
|
|
|
|
| def sha(path): |
| h=hashlib.sha256() |
| with path.open('rb') as f: |
| for b in iter(lambda:f.read(8<<20),b''):h.update(b) |
| return h.hexdigest() |
|
|
|
|
| def status(phase,**kw): |
| row=dict(phase=phase,pid=os.getpid(),updated_unix=time.time(),**kw) |
| write(WORK/'training_status.json',row) |
| print('STATUS',json.dumps(row),flush=True) |
|
|
|
|
| def backbone(path): |
| d=yaml.safe_load((ORIGINAL/'post_training_architecture_config.yaml').read_text()) |
| d.pop('decoder_config',None);d['kind']='encoder_only';d['mace_config']=None |
| cfg=EncoderOnlyArchitectureConfig.model_validate(d) |
| pre,enc=cfg._build_preprocessor(),cfg._build_encoder() |
| |
| pre.atomic_preprocessor.double();pre.geometric_preprocessor.float();enc.float() |
| for name,module in [('encoder.pth',enc),('atomic_preprocessor.pth',pre.atomic_preprocessor),('geometric_preprocessor.pth',pre.geometric_preprocessor)]: |
| module.load_state_dict(torch.load(path/name,map_location='cpu',weights_only=True),strict=True) |
| return pre.cuda(),enc.cuda() |
|
|
|
|
| class COSMORegressor(nn.Module): |
| def __init__(self,pre,enc): |
| super().__init__();self.pre=pre;self.encoder=enc |
| self.head=nn.Sequential(nn.Linear(640,256),nn.SiLU(),nn.Linear(256,155)).cuda() |
|
|
| def forward(self,sample): |
| z=self.encoder(self.pre(sample)) |
| if not torch.is_tensor(z):z=z.flat |
| return self.head(z) |
|
|
|
|
| def batches(indices,lengths,rng=None,max_size=96,atom_cap=4000): |
| order=np.asarray(indices).copy() |
| if rng is not None:rng.shuffle(order) |
| for start in range(0,len(order),1024): |
| chunk=sorted(order[start:start+1024],key=lambda i:lengths[i]);ids=[];atoms=0 |
| for i in chunk: |
| if ids and (atoms+lengths[i]>atom_cap or len(ids)>=max_size): |
| yield ids;ids=[];atoms=0 |
| ids.append(int(i));atoms+=lengths[i] |
| if ids:yield ids |
|
|
|
|
| def inverse(z,mean,scale): |
| t=z*scale+mean |
| |
| return torch.cat([torch.expm1(t[:,:153]).clamp_min(0),torch.exp(t[:,153:])],dim=1) |
|
|
|
|
| def loss_terms(pred,target): |
| e=(pred-target).square() |
| bins=e[:,:153].mean();area=e[:,153].mean();volume=e[:,154].mean() |
| return (bins+area+volume)/3,dict(sigma=float(bins.detach()),area=float(area.detach()),volume=float(volume.detach())) |
|
|
|
|
| def upload_epoch(api,cfg,path): |
| marker=path/'hub_verified.json' |
| if marker.exists():return json.loads(marker.read_text()) |
| names=['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth','cosmo_head.pth', |
| 'training_state.pt','metrics.json','manifest.json','label_transform.json','post_training_architecture_config.yaml','cosmo_decoder_config.json'] |
| ops=[CommitOperationAdd(path_in_repo=f'epoch_{path.name.split("_")[-1]}/{n}',path_or_fileobj=str(path/n)) for n in names] |
| commit=api.create_commit(repo_id=cfg['model_repository'],operations=ops, |
| commit_message=f'Save complete COSMO epoch {path.name}: weights, optimizer, scheduler, RNG and transform') |
| entries={s.rfilename:s for s in api.model_info(cfg['model_repository'],revision=commit.oid,files_metadata=True).siblings} |
| for n in ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth','cosmo_head.pth','training_state.pt']: |
| entry=entries[f'{path.name}/{n}'] |
| remote=entry.lfs.sha256 if entry.lfs else None |
| assert remote==sha(path/n),(n,remote) |
| result=dict(epoch=int(path.name.split('_')[-1]),repository=cfg['model_repository'],revision=commit.oid,verified=True,time=time.time()) |
| write(marker,result) |
| return result |
|
|
|
|
| def main(): |
| parser=argparse.ArgumentParser();parser.add_argument('--smoke',action='store_true') |
| parser.add_argument('--verify-resume',action='store_true');args=parser.parse_args() |
| lock=(WORK/'training.lock').open('a') |
| try:fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) |
| except BlockingIOError: |
| print('TRAINING_ALREADY_RUNNING',flush=True);return |
| cfg=json.loads((WORK/'config.json').read_text()) |
| api=HfApi() |
| status('loading_published_dataset') |
| reference_path=WORK/'dataset_reference.json' |
| reference=json.loads(reference_path.read_text()) if reference_path.exists() else {} |
| revision=reference.get('revision') or api.dataset_info(cfg['dataset_repository']).sha |
| data_root=Path(snapshot_download(cfg['dataset_repository'],repo_type='dataset',revision=revision, |
| allow_patterns=['data/*','manifest.json','label_transform.json','pm6_source.json'])) |
| manifest=json.loads((data_root/'manifest.json').read_text()) |
| assert manifest['rows']==361583 |
| |
| for name,digest in manifest['files'].items():assert sha(data_root/name)==digest,name |
| transform=json.loads((data_root/'label_transform.json').read_text()) |
| frames=[] |
| for split in ['train','validation']: |
| frame=pd.concat([pd.read_parquet(p) for p in sorted((data_root/'data').glob(split+'-*.parquet'))],ignore_index=True) |
| assert len(frame)==manifest['splits'][split] |
| frame['is_validation']=split=='validation';frames.append(frame) |
| df=pd.concat(frames,ignore_index=True) |
| tr=np.flatnonzero(~df.is_validation.values);va=np.flatnonzero(df.is_validation.values) |
| assert not set(df.molecule_group.iloc[tr])&set(df.molecule_group.iloc[va]) |
| target_cols=[f'p_{i:03d}' for i in range(153)]+['area_A2','volume_A3'] |
| labels=df[target_cols].to_numpy(np.float64) |
| transformed=np.concatenate([np.log1p(labels[:,:153]),np.log(labels[:,153:])],axis=1) |
| mean=np.asarray(transform['mean']);scale=np.asarray(transform['scale']) |
| assert np.allclose(transformed[tr].mean(0),mean,atol=1e-10) |
| assert np.allclose(np.maximum(transformed[tr].std(0),1e-3),scale,atol=1e-10) |
| targets=torch.tensor((transformed-mean)/scale,dtype=torch.float32,device='cuda') |
| mean=torch.tensor(mean,dtype=torch.float32,device='cuda');scale=torch.tensor(scale,dtype=torch.float32,device='cuda') |
| positions=[torch.tensor(np.stack(x),dtype=torch.float32) for x in df.positions] |
| numbers=[torch.tensor(np.asarray(x),dtype=torch.long) for x in df.atomic_numbers] |
| lengths=np.array([len(z) for z in numbers]) |
| assert np.isfinite(labels).all() |
| write(WORK/'dataset_reference.json',dict(repository=cfg['dataset_repository'],revision=revision,manifest=manifest)) |
| completed=sorted((WORK/'checkpoints').glob('epoch_[0-9][0-9]')) if (WORK/'checkpoints').exists() else [] |
| if completed and not args.smoke: |
| init=completed[-1] |
| else: |
| init=WORK/'initial_epoch_02';init.mkdir(exist_ok=True) |
| for name in ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth','manifest.json']: |
| src=Path(hf_hub_download(cfg['initial_checkpoint_repository'], |
| filename=cfg['initial_checkpoint_path']+'/'+name,revision=cfg['initial_checkpoint_revision'])) |
| if not (init/name).exists():shutil.copyfile(src,init/name) |
| initial_manifest=json.loads((init/'manifest.json').read_text()) |
| for name in ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth']: |
| assert sha(init/name)==initial_manifest['files'][name]['sha256'] |
| torch.manual_seed(cfg['seed']);torch.cuda.manual_seed_all(cfg['seed']);np.random.seed(cfg['seed']);random.seed(cfg['seed']) |
| pre,enc=backbone(init) |
| model=COSMORegressor(pre,enc) |
| optimizer=torch.optim.AdamW([ |
| {'params':list(pre.parameters())+list(enc.parameters()),'lr':cfg['learning_rate_backbone']}, |
| {'params':model.head.parameters(),'lr':cfg['learning_rate_head']}],weight_decay=cfg['weight_decay']) |
| |
| scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer,lambda e:.1+.9*.5*(1+math.cos(math.pi*e/cfg['epochs']))) |
| mace=build_author_mace_runtime().cuda().eval().requires_grad_(False) |
| captures=[] |
| hook=mace.preprocessor.atomic_preprocessor.register_forward_pre_hook(lambda m,a:captures.append(a[0].detach())) |
| frozen_buffers={n:b.detach().cpu().clone() for n,b in pre.atomic_preprocessor.invariant_normalization.named_buffers()} |
|
|
| def sample(ids): |
| batch=yield_molecules_collate_fn([Sample(atomic_positions=positions[i],atomic_numbers=numbers[i],total_charge=0.,multiplicity=1.) for i in ids]).to_('cuda') |
| captures.clear() |
| with torch.no_grad():mace.preprocessor(batch) |
| assert len(captures)==1 |
| raw=captures.pop().clone() |
| return LegacySample(embeddings=raw,atomic_positions=batch.atomic_positions,padding_mask=batch.padding_mask) |
|
|
| def update(ids): |
| optimizer.zero_grad(set_to_none=True) |
| pred=model(sample(ids));loss,parts=loss_terms(pred,targets[ids]) |
| assert torch.isfinite(loss),'Nonfinite COSMO loss' |
| loss.backward() |
| norm=nn.utils.clip_grad_norm_(model.parameters(),1.,error_if_nonfinite=True) |
| optimizer.step() |
| return float(loss.detach()),parts,float(norm) |
|
|
| if args.smoke: |
| ids=tr[:8].tolist();before=next(enc.parameters()).detach().clone() |
| model.train();result=update(ids) |
| assert not torch.equal(before,next(enc.parameters()).detach()),'Encoder did not update' |
| assert any(p.grad is not None and bool((p.grad!=0).any()) for p in pre.parameters()) |
| assert all(p.grad is None for p in mace.parameters()) |
| model.eval() |
| with torch.no_grad(): |
| s=sample(ids);ref=model(s);physical=inverse(targets[ids],mean,scale) |
| assert np.allclose(physical.cpu().numpy(),labels[ids],rtol=1e-5,atol=1e-4) |
| state={k:v.cpu() for k,v in model.state_dict().items()} |
| path=WORK/'smoke_model.pt';torch.save(state,path) |
| model.load_state_dict(torch.load(path,map_location='cpu',weights_only=True)) |
| with torch.no_grad():assert torch.allclose(ref,model(s),atol=1e-6,rtol=1e-6) |
| for n,b in pre.atomic_preprocessor.invariant_normalization.named_buffers():assert torch.equal(b.cpu(),frozen_buffers[n]) |
| write(WORK/'smoke_validation.json',dict(passed=True,loss=result[0],gradient_norm=result[2], |
| decoder_outputs=155,encoder_updated=True,preprocessor_gradients=True,mace_frozen=True, |
| transform_roundtrip=True,checkpoint_roundtrip=True,dataset_revision=revision)) |
| path.unlink();status('smoke_passed');return |
|
|
| smoke=json.loads((WORK/'smoke_validation.json').read_text()) |
| assert smoke['passed'] and smoke['dataset_revision']==revision,'Smoke validation must match the pinned dataset' |
| import wandb |
| api.create_repo(cfg['model_repository'],repo_type='model',exist_ok=True) |
| wb=json.loads((WORK/'wandb.json').read_text()) if (WORK/'wandb.json').exists() else {'run_id':wandb.util.generate_id()} |
| write(WORK/'wandb.json',wb) |
| run=None |
| if not args.verify_resume: |
| run=wandb.init(entity=cfg['wandb_entity'],project=cfg['wandb_project'],id=wb['run_id'],resume='allow',mode='online', |
| name='rem3di-pcqm-epoch2-cosmo-10epochs',group='rem3di-cosmo',job_type='cosmo-posttraining', |
| config={**cfg,'dataset_revision':revision,'loss':'(mean sigma z-score MSE + A z-score MSE + V z-score MSE)/3', |
| 'train_rows':len(tr),'validation_rows':len(va),'mace':'frozen; original PM6 coordinates; no geometry augmentation'}, |
| dir=str(WORK),settings=wandb.Settings(init_timeout=90)) |
| write(WORK/'wandb.json',dict(run_id=run.id,url=run.url)) |
| run.define_metric('global_step');run.define_metric('train/*',step_metric='global_step') |
| run.define_metric('epoch');run.define_metric('validation/*',step_metric='epoch') |
| start_epoch=1;global_step=0 |
| if completed: |
| last=completed[-1] |
| saved_manifest=json.loads((last/'manifest.json').read_text()) |
| for name in ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth','cosmo_head.pth','training_state.pt']: |
| assert sha(last/name)==saved_manifest['files'][name]['sha256'],name |
| for n,module in [('encoder.pth',enc),('atomic_preprocessor.pth',pre.atomic_preprocessor),('geometric_preprocessor.pth',pre.geometric_preprocessor),('cosmo_head.pth',model.head)]: |
| module.load_state_dict(torch.load(last/n,map_location='cpu',weights_only=True)) |
| state=torch.load(last/'training_state.pt',map_location='cpu',weights_only=False) |
| assert state['dataset_revision']==revision,'Resume must preserve the original dataset revision' |
| assert state['config']==cfg,'Resume configuration differs from the saved run' |
| assert state['epoch']==int(last.name.split('_')[-1]) |
| optimizer.load_state_dict(state['optimizer']);scheduler.load_state_dict(state['scheduler']) |
| 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']) |
| start_epoch=state['epoch']+1;global_step=state['global_step'] |
| upload_epoch(api,cfg,last) |
| write(WORK/'resume_state.json',dict(restored_epoch=state['epoch'],next_epoch=start_epoch, |
| global_step=global_step,dataset_revision=revision,optimizer_restored=True,scheduler_restored=True,rng_restored=True)) |
| print('RESUMED',json.dumps(dict(epoch=state['epoch'],next_epoch=start_epoch,global_step=global_step)),flush=True) |
| if args.verify_resume: |
| assert completed,'No completed checkpoint to verify' |
| model.eval() |
| with torch.no_grad(): |
| output=model(sample(tr[:8].tolist())) |
| loss,_=loss_terms(output,targets[tr[:8]]) |
| physical=inverse(output,mean,scale) |
| assert output.shape==(8,155) and torch.isfinite(output).all() and torch.isfinite(physical).all() |
| assert torch.isfinite(loss) and optimizer.state and scheduler.last_epoch==start_epoch-1 |
| write(WORK/'resume_validation.json',dict(passed=True,restored_epoch=start_epoch-1,next_epoch=start_epoch, |
| global_step=global_step,loss=float(loss),optimizer_state_tensors=len(optimizer.state), |
| scheduler_epoch=scheduler.last_epoch,dataset_revision=revision,finite_predictions=True)) |
| hook.remove();status('resume_verified',next_epoch=start_epoch,global_step=global_step);return |
| try: |
| for epoch in range(start_epoch,cfg['epochs']+1): |
| model.train();started=time.time();loss_sum=0.;count=0 |
| rng=np.random.default_rng(cfg['seed']+epoch) |
| plan=list(batches(tr,lengths,rng,max_size=cfg['max_molecules_per_batch'],atom_cap=cfg['max_atoms_per_batch'])) |
| status('training',epoch=epoch,epochs=cfg['epochs'],batches=len(plan),wandb_url=run.url) |
| for bi,ids in enumerate(plan,1): |
| loss,parts,norm=update(ids) |
| global_step+=1;loss_sum+=loss*len(ids);count+=len(ids) |
| if bi==1 or bi%25==0: |
| row=dict(global_step=global_step,epoch=epoch,batch=bi,batches=len(plan), |
| **{'train/'+k:v for k,v in dict(loss=loss,gradient_norm=norm,backbone_lr=optimizer.param_groups[0]['lr'],head_lr=optimizer.param_groups[1]['lr'],**parts).items()}) |
| run.log(row) |
| status('training',**row,wandb_url=run.url,molecules_seen=count,seconds=time.time()-started) |
| model.eval();val_sum=0.;val_count=0;sse=np.zeros(3);mae=np.zeros(3);sumarea=0. |
| with torch.no_grad(): |
| for ids in batches(va,lengths,max_size=cfg['max_molecules_per_batch'],atom_cap=cfg['max_atoms_per_batch']): |
| pred=model(sample(ids));loss,_=loss_terms(pred,targets[ids]) |
| val_sum+=float(loss)*len(ids);val_count+=len(ids) |
| physical=inverse(pred,mean,scale).cpu().numpy() |
| assert np.isfinite(physical).all() |
| err=physical-labels[ids] |
| sse += [float((err[:,:153]**2).sum()/153),float((err[:,153]**2).sum()),float((err[:,154]**2).sum())] |
| mae += [float(abs(err[:,:153]).sum()/153),float(abs(err[:,153]).sum()),float(abs(err[:,154]).sum())] |
| sumarea+=float(abs(physical[:,:153].sum(1)-physical[:,153]).sum()) |
| assert count==len(tr) and val_count==len(va) |
| for n,b in pre.atomic_preprocessor.invariant_normalization.named_buffers():assert torch.equal(b.cpu(),frozen_buffers[n]) |
| scheduler.step() |
| metrics=dict(epoch=epoch,global_step=global_step,train_loss=loss_sum/count,validation_loss=val_sum/val_count, |
| validation_sigma_rmse=float(np.sqrt(sse[0]/val_count)),validation_area_rmse=float(np.sqrt(sse[1]/val_count)), |
| validation_volume_rmse=float(np.sqrt(sse[2]/val_count)),validation_sigma_mae=float(mae[0]/val_count), |
| validation_area_mae=float(mae[1]/val_count),validation_volume_mae=float(mae[2]/val_count), |
| validation_profile_area_inconsistency=float(sumarea/val_count),seconds=time.time()-started) |
| assert all(math.isfinite(v) for v in metrics.values()) |
| target=WORK/'checkpoints'/f'epoch_{epoch:02d}' |
| pending=target.with_name(target.name+'.partial');pending.mkdir(parents=True,exist_ok=True) |
| for n,module in [('encoder.pth',enc),('atomic_preprocessor.pth',pre.atomic_preprocessor),('geometric_preprocessor.pth',pre.geometric_preprocessor),('cosmo_head.pth',model.head)]: |
| torch.save(module.state_dict(),pending/n) |
| torch.save(dict(epoch=epoch,global_step=global_step,optimizer=optimizer.state_dict(),scheduler=scheduler.state_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(),dataset_revision=revision,config=cfg),pending/'training_state.pt') |
| write(pending/'metrics.json',metrics) |
| shutil.copyfile(data_root/'label_transform.json',pending/'label_transform.json') |
| architecture=yaml.safe_load((ORIGINAL/'post_training_architecture_config.yaml').read_text()) |
| architecture.pop('decoder_config',None);architecture['kind']='encoder_only' |
| architecture['mace_config']['model_path']='MACE-POLAR-1-M.model' |
| (pending/'post_training_architecture_config.yaml').write_text(yaml.safe_dump(architecture,sort_keys=False)) |
| write(pending/'cosmo_decoder_config.json',dict(input_features=640,hidden_features=256,output_features=155,activation='SiLU')) |
| write(pending/'manifest.json',dict(epoch=epoch,dataset_revision=revision,config=cfg, |
| files={p.name:dict(bytes=p.stat().st_size,sha256=sha(p)) for p in pending.iterdir() if p.is_file()})) |
| pending.rename(target) |
| status('uploading_checkpoint',epoch=epoch,metrics=metrics,wandb_url=run.url) |
| for attempt in range(5): |
| try: |
| result=upload_epoch(api,cfg,target);break |
| except Exception: |
| if attempt==4:raise |
| print('UPLOAD_RETRY',traceback.format_exc(),flush=True);time.sleep(min(10*2**attempt,60)) |
| run.log(dict(epoch=epoch,**{'validation/'+k.removeprefix('validation_'):v for k,v in metrics.items() if k.startswith('validation_')}, |
| **{'epoch/'+k:v for k,v in metrics.items()},**{'backup/revision':result['revision'],'backup/verified_epoch':epoch})) |
| write(WORK/'latest_epoch.json',dict(**metrics,hub=result,wandb_url=run.url)) |
| write(WORK/'training_complete.json',dict(epochs=cfg['epochs'],wandb_url=run.url,completed_unix=time.time())) |
| run.summary['training_complete']=True;run.finish();status('complete',epochs=cfg['epochs']) |
| except Exception: |
| status('failed',error=traceback.format_exc());run.finish(exit_code=1);raise |
| finally: |
| hook.remove() |
|
|
|
|
| if __name__=='__main__': |
| try:main() |
| except Exception: |
| status('failed',error=traceback.format_exc());raise |
|
|
