scripts/train.py

"""Non-official from-scratch training baseline for the FuXi 2.1 reconstruction."""

from __future__ import annotations

import argparse
import json
import os
import random
from contextlib import nullcontext
from datetime import datetime
from pathlib import Path

import numpy as np
import torch
import torch.distributed as dist
from onescience.datapipes.climate.era5 import ERA5Dataset
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

from common import load_config, resolve_path
from build_static import ensure_static_resources
from model.FuXi21 import FuXi21
from variables import c85_from_config

try:
    from torch.distributed.fsdp import (
        FullStateDictConfig,
        FullyShardedDataParallel,
        ShardingStrategy,
        StateDictType,
    )
except ImportError:  # pragma: no cover - depends on the installed torch build
    FullyShardedDataParallel = None
    FullStateDictConfig = None
    ShardingStrategy = None
    StateDictType = None


def setup_distributed(device: torch.device) -> tuple[bool, int, int, int]:
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    distributed = world_size > 1
    if distributed:
        dist.init_process_group(backend="nccl" if device.type == "cuda" else "gloo")
    rank = dist.get_rank() if distributed else 0
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    return distributed, rank, local_rank, world_size


def select_device(requested: str, local_rank: int) -> torch.device:
    if requested not in {"auto", "cpu", "cuda"}:
        raise ValueError("training.device must be auto, cpu, or cuda")
    use_accelerator = requested == "cuda" or (requested == "auto" and torch.cuda.is_available())
    if use_accelerator:
        if not torch.cuda.is_available():
            raise RuntimeError("training.device=cuda, but no CUDA/HIP device is available")
        if local_rank >= torch.cuda.device_count():
            raise RuntimeError(
                f"LOCAL_RANK={local_rank} exceeds {torch.cuda.device_count()} visible CUDA/HIP device(s); "
                "use one process per visible device or --device cpu for DDP logic testing"
            )
        torch.cuda.set_device(local_rank)
        return torch.device("cuda", local_rank)
    return torch.device("cpu")


def seed_everything(seed: int, rank: int) -> None:
    seed += rank
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def load_array(path_value: str | None, cfg: dict, expected_shape: tuple[int, ...], name: str) -> torch.Tensor:
    if path_value is None:
        raise ValueError(f"model.{name}_file is required outside the smoke profile")
    path = resolve_path(path_value, cfg)
    if not path.is_file():
        raise FileNotFoundError(f"{name} file not found: {path}")
    value = torch.from_numpy(np.load(path)).float()
    if tuple(value.shape) != expected_shape:
        raise ValueError(f"{name} must have shape {expected_shape}, got {tuple(value.shape)}")
    return value


def build_model(cfg: dict) -> FuXi21:
    model_cfg = cfg["model"]
    profile_name = model_cfg["profile"]
    profile = model_cfg["profiles"][profile_name]
    height, width = profile["grid_size"]
    if profile_name == "smoke":
        static_fields = torch.zeros(6, height, width)
        channel_mask = torch.ones(85, height, width)
    else:
        static_fields = load_array(model_cfg["static_fields_file"], cfg, (6, height, width), "static_fields")
        channel_mask = load_array(model_cfg["channel_mask_file"], cfg, (85, height, width), "channel_mask")
    return FuXi21(
        static_fields=static_fields,
        channel_mask=channel_mask,
        activation_checkpointing=cfg["model"].get("activation_checkpointing", False),
        **profile,
    )


def wrap_distributed_model(model: FuXi21, cfg: dict, device: torch.device, local_rank: int):
    strategy = cfg["training"].get("distributed_strategy", "ddp")
    if strategy == "ddp":
        return DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
    if strategy != "fsdp":
        raise ValueError("training.distributed_strategy must be ddp or fsdp")
    if FullyShardedDataParallel is None:
        raise RuntimeError("This PyTorch build does not provide torch.distributed.fsdp")
    if device.type == "cpu":
        raise RuntimeError("FSDP training requires an accelerator device")
    sharding = cfg["training"].get("fsdp_sharding", "full_shard")
    if sharding != "full_shard":
        raise ValueError("Only fsdp_sharding=full_shard is currently supported")
    return FullyShardedDataParallel(
        model,
        device_id=device,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        use_orig_params=True,
    )


def is_fsdp(model) -> bool:
    return FullyShardedDataParallel is not None and isinstance(model, FullyShardedDataParallel)


def unwrap_model(model):
    if isinstance(model, DistributedDataParallel) or is_fsdp(model):
        return model.module
    return model


def build_loader(cfg: dict, split: str, distributed: bool, train: bool):
    data_cfg = cfg["data"]
    channels, _ = c85_from_config(cfg)
    split_cfg = data_cfg["splits"][split]
    dataset_root = resolve_path(cfg["paths"]["data_root"], cfg)
    dataset = ERA5Dataset(
        dataset_dir=str(dataset_root),
        used_years=split_cfg["years"],
        used_variables=channels,
        input_steps=data_cfg["input_steps"],
        output_steps=data_cfg["output_steps"],
        normalize=True,
    )
    sampler = DistributedSampler(dataset, shuffle=train) if distributed else None
    loader = DataLoader(
        dataset,
        batch_size=cfg["training"]["batch_size"],
        shuffle=train and sampler is None,
        sampler=sampler,
        num_workers=data_cfg["num_workers"],
        pin_memory=torch.cuda.is_available(),
        drop_last=train and distributed,
    )
    return loader, sampler


def crop_batch(inputs: torch.Tensor, targets: torch.Tensor, crop_size: list[int] | None):
    if crop_size is None:
        return inputs, targets
    height, width = crop_size
    return inputs[..., :height, :width], targets[..., :height, :width]


def temporal_features(time_index, device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    target_times = time_index[-1]
    if isinstance(target_times, str):
        target_times = [target_times]
    parsed = [datetime.strptime(value, "%Y%m%d%H") for value in target_times]
    step = torch.zeros(len(parsed), device=device)
    hour = torch.tensor([(value.hour * 60 + value.minute) / 1440 for value in parsed], device=device)
    doy = torch.tensor([min(365, value.timetuple().tm_yday) / 365 for value in parsed], device=device)
    return step, hour, doy


def weighted_mse(prediction: torch.Tensor, target: torch.Tensor, channel_weights: torch.Tensor) -> torch.Tensor:
    latitude = torch.linspace(90, -90, prediction.shape[-2], device=prediction.device, dtype=prediction.dtype)
    area = latitude.deg2rad().cos().clamp_min(0)
    area = area / area.mean()
    weights = channel_weights.to(prediction).view(1, -1, 1, 1) * area.view(1, 1, -1, 1)
    return ((prediction - target).square() * weights).mean()


def autocast_context(device: torch.device, precision: str):
    if precision == "fp32":
        return nullcontext()
    if precision != "bf16":
        raise ValueError("training.precision must be fp32 or bf16")
    return torch.autocast(device_type=device.type, dtype=torch.bfloat16)


def run_epoch(model, loader, optimizer, device, cfg, channel_weights, train: bool) -> float:
    model.train(train)
    total_loss = torch.zeros((), device=device)
    total_samples = torch.zeros((), device=device)
    clip_norm = cfg["training"]["gradient_clip_norm"]
    accumulation_steps = max(1, int(cfg["training"].get("gradient_accumulation_steps", 1)))
    if train:
        optimizer.zero_grad(set_to_none=True)
    for batch_index, (inputs, targets, _, _, time_index) in enumerate(loader):
        inputs, targets = crop_batch(inputs, targets, cfg["data"]["crop_size"])
        inputs = inputs.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)
        if targets.ndim == 5:
            targets = targets[:, 0]
        temporal = temporal_features(time_index, device)
        with torch.set_grad_enabled(train), autocast_context(device, cfg["training"]["precision"]):
            prediction = model(inputs, *temporal)[:, -1]
            loss = weighted_mse(prediction, targets, channel_weights)
        if train:
            (loss / accumulation_steps).backward()
            if (batch_index + 1) % accumulation_steps == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm)
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)
        batch_size = inputs.shape[0]
        total_loss += loss.detach() * batch_size
        total_samples += batch_size
    if train and len(loader) % accumulation_steps:
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
    if dist.is_initialized():
        dist.all_reduce(total_loss)
        dist.all_reduce(total_samples)
    return (total_loss / total_samples).item()


def checkpoint_state(model, optimizer, scheduler, epoch: int, cfg: dict) -> dict:
    if is_fsdp(model):
        state_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
        with FullyShardedDataParallel.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_config):
            model_state = model.state_dict()
        optimizer_state = FullyShardedDataParallel.optim_state_dict(model, optimizer)
    else:
        model_state = unwrap_model(model).state_dict()
        optimizer_state = optimizer.state_dict()
    return {
        "format": "fuxi21_reconstructed_checkpoint_v1",
        "protocol": cfg["protocol"],
        "model": model_state,
        "optimizer": optimizer_state,
        "scheduler": scheduler.state_dict(),
        "epoch": epoch,
        "model_profile": cfg["model"]["profile"],
        "model_config": cfg["model"]["profiles"][cfg["model"]["profile"]],
        "distributed_strategy": cfg["training"].get("distributed_strategy", "ddp"),
    }


def load_checkpoint(path: Path, mode: str, model, optimizer, scheduler, device: torch.device) -> int:
    state = torch.load(path, map_location=device, weights_only=True)
    if state.get("format") != "fuxi21_reconstructed_checkpoint_v1":
        raise ValueError("Only checkpoints produced by this reconstruction can be loaded")
    if is_fsdp(model):
        state_config = FullStateDictConfig(offload_to_cpu=False, rank0_only=False)
        with FullyShardedDataParallel.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_config):
            model.load_state_dict(state["model"])
    else:
        unwrap_model(model).load_state_dict(state["model"])
    if mode == "resume":
        if is_fsdp(model):
            optimizer_state = FullyShardedDataParallel.optim_state_dict_to_load(
                model, optimizer, state["optimizer"]
            )
            optimizer.load_state_dict(optimizer_state)
        else:
            optimizer.load_state_dict(state["optimizer"])
        scheduler.load_state_dict(state["scheduler"])
        return int(state["epoch"]) + 1
    if mode == "initialize":
        return 0
    raise ValueError("checkpoint_mode must be scratch, initialize, or resume")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default="conf/config.yaml")
    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default=None)
    parser.add_argument("--dry-run", action="store_true", help="Run one train and validation batch without saving")
    args = parser.parse_args()
    cfg = load_config(args.config)
    if cfg.get("protocol") != "non_official_protocol":
        raise ValueError("Training config must declare protocol: non_official_protocol")

    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    device = select_device(args.device or cfg["training"]["device"], local_rank)
    distributed, rank, local_rank, _ = setup_distributed(device)
    seed_everything(cfg["seed"], rank)
    if cfg["model"]["profile"] == "full":
        if rank == 0:
            ensure_static_resources(cfg)
        if distributed:
            dist.barrier()
    model = build_model(cfg).to(device)
    if distributed:
        model = wrap_distributed_model(model, cfg, device, local_rank)
    train_years = set(cfg["data"]["splits"]["train"]["years"])
    val_years = set(cfg["data"]["splits"]["val"]["years"])
    if train_years & val_years:
        raise ValueError("train and val years must be disjoint")
    train_loader, train_sampler = build_loader(cfg, "train", distributed, True)
    val_loader, _ = build_loader(cfg, "val", distributed, False)

    train_cfg = cfg["training"]
    if train_cfg["optimizer"] != "AdamW" or train_cfg["scheduler"] != "CosineAnnealingLR":
        raise ValueError("This baseline supports optimizer=AdamW and scheduler=CosineAnnealingLR")
    optimizer = torch.optim.AdamW(model.parameters(), lr=train_cfg["learning_rate"], weight_decay=train_cfg["weight_decay"])
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=train_cfg["epochs"], eta_min=train_cfg["min_learning_rate"]
    )
    channel_weights = train_cfg["channel_weights"]
    channel_weights = torch.ones(85) if channel_weights is None else torch.tensor(channel_weights, dtype=torch.float32)
    if channel_weights.shape != (85,) or torch.any(channel_weights <= 0):
        raise ValueError("training.channel_weights must contain 85 positive values")
    channel_weights = channel_weights / channel_weights.mean()

    start_epoch = 0
    checkpoint = train_cfg["load_checkpoint"]
    mode = train_cfg["checkpoint_mode"]
    if checkpoint is not None:
        path = resolve_path(checkpoint, cfg)
        start_epoch = load_checkpoint(path, mode, model, optimizer, scheduler, device)
    elif mode != "scratch":
        raise ValueError(f"checkpoint is required for checkpoint_mode={mode}")

    checkpoint_path = resolve_path(train_cfg["save_checkpoint"], cfg)
    metrics_path = resolve_path(cfg["paths"]["training_metrics"], cfg)
    if rank == 0 and not args.dry_run:
        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
        metrics_path.parent.mkdir(parents=True, exist_ok=True)
    history = []
    end_epoch = min(train_cfg["epochs"], start_epoch + 1) if args.dry_run else train_cfg["epochs"]
    for epoch in range(start_epoch, end_epoch):
        if train_sampler is not None:
            train_sampler.set_epoch(epoch)
        train_loss = run_epoch(model, train_loader, optimizer, device, cfg, channel_weights, True)
        val_loss = run_epoch(model, val_loader, optimizer, device, cfg, channel_weights, False)
        scheduler.step()
        record = {"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss, "learning_rate": scheduler.get_last_lr()[0]}
        history.append(record)
        state = None
        if not args.dry_run:
            # FSDP state-dict collection is collective; every rank must enter it.
            state = checkpoint_state(model, optimizer, scheduler, epoch, cfg)
        if rank == 0:
            print(json.dumps(record))
            if not args.dry_run:
                torch.save(state, checkpoint_path)
                metrics_path.write_text(json.dumps(history, indent=2) + "\n", encoding="utf-8")
    if distributed:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()