airfRANS-model-exploration/src/airfrans_frontier/training/calibration.py

from __future__ import annotations

import json
import tempfile
import time
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F

from airfrans_frontier.training.config import TrainingConfig


def estimate_forward_flops_per_item(model: nn.Module) -> int:
    total = 0
    for module in model.modules():
        if isinstance(module, nn.Linear):
            total += 2 * module.in_features * module.out_features
            if module.bias is not None:
                total += module.out_features
    return int(total)


def estimate_training_compute(*, steps: int, batch_size: int, forward_flops_per_item: int) -> int:
    return int(steps * batch_size * forward_flops_per_item * 3)


def checkpoint_size_bytes(run_dir: Path, name: str) -> int | None:
    path = run_dir / name
    if not path.is_file():
        return None
    return int(path.stat().st_size)


def gpu_memory_metrics(device: torch.device) -> dict[str, int | None]:
    if device.type != "cuda":
        return {
            "gpu_memory_allocated_mb": None,
            "gpu_memory_reserved_mb": None,
            "gpu_memory_peak_allocated_mb": None,
        }
    index = device.index if device.index is not None else torch.cuda.current_device()
    return {
        "gpu_memory_allocated_mb": int(torch.cuda.memory_allocated(index) // (1024 * 1024)),
        "gpu_memory_reserved_mb": int(torch.cuda.memory_reserved(index) // (1024 * 1024)),
        "gpu_memory_peak_allocated_mb": int(torch.cuda.max_memory_allocated(index) // (1024 * 1024)),
    }


def measure_training_step(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    features: np.ndarray,
    targets: np.ndarray,
    *,
    batch_size: int,
    steps: int,
    device: torch.device,
) -> dict[str, Any]:
    if steps <= 0:
        raise ValueError("Calibration steps must be positive")
    rng = np.random.default_rng(1729)
    model.train()
    if device.type == "cuda":
        torch.cuda.reset_peak_memory_stats(device)
        torch.cuda.synchronize(device)
    started = time.perf_counter()
    last_loss = 0.0
    for _ in range(steps):
        indices = rng.integers(0, features.shape[0], size=batch_size)
        batch_features = torch.from_numpy(np.ascontiguousarray(features[indices], dtype=np.float32)).to(device)
        batch_targets = torch.from_numpy(np.ascontiguousarray(targets[indices], dtype=np.float32)).to(device)
        optimizer.zero_grad(set_to_none=True)
        predictions = model(batch_features)
        loss = F.mse_loss(predictions, batch_targets)
        loss.backward()
        optimizer.step()
        last_loss = float(loss.detach().cpu().item())
    if device.type == "cuda":
        torch.cuda.synchronize(device)
    elapsed = time.perf_counter() - started
    return {
        "calibration_steps": steps,
        "step_time_seconds": elapsed / steps,
        "points_per_sec": steps * batch_size / max(elapsed, 1e-12),
        "last_calibration_loss": last_loss,
        **gpu_memory_metrics(device),
    }


def measure_validation_runtime(
    model: nn.Module,
    features: np.ndarray,
    *,
    batch_size: int,
    device: torch.device,
) -> float:
    model.eval()
    if device.type == "cuda":
        torch.cuda.synchronize(device)
    started = time.perf_counter()
    with torch.no_grad():
        for start in range(0, features.shape[0], batch_size):
            stop = min(start + batch_size, features.shape[0])
            batch_features = torch.from_numpy(np.ascontiguousarray(features[start:stop], dtype=np.float32)).to(device)
            model(batch_features)
    if device.type == "cuda":
        torch.cuda.synchronize(device)
    return time.perf_counter() - started


def measure_checkpoint_size(payload: dict[str, Any]) -> int:
    with tempfile.TemporaryDirectory() as tmp:
        path = Path(tmp) / "checkpoint.pt"
        torch.save(payload, path)
        return int(path.stat().st_size)


def write_calibration_report(path: str | Path, data: dict[str, Any]) -> Path:
    report_path = Path(path)
    report_path.parent.mkdir(parents=True, exist_ok=True)
    report_path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n")
    return report_path


def static_calibration_fields(config: TrainingConfig, model: nn.Module) -> dict[str, Any]:
    forward_flops = estimate_forward_flops_per_item(model)
    return {
        "estimated_forward_flops_per_item": forward_flops,
        "estimated_train_flops": estimate_training_compute(
            steps=config.optim.steps,
            batch_size=config.data.batch_size,
            forward_flops_per_item=forward_flops,
        ),
    }