code_snapshots/WiSER_public_release_code/scripts/train_phase_radiomap.py

"""WiSER radiomap V2.2a training entry point: V1.5 arch + frozen D2.2 backbone.

Inherits V1.5 hierarchical radiomap head. Adds:
  * --load-d22-ckpt PATH: load D2.2 best.pt weights into backbone + tx_proj
                          (strict=False, CIR head keys discarded).
  * --freeze-base: freeze backbone+tx_embed+tx_proj; only radiomap_head trains.

The forward uses `torch.no_grad()` around the base when frozen to save
activation memory, and the `train()` override keeps the frozen part in
`eval()` mode regardless of outer train/eval switches.

Usage:
    python -m torch.distributed.run --standalone --nproc_per_node=8 \
        scripts/v2p2a/train_v2p2a.py \
        --manifest benchmarks/formal100_sionna_radiomap_100train_10heldoutval.json \
        --dataset-kind sionna_radiomap --db-floor -300.0 \
        --gradient-weight 1.0 --corridor-budget 192 \
        --num-cross-layers 4 --num-self-layers 2 \
        --load-d22-ckpt /.../D2p2_.../best.pt --freeze-base \
        --run-root <OUT> --epochs 50 --warmup-epochs 5 --batch-size 4
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import os
import platform
import sys
import time
from contextlib import nullcontext
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

ROOT = Path(__file__).resolve().parent.parent  # WiSER root
sys.path.insert(0, str(ROOT))

from wiser.config import ModelConfig as SharedModelConfig
from wiser.data.radiomap_dataset import (
    RadiomapDataset,
    radiomap_collate,
)
from wiser.alt_models import JointRadiomapCIRModel, FreezeConfig
from wiser.alt_engine.checkpoint import (
    detect_csi_head_arch_from_ckpt,
    load_warm_start_ckpt,
    load_phase_ckpt_into_model,
    save_phase_ckpt,
)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--manifest", required=True)
    p.add_argument("--run-root", required=True, help="absolute output dir for this run")
    p.add_argument("--epochs", type=int, default=50)
    p.add_argument("--warmup-epochs", type=int, default=5)
    p.add_argument("--batch-size", type=int, default=16)
    p.add_argument("--val-batch-size", type=int, default=8)
    p.add_argument("--val-every-epochs", type=int, default=2)
    p.add_argument("--num-workers", type=int, default=4)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--lr-min", type=float, default=1e-6)
    p.add_argument("--weight-decay", type=float, default=0.0)
    p.add_argument("--channels", type=int, default=512)
    p.add_argument("--downsample-stages", type=int, default=3)
    p.add_argument("--blocks-per-stage", type=int, default=3)
    p.add_argument("--spatial-channels", type=int, default=256)
    p.add_argument("--stem-depth", type=int, default=4)
    p.add_argument("--output-depth", type=int, default=4)
    p.add_argument("--grid-h", type=int, default=36)
    p.add_argument("--grid-w", type=int, default=36)
    p.add_argument("--db-floor", type=float, default=-200.0)
    p.add_argument("--dataset-kind", default="csi_path_aggregate",
                   choices=["csi_path_aggregate", "sionna_radiomap"])
    p.add_argument("--huber-beta", type=float, default=2.0)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--backbone-kind", default="trellis2", choices=["trellis2", "dense_fallback"])
    # V1.4 addition
    p.add_argument("--gradient-weight", type=float, default=1.0,
                   help="weight of WiSER spatial decoder-style gradient L1 loss (0 disables, same as V1.3)")
    # V1.5 additions
    p.add_argument("--corridor-budget", type=int, default=192,
                   help="RayCorridorGather top-K tokens per query from level 0 (10cm)")
    p.add_argument("--num-cross-layers", type=int, default=4,
                   help="decoder cross-attn depth (WiSER spatial decoder default 4)")
    p.add_argument("--num-self-layers", type=int, default=2,
                   help="decoder self-attn depth (WiSER spatial decoder default 2)")
    # V2.2a additions (kept for backward-compat)
    p.add_argument("--load-d22-ckpt", default=None,
                   help="path to D2.2 best.pt; loads backbone+tx_proj (ignore CIR head)")
    p.add_argument("--freeze-base", action="store_true",
                   help="(legacy) freeze backbone+tx_embed+tx_proj; only radiomap_head trains")
    # WiSER additions
    p.add_argument("--phase-name", default="P0_warmup_backbone",
                   help="phase name, used for logging / ckpt metadata")
    p.add_argument("--load-v22a-ckpt", default=None,
                   help="V2.2a best.pt path; provides radiomap head + backbone + tx_proj")
    p.add_argument("--load-phase-ckpt", default=None,
                   help="Load full prior-phase ckpt (supersedes --load-v22a/--load-d22)")
    p.add_argument("--freeze-backbone", action="store_true",
                   help="Set backbone.requires_grad=False")
    p.add_argument("--freeze-tx-proj", action="store_true")
    p.add_argument("--freeze-radiomap-head", action="store_true")
    p.add_argument("--freeze-csi-head", action="store_true",
                   help="Set csi_head.requires_grad=False (default for radiomap task!)")
    return p.parse_args()


def gradient_l1_loss(pred: torch.Tensor, gt: torch.Tensor, mask_hw: torch.Tensor) -> torch.Tensor:
    """WiSER spatial decoder style: L1 of finite-difference gradient diff, masked to pairs of
    adjacent cells that are BOTH valid. pred/gt: [B, H, W]. mask_hw: bool [B, H, W]."""
    dx_pred = pred[:, :, 1:] - pred[:, :, :-1]
    dx_gt   = gt[:, :, 1:]   - gt[:, :, :-1]
    dx_mask = mask_hw[:, :, 1:] & mask_hw[:, :, :-1]
    dy_pred = pred[:, 1:, :] - pred[:, :-1, :]
    dy_gt   = gt[:, 1:, :]   - gt[:, :-1, :]
    dy_mask = mask_hw[:, 1:, :] & mask_hw[:, :-1, :]

    out = pred.new_zeros(())
    if dx_mask.any():
        out = out + (dx_pred - dx_gt).abs()[dx_mask].mean()
    if dy_mask.any():
        out = out + (dy_pred - dy_gt).abs()[dy_mask].mean()
    return out


def setup_ddp() -> tuple[int, int, int, torch.device]:
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    local_rank = int(os.environ.get("LOCAL_RANK", rank))
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    return rank, local_rank, world_size, device


def cleanup_ddp() -> None:
    if dist.is_initialized():
        dist.destroy_process_group()


def compute_lr(step: int, total_steps: int, warmup_steps: int, lr_max: float, lr_min: float) -> float:
    if step < warmup_steps:
        return lr_max * (step + 1) / max(warmup_steps, 1)
    progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
    progress = min(max(progress, 0.0), 1.0)
    cos = 0.5 * (1.0 + math.cos(math.pi * progress))
    return lr_min + (lr_max - lr_min) * cos


# RadiomapModelV15 class + helpers moved to alt_models/joint_model.py
# _extract_level_tensors, _pack_levels_to_dense are methods there.


def compute_radiomap_metrics(pred_db: torch.Tensor, gt_db: torch.Tensor, mask_hw: torch.Tensor) -> dict[str, float]:
    """pred_db, gt_db, mask_hw ∈ [B, H, W] (mask bool)."""
    m = mask_hw.bool()
    if m.sum() == 0:
        return {"radiomap_mae_db": 0.0, "radiomap_psnr_db": 0.0, "radiomap_bias_db": 0.0, "valid_cells": 0}
    d = (pred_db[m] - gt_db[m]).float()
    mae = d.abs().mean().item()
    bias = d.mean().item()
    mse = (d * d).mean().item()
    span = 120.0  # typical dB span -120..0 → 120 dB dynamic
    psnr = 20.0 * math.log10(span) - 10.0 * math.log10(max(mse, 1e-8))
    return {
        "radiomap_mae_db": float(mae),
        "radiomap_psnr_db": float(psnr),
        "radiomap_bias_db": float(bias),
        "valid_cells": int(m.sum().item()),
    }


def run_val(
    model: torch.nn.Module,
    loader: DataLoader,
    device: torch.device,
    amp_dtype: torch.dtype | None,
    huber_beta: float,
    world_size: int,
    gradient_weight: float = 0.0,
) -> dict[str, float]:
    model.eval()
    total = {"loss": 0.0, "gain": 0.0, "grad": 0.0, "mae": 0.0, "psnr": 0.0, "bias": 0.0, "cells": 0.0, "n": 0.0}
    with torch.no_grad():
        for batch in loader:
            voxel_feats_list = [t.to(device, non_blocking=True) for t in batch["voxel_feats_list"]]
            voxel_coords_list = [t.to(device, non_blocking=True) for t in batch["voxel_coords_list"]]
            tx_xyz = batch["tx_xyz_norm"].to(device, non_blocking=True)
            rx_grid = batch["rx_grid_xyz_norm"].to(device, non_blocking=True)
            gt = batch["gt_radiomap_db"].to(device, non_blocking=True)
            mask = batch["extent_mask"].to(device, non_blocking=True)
            # V1.5: metric fields
            tx_xyz_m = batch["tx_xyz_metric"].to(device, non_blocking=True)
            rx_grid_m = batch["rx_grid_xyz_metric"].to(device, non_blocking=True)
            scene_ext_m = batch["scene_extent_xyz_m"].to(device, non_blocking=True)
            origin_m_list = [t.to(device, non_blocking=True) for t in batch["voxel_origin_m_list"]]
            cell_size_m = float(batch["voxel_cell_size_m"])
            cm = torch.autocast(device_type="cuda", dtype=amp_dtype) if amp_dtype is not None else nullcontext()
            with cm:
                pred = model(
                    voxel_feats_list, voxel_coords_list,
                    tx_xyz_norm=tx_xyz, tx_xyz_metric=tx_xyz_m,
                    rx_grid_xyz_norm=rx_grid, rx_grid_xyz_metric=rx_grid_m,
                    scene_extent_xyz_m=scene_ext_m,
                    voxel_origin_m_list=origin_m_list,
                    voxel_cell_size_m=cell_size_m,
                    extent_mask=mask,
                )
                mh = mask.squeeze(1).bool()
                if mh.sum() > 0:
                    gain = F.smooth_l1_loss(pred[mh], gt[mh], beta=huber_beta).item()
                else:
                    gain = 0.0
                if gradient_weight > 0 and mh.any():
                    grad = gradient_l1_loss(pred, gt, mh).item()
                else:
                    grad = 0.0
                loss = gain + gradient_weight * grad
            mets = compute_radiomap_metrics(pred.float(), gt, mask.squeeze(1))
            B = gt.shape[0]
            total["loss"] += loss * B
            total["gain"] += gain * B
            total["grad"] += grad * B
            total["mae"] += mets["radiomap_mae_db"] * B
            total["psnr"] += mets["radiomap_psnr_db"] * B
            total["bias"] += mets["radiomap_bias_db"] * B
            total["cells"] += mets["valid_cells"]
            total["n"] += B
    # all-reduce across ranks
    t = torch.tensor([total[k] for k in ("loss", "gain", "grad", "mae", "psnr", "bias", "cells", "n")],
                     dtype=torch.float64, device=device)
    dist.all_reduce(t, op=dist.ReduceOp.SUM)
    loss, gain, grad, mae, psnr, bias, cells, n = t.tolist()
    model.train()
    if n == 0:
        return {"val_loss": 0.0, "val_mae": 0.0, "val_psnr": 0.0, "val_bias": 0.0, "val_cells": 0.0}
    return {
        "val_loss": loss / n,
        "val_gain_loss": gain / n,
        "val_grad_loss": grad / n,
        "val_mae_db": mae / n,
        "val_psnr_db": psnr / n,
        "val_bias_db": bias / n,
        "val_cells_mean": cells / n,
    }


def composite_score(val: dict[str, float]) -> float:
    """Selector for best.pt. Lower MAE is better, higher PSNR is better."""
    return -val.get("val_mae_db", 1e9) + 0.1 * val.get("val_psnr_db", 0.0) - 0.1 * abs(val.get("val_bias_db", 0.0))


def save_ckpt(
    path: Path,
    model: torch.nn.Module,
    opt: torch.optim.Optimizer,
    epoch: int,
    step: int,
    metrics: dict[str, float],
    cli_args: dict[str, Any],
    phase_name: str = "unknown",
) -> None:
    raw = model.module if isinstance(model, DDP) else model
    torch.save(
        {
            "phase_name": phase_name,
            "epoch": int(epoch),
            "step": int(step),
            "metrics": metrics,
            "cli_args": cli_args,
            "model_state_dict": {k: v.detach().cpu() for k, v in raw.state_dict().items()},
            "optimizer_state_dict": opt.state_dict(),
        },
        path,
    )


def main() -> None:
    args = parse_args()
    rank, local_rank, world_size, device = setup_ddp()
    torch.manual_seed(args.seed + rank)

    run_root = Path(args.run_root)
    reports_dir = run_root / "reports"
    artifacts_dir = run_root / "artifacts"
    ckpt_dir = run_root / "checkpoints"
    for d in (reports_dir, reports_dir / "plots" / "best", reports_dir / "plots" / "last", artifacts_dir, ckpt_dir):
        d.mkdir(parents=True, exist_ok=True)

    with open(args.manifest) as f:
        manifest = json.load(f)
    train_records = manifest["train"]
    val_records = manifest["val_heldout"]

    if rank == 0:
        print("============== WiSER radiomap V1 training ==============")
        print(f"world_size={world_size}  per-rank bs={args.batch_size}  global bs={args.batch_size * world_size}")
        print(f"epochs={args.epochs}  warmup={args.warmup_epochs}  lr={args.lr}→{args.lr_min}")
        print(f"manifest: train={len(train_records)}  val={len(val_records)}")
        print(f"run_root: {run_root}")

        snap = {k: (v if not isinstance(v, Path) else str(v)) for k, v in vars(args).items()}
        snap["world_size"] = world_size
        snap["num_train_records"] = len(train_records)
        snap["num_val_records"] = len(val_records)
        with open(artifacts_dir / "config_snapshot.json", "w") as f:
            json.dump(snap, f, indent=2)
        with open(artifacts_dir / "env.txt", "w") as f:
            f.write(f"python={platform.python_version()}\n")
            f.write(f"torch={torch.__version__}\n")
            f.write(f"cuda_available={torch.cuda.is_available()}\n")
            if torch.cuda.is_available():
                f.write(f"cuda={torch.version.cuda}\n")
                f.write(f"gpu0={torch.cuda.get_device_name(0)}\n")
                f.write(f"device_count={torch.cuda.device_count()}\n")

    dist.barrier()

    train_ds = RadiomapDataset(train_records, channels=args.channels, grid_h=args.grid_h, grid_w=args.grid_w, db_floor=args.db_floor, dataset_kind=args.dataset_kind)
    val_ds = RadiomapDataset(val_records, channels=args.channels, grid_h=args.grid_h, grid_w=args.grid_w, db_floor=args.db_floor, dataset_kind=args.dataset_kind)

    train_sampler = DistributedSampler(train_ds, num_replicas=world_size, rank=rank, shuffle=True, drop_last=True, seed=args.seed)
    val_sampler = DistributedSampler(val_ds, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False)

    train_loader = DataLoader(
        train_ds,
        batch_size=args.batch_size,
        sampler=train_sampler,
        collate_fn=radiomap_collate,
        num_workers=args.num_workers,
        pin_memory=True,
        persistent_workers=args.num_workers > 0,
        drop_last=True,
    )
    val_loader = DataLoader(
        val_ds,
        batch_size=args.val_batch_size,
        sampler=val_sampler,
        collate_fn=radiomap_collate,
        num_workers=max(1, args.num_workers // 2),
        pin_memory=True,
        persistent_workers=False,
        drop_last=False,
    )

    shared = SharedModelConfig(
        backbone_kind=args.backbone_kind,
        backbone_channels=args.channels,
        backbone_downsample_stages=args.downsample_stages,
        backbone_blocks_per_stage=args.blocks_per_stage,
    )
    # ---- Build JointRadiomapCIRModel ----
    # Detect CIR head arch from whatever ckpt we're loading (D2.2 or phase ckpt).
    # detect_csi_head_arch_from_ckpt handles both "head.*" and "csi_head.*" prefixes.
    csi_arch_kwargs = {}
    ckpt_for_arch = args.load_phase_ckpt or args.load_d22_ckpt
    if ckpt_for_arch:
        info = detect_csi_head_arch_from_ckpt(ckpt_for_arch)
        csi_arch_kwargs = {k: v for k, v in info.items() if k.startswith("csi_")}
        if "backbone_downsample_stages" in info:
            shared.backbone_downsample_stages = info["backbone_downsample_stages"]
        if "backbone_blocks_per_stage" in info:
            shared.backbone_blocks_per_stage = info["backbone_blocks_per_stage"]
        if rank == 0:
            print(f"[arch-detect] from {ckpt_for_arch}: {csi_arch_kwargs}")

    model = JointRadiomapCIRModel(
        shared=shared,
        channels=args.channels,
        spatial_channels=args.spatial_channels,
        stem_depth=args.stem_depth,
        output_depth=args.output_depth,
        grid_h=args.grid_h,
        grid_w=args.grid_w,
        corridor_budget=args.corridor_budget,
        num_cross_layers=args.num_cross_layers,
        num_self_layers=args.num_self_layers,
        **csi_arch_kwargs,
    ).to(device)

    # ---- Load checkpoints ----
    if args.load_phase_ckpt:
        # Continuing from a previous phase (full model state)
        meta = load_phase_ckpt_into_model(args.load_phase_ckpt, model)
        if rank == 0:
            print(f"[load] phase ckpt from {meta.get('phase_name')} ep={meta.get('epoch')} metrics={meta.get('metrics', {})}")
    elif args.load_v22a_ckpt and args.load_d22_ckpt:
        # Warm start: V2.2a (rm head) + D2.2 (CIR head)
        report = load_warm_start_ckpt(model, args.load_v22a_ckpt, args.load_d22_ckpt, verify=True)
        if rank == 0:
            print(f"[load] warm-start: V2.2a + D2.2 merged: {report}")
    elif args.load_d22_ckpt:
        # Legacy: only D2.2 (no radiomap head init, like V2.2a)
        import torch as _torch
        d22_sd = _torch.load(args.load_d22_ckpt, map_location="cpu", weights_only=False)["model_state_dict"]
        keep = {k: v for k, v in d22_sd.items() if k.startswith("backbone.") or k.startswith("tx_proj.")}
        # Also take head.* → csi_head.*
        csi_remap = {k.replace("head.", "csi_head.", 1): v for k, v in d22_sd.items() if k.startswith("head.")}
        model.load_state_dict(keep, strict=False)
        model.load_state_dict(csi_remap, strict=False)
        if rank == 0:
            print(f"[load] D2.2 only: backbone+tx_proj ({len(keep)}) + csi_head ({len(csi_remap)})")

    # ---- Apply freeze config ----
    freeze_cfg = FreezeConfig(
        backbone=args.freeze_backbone or args.freeze_base,
        tx_proj=args.freeze_tx_proj or args.freeze_base,
        radiomap_head=args.freeze_radiomap_head,
        csi_head=args.freeze_csi_head or True,   # radiomap phase ALWAYS freezes CIR head (it's not used in forward_radiomap)
    )
    model.set_freeze_state(freeze_cfg)

    # DDP wraps .forward, so bind .forward = .forward_radiomap for this phase
    model.forward = model.forward_radiomap

    n_total = sum(p.numel() for p in model.parameters())
    n_train = sum(p.numel() for p in model.parameters() if p.requires_grad)
    if rank == 0:
        print(f"[phase={args.phase_name}] params: total={n_total/1e6:.2f}M trainable={n_train/1e6:.2f}M frozen={(n_total-n_train)/1e6:.2f}M")
        print(f"[phase={args.phase_name}] freeze: backbone={freeze_cfg.backbone} tx_proj={freeze_cfg.tx_proj} rm_head={freeze_cfg.radiomap_head} csi_head={freeze_cfg.csi_head}")

    # find_unused_parameters when any part frozen
    any_frozen = any([freeze_cfg.backbone, freeze_cfg.tx_proj, freeze_cfg.radiomap_head, freeze_cfg.csi_head])
    ddp = DDP(model, device_ids=[local_rank],
              find_unused_parameters=any_frozen,
              gradient_as_bucket_view=True)
    opt = torch.optim.Adam([p for p in ddp.parameters() if p.requires_grad],
                           lr=args.lr, weight_decay=args.weight_decay)

    steps_per_epoch = max(1, len(train_loader))
    total_steps = args.epochs * steps_per_epoch
    warmup_steps = args.warmup_epochs * steps_per_epoch
    amp_dtype = torch.bfloat16

    csv_path = reports_dir / "train_curves.csv"
    json_path = reports_dir / "train_curves.json"
    if rank == 0:
        with open(csv_path, "w", newline="") as f:
            w = csv.writer(f)
            w.writerow([
                "epoch", "step", "lr",
                "train_loss", "train_gain_loss", "train_grad_loss",
                "train_mae_db", "train_psnr_db", "train_bias_db",
                "val_loss", "val_gain_loss", "val_grad_loss",
                "val_mae_db", "val_psnr_db", "val_bias_db", "val_cells_mean",
                "secs_train", "secs_val", "secs_per_epoch", "peak_mem_gb",
            ])

    best_score = -1e9
    history: list[dict] = []
    step = 0

    for epoch in range(args.epochs):
        ddp.train()
        train_sampler.set_epoch(epoch)
        epoch_start = time.perf_counter()
        torch.cuda.reset_peak_memory_stats(device)

        ep_loss = 0.0; ep_gain = 0.0; ep_grad = 0.0
        ep_mae = 0.0; ep_psnr = 0.0; ep_bias = 0.0; ep_n = 0
        for batch in train_loader:
            lr = compute_lr(step, total_steps, warmup_steps, args.lr, args.lr_min)
            for pg in opt.param_groups:
                pg["lr"] = lr

            voxel_feats_list = [t.to(device, non_blocking=True) for t in batch["voxel_feats_list"]]
            voxel_coords_list = [t.to(device, non_blocking=True) for t in batch["voxel_coords_list"]]
            tx_xyz = batch["tx_xyz_norm"].to(device, non_blocking=True)
            rx_grid = batch["rx_grid_xyz_norm"].to(device, non_blocking=True)
            gt = batch["gt_radiomap_db"].to(device, non_blocking=True)
            mask = batch["extent_mask"].to(device, non_blocking=True)
            # V1.5 metric fields
            tx_xyz_m = batch["tx_xyz_metric"].to(device, non_blocking=True)
            rx_grid_m = batch["rx_grid_xyz_metric"].to(device, non_blocking=True)
            scene_ext_m = batch["scene_extent_xyz_m"].to(device, non_blocking=True)
            origin_m_list = [t.to(device, non_blocking=True) for t in batch["voxel_origin_m_list"]]
            cell_size_m = float(batch["voxel_cell_size_m"])

            opt.zero_grad(set_to_none=True)
            with torch.autocast(device_type="cuda", dtype=amp_dtype):
                pred = ddp(
                    voxel_feats_list, voxel_coords_list,
                    tx_xyz_norm=tx_xyz, tx_xyz_metric=tx_xyz_m,
                    rx_grid_xyz_norm=rx_grid, rx_grid_xyz_metric=rx_grid_m,
                    scene_extent_xyz_m=scene_ext_m,
                    voxel_origin_m_list=origin_m_list,
                    voxel_cell_size_m=cell_size_m,
                    extent_mask=mask,
                )
                mh = mask.squeeze(1).bool()
                if mh.sum() > 0:
                    loss_gain = F.smooth_l1_loss(pred[mh], gt[mh], beta=args.huber_beta)
                else:
                    loss_gain = pred.sum() * 0.0
                # V1.4: WiSER spatial decoder style gradient L1 loss
                if args.gradient_weight > 0:
                    loss_grad = gradient_l1_loss(pred, gt, mh)
                else:
                    loss_grad = pred.sum() * 0.0
                loss = loss_gain + args.gradient_weight * loss_grad
            loss.backward()
            opt.step()

            with torch.no_grad():
                m = compute_radiomap_metrics(pred.float(), gt, mask.squeeze(1))
            B = gt.shape[0]
            ep_loss += loss.item() * B
            ep_gain += loss_gain.item() * B
            ep_grad += loss_grad.item() * B
            ep_mae += m["radiomap_mae_db"] * B
            ep_psnr += m["radiomap_psnr_db"] * B
            ep_bias += m["radiomap_bias_db"] * B
            ep_n += B
            step += 1

        # reduce train metrics across ranks
        t_train = torch.tensor([ep_loss, ep_gain, ep_grad, ep_mae, ep_psnr, ep_bias, float(ep_n)], dtype=torch.float64, device=device)
        dist.all_reduce(t_train, op=dist.ReduceOp.SUM)
        tl, tgain, tgrad, tmae, tpsnr, tbias, tn = t_train.tolist()
        train_stats = {
            "train_loss": tl / max(tn, 1),
            "train_gain_loss": tgain / max(tn, 1),
            "train_grad_loss": tgrad / max(tn, 1),
            "train_mae_db": tmae / max(tn, 1),
            "train_psnr_db": tpsnr / max(tn, 1),
            "train_bias_db": tbias / max(tn, 1),
        }

        train_end = time.perf_counter()
        secs_train = train_end - epoch_start
        peak_mem_gb = torch.cuda.max_memory_allocated(device) / 1024**3

        val_stats: dict[str, float] = {}
        secs_val = 0.0
        if (epoch + 1) % args.val_every_epochs == 0 or epoch == args.epochs - 1:
            val_start = time.perf_counter()
            val_stats = run_val(ddp, val_loader, device, amp_dtype, args.huber_beta, world_size,
                                gradient_weight=args.gradient_weight)
            dist.barrier()
            secs_val = time.perf_counter() - val_start
        secs = time.perf_counter() - epoch_start

        if rank == 0:
            ep_entry = {
                "epoch": int(epoch),
                "step": int(step),
                "lr": float(lr),
                **train_stats,
                **val_stats,
                "secs_train": float(secs_train),
                "secs_val": float(secs_val),
                "secs_per_epoch": float(secs),
                "peak_mem_gb": float(peak_mem_gb),
            }
            history.append(ep_entry)
            with open(csv_path, "a", newline="") as f:
                w = csv.writer(f)
                w.writerow([
                    epoch, step, f"{lr:.6e}",
                    f"{train_stats['train_loss']:.4f}",
                    f"{train_stats.get('train_gain_loss', 0):.4f}",
                    f"{train_stats.get('train_grad_loss', 0):.4f}",
                    f"{train_stats['train_mae_db']:.4f}", f"{train_stats['train_psnr_db']:.4f}", f"{train_stats['train_bias_db']:.4f}",
                    f"{val_stats.get('val_loss', ''):.4f}" if val_stats else "",
                    f"{val_stats.get('val_gain_loss', ''):.4f}" if val_stats else "",
                    f"{val_stats.get('val_grad_loss', ''):.4f}" if val_stats else "",
                    f"{val_stats.get('val_mae_db', ''):.4f}" if val_stats else "",
                    f"{val_stats.get('val_psnr_db', ''):.4f}" if val_stats else "",
                    f"{val_stats.get('val_bias_db', ''):.4f}" if val_stats else "",
                    f"{val_stats.get('val_cells_mean', ''):.1f}" if val_stats else "",
                    f"{secs_train:.2f}", f"{secs_val:.2f}", f"{secs:.2f}", f"{peak_mem_gb:.2f}",
                ])
            with open(json_path, "w") as f:
                json.dump({"history": history}, f, indent=2)

            msg = (
                f"[ep {epoch:3d}/{args.epochs}] "
                f"lr={lr:.2e}  "
                f"tr(gain={train_stats.get('train_gain_loss',0):.3f},grad={train_stats.get('train_grad_loss',0):.3f})  "
                f"mae={train_stats['train_mae_db']:.2f}dB  "
                f"psnr={train_stats['train_psnr_db']:.2f}  bias={train_stats['train_bias_db']:+.2f}"
            )
            if val_stats:
                msg += (
                    f"  | val(gain={val_stats.get('val_gain_loss',0):.3f},grad={val_stats.get('val_grad_loss',0):.3f}) "
                    f"mae={val_stats['val_mae_db']:.2f}dB  "
                    f"psnr={val_stats['val_psnr_db']:.2f}  "
                    f"bias={val_stats['val_bias_db']:+.2f}"
                )
            msg += f"  | train={secs_train:.1f}s val={secs_val:.1f}s tot={secs:.1f}s  peak={peak_mem_gb:.1f}GB"
            print(msg, flush=True)

            # best.pt by composite score on val
            if val_stats:
                score = composite_score(val_stats)
                if score > best_score:
                    best_score = score
                    save_ckpt(ckpt_dir / "best.pt", ddp, opt, epoch, step, {**train_stats, **val_stats, "composite_score": score}, vars(args), phase_name=args.phase_name)
                    print(f"  ✓ new best.pt (epoch {epoch}, score {score:.4f}, val_mae_db {val_stats['val_mae_db']:.3f})", flush=True)

            save_ckpt(ckpt_dir / "last.pt", ddp, opt, epoch, step, {**train_stats, **val_stats}, vars(args), phase_name=args.phase_name)

        dist.barrier()

    if rank == 0:
        save_ckpt(ckpt_dir / "final.pt", ddp, opt, args.epochs - 1, step, history[-1] if history else {}, vars(args), phase_name=args.phase_name)
        print("============== V1 training DONE ==============")

    cleanup_ddp()


if __name__ == "__main__":
    main()