scripts/train.py

"""Single-device and torchrun-based distributed training for Samudra v1."""

try:
    from ._bootstrap import ROOT
except ImportError:
    from _bootstrap import ROOT

import argparse
import os
import random
from pathlib import Path

import numpy as np
import torch
import yaml
from torch import nn
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, Dataset, DistributedSampler

from model.samudra import build_model

STATE_CHANNELS = 77
BOUNDARY_CHANNELS = 4


def load_data(path: str | Path) -> tuple[np.ndarray, np.ndarray]:
    """Load and validate native Samudra time-major arrays."""
    with np.load(path) as data:
        prognostic = np.asarray(data["prognostic"], dtype=np.float32)
        boundary = np.asarray(data["boundary"], dtype=np.float32)
    if prognostic.ndim != 4 or prognostic.shape[1] != STATE_CHANNELS:
        raise ValueError("prognostic must have shape [time, 77, lat, lon]")
    if boundary.ndim != 4 or boundary.shape[1] != BOUNDARY_CHANNELS:
        raise ValueError("boundary must have shape [time, 4, lat, lon]")
    if prognostic.shape[0] != boundary.shape[0] or prognostic.shape[2:] != boundary.shape[2:]:
        raise ValueError("prognostic and boundary time/grid dimensions must match")
    return prognostic, boundary


class SamudraDataset(Dataset):
    """Build recurrent training windows from native Samudra arrays."""

    def __init__(self, path: str | Path, recurrent_passes: int):
        self.prognostic, self.boundary = load_data(path)
        self.recurrent_passes = recurrent_passes
        self.end = self.prognostic.shape[0] - 2 * recurrent_passes
        if self.end <= 1:
            raise ValueError("dataset does not contain enough samples for recurrent training")

    def __len__(self) -> int:
        return self.end - 1

    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        t = index + 1
        history = np.stack((self.prognostic[t - 1], self.prognostic[t]))
        forcing = self.boundary[t : t + self.recurrent_passes]
        labels = np.stack(
            [
                np.concatenate(
                    (self.prognostic[t + 2 * step + 1], self.prognostic[t + 2 * step + 2])
                )
                for step in range(self.recurrent_passes)
            ]
        )
        return torch.from_numpy(history), torch.from_numpy(forcing), torch.from_numpy(labels)


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


def load_config(path: str) -> dict:
    with open(path, encoding="utf-8") as handle:
        return yaml.safe_load(handle)


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


def save_checkpoint(path: Path, model: nn.Module, optimizer: torch.optim.Optimizer, scheduler, epoch: int, loss: float) -> None:
    state = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict()
    torch.save({"epoch": epoch, "loss": loss, "model": state, "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict()}, path)


def freeze_batch_norm_stats(module: nn.Module) -> None:
    """Prevent recurrent forwards from mutating BatchNorm buffers in one graph."""
    for child in module.modules():
        if isinstance(child, nn.modules.batchnorm._BatchNorm):
            child.eval()


def train(
    config_path: str,
    data_path: str,
    device_name: str | None = None,
    epochs_override: int | None = None,
    output_dir_override: str | None = None,
) -> None:
    config = load_config(config_path)
    rank = int(os.environ.get("RANK", "0"))
    local_rank = int(os.environ.get("LOCAL_RANK", rank))
    if device_name:
        device = torch.device(device_name)
    elif torch.cuda.is_available():
        device = torch.device(f"cuda:{local_rank}")
    else:
        device = torch.device("cpu")
    rank, world_size, distributed = distributed_setup(device)
    set_seed(int(config["project"].get("seed", 1)) + rank)
    if device.type == "cuda":
        torch.cuda.set_device(device)

    recurrent_passes = int(config["data"].get("recurrent_passes", 1))
    dataset = SamudraDataset(data_path, recurrent_passes)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True) if distributed else None
    loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], shuffle=sampler is None, sampler=sampler, num_workers=config["training"].get("num_workers", 0), pin_memory=device.type == "cuda")
    model = build_model(config).to(device)
    if distributed:
        model = DistributedDataParallel(
            model,
            device_ids=[device.index] if device.type == "cuda" else None,
            broadcast_buffers=False,
        )
    optimizer = torch.optim.Adam(model.parameters(), lr=config["training"]["learning_rate"], weight_decay=config["training"].get("weight_decay", 0.0))
    epochs = epochs_override if epochs_override is not None else int(config["training"]["epochs"])
    if epochs < 1:
        raise ValueError("epochs must be positive")
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    resume = config["training"].get("resume_checkpoint")
    start_epoch = 0
    if resume:
        try:
            checkpoint = torch.load(resume, map_location=device, weights_only=True)
        except TypeError:
            checkpoint = torch.load(resume, map_location=device)
        target = model.module if isinstance(model, DistributedDataParallel) else model
        state = {key: value for key, value in checkpoint["model"].items() if not key.endswith(".cap")}
        target.load_state_dict(state)
        optimizer.load_state_dict(checkpoint["optimizer"])
        scheduler.load_state_dict(checkpoint["scheduler"])
        start_epoch = int(checkpoint["epoch"]) + 1

    output_dir = Path(output_dir_override or config["training"]["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)
    for epoch in range(start_epoch, epochs):
        if sampler is not None:
            sampler.set_epoch(epoch)
        model.train()
        freeze_batch_norm_stats(model)
        total_loss = 0.0
        for batch in loader:
            optimizer.zero_grad(set_to_none=True)
            history, forcing, labels = (item.to(device, non_blocking=True) for item in batch)
            previous, current = history[:, 0], history[:, 1]
            losses = []
            for step in range(recurrent_passes):
                prediction = model(torch.cat((previous, current, forcing[:, step]), dim=1))
                losses.append(nn.functional.mse_loss(prediction, labels[:, step]))
                previous, current = prediction[:, :STATE_CHANNELS], prediction[:, STATE_CHANNELS:]
            loss = torch.stack(losses).mean()
            loss.backward()
            optimizer.step()
            total_loss += float(loss.detach())
        scheduler.step()
        mean_loss = total_loss / max(1, len(loader))
        if rank == 0:
            print(f"epoch={epoch + 1} loss={mean_loss:.6e} lr={scheduler.get_last_lr()[0]:.6e}")
            frequency = int(config["training"].get("save_frequency", 5))
            if (epoch + 1) % frequency == 0 or epoch + 1 == epochs:
                save_checkpoint(output_dir / f"epoch_{epoch + 1:04d}.pt", model, optimizer, scheduler, epoch, mean_loss)
                save_checkpoint(output_dir / "model_bak.pth", model, optimizer, scheduler, epoch, mean_loss)
    if distributed:
        torch.distributed.destroy_process_group()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="./conf/config.yaml")
    parser.add_argument("--data", default="./data/train.npz")
    parser.add_argument("--device", default=None)
    parser.add_argument("--epochs", type=int, default=None)
    parser.add_argument("--output-dir", default=None)
    args = parser.parse_args()
    train(args.config, args.data, args.device, args.epochs, args.output_dir)


if __name__ == "__main__":
    main()