from __future__ import annotations import sys import time from pathlib import Path from typing import Any import torch import torch.nn.functional as F from torch.nn.parallel import DistributedDataParallel as DDP from tqdm import tqdm PROJECT_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(PROJECT_ROOT)) sys.path.insert(0, str(Path(__file__).resolve().parent)) LOCAL_WORKSPACE = PROJECT_ROOT.parents[1] if (LOCAL_WORKSPACE / "onescience" / "src" / "onescience" / "datapipes").is_dir(): sys.path.insert(0, str(LOCAL_WORKSPACE)) from common import ( # noqa: E402 load_config, project_path, relative_l2, resolve_device, rollout, seed_everything, to_attr_dict, to_plain_dict, ) from model.factformer import FactFormer2D # noqa: E402 from onescience.datapipes.cfd import KolmogorovFlow2DDatapipe # noqa: E402 from onescience.distributed.manager import DistributedManager # noqa: E402 DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml" def prepare_config(config: dict[str, Any]) -> None: datapipe = config["datapipe"] data = datapipe["data"] model = config["model"] datapipe["source"]["data_dir"] = str( project_path(datapipe["source"]["data_dir"], PROJECT_ROOT).resolve() ) stats_file = data.get("stats_file") if stats_file: data["stats_file"] = str(project_path(stats_file, PROJECT_ROOT).resolve()) model["in_dim"] = int(data["t_in"]) * int(data["out_dim"]) model["out_dim"] = int(data["out_dim"]) def validate_config(config: dict[str, Any]) -> None: data = config["datapipe"]["data"] loader = config["datapipe"]["dataloader"] model = config["model"] training = config["training"] positive_values = { "train_num": data["train_num"], "test_num": data["test_num"], "resolution": data["resolution"], "interval": data["interval"], "t_in": data["t_in"], "t_out": data["t_out"], "batch_size": loader["batch_size"], "hidden_dim": model["hidden_dim"], "depth": model["depth"], "heads": model["heads"], "mlp_ratio": model["mlp_ratio"], "max_latent_steps": model["max_latent_steps"], "epochs": training["epochs"], "eval_interval": training["eval_interval"], "train_latent_steps": training["train_latent_steps"], } for name, value in positive_values.items(): if int(value) < 1: raise ValueError(f"{name} must be positive, got {value}") if int(loader["num_workers"]) < 0: raise ValueError("num_workers cannot be negative") if int(model["hidden_dim"]) % int(model["heads"]): raise ValueError("hidden_dim must be divisible by heads") if int(training["train_latent_steps"]) > int(model["max_latent_steps"]): raise ValueError("train_latent_steps cannot exceed model.max_latent_steps") if int(training["train_latent_steps"]) > int(data["t_out"]): raise ValueError("train_latent_steps cannot exceed datapipe.data.t_out") for name in ("max_train_batches", "max_eval_batches"): value = training.get(name) if value is not None and int(value) < 1: raise ValueError(f"{name} must be positive when set") def build_model( model_config: dict[str, Any], spatial_shape: tuple[int, int] ) -> FactFormer2D: return FactFormer2D( in_dim=int(model_config["in_dim"]), out_dim=int(model_config["out_dim"]), spatial_shape=spatial_shape, hidden_dim=int(model_config["hidden_dim"]), depth=int(model_config["depth"]), heads=int(model_config["heads"]), mlp_ratio=int(model_config["mlp_ratio"]), dropout=float(model_config["dropout"]), activation=str(model_config["activation"]), include_pos=bool(model_config["include_pos"]), space_dim=int(model_config["space_dim"]), latent_multiplier=float(model_config["latent_multiplier"]), max_latent_steps=int(model_config["max_latent_steps"]), ) def evaluate( model: torch.nn.Module, loader: torch.utils.data.DataLoader, device: torch.device, datapipe: KolmogorovFlow2DDatapipe, t_out: int, out_dim: int, max_latent_steps: int, max_batches: int | None, distributed: bool, rank: int, ) -> tuple[float, float]: model.eval() totals = torch.zeros(4, dtype=torch.float64, device=device) with torch.no_grad(): iterator = tqdm(loader, desc="Evaluating", disable=rank != 0) for batch_index, batch in enumerate(iterator): if max_batches is not None and batch_index >= max_batches: break pos = batch["pos"].to(device) state = batch["x"].to(device) target = batch["y"].to(device) prediction = rollout( model, pos, state, t_out, out_dim, max_latent_steps ) prediction = datapipe.decode_solution(prediction) target = datapipe.decode_solution(target) totals[0] += F.mse_loss(prediction, target, reduction="sum") totals[1] += target.numel() totals[2] += relative_l2(prediction, target).sum() totals[3] += target.shape[0] if distributed: torch.distributed.all_reduce(totals) return ( (totals[0] / totals[1].clamp_min(1)).item(), (totals[2] / totals[3].clamp_min(1)).item(), ) def main() -> None: config_path = DEFAULT_CONFIG.resolve() config = load_config(config_path) prepare_config(config) validate_config(config) common = config["common"] datapipe_config = config["datapipe"] model_config = config["model"] training = config["training"] seed_everything(int(common["seed"])) DistributedManager.initialize() dist = DistributedManager() distributed = dist.world_size > 1 device = dist.device if str(common["device"]) == "auto" else resolve_device( str(common["device"]) ) weight_dir = project_path(training["weight_dir"], PROJECT_ROOT).resolve() checkpoint_path = weight_dir / str(training["checkpoint_name"]) if dist.rank == 0: print(f"Config: {config_path}") print( "Data: " f"{Path(datapipe_config['source']['data_dir']) / datapipe_config['source']['file_name']}" ) print(f"Device: {device}") print( f"Samples: train={datapipe_config['data']['train_num']} " f"test={datapipe_config['data']['test_num']} " f"t_in={datapipe_config['data']['t_in']} " f"t_out={datapipe_config['data']['t_out']}" ) started = time.time() try: datapipe = KolmogorovFlow2DDatapipe( to_attr_dict(datapipe_config), distributed=distributed ) train_loader, train_sampler = datapipe.train_dataloader() test_loader, _ = datapipe.test_dataloader() spatial_shape = tuple(datapipe.spatial_shape) model = build_model(model_config, spatial_shape).to(device) if distributed: device_ids = [dist.local_rank] if device.type == "cuda" else None model = DDP(model, device_ids=device_ids) if dist.rank == 0: parameter_count = sum(parameter.numel() for parameter in model.parameters()) print(f"Spatial shape: {spatial_shape}") print(f"Parameters: {parameter_count:,}") optimizer = torch.optim.AdamW( model.parameters(), lr=float(training["lr"]), weight_decay=float(training["weight_decay"]), ) scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=int(training["step_size"]), gamma=float(training["gamma"]), ) data = datapipe_config["data"] t_out = int(data["t_out"]) out_dim = int(data["out_dim"]) max_latent_steps = int(model_config["max_latent_steps"]) train_latent_steps = int(training["train_latent_steps"]) max_train_batches = training.get("max_train_batches") max_train_batches = None if max_train_batches is None else int(max_train_batches) max_eval_batches = training.get("max_eval_batches") max_eval_batches = None if max_eval_batches is None else int(max_eval_batches) best_relative_l2 = float("inf") stale_evaluations = 0 for epoch in range(1, int(training["epochs"]) + 1): if train_sampler is not None: train_sampler.set_epoch(epoch) model.train() epoch_loss = 0.0 batch_count = 0 iterator = tqdm(train_loader, desc=f"Epoch {epoch}", disable=dist.rank != 0) for batch_index, batch in enumerate(iterator): if max_train_batches is not None and batch_index >= max_train_batches: break pos = batch["pos"].to(device) state = batch["x"].to(device) target = batch["y"].to(device)[..., : train_latent_steps * out_dim] optimizer.zero_grad(set_to_none=True) prediction = model(pos, state, latent_steps=train_latent_steps) loss = F.mse_loss(prediction, target) if not torch.isfinite(loss): raise FloatingPointError(f"Non-finite training loss at epoch {epoch}") loss.backward() if training.get("max_grad_norm") is not None: torch.nn.utils.clip_grad_norm_( model.parameters(), float(training["max_grad_norm"]) ) optimizer.step() epoch_loss += loss.item() batch_count += 1 if dist.rank == 0: iterator.set_postfix(loss=f"{loss.item():.3e}") scheduler.step() if batch_count == 0: raise RuntimeError("No training batches were processed") should_evaluate = ( epoch % int(training["eval_interval"]) == 0 or epoch == int(training["epochs"]) ) if not should_evaluate: continue validation_mse, validation_relative_l2 = evaluate( model, test_loader, device, datapipe, t_out, out_dim, max_latent_steps, max_eval_batches, distributed, dist.rank, ) if dist.rank == 0: print( f"epoch={epoch:4d} train_mse={epoch_loss / batch_count:.6e} " f"val_mse={validation_mse:.6e} " f"val_relative_l2={validation_relative_l2:.6e}" ) if validation_relative_l2 < best_relative_l2: best_relative_l2 = validation_relative_l2 stale_evaluations = 0 weight_dir.mkdir(parents=True, exist_ok=True) model_to_save = model.module if isinstance(model, DDP) else model torch.save( { "epoch": epoch, "model_state": model_to_save.state_dict(), "model_config": to_plain_dict(model_config), "datapipe_config": to_plain_dict(datapipe_config), "training_config": to_plain_dict(training), "spatial_shape": spatial_shape, "normalizer": datapipe.get_normalizer_state(), "best_relative_l2": best_relative_l2, }, checkpoint_path, ) print(f"Saved checkpoint: {checkpoint_path}") else: stale_evaluations += 1 stop = torch.tensor( [int(stale_evaluations >= int(training["patience"]))], device=device ) if distributed: torch.distributed.broadcast(stop, src=0) if stop.item(): if dist.rank == 0: print("Early stopping triggered") break if dist.rank == 0: print(f"Elapsed: {time.time() - started:.1f}s") finally: DistributedManager.cleanup() if __name__ == "__main__": main()