"""Shared training utilities for pretraining and fine-tuning. All training loops import from here. No duplicated optimizer construction, scheduler logic, or checkpoint I/O across pretrain.py and finetune.py. Functions: create_optimizer — AdamW, weight decay on 2D+ params only create_scheduler — Linear warmup + cosine decay to min_lr_fraction save_checkpoint — Atomic write (.tmp then rename) load_checkpoint — Restore with optional fingerprint verification setup_deterministic — Seed all RNGs, deterministic algorithms nan_guard — Per-step finiteness check with debug dump Class: MetricsLogger — Tensorboard + optional wandb """ from __future__ import annotations import logging import math import os import random from pathlib import Path from typing import Any import numpy as np import torch import torch.nn as nn from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR log = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Optimizer / Scheduler # --------------------------------------------------------------------------- def create_optimizer( model: nn.Module, lr: float = 3e-4, betas: tuple[float, float] = (0.9, 0.95), weight_decay: float = 0.1, ) -> AdamW: """AdamW with weight decay on matrices only (not norms, not biases).""" decay, no_decay = [], [] for param in model.parameters(): if not param.requires_grad: continue (decay if param.dim() >= 2 else no_decay).append(param) return AdamW( [{"params": decay, "weight_decay": weight_decay}, {"params": no_decay, "weight_decay": 0.0}], lr=lr, betas=betas, ) def create_scheduler( optimizer: AdamW, warmup_steps: int, total_steps: int, min_lr_fraction: float = 0.1, ) -> LambdaLR: """Linear warmup then cosine decay to min_lr_fraction * peak LR.""" def lr_lambda(step: int) -> float: if step < warmup_steps: return step / max(1, warmup_steps) progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) return min_lr_fraction + (1 - min_lr_fraction) * 0.5 * (1 + math.cos(math.pi * progress)) return LambdaLR(optimizer, lr_lambda) # --------------------------------------------------------------------------- # Checkpointing # --------------------------------------------------------------------------- def save_checkpoint( path: str | Path, model: nn.Module, optimizer: AdamW, scheduler: LambdaLR, step: int, config: dict[str, Any], fingerprint: str = "", ) -> None: """Atomic save: writes to .tmp then os.replace for crash safety.""" path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) tmp = path.with_suffix(".tmp") torch.save({ "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "step": step, "config": config, "fingerprint": fingerprint, }, tmp) os.replace(tmp, path) log.info("Checkpoint saved: %s (step %d)", path, step) def load_checkpoint( path: str | Path, model: nn.Module, optimizer: AdamW | None = None, scheduler: LambdaLR | None = None, expected_fingerprint: str | None = None, strict: bool = True, ) -> dict[str, Any]: """Load checkpoint. Pass optimizer/scheduler=None to skip restoring them. strict=False allows loading pretrained weights into a model with extra parameters (e.g. a fraud head added for fine-tuning). Raises ValueError on fingerprint mismatch. """ ckpt = torch.load(path, map_location="cpu", weights_only=False) if expected_fingerprint is not None: saved = ckpt.get("fingerprint", "") if saved != expected_fingerprint: raise ValueError( f"Fingerprint mismatch: checkpoint='{saved}', expected='{expected_fingerprint}'" ) model.load_state_dict(ckpt["model_state_dict"], strict=strict) if optimizer is not None and "optimizer_state_dict" in ckpt: optimizer.load_state_dict(ckpt["optimizer_state_dict"]) if scheduler is not None and "scheduler_state_dict" in ckpt: scheduler.load_state_dict(ckpt["scheduler_state_dict"]) log.info("Checkpoint loaded: %s (step %s)", path, ckpt.get("step", "?")) return ckpt # --------------------------------------------------------------------------- # Reproducibility # --------------------------------------------------------------------------- def setup_deterministic( seed: int = 42, warn_only: bool = True, cublas_workspace_config: str = ":4096:8", ) -> None: """Seed all RNGs and enable deterministic CUDA kernels. warn_only=True because depthwise Conv1d backward may lack a deterministic CUDA kernel — logs a warning instead of crashing (D19e). """ random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) torch.use_deterministic_algorithms(True, warn_only=warn_only) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False os.environ["CUBLAS_WORKSPACE_CONFIG"] = cublas_workspace_config log.info("Deterministic mode: seed=%d, warn_only=%s", seed, warn_only) # --------------------------------------------------------------------------- # NaN guard # --------------------------------------------------------------------------- class NaNError(RuntimeError): """Raised when loss contains NaN or Inf.""" pass def nan_guard( loss: torch.Tensor, step: int, output_dir: str | Path, model: nn.Module | None = None, ) -> None: """Check loss is finite. Dumps debug state and raises NaNError on failure.""" if torch.isfinite(loss): return output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) debug_path = output_dir / f"nan_debug_step{step}.pt" state: dict[str, Any] = {"step": step, "loss": loss.detach().cpu()} if model is not None: state["model_state_dict"] = { k: v.detach().cpu() for k, v in model.state_dict().items() } torch.save(state, debug_path) log.error("NaN/Inf loss at step %d. Debug state: %s", step, debug_path) raise NaNError(f"Loss is {loss.item()} at step {step}. Debug: {debug_path}") # --------------------------------------------------------------------------- # Metrics logging # --------------------------------------------------------------------------- class MetricsLogger: """Tensorboard writer + optional wandb. Wandb activates on WANDB_API_KEY.""" def __init__( self, log_dir: str | Path, experiment_name: str = "run", use_wandb: str = "auto", wandb_project: str = "lfm2-transactions", ) -> None: from torch.utils.tensorboard import SummaryWriter self.tb = SummaryWriter(log_dir=str(log_dir)) self.wandb_run = None if use_wandb == "auto" and os.environ.get("WANDB_API_KEY"): try: import wandb self.wandb_run = wandb.init( project=wandb_project, name=experiment_name, dir=str(log_dir), ) except ImportError: log.info("wandb not installed, tensorboard only") def log_scalar(self, tag: str, value: float, step: int) -> None: self.tb.add_scalar(tag, value, step) if self.wandb_run is not None: import wandb wandb.log({tag: value}, step=step) def log_per_feature_losses( self, total_loss: float, per_feature_losses: dict[str, float], step: int, ) -> None: self.log_scalar("loss/total", total_loss, step) for name, val in per_feature_losses.items(): self.log_scalar(f"loss/{name}", val, step) def log_config(self, config: dict[str, Any]) -> None: self.tb.add_text("config", str(config), 0) if self.wandb_run is not None: import wandb wandb.config.update(config) def close(self) -> None: self.tb.close() if self.wandb_run is not None: import wandb wandb.finish()