scripts/train.py

import logging
import time
from pathlib import Path

import hydra
import torch
from omegaconf import OmegaConf
from torch.amp import GradScaler, autocast
from torch.nn.parallel import DistributedDataParallel

from common import load_config
from onescience.datapipes.cfd import DeepMindLagrangianDatapipe
from onescience.distributed.manager import DistributedManager
from onescience.launch.utils import load_checkpoint, save_checkpoint


def setup_logging(rank: int):
    level = logging.INFO if rank == 0 else logging.WARNING
    logging.basicConfig(level=level, format="%(asctime)s - %(levelname)s - %(message)s")
    return logging.getLogger("lagrangian_mgn.train")


def resolve_device(requested: str, manager: DistributedManager):
    if manager.distributed:
        return manager.device
    if requested == "cpu":
        return torch.device("cpu")
    if requested in ("cuda", "gpu"):
        if not torch.cuda.is_available():
            raise RuntimeError("Config requested CUDA, but torch.cuda.is_available() is false.")
        return torch.device("cuda:0")
    return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def train_step(model, graph, dataset, criterion, optimizer, scaler, scheduler, device, amp_enabled):
    graph = graph.to(device)
    optimizer.zero_grad(set_to_none=True)

    with autocast(device_type=device.type, enabled=amp_enabled):
        gt_pos, gt_vel, gt_acc = dataset.unpack_targets(graph)
        pred_acc = model(graph.ndata["x"], graph.edata["x"], graph)

        mask = graph.ndata["mask"].unsqueeze(-1).to(pred_acc.dtype)
        num_nz = torch.clamp(mask.sum() * dataset.dim, min=1.0)
        loss_acc_norm = (mask * criterion(pred_acc, gt_acc)).sum() / num_nz

        with torch.no_grad():
            pos, vel, _ = dataset.unpack_inputs(graph)
            pred_pos, pred_vel = dataset.time_integrator(
                position=pos,
                velocity=vel[-1],
                acceleration=pred_acc,
                dt=dataset.dt,
                denormalize=True,
            )
            loss_pos = (mask * criterion(pred_pos, gt_pos)).sum() / num_nz
            loss_vel = (
                mask * criterion(pred_vel, dataset.denormalize_velocity(gt_vel))
            ).sum() / num_nz
            loss_acc = (
                mask
                * criterion(
                    dataset.denormalize_acceleration(pred_acc),
                    dataset.denormalize_acceleration(gt_acc),
                )
            ).sum() / num_nz

    scaler.scale(loss_acc_norm).backward()
    scaler.step(optimizer)
    scaler.update()
    scheduler.step()

    return {
        "loss": loss_acc_norm.item() + loss_pos.item() + loss_vel.item(),
        "loss_acc_norm": loss_acc_norm.item(),
        "loss_pos": loss_pos.item(),
        "loss_vel": loss_vel.item(),
        "loss_acc": loss_acc.item(),
    }


def main():
    cfg = load_config()
    Path(cfg.output).mkdir(parents=True, exist_ok=True)
    Path(cfg.resume_dir).mkdir(parents=True, exist_ok=True)

    DistributedManager.initialize()
    manager = DistributedManager()
    logger = setup_logging(manager.rank)

    if manager.rank == 0:
        logger.info("Config:\n%s", OmegaConf.to_yaml(cfg))

    datapipe = DeepMindLagrangianDatapipe(cfg, distributed=manager.distributed)
    train_loader = datapipe.train_dataloader()
    val_loader = datapipe.val_dataloader()
    dataset = datapipe.train_dataset

    if cfg.model.recompute_activation and str(cfg.model.mlp_activation_fn).lower() != "silu":
        raise ValueError("recompute_activation only supports SiLU.")

    device = resolve_device(str(cfg.train.device), manager)
    model = hydra.utils.instantiate(cfg.model).to(device)
    if cfg.compile.enabled:
        model = torch.compile(model, **cfg.compile.args)
    if manager.distributed:
        model = DistributedDataParallel(
            model,
            device_ids=[manager.local_rank],
            output_device=manager.local_rank,
            find_unused_parameters=False,
        )

    criterion = hydra.utils.instantiate(cfg.loss)
    optimizer = hydra.utils.instantiate(cfg.optimizer, model.parameters())

    num_iterations = max(int(cfg.train.epochs) * len(train_loader), 1)
    if cfg.lr_scheduler._target_ == "torch.optim.lr_scheduler.CosineAnnealingLR":
        if cfg.lr_scheduler.T_max is None:
            cfg.lr_scheduler.T_max = num_iterations
    elif cfg.lr_scheduler._target_ == "torch.optim.lr_scheduler.OneCycleLR":
        if cfg.lr_scheduler.total_steps is None:
            cfg.lr_scheduler.total_steps = num_iterations
    scheduler = hydra.utils.instantiate(cfg.lr_scheduler, optimizer)

    amp_enabled = bool(cfg.amp.enabled and device.type == "cuda")
    scaler = GradScaler(enabled=amp_enabled)

    epoch_init = load_checkpoint(
        cfg.resume_dir,
        models=model,
        optimizer=optimizer,
        scheduler=scheduler,
        scaler=scaler,
        device=device,
    )

    best_valid_loss = float("inf")
    best_loss_epoch = epoch_init
    logger.info("Starting training on %s", device)

    for epoch in range(epoch_init, int(cfg.train.epochs)):
        if manager.distributed and hasattr(train_loader.sampler, "set_epoch"):
            train_loader.sampler.set_epoch(epoch)

        start = time.time()
        model.train()
        epoch_losses = {}

        for step, graph in enumerate(train_loader, start=1):
            losses = train_step(
                model,
                graph,
                dataset,
                criterion,
                optimizer,
                scaler,
                scheduler,
                device,
                amp_enabled,
            )
            for key, value in losses.items():
                epoch_losses.setdefault(key, []).append(value)
            if manager.rank == 0 and step % int(cfg.train.log_interval) == 0:
                logger.info(
                    "Epoch %s/%s step %s/%s loss %.4e acc_norm %.4e",
                    epoch + 1,
                    cfg.train.epochs,
                    step,
                    len(train_loader),
                    losses["loss"],
                    losses["loss_acc_norm"],
                )

        model.eval()
        valid_loss = 0.0
        with torch.no_grad():
            for graph in val_loader:
                graph = graph.to(device)
                gt_pos, _, gt_acc = datapipe.valid_dataset.unpack_targets(graph)
                pred_acc = model(graph.ndata["x"], graph.edata["x"], graph)
                mask = graph.ndata["mask"].unsqueeze(-1).to(pred_acc.dtype)
                num_nz = torch.clamp(mask.sum() * datapipe.valid_dataset.dim, min=1.0)
                loss_acc = (mask * criterion(pred_acc, gt_acc)).sum() / num_nz
                pos, vel, _ = datapipe.valid_dataset.unpack_inputs(graph)
                pred_pos, _ = datapipe.valid_dataset.time_integrator(
                    position=pos,
                    velocity=vel[-1],
                    acceleration=pred_acc,
                    dt=datapipe.valid_dataset.dt,
                    denormalize=True,
                )
                loss_pos = (mask * criterion(pred_pos, gt_pos)).sum() / num_nz
                valid_loss += (loss_acc + loss_pos).item()

        valid_loss /= max(len(val_loader), 1)

        if manager.rank == 0:
            mean_losses = {k: sum(v) / len(v) for k, v in epoch_losses.items()}
            logger.info(
                "Epoch %s finished in %.2fs train_loss %.4e valid_loss %.4e lr %.2e",
                epoch + 1,
                time.time() - start,
                mean_losses.get("loss", 0.0),
                valid_loss,
                scheduler.get_last_lr()[0],
            )
            if valid_loss < best_valid_loss:
                best_valid_loss = valid_loss
                best_loss_epoch = epoch
                save_checkpoint(
                    cfg.resume_dir,
                    models=model,
                    optimizer=optimizer,
                    scheduler=scheduler,
                    scaler=scaler,
                    epoch=epoch + 1,
                )
                logger.info("Checkpoint saved to %s", cfg.resume_dir)
            if epoch - best_loss_epoch >= int(cfg.train.patience):
                logger.warning("Early stopping after %s stale epochs", cfg.train.patience)
                break

        if manager.distributed:
            torch.distributed.barrier()

    manager.cleanup()
    logger.info("Training finished")


if __name__ == "__main__":
    main()