scripts/train.py

"""OneForecast training entry point with integrated data checking."""

from __future__ import annotations

import argparse
import os
from pathlib import Path
import sys
import random

import numpy as np
import torch
import torch.distributed as dist
from torch.nn import functional as F
from torch.autograd import Function
from torch.nn.parallel import DistributedDataParallel

import yaml

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from model.era5_adapter import OFFICIAL_VARIABLES, OneForecastERA5Adapter
from model.oneforecast import build_model, check_checkpoint_compatibility


def _resolve_path(value: str | Path, config_path: Path) -> Path:
    path = Path(value).expanduser()
    return path if path.is_absolute() else (config_path.parent.parent / path).resolve()


def _load_config(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as handle:
        config = yaml.safe_load(handle)
    config["datapipe"]["dataset_dir"] = str(_resolve_path(config["datapipe"]["dataset_dir"], path))
    config["model"]["official_checkpoint_path"] = str(
        _resolve_path(config["model"]["official_checkpoint_path"], path)
    )
    config["model"]["checkpoint_path"] = config["model"]["official_checkpoint_path"]
    config["training"]["checkpoint_dir"] = str(_resolve_path(config["training"]["checkpoint_dir"], path))
    return config


def _set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def _resolve_device(name: str) -> torch.device:
    """Map the logical DCU name to the backend exposed by this PyTorch build."""
    requested = str(name).lower()
    if requested == "dcu":
        if torch.cuda.is_available():
            return torch.device("cuda")
        privateuse = torch._C._get_privateuse1_backend_name()
        if privateuse != "privateuseone":
            return torch.device(privateuse)
        raise RuntimeError("runtime.device=dcu, but this PyTorch build exposes no usable accelerator")
    if requested == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    device = torch.device(requested)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("runtime.device=cuda, but torch.cuda.is_available() is False")
    return device


def _setup_distributed(device_name: str, backend: str = "nccl") -> tuple[torch.device, int, int, bool]:
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    distributed = world_size > 1
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if distributed:
        device = _resolve_device(device_name)
        if device.type == "cuda":
            torch.cuda.set_device(local_rank)
            device = torch.device("cuda", local_rank)
        dist.init_process_group(backend=backend, init_method="env://")
        return device, dist.get_rank(), world_size, True
    return _resolve_device(device_name), 0, 1, False


def _reduce_metrics(total: float, count: int, device: torch.device, distributed: bool) -> float:
    metrics = torch.tensor([total, count], dtype=torch.float64, device=device)
    if distributed:
        dist.all_reduce(metrics, op=dist.ReduceOp.SUM)
    return float(metrics[0] / metrics[1].clamp_min(1))


def _loader_batch(batch: tuple) -> tuple[torch.Tensor, torch.Tensor]:
    inputs, targets = batch[0], batch[1]
    if inputs.ndim == 5 or targets.ndim == 5:
        raise ValueError("OneForecast currently supports input_steps=1 and output_steps=1 only")
    if inputs.ndim != 4 or targets.ndim != 4:
        raise ValueError(f"Expected batched fields with four dimensions, got {inputs.shape} and {targets.shape}")
    if inputs.shape[-2] == 121:
        inputs = inputs[..., :120, :]
    if targets.shape[-2] == 121:
        targets = targets[..., :120, :]
    if inputs.shape[-2:] != (120, 240) or targets.shape[-2:] != (120, 240):
        raise ValueError(f"Expected official model grid 120x240, got {inputs.shape} and {targets.shape}")
    return torch.nan_to_num(inputs.float()), torch.nan_to_num(targets.float())


class _LossScaleFunction(Function):
    @staticmethod
    def forward(ctx, values: torch.Tensor, eps: float) -> torch.Tensor:
        ctx.eps = eps
        return values

    @staticmethod
    def backward(ctx, gradients: torch.Tensor) -> tuple[torch.Tensor, None]:
        channels = gradients.shape[1]
        weights = 1.0 / gradients.norm(p=2, dim=(-1, -2), keepdim=True).clamp_min(ctx.eps)
        weights = weights / weights.sum(dim=1, keepdim=True).clamp_min(ctx.eps)
        return channels * weights * gradients, None


def _relative_channel_l2(prediction: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    difference = (prediction - target).flatten(2).norm(p=2, dim=2)
    target_norm = target.flatten(2).norm(p=2, dim=2).clamp_min(1e-10)
    channel_loss = (difference / target_norm).mean(dim=0)
    return channel_loss.mean(), channel_loss


def check_data(config: dict) -> dict:
    settings = config["datapipe"]
    adapter = OneForecastERA5Adapter(
        settings["dataset_dir"], settings["train_years"],
        batch_size=settings["batch_size"], input_steps=settings["input_steps"],
        output_steps=settings["output_steps"], normalize=settings["normalize"],
        num_workers=settings["num_workers"],
    )
    report = adapter.inspect()
    print(report)
    return report


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, default=Path("conf/config.yaml"))
    parser.add_argument("--check-data", action="store_true")
    parser.add_argument("--check-model", action="store_true")
    parser.add_argument("--check-checkpoint", action="store_true")
    parser.add_argument("--check-distributed", action="store_true")
    parser.add_argument("--device", default=None)
    parser.add_argument("--distributed-backend", default=None)
    parser.add_argument("--max-epochs", type=int, default=None)
    parser.add_argument("--max-batches", type=int, default=None)
    parser.add_argument("--weight-init", choices=("scratch", "official"), default=None)
    args = parser.parse_args()
    config = _load_config(args.config.resolve())
    if args.device is not None:
        config["runtime"]["device"] = args.device
    if args.distributed_backend is not None:
        config["runtime"]["distributed_backend"] = args.distributed_backend
    if args.max_epochs is not None:
        config["training"]["max_epoch"] = args.max_epochs
    if args.max_batches is not None:
        config["training"]["max_batches"] = args.max_batches
    if tuple(config["datapipe"]["variables"]) != OFFICIAL_VARIABLES:
        raise ValueError("datapipe.variables must exactly match the official 69-channel order")
    if args.weight_init is not None:
        config["model"]["weight_init"] = args.weight_init
    if config["model"].get("weight_init") == "official":
        config["model"]["checkpoint_path"] = config["model"]["official_checkpoint_path"]
    if args.check_data:
        check_data(config)
        return
    if args.check_model:
        configured_init = config["model"].get("weight_init", "scratch")
        config["model"]["weight_init"] = "scratch"
        with __import__("torch").device("meta"):
            model = build_model(config, build_graph=False)
        print({"model": type(model).__name__, "parameters": sum(p.numel() for p in model.parameters()),
               "configured_weight_init": configured_init})
        return
    if args.check_checkpoint:
        with __import__("torch").device("meta"):
            model = build_model(config, build_graph=False)
        report = check_checkpoint_compatibility(
            model, config["model"]["official_checkpoint_path"]
        )
        print(report)
        if not report.compatible:
            raise SystemExit(1)
        return
    if args.check_distributed:
        device, rank, world_size, distributed = _setup_distributed(
            config["runtime"].get("device", "cpu"), config["runtime"].get("distributed_backend", "nccl")
        )
        settings = config["datapipe"]
        adapter = OneForecastERA5Adapter(
            settings["dataset_dir"], settings["train_years"], batch_size=settings["batch_size"],
            input_steps=settings["input_steps"], output_steps=settings["output_steps"],
            normalize=settings["normalize"], num_workers=0, distributed=distributed,
        )
        loader, sampler = adapter.get_dataloader("train")
        sample_indices = list(iter(sampler)) if sampler is not None else list(range(len(loader.dataset)))
        print({"rank": rank, "world_size": world_size, "distributed": distributed,
               "backend": dist.get_backend() if distributed else None, "device": str(device),
               "sampler": type(sampler).__name__ if sampler is not None else None,
               "sample_indices": sample_indices})
        if distributed:
            dist.barrier()
            dist.destroy_process_group()
        return
    settings = config["datapipe"]
    if settings["input_steps"] != 1 or settings["output_steps"] != 1:
        raise SystemExit("OneForecast training currently requires input_steps=1 and output_steps=1")
    device, rank, world_size, distributed = _setup_distributed(
        config["runtime"].get("device", "cpu"), config["runtime"].get("distributed_backend", "nccl")
    )
    _set_seed(int(config["runtime"].get("seed", 42)))
    model = build_model(config).to(device)
    if distributed:
        ddp_devices = {"device_ids": [device.index], "output_device": device.index} if device.type == "cuda" else {}
        model = DistributedDataParallel(model, broadcast_buffers=False, **ddp_devices)
    optimizer = torch.optim.Adam(
        model.parameters(), lr=float(config["training"]["learning_rate"]),
    )
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=max(1, int(config["training"]["max_epoch"])),
    )
    train_adapter = OneForecastERA5Adapter(
        _resolve_path(settings["dataset_dir"], args.config), settings["train_years"],
        batch_size=settings["batch_size"], input_steps=1, output_steps=1,
        normalize=settings["normalize"], num_workers=settings["num_workers"], distributed=distributed,
    )
    valid_adapter = OneForecastERA5Adapter(
        _resolve_path(settings["dataset_dir"], args.config), settings["valid_years"],
        batch_size=settings["batch_size"], input_steps=1, output_steps=1,
        normalize=settings["normalize"], num_workers=settings["num_workers"], distributed=distributed,
    )
    train_loader, train_sampler = train_adapter.get_dataloader("train")
    valid_loader, valid_sampler = valid_adapter.get_dataloader("val")
    checkpoint_dir = Path(config["training"]["checkpoint_dir"])
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    max_batches = int(config["training"].get("max_batches", -1))
    for epoch in range(int(config["training"]["start_epoch"]), int(config["training"]["max_epoch"])):
        if train_sampler is not None:
            train_sampler.set_epoch(epoch)
        if valid_sampler is not None:
            valid_sampler.set_epoch(epoch)
        model.train()
        train_loss = 0.0
        train_batches = 0
        for batch in train_loader:
            inputs, targets = _loader_batch(batch)
            optimizer.zero_grad(set_to_none=True)
            prediction = _LossScaleFunction.apply(model(inputs.to(device)), 1e-5)
            loss, _ = _relative_channel_l2(prediction, targets.to(device))
            loss.backward()
            optimizer.step()
            train_loss += float(loss.detach())
            train_batches += 1
            if max_batches >= 0 and train_batches >= max_batches:
                break
        model.eval()
        valid_loss = 0.0
        with torch.no_grad():
            valid_batches = 0
            for batch in valid_loader:
                inputs, targets = _loader_batch(batch)
                prediction = model(inputs.to(device))
                valid_loss += float(F.mse_loss(prediction, targets.to(device)))
                valid_batches += 1
                if max_batches >= 0 and valid_batches >= max_batches:
                    break
        train_mean = _reduce_metrics(train_loss, train_batches, device, distributed)
        valid_mean = _reduce_metrics(valid_loss, valid_batches, device, distributed)
        if rank == 0:
            print({"epoch": epoch + 1, "train_loss": train_mean, "valid_loss": valid_mean,
                   "world_size": world_size})
        if rank == 0 and (epoch + 1) % int(config["training"].get("save_every_epoch", 1)) == 0:
            model_name = config["training"].get("model_name", "model_bak")
            state = model.module.state_dict() if distributed else model.state_dict()
            torch.save({"model_state": state, "epoch": epoch + 1, "world_size": world_size},
                       checkpoint_dir / f"{model_name}.tar")
        scheduler.step()
    if distributed:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()