scripts/train.py

"""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', '<unset>')})"
                )
            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())