"""Train CorrDiff's conditional mean, then its frozen-mean residual EDM.""" import argparse import json import os import random import sys from contextlib import nullcontext from pathlib import Path import numpy as np import torch import yaml from torch import distributed as dist from torch.nn import functional as F from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, DistributedSampler, TensorDataset ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from model.corrdiff import CorrDiff def scalar(archive, key, default=None): if key not in archive: if default is not None: return default raise ValueError(f"NPZ is missing required metadata: {key}") value = archive[key] if value.ndim != 0: raise ValueError(f"NPZ metadata {key} must be a scalar") return str(value.item()) def reduced_average(total, count, device, distributed): values = torch.tensor([total, count], dtype=torch.float64, device=device) if distributed: dist.all_reduce(values, op=dist.ReduceOp.SUM) if values[1].item() == 0: raise RuntimeError("Training stage processed no batches") return (values[0] / values[1]).item() def main(): parser = argparse.ArgumentParser() parser.add_argument("--config", default=str(ROOT / "conf/config.yaml")) args = parser.parse_args() config = yaml.safe_load(Path(args.config).read_text(encoding="utf-8")) global_rank = int(os.getenv("RANK", 0)) local_rank = int(os.getenv("LOCAL_RANK", 0)) world = int(os.getenv("WORLD_SIZE", 1)) distributed = world > 1 if distributed: dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo") use_cuda = torch.cuda.is_available() and config["runtime"]["device"] != "cpu" device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu") if use_cuda: torch.cuda.set_device(local_rank) seed = config["seed"] + global_rank random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) archive = np.load(ROOT / config["data"]["path"]) protocol = scalar(archive, "protocol") data_source = scalar(archive, "data_source") if protocol != config["data"]["protocol"]: raise ValueError(f"Expected protocol {config['data']['protocol']}, got {protocol}") if not data_source: raise ValueError("data_source must be a non-empty scalar") coarse = torch.from_numpy(archive[config["data"]["input_key"]]) target = torch.from_numpy(archive[config["data"]["target_key"]]) if coarse.ndim != 4 or target.ndim != 4: raise ValueError("CorrDiff input and target must be NCHW tensors") if len(coarse) != len(target) or tuple(coarse.shape[1:]) != tuple(config["data"]["input_shape"]) or tuple(target.shape[1:]) != tuple(config["data"]["target_shape"]): raise ValueError("NPZ tensor shapes do not match config") dataset = TensorDataset(coarse, target) sampler = DistributedSampler(dataset, shuffle=True) if distributed else None loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], sampler=sampler, shuffle=sampler is None, num_workers=config["training"]["num_workers"]) model = CorrDiff(**config["model"]).to(device) if distributed: model = DDP(model, device_ids=[local_rank] if use_cuda else None, find_unused_parameters=True) base = model.module if distributed else model reg_opt = torch.optim.AdamW(base.regression.parameters(), lr=config["training"]["learning_rate"]) diff_opt = torch.optim.AdamW(base.diffusion.parameters(), lr=config["training"]["learning_rate"]) amp = bool(config["training"]["amp"] and use_cuda) scaler = torch.amp.GradScaler("cuda", enabled=amp) autocast = (lambda: torch.amp.autocast("cuda", enabled=True)) if amp else nullcontext history = [] # Stage 1 is completed in full before any residual-EDM update occurs. for epoch in range(config["training"]["regression_epochs"]): if sampler is not None: sampler.set_epoch(epoch) model.train() total = count = 0 for batch_index, (coarse_batch, target_batch) in enumerate(loader): coarse_batch, target_batch = coarse_batch.to(device), target_batch.to(device) reg_opt.zero_grad(set_to_none=True) with autocast(): loss = F.mse_loss(model(coarse_batch, mode="mean"), target_batch) scaler.scale(loss).backward() scaler.step(reg_opt) scaler.update() total += loss.item() count += 1 if batch_index + 1 >= config["training"]["max_batches_per_epoch"]: break value = reduced_average(total, count, device, distributed) record = {"stage": "regression", "epoch": epoch + 1, "regression_mse": value} history.append(record) if global_rank == 0: print(json.dumps(record)) base.regression.eval() for parameter in base.regression.parameters(): parameter.requires_grad_(False) for epoch in range(config["training"]["diffusion_epochs"]): if sampler is not None: sampler.set_epoch(config["training"]["regression_epochs"] + epoch) base.diffusion.train() total = count = 0 for batch_index, (coarse_batch, target_batch) in enumerate(loader): coarse_batch, target_batch = coarse_batch.to(device), target_batch.to(device) with torch.no_grad(): mean = base.mean(coarse_batch) residual = target_batch - mean sigma = (torch.randn(len(coarse_batch), device=device) * config["training"]["p_std"] + config["training"]["p_mean"]).exp() noisy = residual + sigma[:, None, None, None] * torch.randn_like(residual) diff_opt.zero_grad(set_to_none=True) with autocast(): denoised = model(coarse_batch, mode="denoise", mean=mean, noisy=noisy, sigma=sigma) weight = (sigma.square() + base.sigma_data**2) / (sigma * base.sigma_data).square() loss = (weight[:, None, None, None] * (denoised - residual).square()).mean() scaler.scale(loss).backward() scaler.step(diff_opt) scaler.update() total += loss.item() count += 1 if batch_index + 1 >= config["training"]["max_batches_per_epoch"]: break value = reduced_average(total, count, device, distributed) record = {"stage": "diffusion", "epoch": epoch + 1, "edm_loss": value} history.append(record) if global_rank == 0: print(json.dumps(record)) if global_rank == 0: checkpoint = ROOT / config["paths"]["checkpoint"] checkpoint.parent.mkdir(parents=True, exist_ok=True) torch.save({"model": base.state_dict(), "config": config, "format": "corrdiff-edm-v3", "protocol": protocol, "data_source": data_source}, checkpoint) metrics = ROOT / config["paths"]["training_metrics"] metrics.parent.mkdir(parents=True, exist_ok=True) metrics.write_text(json.dumps({"history": history, "protocol": protocol, "data_source": data_source}, indent=2) + "\n") print(f"checkpoint={checkpoint}") if distributed: dist.destroy_process_group() if __name__ == "__main__": main()