| """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__) |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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 |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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}") |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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() |
|
|
