| """Run the full Table 9 Darcy schedule from the package root.""" |
| from __future__ import annotations |
|
|
| import csv |
| import json |
| import os |
| from pathlib import Path |
| import random |
| import time |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader, TensorDataset |
|
|
| from common import (HNO, atomic_checkpoint, atomic_json, compute_device, configuration, |
| decode_output, encode_input, load_checkpoint, objective, read_config, |
| read_split, relative_l2, set_seed, source_hashes, training_statistics) |
|
|
|
|
| def optimizer_and_schedule(model, c, steps): |
| t = c['training'] |
| optimizer = torch.optim.AdamW(model.parameters(), lr=t['learning_rate'], |
| weight_decay=t['weight_decay'], betas=tuple(t['adam_betas']), eps=t['adam_eps']) |
| scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=t['learning_rate'], |
| epochs=t['epochs'], steps_per_epoch=steps, **t['onecycle']) |
| return optimizer, scheduler |
|
|
|
|
| def evaluate(model, loader, stats, device): |
| model.eval() |
| accumulated = torch.zeros((), device=device, dtype=stats['target_mean'].dtype) |
| count = 0 |
| with torch.no_grad(): |
| for x, target in loader: |
| x, target = x.to(device), target.to(device) |
| prediction = decode_output(model(x), stats) |
| accumulated += relative_l2(prediction, target).sum() |
| count += len(x) |
| return accumulated.item() / count |
|
|
|
|
| def epoch_train(model, loader, stats, device, optimizer, scheduler, c, epoch, smoke=False): |
| model.train() |
| totals = torch.zeros(3, device=device, dtype=stats['target_mean'].dtype) |
| count = 0 |
| for batch, (x, target) in enumerate(loader, start=1): |
| x, target = x.to(device), target.to(device) |
| optimizer.zero_grad(set_to_none=True) |
| prediction = decode_output(model(x), stats) |
| loss, primary, auxiliary = objective(prediction, target, c['training']) |
| loss.backward() |
| norm = torch.nn.utils.clip_grad_norm_(model.parameters(), c['training']['gradient_clip'], error_if_nonfinite=True) |
| optimizer.step() |
| scheduler.step() |
| totals += torch.stack((loss.detach(), primary.detach(), auxiliary.detach())) |
| count += len(x) |
| if batch % c['training']['print_every_batches'] == 0 or batch == len(loader): |
| print(f'{"SMOKE " if smoke else ""}epoch={epoch:03d} batch={batch:03d}/{len(loader)} ' |
| f'loss={loss.item()/len(x):.8f} rel_l2={primary.item()/len(x):.8f} ' |
| f'derivative={auxiliary.item()/len(x):.8f} grad_norm_before_clip={float(norm):.6g}', flush=True) |
| values = (totals / count).cpu().tolist() |
| if not all(np.isfinite(values)): |
| raise FloatingPointError('Nonfinite epoch losses') |
| return values |
|
|
|
|
| def rng_snapshot(generator): |
| ns = np.random.get_state() |
| return {'python': random.getstate(), 'numpy': (ns[0], ns[1].tolist(), ns[2], ns[3], ns[4]), |
| 'torch': torch.get_rng_state(), 'accelerator': torch.cuda.get_rng_state_all(), |
| 'loader': generator.get_state()} |
|
|
|
|
| def restore_rng(state, generator): |
| random.setstate(state['python']) |
| ns = state['numpy'] |
| np.random.set_state((ns[0], np.asarray(ns[1], dtype=np.uint32), ns[2], ns[3], ns[4])) |
| torch.set_rng_state(state['torch']) |
| torch.cuda.set_rng_state_all(state['accelerator']) |
| generator.set_state(state['loader']) |
|
|
|
|
| def main(): |
| parser = configuration('HNO Darcy Table 9 training') |
| parser.add_argument('--resume', help='Explicit last checkpoint to continue the same schedule') |
| parser.add_argument('--smoke', action='store_true', help='Six diagnostic checks; separate outputs, not a formal run') |
| args = parser.parse_args() |
| c = read_config(args.config) |
| device = compute_device(args.device) |
| torch.set_num_threads(c['training']['cpu_threads']) |
| set_seed(c['seed']) |
| torch.backends.cudnn.benchmark = True |
| results = Path(c['output']['results_dir']) |
| results.mkdir(parents=True, exist_ok=True) |
| print(f'SLURM_JOB_ID={os.environ.get("SLURM_JOB_ID")} node={os.uname().nodename} ' |
| f'device={torch.cuda.get_device_name()} torch={torch.__version__} hip={torch.version.hip}', flush=True) |
| start = time.monotonic() |
| a_train, u_train = read_split(c, 'train') |
| a_eval, u_eval = read_split(c, 'eval') |
| stats_cpu = training_statistics(a_train, u_train) |
| if any(torch.linalg.vector_norm(u.flatten(1), dim=1).min() <= 1e-12 for u in [u_train, u_eval]): |
| raise ValueError('Zero-norm solution sample') |
| stats = {k: v.to(device) for k, v in stats_cpu.items()} |
| generator = torch.Generator().manual_seed(c['seed']) |
| train_loader = DataLoader(TensorDataset(encode_input(a_train, stats_cpu), u_train), |
| batch_size=c['training']['batch_size'], shuffle=True, generator=generator, |
| num_workers=c['training']['num_workers'], pin_memory=True) |
| eval_loader = DataLoader(TensorDataset(encode_input(a_eval, stats_cpu), u_eval), |
| batch_size=c['training']['batch_size'], shuffle=False, |
| num_workers=c['training']['num_workers'], pin_memory=True) |
| model = HNO(**c['model']).to(device) |
| parameters = sum(p.numel() for p in model.parameters()) |
| print(f'Parameters={parameters:,}; train={len(a_train)} eval={len(a_eval)} ' |
| f'grid={tuple(a_train.shape[1:])} tokens={model.patch_grid**2} ' |
| f'target_dtype={u_train.dtype}; epochs={c["training"]["epochs"]}', flush=True) |
| optimizer, scheduler = optimizer_and_schedule(model, c, len(train_loader)) |
| best, best_epoch, start_epoch = float('inf'), 0, 1 |
| if args.resume: |
| model, stats, checkpoint = load_checkpoint(args.resume, device) |
| if checkpoint['training_config'] != c['training'] or checkpoint['model_config'] != c['model']: |
| raise ValueError('Resume requires identical architecture and training schedule') |
| optimizer, scheduler = optimizer_and_schedule(model, c, len(train_loader)) |
| optimizer.load_state_dict(checkpoint['optimizer']) |
| scheduler.load_state_dict(checkpoint['scheduler']) |
| restore_rng(checkpoint['rng'], generator) |
| stats_cpu = checkpoint['normalization'] |
| |
| train_loader.dataset.tensors = (encode_input(a_train, stats_cpu), u_train) |
| eval_loader.dataset.tensors = (encode_input(a_eval, stats_cpu), u_eval) |
| best, best_epoch = checkpoint['best_eval_relative_l2'], checkpoint['best_epoch'] |
| start_epoch = checkpoint['epoch'] + 1 |
| if args.smoke: |
| checks = {} |
| x, y = next(iter(eval_loader)) |
| x, y = x.to(device), y.to(device) |
| prediction = decode_output(model(x), stats) |
| assert prediction.shape == y.shape and torch.isfinite(prediction).all() |
| checks['forward_pass'] = 'passed' |
| loss, _, _ = objective(prediction, y, c['training']) |
| loss.backward() |
| assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters()) |
| checks['backward_pass'] = 'passed' |
| losses = epoch_train(model, train_loader, stats, device, optimizer, scheduler, c, 1, smoke=True) |
| checks['train_loop_dry_run'] = {'status': 'passed', 'n_samples':len(a_train), 'losses':losses} |
| score = evaluate(model, eval_loader, stats, device) |
| assert np.isfinite(score) |
| checks['validation_loop_dry_run'] = {'status': 'passed','n_samples':len(a_eval),'relative_l2':score} |
| assert torch.equal(relative_l2(y, y), torch.zeros(len(y), device=device, dtype=y.dtype)) |
| assert torch.allclose(relative_l2(2*y, y), torch.ones(len(y), device=device,dtype=y.dtype)) |
| checks['domain_metric_dry_run'] = 'passed: exact identity and scaling checks' |
| assert parameters == 815873, f'Unexpected parameter count: {parameters}' |
| assert len(train_loader) == 250 and len(eval_loader) == 50 |
| checks['config_consistency'] = {'status':'passed','parameters':parameters,'training_batches':250,'evaluation_batches':50} |
| atomic_json(results/'smoke_checks.json', checks) |
| print('SMOKE_ALL_SIX_CHECKS_PASSED', flush=True) |
| return |
| manifest = {'job_id':os.environ.get('SLURM_JOB_ID'), 'node':os.uname().nodename, |
| 'device':torch.cuda.get_device_name(), 'torch':str(torch.__version__), 'hip':torch.version.hip, |
| 'parameters':parameters,'config':c,'source_sha256':source_hashes(), |
| 'n_train':len(a_train),'n_eval':len(a_eval),'steps_per_epoch':len(train_loader), |
| 'target_dtype':str(u_train.dtype), 'checkpoint_selection':'minimum evaluation relative L2; no separate validation set', |
| 'normalization':{k:v.tolist() for k,v in stats_cpu.items()},'status':'running'} |
| atomic_json(results/'run_manifest.json',manifest) |
| fields=['epoch','global_step','train_loss','train_relative_l2','train_derivative','eval_relative_l2','learning_rate','seconds','best_epoch','best_relative_l2'] |
| history=results/'training_history.csv' |
| with history.open('a' if args.resume else 'w', newline='') as f: |
| writer=csv.DictWriter(f,fieldnames=fields) |
| if not args.resume: writer.writeheader() |
| for epoch in range(start_epoch,c['training']['epochs']+1): |
| tic=time.monotonic() |
| losses=epoch_train(model,train_loader,stats,device,optimizer,scheduler,c,epoch) |
| score=evaluate(model,eval_loader,stats,device) |
| if not np.isfinite(score): raise FloatingPointError('Nonfinite evaluation metric') |
| improved=score<best |
| if improved: best,best_epoch=score,epoch |
| row=dict(zip(fields,[epoch,epoch*len(train_loader),*losses,score,optimizer.param_groups[0]['lr'],time.monotonic()-tic,best_epoch,best])) |
| writer.writerow(row); f.flush() |
| print(f'epoch={epoch:03d}/{c["training"]["epochs"]} step={row["global_step"]} ' |
| f'train_loss={losses[0]:.8f} train_rel={losses[1]:.8f} derivative={losses[2]:.8f} ' |
| f'eval_rel={score:.8f} best={best:.8f}@{best_epoch} ' |
| f'lr={row["learning_rate"]:.8g} seconds={row["seconds"]:.2f}',flush=True) |
| saved={'format_version':1,'model_config':c['model'],'model_state':model.state_dict(), |
| 'normalization':{k:v.detach().cpu() for k,v in stats.items()},'epoch':epoch, |
| 'global_step':epoch*len(train_loader),'eval_relative_l2':score, |
| 'best_eval_relative_l2':best,'best_epoch':best_epoch,'seed':c['seed'], |
| 'training_config':c['training'],'config_at_training':c} |
| if improved: atomic_checkpoint(c['output']['checkpoint'],saved) |
| saved.update(optimizer=optimizer.state_dict(),scheduler=scheduler.state_dict(),rng=rng_snapshot(generator)) |
| atomic_checkpoint(c['output']['last_checkpoint'],saved) |
| atomic_json(results/'training_progress.json',row) |
| manifest.update(status='completed',completed_epochs=c['training']['epochs'], |
| global_step=c['training']['epochs']*len(train_loader),best_epoch=best_epoch, |
| best_eval_relative_l2=best,elapsed_seconds=time.monotonic()-start) |
| atomic_json(results/'run_manifest.json',manifest) |
| atomic_json(results/'training_summary.json',{k:manifest[k] for k in ['status','completed_epochs','global_step','best_epoch','best_eval_relative_l2','elapsed_seconds']}) |
| print(f'FORMAL_TRAINING_COMPLETED epochs={c["training"]["epochs"]} best={best:.9f} checkpoint={c["output"]["checkpoint"]}',flush=True) |
|
|
|
|
| if __name__=='__main__': |
| main() |
|
|
