"""Train ACE for one-step normalized field prediction.""" from __future__ import annotations import argparse import json import os import random import sys from pathlib import Path import numpy as np import torch import yaml from torch import nn import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler if __package__ in (None, ""): sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from ACE.model.data import ArrayPairDataset, load_npz, make_fake_pairs, save_fake_pairs from ACE.model.ace import ACEModel, ACEModelConfig from ACE.model.normalization import ACEDataNormalizer from ACE.model.paths import CHECKPOINT_DIR, GENERATED_DATA_PATH, configured_path class EMA: def __init__(self, model: nn.Module, decay: float) -> None: self.decay = float(decay) self.shadow = {name: value.detach().clone() for name, value in model.state_dict().items()} def update(self, model: nn.Module) -> None: with torch.no_grad(): for name, value in model.state_dict().items(): self.shadow[name].mul_(self.decay).add_(value.detach(), alpha=1.0 - self.decay) def copy_to(self, model: nn.Module) -> None: model.load_state_dict(self.shadow, strict=True) def parameter_statistics(model: nn.Module) -> tuple[int, int]: """Count real scalar parameters, expanding complex values to real/imag parts.""" total = 0 nonzero = 0 for parameter in model.parameters(): values = torch.view_as_real(parameter.detach()) if parameter.is_complex() else parameter.detach() total += values.numel() nonzero += int(torch.count_nonzero(values).item()) return total, nonzero def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", type=Path, default=Path(__file__).resolve().parents[1] / "conf" / "config.yaml") parser.add_argument("--data-path", type=Path, default=None, help="NPZ with inputs[N,40,H,W] and targets[N,44,H,W]") parser.add_argument("--fake-data", action="store_true", help="Use fake fields for smoke only") parser.add_argument("--num-samples", type=int, default=8) parser.add_argument("--height", type=int, default=180) parser.add_argument("--width", type=int, default=360) parser.add_argument("--output-dir", type=Path, default=None, help="Checkpoint directory (default: ACE/data/checkpoint)") parser.add_argument("--epochs", type=int, default=None) parser.add_argument("--batch-size", type=int, default=None) parser.add_argument("--learning-rate", type=float, default=None) parser.add_argument("--embed-dim", type=int, default=None) parser.add_argument("--num-layers", type=int, default=None) parser.add_argument("--spectral-layers", type=int, default=None) parser.add_argument("--seed", type=int, default=None) parser.add_argument("--device", default="auto") return parser.parse_args() def load_config(path: Path) -> dict: with path.open("r", encoding="utf-8") as handle: if path.suffix.lower() in {".yaml", ".yml"}: return yaml.safe_load(handle) return json.load(handle) def initialize_distributed(requested_device: str, backend: str | None) -> tuple[torch.device, int, int, int]: """Initialize torchrun/Slurm process groups and select the local device.""" world_size = int(os.environ.get("WORLD_SIZE", "1")) local_rank = int(os.environ.get("LOCAL_RANK", "0")) distributed = world_size > 1 use_cuda = torch.cuda.is_available() and requested_device != "cpu" if use_cuda: device_count = torch.cuda.device_count() if distributed: if not 0 <= local_rank < device_count: raise RuntimeError( f"LOCAL_RANK={local_rank} is unavailable; this process sees " f"{device_count} CUDA devices " f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')})" ) device_index = local_rank elif requested_device == "auto" or requested_device == "cuda": device_index = 0 else: device_index = torch.device(requested_device).index if device_index is None: device_index = 0 if not 0 <= device_index < device_count: raise RuntimeError(f"Requested {requested_device}, but only {device_count} CUDA devices are visible") # Set the rank-local device before NCCL initialization. torch.cuda.set_device(device_index) device = torch.device("cuda", device_index) else: device = torch.device("cpu") if distributed: selected_backend = backend or ("nccl" if device.type == "cuda" else "gloo") if selected_backend == "nccl" and device.type != "cuda": raise RuntimeError("NCCL distributed training requires CUDA; use distributed.backend=gloo for CPU") if not dist.is_initialized(): dist.init_process_group(backend=selected_backend, init_method="env://") rank = dist.get_rank() else: rank = 0 return device, rank, local_rank, world_size def reduce_epoch_loss(total: float, count: int, device: torch.device) -> float: """Return the sample-weighted loss across all distributed ranks.""" values = torch.tensor([total, float(count)], dtype=torch.float64, device=device) if dist.is_initialized(): dist.all_reduce(values, op=dist.ReduceOp.SUM) return float((values[0] / values[1].clamp_min(1.0)).item()) def main() -> int: args = parse_args() config = load_config(args.config) data_path = args.data_path or configured_path(config, "data_path", GENERATED_DATA_PATH) output_dir = args.output_dir or configured_path(config, "checkpoint_dir", CHECKPOINT_DIR) train_cfg = config["training"] model_cfg = config["model"] seed = train_cfg["seed"] if args.seed is None else args.seed random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) distributed_cfg = config.get("distributed", {}) device, rank, local_rank, world_size = initialize_distributed( args.device, distributed_cfg.get("backend"), ) print( json.dumps( { "rank": rank, "world_size": world_size, "local_rank": local_rank, "device": str(device), "visible_devices": torch.cuda.device_count(), } ), flush=True, ) if args.fake_data: inputs, targets = make_fake_pairs(args.num_samples, args.height, args.width, seed) dataset = ArrayPairDataset(inputs, targets) data_source = "fake-data (smoke only)" else: if not data_path.exists() and rank == 0: data_cfg = config.get("data", {}) save_fake_pairs( data_path, num_samples=int(data_cfg.get("synthetic_num_samples", args.num_samples)), height=int(data_cfg.get("synthetic_height", args.height)), width=int(data_cfg.get("synthetic_width", args.width)), seed=seed, ) print(json.dumps({"status": "generated_data", "path": str(data_path)}), flush=True) if dist.is_initialized(): dist.barrier() if not data_path.exists(): raise FileNotFoundError(f"training data was not created: {data_path}") dataset = load_npz(data_path) data_source = str(data_path) normalizer = ACEDataNormalizer().fit(dataset.inputs, dataset.targets) dataset = ArrayPairDataset( normalizer.transform_inputs(dataset.inputs), normalizer.transform_targets(dataset.targets), ) nlat, nlon = dataset.inputs.shape[-2:] model_config = ACEModelConfig( nlat=nlat, nlon=nlon, embed_dim=model_cfg["embed_dim"] if args.embed_dim is None else args.embed_dim, num_layers=model_cfg["num_layers"] if args.num_layers is None else args.num_layers, spectral_layers=model_cfg["spectral_layers"] if args.spectral_layers is None else args.spectral_layers, filter_type=model_cfg["filter_type"], operator_type=model_cfg["operator_type"], scale_factor=model_cfg["scale_factor"], grid=model_cfg.get("grid", "legendre-gauss"), grid_internal=model_cfg.get("grid_internal", "legendre-gauss"), mlp_ratio=float(model_cfg.get("mlp_ratio", 2.0)), fallback=False, ) model = ACEModel(model_config).to(device) raw_model = model optimizer = torch.optim.Adam(model.parameters(), lr=train_cfg["learning_rate"] if args.learning_rate is None else args.learning_rate) epochs = train_cfg["epochs"] if args.epochs is None else args.epochs batch_size = train_cfg["batch_size"] if args.batch_size is None else args.batch_size scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(epochs, 1)) ema = EMA(raw_model, train_cfg["ema_decay"]) sampler = DistributedSampler(dataset, shuffle=True) if world_size > 1 else None loader = DataLoader( dataset, batch_size=batch_size, shuffle=sampler is None, sampler=sampler, ) if world_size > 1: model = DistributedDataParallel( model, device_ids=[device.index] if device.type == "cuda" else None, output_device=device.index if device.type == "cuda" else None, # SFNO transform coefficients are static buffers. Broadcasting # them before every forward mutates their version in-place and # can invalidate an autoregressive backward graph. broadcast_buffers=False, ) history = [] for epoch in range(epochs): if sampler is not None: sampler.set_epoch(epoch) model.train() total = 0.0 count = 0 for inputs, targets in loader: inputs, targets = inputs.to(device), targets.to(device) optimizer.zero_grad(set_to_none=True) prediction = model(inputs) loss = torch.mean((prediction - targets) ** 2) loss.backward() optimizer.step() ema.update(raw_model) total += float(loss.detach()) * inputs.shape[0] count += inputs.shape[0] scheduler.step() epoch_loss = reduce_epoch_loss(total, count, device) history_entry = {"epoch": epoch + 1, "loss": epoch_loss, "lr": scheduler.get_last_lr()[0]} if rank == 0: history.append(history_entry) print(json.dumps(history_entry), flush=True) if rank == 0: output_dir.mkdir(parents=True, exist_ok=True) checkpoint = output_dir / "model_bak.pt" parameter_count, nonzero_parameter_count = parameter_statistics(raw_model) torch.save( { "model_config": model_config.to_dict(), "model_state": raw_model.state_dict(), "ema_state": ema.shadow, "normalizer": normalizer.to_dict(), "model_implementation": "spherical_sfno_gauss_legendre", "parameter_count": parameter_count, "nonzero_parameter_count": nonzero_parameter_count, "history": history, "data_source": data_source, "world_size": world_size, "paper_reproduction": "ACE arXiv:2310.02074; torch_harmonics spherical SFNO", }, checkpoint, ) (output_dir / "train_history.json").write_text(json.dumps(history, indent=2), encoding="utf-8") print(json.dumps({"status": "success", "checkpoint": str(checkpoint), "data_source": data_source, "world_size": world_size}), flush=True) else: checkpoint = output_dir / "model_bak.pt" if dist.is_initialized(): dist.barrier() dist.destroy_process_group() return 0 if __name__ == "__main__": raise SystemExit(main())