giant/giant/config.py

import random
import subprocess
import sys
import tomllib
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import torch

DEFAULT_CONFIG: dict = {
    "train": {
        "mode": "flow",
        "epochs": 100,
        "batch_size": 4096,
        "lr": 3e-4,
        "weight_decay": 0.01,  # AdamW default — exposed so it can be tuned
        "ema_decay": 0.9999,  # EMA of model weights for sampling; 0 disables
        # per-epoch val loss (not the marginal/KL validate_every pass) is
        # capped to this many batches; 0 = full val set every epoch
        "max_val_batches": 200,
        "val_fraction": 0.1,
        "num_workers": 4,
        "seed": 0,
        "validate_every": 10,
        "validate_steps": 10,
        "warmup_epochs": 5,
        "lambda_nsec": 0.1,
        "lambda_s2": 1.0,
        # WGAN-GP-only knobs (mode == "wgan"; ignored by flow/ddpm). n_critic:
        # critic updates per generator update. gp_weight: gradient-penalty
        # coefficient (Gulrajani et al. 2017). critic_lr: 0.0 means "use
        # `lr`" — not None, since save_config's TOML writer has no null
        # literal to round-trip.
        "n_critic": 5,
        "gp_weight": 10.0,
        "critic_lr": 0.0,
    },
    "model": {
        "hidden_dim": 256,
        "n_blocks": 6,
        "emb_dim": 16,
        "dropout": 0.1,
        # WGAN generator noise-vector width (mode == "wgan" only).
        "noise_dim": 64,
        # "physical" conditions on material/particle physical properties via
        # a small MLP (giant.model.network.ConditionEncoder); "embedding"
        # keeps the original learned pdg/material embedding tables — kept
        # available as the generalization-comparison baseline. Checkpoints
        # from before this option existed have no "conditioning" key and
        # load as "embedding" (see giant.model.network.build_models).
        "conditioning": "physical",
        "router": {
            "enabled": False,
            "type": "energy",  # selects the Router impl from ROUTER_REGISTRY
            "n_experts": 4,
            "expert_hidden_dim": 128,
            "expert_n_blocks": 3,
            "temperature": 0.5,  # energy/pdg-router kwarg
            "learn_centers": True,  # energy/pdg-router kwarg
            "lambda_balance": 0.0,  # optional load-balance aux loss weight
            "emb_dim": 8,  # process/pdg-router kwarg: own pdg(/mat) embedding width
            "hidden_dim": 64,  # process-router kwarg: its classifier's hidden width
            "lambda_proc": 0.0,  # process-router kwarg: supervised process-CE weight
            # (0.0 still trains a working router — the gate gets gradient
            # through the downstream flow loss like EnergyRouter's centers —
            # but only lambda_proc > 0 grounds it in the true `process` label)
            # type = "composed" routes on multiple axes at once (e.g. energy x
            # pdg), each with its own expert count/hyperparameters. Axes are
            # NOT in these defaults (there's no meaningful default axis list)
            # — set them as flat axis{i}_{field} keys instead of "n_experts",
            # e.g. axis0_type = "energy", axis0_n_experts = 4, axis1_type =
            # "pdg", axis1_n_experts = 3, axis1_emb_dim = 8. See
            # giant.model.network._parse_composed_axes / `--router-axis`.
        },
    },
}


def git_hash() -> str:
    try:
        return (
            subprocess.check_output(
                ["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
            )
            .decode()
            .strip()
        )
    except Exception:
        return "unknown"


def auto_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


# Calibration point for estimate_batch_size(training=True): hidden_dim=1024,
# n_blocks=6, batch_size=29696 measured at ~7683 MiB VRAM (post-Phase-2
# architecture, including the Stage-2 secondary decoder and n_sec head).
# Activation memory is assumed to scale linearly with
# batch_size * hidden_dim * n_blocks (the ResBlock stack dominates), so this
# is a rough estimate rather than a guaranteed bound.
_REF_BYTES = 7683 * 1024**2
_REF_BATCH_SIZE = 29696
_REF_HIDDEN_DIM = 1024
_REF_N_BLOCKS = 6

# Calibration point for estimate_batch_size(training=False): inference has no
# backward graph or optimizer state, so its memory footprint is much smaller
# per sample. hidden_dim=1024, n_blocks=8, batch_size=65536 measured at ~2037
# MiB VRAM.
_REF_BYTES_PREDICT = 2037 * 1024**2
_REF_BATCH_SIZE_PREDICT = 65536
_REF_HIDDEN_DIM_PREDICT = 1024
_REF_N_BLOCKS_PREDICT = 8


def estimate_batch_size(
    hidden_dim: int,
    n_blocks: int,
    device: torch.device,
    safety_factor: float = 0.8,
    min_batch_size: int = 1024,
    training: bool = True,
) -> int:
    """Estimate a batch size that fits in the free memory on `device`.

    Only supported on CUDA devices, which expose a free/total memory query;
    other backends (cpu, mps) raise ValueError. Pass `training=False` for
    inference (e.g. `predict`), which uses a much lower per-sample memory
    calibration since there's no backward graph or optimizer state.
    """
    if device.type != "cuda":
        raise ValueError(
            f"--batch-size auto is only supported on cuda devices, got {device.type!r}"
        )
    device_index = (
        device.index if device.index is not None else torch.cuda.current_device()
    )
    free_bytes, _total_bytes = torch.cuda.mem_get_info(device_index)
    if training:
        ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
            _REF_BYTES,
            _REF_BATCH_SIZE,
            _REF_HIDDEN_DIM,
            _REF_N_BLOCKS,
        )
    else:
        ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
            _REF_BYTES_PREDICT,
            _REF_BATCH_SIZE_PREDICT,
            _REF_HIDDEN_DIM_PREDICT,
            _REF_N_BLOCKS_PREDICT,
        )
    bytes_per_unit = ref_bytes / (ref_batch_size * ref_hidden_dim * ref_n_blocks)
    bytes_per_sample = bytes_per_unit * hidden_dim * n_blocks
    batch_size = int(free_bytes * safety_factor / bytes_per_sample)
    batch_size = max(min_batch_size, (batch_size // 1024) * 1024)
    return batch_size


def load_toml(path: Path) -> dict:
    with open(path, "rb") as f:
        return tomllib.load(f)


def warn_if_git_hash_mismatch(file_cfg: dict, config_path: Path) -> None:
    """Warn (don't fail) if a config.toml's [meta].git_hash predates the current checkout.

    A config saved by a previous run may have been produced by code that has
    since changed, so its hyperparameters might not mean what they used to —
    surface that as a heads-up rather than blocking the rerun.
    """
    file_hash = file_cfg.get("meta", {}).get("git_hash")
    current_hash = git_hash()
    if not file_hash or file_hash == "unknown" or current_hash == "unknown":
        return
    if file_hash != current_hash:
        print(
            f"warning: {config_path} was generated at git commit {file_hash}, "
            f"but the current checkout is at {current_hash} — hyperparameters "
            "may not match the code that originally produced this config",
            file=sys.stderr,
        )


def warn_if_checkpoint_config_mismatch(ckpt_path: str | Path) -> None:
    """Look for a config.toml next to a checkpoint and warn on a git_hash mismatch.

    Training writes config.toml into the same out_dir as its checkpoints, so a
    checkpoint loaded later (for `predict` or `giant.analysis`) can be
    cross-checked the same way `--config` loading is, without the caller having
    to pass the toml path explicitly. Silently does nothing if no config.toml
    is found alongside the checkpoint.
    """
    config_path = Path(ckpt_path).parent / "config.toml"
    if not config_path.exists():
        return
    warn_if_git_hash_mismatch(load_toml(config_path), config_path)


def merge_cli_overrides(
    defaults: dict,
    config_path: Path | None,
    train_overrides: dict,
    model_overrides: dict,
) -> dict:
    """Resolve config as defaults -> TOML file -> explicit CLI flags.

    `model.router` is deep-merged one level (rather than replaced wholesale)
    at each stage, so a TOML file or CLI flag only overriding e.g.
    `router.enabled` doesn't drop the rest of the router defaults.
    """
    cfg = {"train": dict(defaults["train"]), "model": dict(defaults["model"])}
    cfg["model"]["router"] = dict(defaults["model"]["router"])
    if config_path is not None:
        file_cfg = load_toml(config_path)
        cfg["train"].update(file_cfg.get("train", {}))
        file_model = dict(file_cfg.get("model", {}))
        file_router = file_model.pop("router", None)
        cfg["model"].update(file_model)
        if file_router:
            cfg["model"]["router"].update(file_router)
        warn_if_git_hash_mismatch(file_cfg, config_path)
    model_overrides = dict(model_overrides)
    router_overrides = model_overrides.pop("router", None)
    cfg["train"].update(train_overrides)
    cfg["model"].update(model_overrides)
    if router_overrides:
        cfg["model"]["router"].update(router_overrides)
    return cfg


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _toml_value(v) -> str:
    if isinstance(v, bool):
        return "true" if v else "false"
    if isinstance(v, str):
        return repr(v)
    return str(v)


def save_config(cfg: dict, out_dir: Path, meta: dict) -> None:
    lines = []
    # One-level-nested dict values (e.g. model.router) are rendered as their
    # own [section.subsection] table after the parent section, since TOML
    # doesn't accept a bare dict as a `key = value` scalar line.
    nested_sections: list[tuple[str, dict]] = []
    for section, values in cfg.items():
        lines.append(f"[{section}]")
        for k, v in values.items():
            if isinstance(v, dict):
                nested_sections.append((f"{section}.{k}", v))
                continue
            lines.append(f"{k:<14} = {_toml_value(v)}")
        lines.append("")

    for name, values in nested_sections:
        lines.append(f"[{name}]")
        for k, v in values.items():
            lines.append(f"{k:<14} = {_toml_value(v)}")
        lines.append("")

    lines.append("[meta]")
    for k, v in meta.items():
        lines.append(f"{k:<14} = {_toml_value(v)}")

    (out_dir / "config.toml").write_text("\n".join(lines))


def build_run_meta(
    data: Path,
    seed: int,
    n_pdg_codes: int,
    n_materials: int,
    n_train_events: int,
    n_val_events: int,
    n_train_steps: int,
) -> dict:
    return {
        "git_hash": git_hash(),
        "seed": seed,
        "timestamp_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "python_version": sys.version.split()[0],
        "torch_version": torch.__version__,
        "command": " ".join(sys.argv),
        "data_path": str(data),
        "n_pdg_codes": n_pdg_codes,
        "n_materials": n_materials,
        "n_train_events": n_train_events,
        "n_val_events": n_val_events,
        "n_train_steps": n_train_steps,
    }