giant/giant/train.py

import copy
import csv
import math
import os
import signal
import time
from pathlib import Path
from types import FrameType
from typing import Callable

import numpy as np
import torch
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm

from giant.constants import K_MAX, SEC_SLOT_DIM
from giant.model.schedule import (
    CosineSchedule,
    flow_matching_loss,
    flow_matching_loss_secondary,
)
from giant.model.wgan import gradient_penalty, generator_loss
from giant.validate import validate_marginals

_METRICS_FIELDS = [
    "epoch",
    "train_loss",
    "train_loss_s1",
    "train_loss_nsec",
    "train_loss_s2",
    "train_loss_balance",
    "train_loss_proc",
    "train_loss_entropy",
    "train_nsec_acc",
    "d_loss",
    "g_loss",
    "wasserstein_estimate",
    "gp_loss",
    "val_loss",
    "val_loss_s1",
    "val_loss_nsec",
    "val_loss_s2",
    "val_loss_balance",
    "val_loss_proc",
    "val_loss_entropy",
    "val_nsec_acc",
    "val_marginal_kl",
    "router_s1_entropy",
    "router_s1_util_min",
    "router_s1_util_max",
    "router_s1_util_std",
    "router_s2_entropy",
    "router_s2_util_min",
    "router_s2_util_max",
    "router_s2_util_std",
    "lr",
    "critic_lr",
    "grad_norm",
    "grad_norm_d",
    "grad_norm_g",
    "gpu_mem_mb",
    "samples_per_sec",
    "is_best",
    "epoch_time_s",
]

_CATCHABLE_SIGNALS = (signal.SIGINT, signal.SIGTERM)


class _GracefulShutdown:
    """Turns SIGINT/SIGTERM into a flag check instead of an immediate crash.

    A second signal while already shutting down restores the default
    handler and re-sends the signal, so an unresponsive run can still be
    force-killed.
    """

    def __init__(self) -> None:
        self.requested = False
        self._previous: dict[
            int,
            Callable[[int, FrameType | None], object] | signal.Handlers | int | None,
        ] = {}

    def __enter__(self) -> "_GracefulShutdown":
        for sig in _CATCHABLE_SIGNALS:
            self._previous[sig] = signal.getsignal(sig)
            signal.signal(sig, self._handle)
        return self

    def __exit__(self, *exc_info) -> None:
        for sig, handler in self._previous.items():
            signal.signal(sig, handler)

    def _handle(self, signum: int, frame) -> None:
        if self.requested:
            signal.signal(signum, self._previous[signum])
            os.kill(os.getpid(), signum)
            return
        self.requested = True
        print(
            f"\nreceived {signal.Signals(signum).name} — finishing the current "
            "batch, then saving a checkpoint and exiting (send again to force-quit)"
        )


@torch.no_grad()
def _update_ema(
    ema_model: torch.nn.Module, model: torch.nn.Module, decay: float
) -> None:
    for ema_p, p in zip(ema_model.parameters(), model.parameters()):
        ema_p.mul_(decay).add_(p, alpha=1 - decay)


def _compute_losses(
    stage1_model: torch.nn.Module,
    sec_decoder: torch.nn.Module,
    batch: tuple,
    mode: str,
    ddpm_schedule,
    device: torch.device,
    lambda_nsec: float,
    lambda_s2: float,
    lambda_balance: float = 0.0,
    lambda_proc: float = 0.0,
    lambda_entropy: float = 0.0,
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
]:
    """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc, L_entropy, nsec_acc) for one batch."""
    cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx = batch
    cond_cont = cond_cont.to(device)
    cond_cat = cond_cat.to(device)
    x1_s1 = x1_s1.to(device)
    n_sec = n_sec.to(device)
    sec_cont = sec_cont.to(device)
    proc_idx = proc_idx.to(device)

    # Stage-1 flow loss
    if mode == "flow":
        l_s1 = flow_matching_loss(stage1_model, x1_s1, cond_cont, cond_cat)
    else:
        assert ddpm_schedule is not None
        l_s1 = ddpm_schedule.loss(stage1_model, x1_s1, cond_cont, cond_cat)

    # n_sec classification loss
    n_sec_logits = stage1_model.predict_n_sec(cond_cont, cond_cat)
    l_nsec = F.cross_entropy(n_sec_logits, n_sec)
    nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean()

    # Stage-2 secondary flow loss
    # Use a noiseless Stage-1 target as context (detach to avoid back-prop
    # coupling between the two flow paths). sec_cont's log_mass/charge
    # columns are already a fixed physics-derived regression target (see
    # giant.data.transforms.encode_secondaries) rather than a learned/moving
    # one, so — unlike the embedding-table target this replaced — no
    # detaching is needed to keep the target from chasing the decoder.
    x1_s2 = sec_cont.flatten(1)  # (B, SEC_DIM)

    sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
    l_s2 = flow_matching_loss_secondary(
        sec_decoder,
        x1_s2,
        cond_cont,
        cond_cat,
        x1_s1.detach(),
        sec_mask,
    )

    # Optional MoE load-balance auxiliary loss: only present when both stages
    # are routed (RoutedDenoisingMLP/RoutedSecondaryDecoder carry `.router`,
    # the monolith models don't), computed on cond_cont alone (cheap — no
    # trunk compute) so it's reported even when lambda_balance == 0.
    if hasattr(stage1_model, "router") and hasattr(sec_decoder, "router"):
        l_balance = stage1_model.router.balance_loss(
            cond_cont, cond_cat
        ) + sec_decoder.router.balance_loss(cond_cont, cond_cat)
        # Supervised router auxiliary loss (e.g. ProcessRouter's process
        # classifier); a scalar 0 for routers with no such loss (EnergyRouter).
        l_proc = stage1_model.router.classify_loss(
            cond_cont, cond_cat, proc_idx
        ) + sec_decoder.router.classify_loss(cond_cont, cond_cat, proc_idx)
        # Optional entropy-regularization aux loss (see Router.entropy_loss):
        # penalizes uniform/collapsed gating, a secondary guard against
        # gate-sharpness collapse that lambda_balance alone can't see.
        l_entropy = stage1_model.router.entropy_loss(
            cond_cont, cond_cat
        ) + sec_decoder.router.entropy_loss(cond_cont, cond_cat)
    else:
        l_balance = torch.zeros((), device=device)
        l_proc = torch.zeros((), device=device)
        l_entropy = torch.zeros((), device=device)

    total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2
    if lambda_balance > 0:
        total = total + lambda_balance * l_balance
    if lambda_proc > 0:
        total = total + lambda_proc * l_proc
    if lambda_entropy > 0:
        total = total + lambda_entropy * l_entropy
    return total, l_s1, l_nsec, l_s2, l_balance, l_proc, l_entropy, nsec_acc


def _wgan_train_step(
    generator: torch.nn.Module,
    sec_generator: torch.nn.Module,
    critic: torch.nn.Module,
    sec_critic: torch.nn.Module,
    batch: tuple,
    device: torch.device,
    optimizer_g: optim.Optimizer,
    optimizer_d: optim.Optimizer,
    g_params: list,
    d_params: list,
    step_count: int,
    n_critic: int,
    gp_weight: float,
    lambda_nsec: float,
    lambda_s2: float,
) -> dict:
    """One WGAN-GP training step, both stages (see giant/model/wgan.py for the losses).

    Both critics update every batch. Every `n_critic`-th batch additionally
    updates both generators. The (non-adversarial) n_sec classifier updates
    every batch regardless — folded into whichever `optimizer_g` step happens
    this batch (full adversarial g_loss on generator batches, n_sec-only in
    between) rather than throttled to the generator's cadence, since n_sec
    accuracy is a headline flow-vs-wgan comparison metric and shares the
    generator's ConditionEncoder.

    Stage 2's real/fake target is a flattened (B, SEC_DIM) vector with
    `K_MAX - n_sec` padded slots per row; both critic's input and its
    gradient-penalty gradient are masked to the valid slots (see
    `giant.model.wgan.gradient_penalty`) so the critic can't key on padding
    instead of genuine content. Stage 2 is conditioned on the *real*
    ground-truth Stage-1 target (`x1_s1`, detached) rather than the
    generator's own fake Stage-1 output — same precedent as the flow-matching
    path's `flow_matching_loss_secondary` call, avoiding compounding errors
    during training.
    """
    cond_cont, cond_cat, x1_s1, n_sec, sec_cont, _proc_idx = batch
    cond_cont = cond_cont.to(device)
    cond_cat = cond_cat.to(device)
    x1_s1 = x1_s1.to(device)
    n_sec = n_sec.to(device)
    sec_cont = sec_cont.to(device)

    B = x1_s1.size(0)
    x1_s2 = sec_cont.flatten(1)  # (B, SEC_DIM)
    sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
    mask_flat = (
        sec_mask.unsqueeze(-1).expand(-1, -1, SEC_SLOT_DIM).reshape(B, -1).float()
    )
    stage1_ctx = x1_s1.detach()

    def critic_fn1(x: torch.Tensor) -> torch.Tensor:
        return critic(x, cond_cont, cond_cat)

    def critic_fn2(x: torch.Tensor) -> torch.Tensor:
        return sec_critic(x, cond_cont, cond_cat, stage1_ctx)

    z1 = torch.randn(B, generator.noise_dim, device=device)
    fake1 = generator(z1, cond_cont, cond_cat)
    z2 = torch.randn(B, sec_generator.noise_dim, device=device)
    fake2 = sec_generator(z2, cond_cont, cond_cat, stage1_ctx)
    fake2_masked = fake2 * mask_flat
    real2_masked = x1_s2 * mask_flat

    # --- Critic step (every batch) ---
    fake1_detached = fake1.detach()
    real1_score = critic_fn1(x1_s1)
    fake1_score = critic_fn1(fake1_detached)
    gp1 = gradient_penalty(critic_fn1, x1_s1, fake1_detached)
    d1 = fake1_score.mean() - real1_score.mean() + gp_weight * gp1
    wasserstein_estimate = (real1_score.mean() - fake1_score.mean()).detach()

    fake2_detached_masked = fake2_masked.detach()
    real2_score = critic_fn2(real2_masked)
    fake2_score = critic_fn2(fake2_detached_masked)
    gp2 = gradient_penalty(
        critic_fn2, real2_masked, fake2_detached_masked, mask=mask_flat
    )
    d2 = fake2_score.mean() - real2_score.mean() + gp_weight * gp2

    d_loss = d1 + lambda_s2 * d2
    optimizer_d.zero_grad()
    d_loss.backward()
    grad_norm_d = torch.nn.utils.clip_grad_norm_(d_params, 1.0)
    optimizer_d.step()

    # --- Generator (+ n_sec) step ---
    did_g_step = step_count % n_critic == 0
    n_sec_logits = generator.predict_n_sec(cond_cont, cond_cat)
    l_nsec = F.cross_entropy(n_sec_logits, n_sec)
    nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean()
    optimizer_g.zero_grad()
    if did_g_step:
        g1 = generator_loss(critic_fn1, fake1)
        g2 = generator_loss(critic_fn2, fake2_masked)
        g_loss = g1 + lambda_nsec * l_nsec + lambda_s2 * g2
    else:
        g1 = torch.zeros((), device=device)
        g2 = torch.zeros((), device=device)
        g_loss = lambda_nsec * l_nsec
    g_loss.backward()
    grad_norm_g = torch.nn.utils.clip_grad_norm_(g_params, 1.0)
    optimizer_g.step()

    return {
        "d_loss": d_loss.detach(),
        "g_loss": (g1 + lambda_s2 * g2).detach(),
        "wasserstein_estimate": wasserstein_estimate,
        "gp_loss": (gp1 + lambda_s2 * gp2).detach(),
        "l_nsec": l_nsec.detach(),
        "nsec_acc": nsec_acc.detach(),
        "did_g_step": did_g_step,
        "grad_norm": grad_norm_d.item() + grad_norm_g.item(),
        "grad_norm_d": grad_norm_d.item(),
        "grad_norm_g": grad_norm_g.item(),
    }


def train(
    stage1_model: torch.nn.Module,
    sec_decoder: torch.nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    mode: str,
    epochs: int,
    lr: float,
    warmup_epochs: int,
    device: torch.device,
    out_dir: str | Path,
    weight_decay: float = 0.01,
    ema_decay: float = 0.9999,
    lambda_nsec: float = 0.1,
    lambda_s2: float = 1.0,
    lambda_balance: float = 0.0,
    lambda_proc: float = 0.0,
    lambda_entropy: float = 0.0,
    normalizer_dict: dict | None = None,
    pdg_map: dict | None = None,
    mat_map: dict | None = None,
    proc_map: dict | None = None,
    model_config: dict | None = None,
    resume_path: str | Path | None = None,
    validate_every: int = 0,
    validate_steps: int = 10,
    max_val_batches: int = 0,
    total_train_batches: int = 0,
    critic: torch.nn.Module | None = None,
    sec_critic: torch.nn.Module | None = None,
    n_critic: int = 5,
    gp_weight: float = 10.0,
    critic_lr: float | None = None,
    use_wandb: bool = False,
    wandb_project: str = "giant",
    wandb_run_name: str = "",
    wandb_log_every: int = 50,
) -> None:
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    stage1_params = sum(p.numel() for p in stage1_model.parameters())
    sec_decoder_params = sum(p.numel() for p in sec_decoder.parameters())
    critic_params = (
        sum(p.numel() for p in critic.parameters()) if critic is not None else 0
    )
    sec_critic_params = (
        sum(p.numel() for p in sec_critic.parameters()) if sec_critic is not None else 0
    )
    total_params = (
        stage1_params + sec_decoder_params + critic_params + sec_critic_params
    )

    wandb_run = None
    if use_wandb:
        try:
            import wandb
        except ImportError as exc:
            raise RuntimeError(
                "train.wandb = true (--wandb) requires the 'wandb' package — "
                "install it via `uv sync --extra wandb`"
            ) from exc
        # `id` is derived from out_dir so resuming a run (--resume) reattaches
        # to the same wandb run instead of starting a new one.
        wandb_run = wandb.init(
            project=wandb_project,
            name=wandb_run_name or out_dir.name,
            id=out_dir.name,
            resume="allow",
            config={
                "mode": mode,
                "epochs": epochs,
                "lr": lr,
                "warmup_epochs": warmup_epochs,
                "weight_decay": weight_decay,
                "ema_decay": ema_decay,
                "lambda_nsec": lambda_nsec,
                "lambda_s2": lambda_s2,
                "lambda_balance": lambda_balance,
                "lambda_proc": lambda_proc,
                "lambda_entropy": lambda_entropy,
                "n_critic": n_critic,
                "gp_weight": gp_weight,
                "model": model_config or {},
                "stage1_params": stage1_params,
                "sec_decoder_params": sec_decoder_params,
                "critic_params": critic_params,
                "sec_critic_params": sec_critic_params,
                "total_params": total_params,
            },
        )

    stage1_model = stage1_model.to(device)
    sec_decoder = sec_decoder.to(device)
    if mode == "wgan":
        assert critic is not None and sec_critic is not None, (
            "mode='wgan' requires critic/sec_critic (see giant.model.network.build_critics)"
        )
        critic = critic.to(device)
        sec_critic = sec_critic.to(device)

    # MoE routing trunk (RoutedDenoisingMLP/RoutedSecondaryDecoder) is
    # optional and orthogonal to `mode` — both stages carry a `.router`
    # when enabled. Each router is an independent instance (their
    # `n_experts` need not match), used both for the batch-level gate
    # entropy snapshot below and the val-level gate stats further down.
    has_router = hasattr(stage1_model, "router") and hasattr(sec_decoder, "router")

    # Flow-matching/diffusion models sample noticeably better from an EMA of
    # the weights than from the raw SGD-noisy ones — buffers (e.g. the fixed
    # sinusoidal-embedding freqs, or non-learned router centers) never change
    # after this initial copy, so only parameters need the running average.
    ema_stage1_model: torch.nn.Module | None = None
    ema_sec_decoder: torch.nn.Module | None = None
    if ema_decay > 0:
        ema_stage1_model = copy.deepcopy(stage1_model).eval()
        ema_sec_decoder = copy.deepcopy(sec_decoder).eval()
        for p in ema_stage1_model.parameters():
            p.requires_grad_(False)
        for p in ema_sec_decoder.parameters():
            p.requires_grad_(False)

    all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters())
    all_params_d: list = []
    optimizer_d: optim.Optimizer | None = None
    if mode == "wgan":
        assert critic is not None and sec_critic is not None
        # Standard WGAN-GP recipe (Gulrajani et al. 2017): Adam with
        # beta1=0 (momentum destabilizes critic training) and no weight
        # decay, rather than the AdamW(weight_decay=...) used for flow/ddpm.
        optimizer = optim.Adam(all_params, lr=lr, betas=(0.0, 0.9))
        all_params_d = list(critic.parameters()) + list(sec_critic.parameters())
        optimizer_d = optim.Adam(
            all_params_d,
            lr=critic_lr if critic_lr is not None else lr,
            betas=(0.0, 0.9),
        )
    else:
        optimizer = optim.AdamW(all_params, lr=lr, weight_decay=weight_decay)

    # Warmup/decay in units of optimizer steps rather than epochs: at large
    # dataset sizes a single epoch can be tens of thousands of steps, and an
    # epoch-granularity schedule would leave warmup/cosine decay unable to
    # move within it. Requires an accurate `total_train_batches` (steps per
    # epoch); the only caller, run_train_job, always supplies one.
    #
    # In wgan mode, `lr_sched.step()`/EMA only fire on generator steps (see
    # the per-batch loop below) — 1 in every `n_critic` batches — so the
    # schedule's own step-counting must be in those same units, or warmup
    # would never finish and cosine decay would barely move.
    steps_per_epoch = max(total_train_batches, 1)
    if mode == "wgan":
        steps_per_epoch = max(total_train_batches // (n_critic + 1), 1)
    warmup_steps = warmup_epochs * steps_per_epoch
    total_steps = max(epochs * steps_per_epoch, 1)

    def _lr_lambda(step: int) -> float:
        if warmup_steps > 0 and step < warmup_steps:
            return (step + 1) / warmup_steps
        t = step - warmup_steps
        T = max(total_steps - warmup_steps, 1)
        return 0.5 * (1.0 + math.cos(math.pi * min(t, T) / T))

    lr_sched = optim.lr_scheduler.LambdaLR(optimizer, _lr_lambda)

    ddpm_schedule = CosineSchedule().to(device) if mode == "ddpm" else None

    start_epoch = 1
    best_val_loss = float("inf")
    resumed_global_step = 0
    if resume_path is not None:
        ckpt = torch.load(resume_path, map_location=device, weights_only=False)
        stage1_model.load_state_dict(ckpt["model"])
        sec_decoder.load_state_dict(ckpt["sec_decoder"])
        if ema_decay > 0:
            assert ema_stage1_model is not None and ema_sec_decoder is not None
            ema_stage1_model.load_state_dict(ckpt.get("model_ema", ckpt["model"]))
            ema_sec_decoder.load_state_dict(
                ckpt.get("sec_decoder_ema", ckpt["sec_decoder"])
            )
        if mode == "wgan":
            assert (
                critic is not None
                and sec_critic is not None
                and optimizer_d is not None
            )
            critic.load_state_dict(ckpt["critic"])
            sec_critic.load_state_dict(ckpt["sec_critic"])
            optimizer_d.load_state_dict(ckpt["optimizer_d"])
        optimizer.load_state_dict(ckpt["optimizer"])
        lr_sched.load_state_dict(ckpt["lr_sched"])
        start_epoch = ckpt.get("epoch", 0) + 1
        best_val_loss = ckpt.get("best_val_loss", float("inf"))
        resumed_global_step = ckpt.get("global_step", 0)

        # optimizer/lr_sched.load_state_dict() above restore the checkpoint's
        # own base LR, which would otherwise silently override an explicit
        # `lr` argument. Make `lr` authoritative again, applied at whatever
        # point the cosine/warmup schedule has already reached.
        lr_sched.base_lrs = [lr for _ in lr_sched.base_lrs]
        resumed_lr = lr * _lr_lambda(lr_sched.last_epoch)
        for group in optimizer.param_groups:
            group["lr"] = resumed_lr

        if start_epoch > epochs:
            print(
                f"checkpoint already completed epoch {start_epoch - 1} "
                f"(>= --epochs {epochs}) — nothing to train"
            )
            return

    metrics_path = out_dir / "metrics.csv"
    resuming_existing_metrics = resume_path is not None and metrics_path.exists()
    write_header = not resuming_existing_metrics
    metrics_file = open(
        metrics_path, "a" if resuming_existing_metrics else "w", newline=""
    )
    metrics_writer = csv.DictWriter(metrics_file, fieldnames=_METRICS_FIELDS)
    if write_header:
        metrics_writer.writeheader()

    epoch_w = len(str(epochs))
    last_completed_epoch = start_epoch - 1
    # Restored from the checkpoint on --resume so wandb_run.log(..., step=...)
    # keeps advancing monotonically instead of restarting at 0 mid-run (a
    # reattached wandb run — see wandb.init(id=..., resume="allow") below —
    # would otherwise silently drop every post-resume point).
    global_step = resumed_global_step
    with _GracefulShutdown() as shutdown:
        for epoch in range(start_epoch, epochs + 1):
            epoch_start = time.monotonic()
            if device.type == "cuda":
                torch.cuda.reset_peak_memory_stats(device)
            stage1_model.train()
            sec_decoder.train()
            if mode == "wgan":
                assert critic is not None and sec_critic is not None
                critic.train()
                sec_critic.train()
            train_loss_sum = 0.0
            train_s1_sum = 0.0
            train_nsec_sum = 0.0
            train_s2_sum = 0.0
            train_balance_sum = 0.0
            train_proc_sum = 0.0
            train_entropy_sum = 0.0
            train_d_sum = 0.0
            train_g_sum = 0.0
            train_wasserstein_sum = 0.0
            train_gp_sum = 0.0
            train_nsec_acc_sum = 0.0
            train_grad_norm_d_sum = 0.0
            train_grad_norm_g_sum = 0.0
            train_n = 0
            train_batches = 0
            grad_norm_sum = 0.0
            ema_loss = 0.0
            ema_grad_norm = 0.0
            bar = tqdm(
                train_loader,
                desc=f"  epoch {epoch:{epoch_w}d}/{epochs}",
                total=total_train_batches or None,
                leave=False,
                unit="batch",
                dynamic_ncols=True,
            )
            for batch in bar:
                if mode == "wgan":
                    assert (
                        critic is not None
                        and sec_critic is not None
                        and optimizer_d is not None
                    )
                    stats = _wgan_train_step(
                        stage1_model,
                        sec_decoder,
                        critic,
                        sec_critic,
                        batch,
                        device,
                        optimizer,
                        optimizer_d,
                        all_params,
                        all_params_d,
                        global_step,
                        n_critic,
                        gp_weight,
                        lambda_nsec,
                        lambda_s2,
                    )
                    if stats["did_g_step"]:
                        lr_sched.step()
                        if ema_decay > 0:
                            assert (
                                ema_stage1_model is not None
                                and ema_sec_decoder is not None
                            )
                            _update_ema(ema_stage1_model, stage1_model, ema_decay)
                            _update_ema(ema_sec_decoder, sec_decoder, ema_decay)

                    B = batch[0].size(0)
                    batch_loss = stats["d_loss"].item() + stats["g_loss"].item()
                    batch_grad_norm = stats["grad_norm"]
                    train_loss_sum += batch_loss * B
                    train_nsec_sum += stats["l_nsec"].item() * B
                    train_d_sum += stats["d_loss"].item() * B
                    train_g_sum += stats["g_loss"].item() * B
                    train_wasserstein_sum += stats["wasserstein_estimate"].item() * B
                    train_gp_sum += stats["gp_loss"].item() * B
                    train_nsec_acc_sum += stats["nsec_acc"].item() * B
                    train_grad_norm_d_sum += stats["grad_norm_d"] * B
                    train_grad_norm_g_sum += stats["grad_norm_g"] * B
                else:
                    loss, l_s1, l_nsec, l_s2, l_balance, l_proc, l_entropy, nsec_acc = (
                        _compute_losses(
                            stage1_model,
                            sec_decoder,
                            batch,
                            mode,
                            ddpm_schedule,
                            device,
                            lambda_nsec,
                            lambda_s2,
                            lambda_balance,
                            lambda_proc,
                            lambda_entropy,
                        )
                    )
                    optimizer.zero_grad()
                    loss.backward()
                    grad_norm = torch.nn.utils.clip_grad_norm_(all_params, 1.0)
                    optimizer.step()
                    lr_sched.step()
                    if ema_decay > 0:
                        assert (
                            ema_stage1_model is not None and ema_sec_decoder is not None
                        )
                        _update_ema(ema_stage1_model, stage1_model, ema_decay)
                        _update_ema(ema_sec_decoder, sec_decoder, ema_decay)

                    B = batch[0].size(0)
                    batch_loss = loss.item()
                    batch_grad_norm = grad_norm.item()
                    train_loss_sum += batch_loss * B
                    train_s1_sum += l_s1.item() * B
                    train_nsec_sum += l_nsec.item() * B
                    train_s2_sum += l_s2.item() * B
                    train_balance_sum += l_balance.item() * B
                    train_proc_sum += l_proc.item() * B
                    train_entropy_sum += l_entropy.item() * B
                    train_nsec_acc_sum += nsec_acc.item() * B

                train_n += B
                train_batches += 1
                grad_norm_sum += batch_grad_norm
                ema_loss = (
                    batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss
                )
                ema_grad_norm = (
                    batch_grad_norm
                    if train_batches == 1
                    else 0.95 * ema_grad_norm + 0.05 * batch_grad_norm
                )
                bar.set_postfix_str(
                    f"loss={ema_loss:.4f} gnorm={ema_grad_norm:.3f}", refresh=False
                )

                global_step += 1
                if (
                    wandb_run is not None
                    and wandb_log_every > 0
                    and global_step % wandb_log_every == 0
                ):
                    log_payload = {
                        "batch/epoch": epoch,
                        "batch/loss": batch_loss,
                        "batch/loss_ema": ema_loss,
                        "batch/grad_norm": batch_grad_norm,
                        "batch/lr": optimizer.param_groups[0]["lr"],
                        "batch/critic_lr": (
                            optimizer_d.param_groups[0]["lr"]
                            if optimizer_d is not None
                            else 0.0
                        ),
                    }
                    if has_router:
                        # Cheap re-use of the batch already in hand — no
                        # extra data loading, just a small forward through
                        # each router's own gate function. Only entropy is
                        # logged at this granularity (not per-expert
                        # utilization): a single batch's importance sum is
                        # too noisy as a "global share" estimate, whereas
                        # the val-loop aggregate (below) sums over the
                        # whole val set for that. Batch-level entropy alone
                        # is still enough to see a router collapsing in
                        # real time, mid-epoch, rather than only at the
                        # next validation pass.
                        with torch.no_grad():
                            cond_cont_b = batch[0].to(device)
                            cond_cat_b = batch[1].to(device)
                            s1_entropy, _ = stage1_model.router.gate_stats(
                                cond_cont_b, cond_cat_b
                            )
                            s2_entropy, _ = sec_decoder.router.gate_stats(
                                cond_cont_b, cond_cat_b
                            )
                        log_payload["batch/router_s1_entropy"] = s1_entropy.item()
                        log_payload["batch/router_s2_entropy"] = s2_entropy.item()
                    wandb_run.log(log_payload, step=global_step)

                if shutdown.requested:
                    break
            bar.close()

            if shutdown.requested:
                break

            train_loss = train_loss_sum / max(train_n, 1)
            train_grad_norm = grad_norm_sum / max(train_batches, 1)
            train_nsec_acc = train_nsec_acc_sum / max(train_n, 1)
            train_grad_norm_d = train_grad_norm_d_sum / max(train_n, 1)
            train_grad_norm_g = train_grad_norm_g_sum / max(train_n, 1)
            current_lr = optimizer.param_groups[0]["lr"]
            critic_lr_value = (
                optimizer_d.param_groups[0]["lr"] if optimizer_d is not None else 0.0
            )

            stage1_model.eval()
            sec_decoder.eval()
            if mode == "wgan":
                assert critic is not None and sec_critic is not None
                critic.eval()
                sec_critic.eval()

            val_marginal_kl = float("nan")
            if mode == "wgan":
                # WGANGenerator.forward(z, cond_cont, cond_cat) has no
                # diffusion/flow `t` argument, so the usual _compute_losses
                # val loop below (which calls flow_matching_loss ->
                # stage1_model(x_t, t, ...)) doesn't apply — and a WGAN
                # critic loss isn't a monotone quality signal fit for
                # best-checkpoint selection anyway. Select on marginal KL
                # against the EMA generators instead (matches what
                # predict/rollout sample from by default, --weights ema).
                eval_stage1 = (
                    ema_stage1_model if ema_stage1_model is not None else stage1_model
                )
                eval_sec_decoder = (
                    ema_sec_decoder if ema_sec_decoder is not None else sec_decoder
                )
                marginal_result = validate_marginals(
                    eval_stage1,
                    val_loader,
                    mode=mode,
                    device=device,
                    sec_decoder=eval_sec_decoder,
                )
                val_marginal_kl = float(np.mean(marginal_result["kl_divergence"]))
                val_loss = val_marginal_kl
                val_s1_sum = val_nsec_sum = val_s2_sum = val_balance_sum = (
                    val_proc_sum
                ) = val_entropy_sum = val_nsec_acc_sum = 0.0
                val_n = 1
                val_nsec_acc = 0.0
                router_s1_entropy = router_s2_entropy = 0.0
                router_s1_util_min = router_s1_util_max = router_s1_util_std = 0.0
                router_s2_util_min = router_s2_util_max = router_s2_util_std = 0.0
            else:
                val_loss_sum = 0.0
                val_s1_sum = 0.0
                val_nsec_sum = 0.0
                val_s2_sum = 0.0
                val_balance_sum = 0.0
                val_proc_sum = 0.0
                val_entropy_sum = 0.0
                val_nsec_acc_sum = 0.0
                val_n = 0
                if has_router:
                    n_experts_s1 = stage1_model.router.n_experts
                    n_experts_s2 = sec_decoder.router.n_experts
                    val_router_s1_entropy_sum = 0.0
                    val_router_s2_entropy_sum = 0.0
                    val_router_s1_importance_sum = torch.zeros(
                        n_experts_s1, device=device
                    )
                    val_router_s2_importance_sum = torch.zeros(
                        n_experts_s2, device=device
                    )
                with torch.no_grad():
                    for val_batch_idx, batch in enumerate(val_loader):
                        if max_val_batches > 0 and val_batch_idx >= max_val_batches:
                            break
                        (
                            loss,
                            l_s1,
                            l_nsec,
                            l_s2,
                            l_balance,
                            l_proc,
                            l_entropy,
                            nsec_acc,
                        ) = _compute_losses(
                            stage1_model,
                            sec_decoder,
                            batch,
                            mode,
                            ddpm_schedule,
                            device,
                            lambda_nsec,
                            lambda_s2,
                            lambda_balance,
                            lambda_proc,
                            lambda_entropy,
                        )
                        B = batch[0].size(0)
                        val_loss_sum += loss.item() * B
                        val_s1_sum += l_s1.item() * B
                        val_nsec_sum += l_nsec.item() * B
                        val_s2_sum += l_s2.item() * B
                        val_balance_sum += l_balance.item() * B
                        val_proc_sum += l_proc.item() * B
                        val_entropy_sum += l_entropy.item() * B
                        val_nsec_acc_sum += nsec_acc.item() * B
                        if has_router:
                            cond_cont = batch[0].to(device)
                            cond_cat = batch[1].to(device)
                            s1_entropy, s1_importance = stage1_model.router.gate_stats(
                                cond_cont, cond_cat
                            )
                            s2_entropy, s2_importance = sec_decoder.router.gate_stats(
                                cond_cont, cond_cat
                            )
                            val_router_s1_entropy_sum += s1_entropy.item() * B
                            val_router_s2_entropy_sum += s2_entropy.item() * B
                            val_router_s1_importance_sum += s1_importance
                            val_router_s2_importance_sum += s2_importance
                        val_n += B
                val_loss = val_loss_sum / max(val_n, 1)
                val_nsec_acc = val_nsec_acc_sum / max(val_n, 1)

                if has_router:
                    router_s1_entropy = val_router_s1_entropy_sum / max(val_n, 1)
                    router_s2_entropy = val_router_s2_entropy_sum / max(val_n, 1)
                    s1_util = val_router_s1_importance_sum / (
                        val_router_s1_importance_sum.sum().clamp_min(1e-8)
                    )
                    s2_util = val_router_s2_importance_sum / (
                        val_router_s2_importance_sum.sum().clamp_min(1e-8)
                    )
                    router_s1_util_min = s1_util.min().item()
                    router_s1_util_max = s1_util.max().item()
                    router_s1_util_std = (
                        s1_util.std().item() if n_experts_s1 > 1 else 0.0
                    )
                    router_s2_util_min = s2_util.min().item()
                    router_s2_util_max = s2_util.max().item()
                    router_s2_util_std = (
                        s2_util.std().item() if n_experts_s2 > 1 else 0.0
                    )
                else:
                    router_s1_entropy = router_s2_entropy = 0.0
                    router_s1_util_min = router_s1_util_max = router_s1_util_std = 0.0
                    router_s2_util_min = router_s2_util_max = router_s2_util_std = 0.0

                if validate_every > 0 and epoch % validate_every == 0:
                    print(f"[epoch {epoch}] marginal validation:")
                    marginal_result = validate_marginals(
                        stage1_model,
                        val_loader,
                        mode=mode,
                        schedule=ddpm_schedule,
                        device=device,
                        steps=validate_steps,
                        sec_decoder=sec_decoder,
                    )
                    val_marginal_kl = float(np.mean(marginal_result["kl_divergence"]))

            epoch_time = time.monotonic() - epoch_start
            gpu_mem_mb = (
                torch.cuda.max_memory_allocated(device) / (1024 * 1024)
                if device.type == "cuda"
                else 0.0
            )

            is_best = val_loss < best_val_loss
            marker = "  [best]" if is_best else ""
            print(
                f"epoch {epoch:{epoch_w}d}/{epochs}"
                f"  train {train_loss:.4f}"
                f"  (s1={train_s1_sum / max(train_n, 1):.3f}"
                f"  nsec={train_nsec_sum / max(train_n, 1):.3f}"
                f"  s2={train_s2_sum / max(train_n, 1):.3f}"
                f"  bal={train_balance_sum / max(train_n, 1):.3f}"
                f"  proc={train_proc_sum / max(train_n, 1):.3f}"
                f"  entropy={train_entropy_sum / max(train_n, 1):.3f}"
                f"  d={train_d_sum / max(train_n, 1):.3f}"
                f"  g={train_g_sum / max(train_n, 1):.3f})"
                f"  val {val_loss:.4f}"
                f"  lr {current_lr:.2e}  gnorm {train_grad_norm:.3f}"
                f"  {epoch_time:.1f}s{marker}"
            )
            metrics_row = {
                "epoch": epoch,
                "train_loss": train_loss,
                "train_loss_s1": train_s1_sum / max(train_n, 1),
                "train_loss_nsec": train_nsec_sum / max(train_n, 1),
                "train_loss_s2": train_s2_sum / max(train_n, 1),
                "train_loss_balance": train_balance_sum / max(train_n, 1),
                "train_loss_proc": train_proc_sum / max(train_n, 1),
                "train_loss_entropy": train_entropy_sum / max(train_n, 1),
                "train_nsec_acc": train_nsec_acc,
                "d_loss": train_d_sum / max(train_n, 1),
                "g_loss": train_g_sum / max(train_n, 1),
                "wasserstein_estimate": train_wasserstein_sum / max(train_n, 1),
                "gp_loss": train_gp_sum / max(train_n, 1),
                "val_loss": val_loss,
                "val_loss_s1": val_s1_sum / max(val_n, 1),
                "val_loss_nsec": val_nsec_sum / max(val_n, 1),
                "val_loss_s2": val_s2_sum / max(val_n, 1),
                "val_loss_balance": val_balance_sum / max(val_n, 1),
                "val_loss_proc": val_proc_sum / max(val_n, 1),
                "val_loss_entropy": val_entropy_sum / max(val_n, 1),
                "val_nsec_acc": val_nsec_acc,
                "val_marginal_kl": val_marginal_kl,
                "router_s1_entropy": router_s1_entropy,
                "router_s1_util_min": router_s1_util_min,
                "router_s1_util_max": router_s1_util_max,
                "router_s1_util_std": router_s1_util_std,
                "router_s2_entropy": router_s2_entropy,
                "router_s2_util_min": router_s2_util_min,
                "router_s2_util_max": router_s2_util_max,
                "router_s2_util_std": router_s2_util_std,
                "lr": current_lr,
                "critic_lr": critic_lr_value,
                "grad_norm": train_grad_norm,
                "grad_norm_d": train_grad_norm_d,
                "grad_norm_g": train_grad_norm_g,
                "gpu_mem_mb": gpu_mem_mb,
                "samples_per_sec": train_n / max(epoch_time, 1e-8),
                "is_best": int(is_best),
                "epoch_time_s": epoch_time,
            }
            metrics_writer.writerow(metrics_row)
            metrics_file.flush()
            if wandb_run is not None:
                # Shares the same monotonic step axis as the per-batch
                # `batch/*` logs above (global_step) rather than `epoch`,
                # since a wandb run's `step` argument across `log()` calls
                # must never decrease.
                wandb_run.log(metrics_row, step=global_step)

            ckpt: dict = {
                "model": stage1_model.state_dict(),
                "sec_decoder": sec_decoder.state_dict(),
                "optimizer": optimizer.state_dict(),
                "lr_sched": lr_sched.state_dict(),
                "epoch": epoch,
                "best_val_loss": best_val_loss,
                "global_step": global_step,
            }
            if mode == "wgan":
                assert (
                    critic is not None
                    and sec_critic is not None
                    and optimizer_d is not None
                )
                ckpt["critic"] = critic.state_dict()
                ckpt["sec_critic"] = sec_critic.state_dict()
                ckpt["optimizer_d"] = optimizer_d.state_dict()
            if ema_decay > 0:
                assert ema_stage1_model is not None and ema_sec_decoder is not None
                ckpt["model_ema"] = ema_stage1_model.state_dict()
                ckpt["sec_decoder_ema"] = ema_sec_decoder.state_dict()
            if normalizer_dict is not None:
                ckpt["normalizer"] = normalizer_dict
            if pdg_map is not None:
                ckpt["pdg_map"] = pdg_map
            if mat_map is not None:
                ckpt["mat_map"] = mat_map
            if proc_map is not None:
                ckpt["proc_map"] = proc_map
            if model_config is not None:
                ckpt["model_config"] = model_config

            if val_loss < best_val_loss:
                best_val_loss = val_loss
                ckpt["best_val_loss"] = best_val_loss
                torch.save(ckpt, out_dir / "best.pt")

            torch.save(ckpt, out_dir / "last.pt")
            last_completed_epoch = epoch

            if shutdown.requested:
                break

    metrics_file.close()
    if wandb_run is not None:
        wandb_run.finish()

    if shutdown.requested:
        print(
            f"stopped after epoch {last_completed_epoch} due to shutdown signal — "
            f"resume with --resume {out_dir / 'last.pt'}"
        )