scripts/train.py

from __future__ import annotations

import argparse
import os
from pathlib import Path
import sys
import traceback
from typing import Any

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
import yaml


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))
SCRIPTS_DIR = PROJECT_ROOT / "scripts"
if str(SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPTS_DIR))

from model.checkpoint import load_checkpoint
from model.w_mae import w_mae_base
from era5_adapter import WMAEERA5Dataset


def load_config(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        config = yaml.safe_load(handle)
    if not isinstance(config, dict):
        raise ValueError(f"Configuration must contain a mapping: {path}")
    return config


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


def validate_config(config: dict[str, Any]) -> None:
    data = config["data"]
    model = config["model"]
    if int(data["expected_channels"]) != int(model["input_channels"]):
        raise ValueError("data.expected_channels must equal model.input_channels.")
    if tuple(data["model_size"]) != tuple(model["image_size"]):
        raise ValueError("data.model_size must equal model.image_size.")
    if tuple(model["patch_size"])[0] <= 0 or tuple(model["patch_size"])[1] <= 0:
        raise ValueError("model.patch_size values must be positive.")
    if float(model["mask_ratio"]) not in {0.0, 0.75}:
        raise ValueError("W-MAE supports mask_ratio 0.0 or 0.75 only.")


def build_model(config: dict[str, Any]) -> torch.nn.Module:
    model_config = config["model"]
    return w_mae_base(
        img_size=tuple(model_config["image_size"]),
        patch_size=tuple(model_config["patch_size"]),
        in_chans=int(model_config["input_channels"]),
        embed_dim=int(model_config["embed_dim"]),
        depth=int(model_config["encoder_depth"]),
        decoder_embed_dim=int(model_config["decoder_embed_dim"]),
        decoder_depth=int(model_config["decoder_depth"]),
        mlp_ratio=float(model_config.get("mlp_ratio", 4.0)),
        norm_pix_loss=bool(model_config["norm_pixel_loss"]),
    )


def build_dataset(config: dict[str, Any], config_path: Path, split: str) -> WMAEERA5Dataset:
    data = config["data"]
    years = data[f"{split}_years"]
    if not data["variables"]:
        raise ValueError(
            "data.variables is empty. Supply the verified, explicitly ordered 20-channel list before training."
        )
    return WMAEERA5Dataset(
        dataset_dir=resolve_project_path(data["dataset_dir"], config_path),
        years=years,
        variables=data["variables"],
        task="pretrain",
        input_steps=int(data["input_steps"]),
        output_steps=int(data["output_steps"]),
        normalize=bool(data["normalize"]),
    )


def validate_architecture_reference(
    model: torch.nn.Module,
    reference_path: Path,
    config_path: Path,
    source_root: Path | None,
) -> None:
    reference_model = build_model(load_config(config_path))
    load_checkpoint(reference_model, reference_path, source_root=source_root)
    model_state = model.state_dict()
    reference_state = reference_model.state_dict()
    if set(model_state) != set(reference_state):
        raise RuntimeError("Model parameter names do not match the official architecture reference.")
    mismatches = {
        key: (tuple(model_state[key].shape), tuple(reference_state[key].shape))
        for key in model_state
        if tuple(model_state[key].shape) != tuple(reference_state[key].shape)
    }
    if mismatches:
        raise RuntimeError(f"Model parameter shapes do not match the official architecture: {mismatches}")


def save_checkpoint(
    path: Path,
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    epoch: int,
    loss: float,
    config: dict[str, Any],
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict(),
            "epoch": epoch,
            "loss": loss,
            "config": config,
        },
        path,
    )


def resolve_device(requested: str) -> torch.device:
    if requested == "auto":
        requested = "cuda" if torch.cuda.is_available() else "cpu"
    device = torch.device(requested)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("The requested cuda device is unavailable. Use --device cpu or check the DCU runtime.")
    return device


def device_summary(device: torch.device) -> str:
    if device.type != "cuda":
        return f"device={device}"
    backend = f"HIP {torch.version.hip}" if torch.version.hip else f"CUDA {torch.version.cuda}"
    return f"device={device} name={torch.cuda.get_device_name(device)} backend={backend}"


def memory_summary(device: torch.device) -> str:
    if device.type != "cuda":
        return "memory=unavailable(cpu)"
    allocated = torch.cuda.memory_allocated(device) / (1024**3)
    reserved = torch.cuda.memory_reserved(device) / (1024**3)
    peak = torch.cuda.max_memory_allocated(device) / (1024**3)
    return f"memory_allocated={allocated:.2f}GB memory_reserved={reserved:.2f}GB peak_allocated={peak:.2f}GB"


def setup_distributed(requested_device: str) -> tuple[torch.device, int, int, int]:
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    rank = int(os.environ.get("RANK", "0"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if world_size == 1:
        return resolve_device(requested_device), rank, world_size, local_rank

    if requested_device == "cpu":
        device = torch.device("cpu")
        backend = "gloo"
    else:
        if not torch.cuda.is_available():
            raise RuntimeError("torchrun requires one accessible DCU/GPU per local rank.")
        if local_rank >= torch.cuda.device_count():
            raise RuntimeError(
                f"LOCAL_RANK={local_rank} exceeds the {torch.cuda.device_count()} visible accelerator devices."
            )
        torch.cuda.set_device(local_rank)
        device = torch.device("cuda", local_rank)
        backend = "nccl"
    dist.init_process_group(backend=backend, init_method="env://")
    return device, rank, world_size, local_rank


def rank_print(message: str, rank: int) -> None:
    if rank == 0:
        print(message, flush=True)


def all_ranks_finite(value: torch.Tensor, world_size: int) -> bool:
    finite = torch.isfinite(value).to(dtype=torch.int32)
    if world_size > 1:
        dist.all_reduce(finite, op=dist.ReduceOp.MIN)
    return bool(finite.item())


def unwrap_model(model: torch.nn.Module) -> torch.nn.Module:
    return model.module if isinstance(model, DistributedDataParallel) else model


def train_one_epoch(
    model: torch.nn.Module,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    mask_ratio: float,
    device: torch.device,
    epoch: int,
    log_interval: int,
    rank: int,
    world_size: int,
) -> float:
    model.train()
    total_loss = 0.0
    batches = 0
    total_batches = len(loader)
    if device.type == "cuda":
        torch.cuda.reset_peak_memory_stats(device)
    for batch_index, (inputs, _, _, _, _) in enumerate(loader, start=1):
        inputs = inputs.to(device)
        output = model(inputs, mask_ratio=mask_ratio)
        if not all_ranks_finite(output.loss.detach(), world_size):
            raise FloatingPointError("Training loss is NaN or Inf on at least one rank.")
        optimizer.zero_grad(set_to_none=True)
        output.loss.backward()
        optimizer.step()
        total_loss += float(output.loss.detach().cpu())
        batches += 1
        if batch_index == 1 or batch_index % log_interval == 0 or batch_index == total_batches:
            display_loss = output.loss.detach()
            if world_size > 1:
                dist.all_reduce(display_loss, op=dist.ReduceOp.SUM)
                display_loss /= world_size
            rank_print(
                f"epoch={epoch} batch={batch_index}/{total_batches} loss={float(display_loss):.6f} "
                f"{memory_summary(device)}",
                rank,
            )
    if batches == 0:
        raise ValueError("Training DataLoader produced no batches.")
    totals = torch.tensor([total_loss, batches], dtype=torch.float64, device=device)
    if world_size > 1:
        dist.all_reduce(totals, op=dist.ReduceOp.SUM)
    return float((totals[0] / totals[1]).cpu())


def main() -> None:
    parser = argparse.ArgumentParser(description="Train W-MAE using the configured ERA5 adapter.")
    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "conf" / "config.yaml")
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--learning-rate", type=float, default=1e-4)
    parser.add_argument(
        "--device",
        default="auto",
        help="Training device: auto (default), cuda (AMD DCU through HIP), or cpu.",
    )
    parser.add_argument("--checkpoint", type=Path, default=None)
    parser.add_argument("--checkpoint-source-root", type=Path, default=None)
    parser.add_argument("--non-strict-checkpoint", action="store_true")
    parser.add_argument(
        "--output-checkpoint",
        type=Path,
        default=PROJECT_ROOT / "data" / "checkpoint" / "model_bak.pth",
    )
    parser.add_argument("--architecture-reference-checkpoint", type=Path, default=None)
    parser.add_argument("--log-interval", type=int, default=1)
    args = parser.parse_args()
    if args.epochs <= 0:
        raise ValueError("--epochs must be positive.")
    if args.log_interval <= 0:
        raise ValueError("--log-interval must be positive.")

    config_path = args.config.resolve()
    device, rank, world_size, local_rank = setup_distributed(args.device)
    rank_print(
        f"training started: config={config_path} epochs={args.epochs} device={args.device}", rank
    )
    rank_print(
        f"distributed: enabled={world_size > 1} world_size={world_size} "
        f"rank={rank} local_rank={local_rank}",
        rank,
    )
    config = load_config(config_path)
    validate_config(config)
    data_config = config["data"]
    print(f"rank={rank} local_rank={local_rank} runtime: {device_summary(device)}", flush=True)
    rank_print("building model", rank)
    model = build_model(config).to(device)
    print(
        f"rank={rank} local_rank={local_rank} model_parameters_device="
        f"{next(model.parameters()).device}",
        flush=True,
    )

    if args.checkpoint is not None:
        checkpoint_path = resolve_project_path(args.checkpoint, config_path)
        source_root = (
            resolve_project_path(args.checkpoint_source_root, config_path)
            if args.checkpoint_source_root
            else None
        )
        report = load_checkpoint(
            model,
            checkpoint_path,
            strict=not args.non_strict_checkpoint,
            map_location=device,
            source_root=source_root,
        )
        rank_print(f"loaded checkpoint: {report}", rank)
    else:
        rank_print("training from randomly initialized weights", rank)

    if world_size > 1:
        if device.type == "cuda":
            model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
        else:
            model = DistributedDataParallel(model)
    rank_print("building dataset", rank)
    dataset = build_dataset(config, config_path, "train")
    sampler = DistributedSampler(dataset, shuffle=True) if world_size > 1 else None
    loader = DataLoader(
        dataset,
        batch_size=int(data_config["batch_size"]),
        shuffle=sampler is None,
        sampler=sampler,
        num_workers=int(data_config["num_workers"]),
    )
    rank_print(
        f"dataset ready: samples={len(dataset)} batches_per_rank={len(loader)} "
        f"batch_size_per_rank={data_config['batch_size']} "
        f"global_batch_size={int(data_config['batch_size']) * world_size}",
        rank,
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate)
    for epoch in range(args.epochs):
        if sampler is not None:
            sampler.set_epoch(epoch)
        rank_print(f"epoch={epoch + 1}/{args.epochs} started", rank)
        loss = train_one_epoch(
            model,
            loader,
            optimizer,
            float(config["model"]["mask_ratio"]),
            device,
            epoch + 1,
            args.log_interval,
            rank,
            world_size,
        )
        rank_print(f"epoch={epoch + 1}/{args.epochs} finished loss={loss:.6f}", rank)

    if world_size > 1:
        dist.barrier()

    if rank != 0:
        return

    checkpoint_model = unwrap_model(model)

    reference_path = (
        resolve_project_path(args.architecture_reference_checkpoint, config_path)
        if args.architecture_reference_checkpoint
        else None
    )
    reference_source_root = (
        resolve_project_path(args.checkpoint_source_root, config_path)
        if args.checkpoint_source_root
        else None
    )
    if reference_path is not None:
        print(f"checking architecture against {reference_path}", flush=True)
        validate_architecture_reference(checkpoint_model, reference_path, config_path, reference_source_root)
        print("architecture reference check passed", flush=True)
    else:
        print("architecture reference check skipped: no reference checkpoint supplied", flush=True)
    output_path = resolve_project_path(args.output_checkpoint, config_path)
    save_checkpoint(output_path, checkpoint_model, optimizer, args.epochs, loss, config)
    print(f"checkpoint saved: {output_path}", flush=True)
    print("training completed successfully", flush=True)


if __name__ == "__main__":
    try:
        main()
    except Exception:
        traceback.print_exc()
        raise
    finally:
        if dist.is_available() and dist.is_initialized():
            dist.destroy_process_group()