"""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'] # Resume inputs must use saved statistics, including when data path changes. 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