"""Phase 2: fine-tune ``EncoderForClassification`` -> legitimate / phishing / spam. The instructor's "experimentation dials" (number of layers, epochs, loss function, plus learning rate, batch size, frozen layers, synthetic data) are first-class CLI flags, and every run leaves ONE JSON line in reports/experiments.jsonl so the experimentation table (``finetune.py list``) is generated from evidence rather than typed up afterwards. Subcommands:: train [--run-name R] [--init pretrained:|scratch] [--n-layers N] [--freeze-layers K] [--loss cross_entropy|weighted_cross_entropy|focal] [--focal-gamma G] [--lr LR] [--batch-size B] [--micro-batch M] [--epochs E] [--max-steps S] [--eval-every S] [--log-every S] [--synthetic on|off] [--limit-rows N] [--num-labels 2|3] [--cpu] [--resume] [--overwrite] plot --run-name R [--out PNG] reports/finetune_curves.png list experiments.jsonl as a markdown table ML decisions, in brief: * **Three classes, partial labels (RUN_LOG entry 6).** ``finetune.num_labels`` (3) sizes the classification head; config ``model:`` is never touched because the live pretraining run refuses to resume if its model snapshot changes. Every row carries ``label3`` in {0 legitimate, 1 phishing, 2 spam, -1 unwanted-but-unsplit} and ``label_set`` in {legit, phish, spam, unwanted}; the binary ``label`` column (1 = phishing OR spam) stays. The TREC-05/07 and CEAS-08 positives never separated spam from phishing, so they train on the probability MASS of {phishing, spam}: per row an allowed-class mask ``[B, 3]`` built from label_set, and the loss is the partial-label negative log-likelihood ``logsumexp(all logits) - logsumexp(logits[allowed])`` = -log sum_{c in allowed} p_c (Jin & Ghahramani 2002; Cour et al. 2011). For a single allowed class this IS ``F.cross_entropy``; for an "unwanted" row it is ``-log(p_phish + p_spam)``, which still shapes the legitimate-vs-unwanted decision without asserting which kind it was. With ``num_labels: 2`` the old binary behaviour is reproduced exactly (label3 ignored; the mask is the one-hot label) so the binary baseline runs remain reproducible. * **Two metric sets.** (a) BINARY roll-up on all val rows: P(unwanted) = p1 + p2 vs the binary label at threshold 0.5 -> accuracy/precision/recall/F1/ROC-AUC/PR-AUC with the same definitions as the binary runs (so results stay comparable). (b) THREE-WAY on rows with label3 in {0, 1, 2}: argmax over 3 classes -> 3x3 confusion, per-class P/R/F1, macro-F1 over all three classes, plus a phishing-vs-spam confusion restricted to true unwanted rows (argmax over {phishing, spam} only: "given it is unwanted, which kind?"). Display names are {0: Legitimate, 1: Phishing, 2: Spam}; the roll-up is "Unwanted (phishing or spam)". Selection / early stopping follow ``finetune.early_stopping.metric``: ``val_f1_unwanted`` (default; = today's binary F1) or ``val_macro_f1_3way``. * **Data path.** Train = finetune_train.parquet (+ optional synthetic rows), selection = finetune_val.parquet. test_holdout.parquet (TREC-06) is refused by file name AND by the ``source`` column of anything loaded, so a misconfigured path cannot leak the holdout. Texts are encoded once with ``tokenizer.encode_batch`` (truncate to model.max_seq_len, pad to a multiple of finetune.dynamic_padding_multiple), stored trimmed, and re-padded per batch with the identical rule; the first batch is checked tensor-for-tensor against a direct ``encode_batch`` call so the app / evaluate.py path (which calls ``encode_batch`` per request) sees exactly the inputs the classifier was trained on. No sequence packing: the bool key-padding mask takes SDPA's memory-efficient kernel (research_model-math.md sec. 5). * **DataLoader.** Every batch goes through ``torch.utils.data.DataLoader`` (``make_loader``): a row-index ``Dataset``, the seeded length-bucketed batch order as ``batch_sampler`` and a collate function that pads per micro-batch. Shuffling, batching and preprocessing therefore live in the standard PyTorch data pipeline the course reference uses. * **Length bucketing.** Each epoch: seeded permutation -> mega-batches of 50 x batch_size -> sort by token length inside each -> cut into batches -> shuffle the batch order. Padding waste drops from ~50% (random batches over a 20..512 length spread) to a few percent while the batch composition stays random at the mega-batch level. The order is a pure function of (seed, epoch) so ``--resume`` can skip exactly the batches already consumed. * **Init.** ``pretrained:`` loads ``encoder.*`` from a Phase 1 checkpoint with strict key matching (``EncoderForClassification.from_pretrained_encoder``); ``scratch`` is the ablation that shows what MLM pretraining buys. ``--n-layers N`` keeps the FIRST N encoder layers (``encoder.layers.0..N-1``) of the pretrained stack and drops the rest; the final LayerNorm is kept although it was trained on layer-24 output (fine-tuning re-fits it). ``--freeze-layers K`` freezes embeddings + layers 0..K-1: cheaper steps (autograd stops at layer K) and less catastrophic forgetting on a small train set. * **Loss dial.** All three variants are defined over the partial-label NLL above (``nll``): cross_entropy = mean(nll); weighted_cross_entropy = sum(w_r nll_r) / sum(w_r) with inverse-frequency class weights w_c = N / (C * n_c) from the fully-labelled rows of the train split actually used and w_r = mean of w over the row's allowed classes (so an "unwanted" row gets mean(w_phish, w_spam)) -- the sum-of-weights normalisation is ``F.cross_entropy(weight=)``'s, so the two agree on fully-labelled batches; focal loss (Lin et al. 2017) ``(1 - exp(-nll))^gamma * nll`` = ``-(1 - p_t)^gamma log p_t`` without alpha, where p_t is the allowed-mass. Loss is computed in fp32 on bf16-autocast logits. * **Optimizer / schedule.** The same ``build_optimizer`` (AdamW, no weight decay on biases and LayerNorm, fused) and the same fwd/bwd -> clip 1.0 -> step -> zero_grad(set_to_none) pattern the smoke test measured. Linear warmup over finetune.warmup_fraction of the total optimizer steps, then linear decay to 0 (BERT fine-tuning schedule). Gradient accumulation: a logical batch of ``--batch-size`` rows is split into ``--micro-batch`` chunks (default 16, the smoke test's 13.65 GiB point for 16 x 512 with the heavier MLM head); each chunk's loss is scaled by its row share so the step gradient is the mean over the logical batch even when the last batch of an epoch is short. * **Selection / early stopping.** Full validation every ``--eval-every`` steps and at every epoch end; best/ is rewritten whenever the selection metric (val F1 of the unwanted roll-up at threshold 0.5 by default, or 3-way macro-F1) improves; finetune.early_stopping.patience evaluations without improvement stop the run. The F1-optimal threshold on val is recorded next to the 0.5 decision threshold so evaluate.py / the app can show both; selection itself uses 0.5 to stay comparable across runs. best/config.json carries num_labels, the label names, the selection metric, both metric sets and the thresholds. * **--cpu (dev only).** Runs the identical code path on the CPU with whatever config it is given (tiny encoder + fixture parquet files) so the whole script can be exercised while the GPU is busy; autocast and the VRAM guard are disabled and torch.cuda is never initialised: device info, allocator stats and the CUDA RNG snapshot are gated on the chosen device, and CUDA_VISIBLE_DEVICES is cleared for the process before the runtime is first queried (torch 2.14's ``AdamW.step`` health check would otherwise create a context on any visible GPU). If a device is still visible after that, the run refuses to start. The default path still requires CUDA. * **Resume.** last/ holds model + trainer_state.pt (optimizer, scheduler, step, epoch, batch-in-epoch, Python/NumPy/torch CPU/torch CUDA RNG states, best score, patience counter, cumulative seconds, the resolved dials). ``--resume`` restores all of it and skips the already-consumed batches of the current epoch. Dials that shape the optimizer state or the batch order (init, n_layers, freeze_layers, batch_size, seed) are refused on resume; ``--lr`` is re-applied to the restored optimizer/scheduler so the record never claims an lr that was not used; metrics.csv / epochs.csv rows logged after the checkpoint (a crash between checkpoints) are dropped before appending; a run already at its --max-steps / --epochs is reported as "nothing to resume" without training an extra step. SIGINT/SIGTERM save last/ and exit 0. * **VRAM guard.** ``PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`` before CUDA starts; after every optimizer step the reserved peak is compared with config.vram.hard_cap_gb (GiB, the smoke test's convention) and the run aborts (after saving last/) on a breach. """ from __future__ import annotations import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import argparse # noqa: E402 import contextlib # noqa: E402 import csv # noqa: E402 import json # noqa: E402 import logging # noqa: E402 import math # noqa: E402 import random # noqa: E402 import signal # noqa: E402 import sys # noqa: E402 import time # noqa: E402 import zlib # noqa: E402 from dataclasses import asdict, dataclass, field, fields # noqa: E402 from pathlib import Path # noqa: E402 from typing import Any # noqa: E402 sys.path.insert(0, str(Path(__file__).parent)) import numpy as np # noqa: E402 import pyarrow.parquet as pq # noqa: E402 import torch # noqa: E402 import torch.nn.functional as F # noqa: E402 from safetensors import safe_open # noqa: E402 from safetensors.torch import load_file, save_file # noqa: E402 from torch import Tensor # noqa: E402 from torch.utils.data import DataLoader, Dataset # noqa: E402 import common # noqa: E402 from model import EncoderConfig, EncoderForClassification, count_parameters # noqa: E402 from smoke_test import build_optimizer # noqa: E402 (the measured optimizer code path) from tokenizer import encode_batch, load_tokenizer # noqa: E402 LOG = logging.getLogger("finetune") GIB = float(2**30) TRAIN_FILE = "finetune_train.parquet" VAL_FILE = "finetune_val.parquet" SYNTHETIC_FILE = "synthetic.parquet" HOLDOUT_FILE = "test_holdout.parquet" LOSSES = ("cross_entropy", "weighted_cross_entropy", "focal") MEGA_BATCH_FACTOR = 50 ENCODE_CHUNK = 2048 DECISION_THRESHOLD = 0.5 METRIC_KEYS = ("accuracy", "precision", "recall", "f1", "roc_auc", "pr_auc") RUN_LOG_MIN_STEPS = 100 # Dials whose change on --resume would leave the restored optimizer state / batch order / head shape inconsistent. RESUME_FIXED_DIALS = ("init", "n_layers", "freeze_layers", "batch_size", "seed", "num_labels") # The only implemented values of the config-only dials; anything else in config.yaml is an error, not a no-op. IMPLEMENTED = {"schedule": "linear_decay", "pooling": "cls"} CONFIG_PATH = {"schedule": "finetune.schedule", "pooling": "finetune.pooling", "selection_metric": "finetune.early_stopping.metric", "num_labels": "finetune.num_labels"} # ---- three-way label contract (RUN_LOG entry 6) ----------------------------------------------- LABEL_NAMES = {0: "Legitimate", 1: "Phishing", 2: "Spam"} UNWANTED_NAME = "Unwanted (phishing or spam)" LABEL_SETS: dict[str, tuple[int, ...]] = {"legit": (0,), "phish": (1,), "spam": (2,), "unwanted": (1, 2)} LABEL3_OF_SET = {"legit": 0, "phish": 1, "spam": 2, "unwanted": -1} NUM_LABELS_ALLOWED = (2, 3) # Selection metrics. val_f1_unwanted is today's binary F1 (p_unwanted = p1 + p2 at 0.5) so results stay # comparable with the binary runs; val_f1_phishing is the pre-entry-6 name of the same number and is # accepted so the recorded binary baselines resolve. val_macro_f1_3way needs num_labels == 3. SELECTION_METRICS = ("val_f1_unwanted", "val_macro_f1_3way", "val_f1_phishing") SELECTION_KEY = {"val_f1_unwanted": "f1", "val_f1_phishing": "f1", "val_macro_f1_3way": "macro_f1_3way"} CLASS_SLUGS = ("legitimate", "phishing", "spam") CM3_KEYS = tuple(f"cm3_{a[0]}{b[0]}" for a in CLASS_SLUGS for b in CLASS_SLUGS) # cm3_ll ... cm3_ss (true, pred) PVS_KEYS = ("pvs_phish_as_phish", "pvs_phish_as_spam", "pvs_spam_as_phish", "pvs_spam_as_spam") THREE_WAY_KEYS = ("n_3way", "accuracy_3way", "macro_f1_3way", *(f"{m}_{s}" for s in CLASS_SLUGS for m in ("precision", "recall", "f1")), *CM3_KEYS, "n_unwanted_3way", *PVS_KEYS) # validated categorical palette (dataviz skill): blue, orange, aqua, yellow PALETTE = ("#2a78d6", "#eb6834", "#1baf7a", "#eda100") # ---------------------------------------------------------------------------------- dials @dataclass class TrainArgs: """Every dial of a run, resolved from CLI > saved state (on --resume) > config.yaml. Stored verbatim in trainer_state.pt and experiments.jsonl so a run is reproducible from its record alone. """ run_name: str = "finetune" init: str = "pretrained" # "pretrained" | "pretrained:" | "scratch" n_layers: int = 24 freeze_layers: int = 0 loss: str = "cross_entropy" focal_gamma: float = 2.0 lr: float = 3e-5 batch_size: int = 32 micro_batch: int = 16 epochs: int = 3 max_steps: int | None = None eval_every: int = 500 log_every: int = 20 eval_batch_size: int = 64 warmup_fraction: float = 0.06 weight_decay: float = 0.01 betas: tuple[float, float] = (0.9, 0.98) eps: float = 1e-6 grad_clip: float = 1.0 patience: int = 2 schedule: str = "linear_decay" # config-only: finetune.schedule selection_metric: str = "val_f1_unwanted" # config: finetune.early_stopping.metric (val_f1_unwanted | val_macro_f1_3way) pooling: str = "cls" # config-only: finetune.pooling (EncoderForClassification reads token 0) num_labels: int = 3 # finetune.num_labels (NOT model.num_labels): 3 = legit/phish/spam, 2 = binary baseline synthetic: str = "off" limit_rows: int | None = None seed: int = 42 cpu: bool = False # dev only: run the whole script on the CPU @classmethod def resolve(cls, ns: argparse.Namespace, cfg: dict, saved: dict | None) -> "TrainArgs": fcfg, mcfg = cfg["finetune"], cfg["model"] pcfg = cfg.get("pretrain", {}) base = cls( n_layers=int(mcfg["n_layers"]), # Absent key -> 2 so a config.yaml predating entry 6 still reproduces the binary baseline. num_labels=int(fcfg.get("num_labels", 2)), freeze_layers=int(fcfg.get("freeze_encoder_layers", 0)), loss=str(fcfg.get("loss", "cross_entropy")), focal_gamma=float(fcfg.get("focal_gamma", 2.0)), lr=float(fcfg["lr"]), batch_size=int(fcfg["batch_size"]), micro_batch=int(fcfg.get("micro_batch", 16)), epochs=int(fcfg["epochs"]), eval_every=int(fcfg.get("eval_every_steps", fcfg.get("checkpoint_every_steps", 500))), log_every=int(fcfg.get("log_every_steps", pcfg.get("log_every_steps", 20))), eval_batch_size=int(fcfg.get("eval_batch_size", 64)), warmup_fraction=float(fcfg["warmup_fraction"]), weight_decay=float(fcfg.get("weight_decay", pcfg.get("weight_decay", 0.01))), betas=tuple(fcfg.get("betas", pcfg.get("betas", (0.9, 0.98)))), eps=float(fcfg.get("eps", pcfg.get("eps", 1e-6))), grad_clip=float(fcfg.get("grad_clip", pcfg.get("grad_clip", 1.0))), patience=int(fcfg["early_stopping"]["patience"]), schedule=str(fcfg.get("schedule", IMPLEMENTED["schedule"])), selection_metric=str(fcfg["early_stopping"].get("metric", "val_f1_unwanted")), pooling=str(fcfg.get("pooling", IMPLEMENTED["pooling"])), synthetic="on" if cfg["data"].get("synthetic", {}).get("enabled") else "off", seed=int(cfg["seed"]), ) if saved: for k, v in saved.items(): if k in {f.name for f in fields(cls)}: setattr(base, k, tuple(v) if k == "betas" else v) for k, v in vars(ns).items(): if k in {f.name for f in fields(cls)} and v is not None and not (k == "cpu" and v is False): setattr(base, k, v) if getattr(ns, "micro_batch", None) is None: # A smaller --batch-size without an explicit --micro-batch means "one micro-batch per step". base.micro_batch = min(base.micro_batch, base.batch_size) for key, allowed in IMPLEMENTED.items(): if getattr(base, key) != allowed: raise ValueError(f"config {CONFIG_PATH[key]}={getattr(base, key)!r} is not implemented; only {allowed!r} is " f"(editing it would otherwise silently change nothing)") if base.num_labels not in NUM_LABELS_ALLOWED: raise ValueError(f"{CONFIG_PATH['num_labels']} must be one of {NUM_LABELS_ALLOWED}, got {base.num_labels!r}") if base.selection_metric not in SELECTION_METRICS: raise ValueError(f"{CONFIG_PATH['selection_metric']}={base.selection_metric!r} is not implemented; " f"one of {SELECTION_METRICS[:2]} (val_f1_phishing is the legacy alias of val_f1_unwanted)") if base.selection_metric == "val_macro_f1_3way" and base.num_labels != 3: raise ValueError("selection metric val_macro_f1_3way needs finetune.num_labels: 3") if base.loss not in LOSSES: raise ValueError(f"--loss must be one of {LOSSES}, got {base.loss!r}") if base.batch_size % base.micro_batch: raise ValueError(f"--batch-size {base.batch_size} must be a multiple of --micro-batch {base.micro_batch}") if not 0 <= base.freeze_layers <= base.n_layers: raise ValueError(f"--freeze-layers {base.freeze_layers} outside [0, {base.n_layers}]") return base def to_record(self) -> dict[str, Any]: d = asdict(self) d["betas"] = list(self.betas) return d # ----------------------------------------------------------------------------------- data @dataclass class Labelled: """The label columns of one parquet split after validation against the three-way contract.""" texts: list[str] labels: np.ndarray # int64 [n] binary: 1 = phishing OR spam (unchanged from the binary build) label3: np.ndarray # int64 [n] in {0, 1, 2, -1}; with num_labels == 2 this is a copy of ``labels`` allowed: np.ndarray # bool [n, C] allowed-class mask built from label_set (one-hot for full labels) sources: list[str] def __len__(self) -> int: return len(self.texts) @dataclass class EncodedSplit: """Token ids for one split, truncated to max_len and stored trimmed (flat int32 + offsets).""" name: str flat: np.ndarray # int32 concatenation of every row's ids offsets: np.ndarray # int64 [n + 1] labels: np.ndarray # int64 [n] binary roll-up target (metrics set a) label3: np.ndarray # int64 [n] three-way target, -1 = unwanted-but-unsplit (metrics set b uses >= 0) allowed: np.ndarray # bool [n, C] partial-label mask (the loss target) sources: list[str] texts: list[str] # kept for the encode_batch consistency check only def __len__(self) -> int: return int(self.labels.shape[0]) @property def lengths(self) -> np.ndarray: return np.diff(self.offsets) @property def num_labels(self) -> int: return int(self.allowed.shape[1]) def row(self, i: int) -> np.ndarray: return self.flat[self.offsets[i]:self.offsets[i + 1]] def collate(self, idx: np.ndarray, multiple: int, device: torch.device) -> tuple[Tensor, Tensor, Tensor, Tensor]: """Pad to the longest row in the batch rounded up to ``multiple`` (encode_batch's rule). Returns (ids int64 [B, L], key-padding mask bool [B, L], binary labels int64 [B], allowed-class mask bool [B, C]). """ lens = self.lengths[idx] L = int(lens.max()) if multiple: L = int(math.ceil(L / multiple) * multiple) ids = np.zeros((len(idx), L), dtype=np.int64) mask = np.zeros((len(idx), L), dtype=np.bool_) for r, i in enumerate(idx): n = int(lens[r]) ids[r, :n] = self.row(int(i)) mask[r, :n] = True return (torch.from_numpy(ids).to(device, non_blocking=True), torch.from_numpy(mask).to(device, non_blocking=True), torch.from_numpy(self.labels[idx]).to(device, non_blocking=True), torch.from_numpy(self.allowed[idx]).to(device, non_blocking=True)) def allowed_mask(label3: np.ndarray, label_set: list[str], num_labels: int) -> np.ndarray: """Per-row allowed-class mask ``[n, C]`` from label_set (3 classes) or the binary label (2 classes).""" out = np.zeros((len(label3), num_labels), dtype=np.bool_) if num_labels == 2: out[np.arange(len(label3)), label3] = True return out for r, s in enumerate(label_set): out[r, list(LABEL_SETS[s])] = True return out def _assert_not_holdout(path: Path, holdout_source: str) -> None: """Hard rule: TREC-06 is touched only by evaluate.py. Refuse by file name and by source value.""" if path.name == HOLDOUT_FILE: raise RuntimeError(f"refusing to load {path.name}: the test holdout is only read by evaluate.py") names = set(pq.read_schema(path).names) if "source" in names: values = {str(v) for v in pq.read_table(path, columns=["source"]).column("source").unique().to_pylist()} norm = holdout_source.lower().replace("_", "-") hits = {v for v in values if v.strip().lower().replace("_", "-") == norm} if hits: raise RuntimeError(f"{path.name} contains holdout source {hits}: test-set leakage") def read_split(path: Path, holdout_source: str, limit: int | None, seed: int, num_labels: int, force_source: str | None = None) -> Labelled: """Read text + labels of one split and validate the three-way contract. With ``num_labels == 3`` the ``label3`` / ``label_set`` columns are REQUIRED (a split built before RUN_LOG entry 6 is an error, not a silent all-"unwanted" run) and checked for internal consistency (label_set <-> label3 <-> binary label). With ``num_labels == 2`` they are ignored and the binary label is the target, reproducing the binary baseline. """ _assert_not_holdout(path, holdout_source) names = set(pq.read_schema(path).names) three = num_labels == 3 if three and not {"label3", "label_set"} <= names: raise ValueError(f"{path.name}: finetune.num_labels is 3 but the columns label3/label_set are missing " f"(have {sorted(names)}); rebuild the split with the three-way data_prep") cols = ["text", "label"] + (["source"] if "source" in names else []) + (["label3", "label_set"] if three else []) table = pq.read_table(path, columns=cols) if limit is not None and table.num_rows > limit: # Seeded random subsample (not a head slice) so a dev run still sees every source. idx = np.sort(np.random.default_rng([seed, zlib.crc32(path.name.encode())]).choice(table.num_rows, size=limit, replace=False)) table = table.take(idx) texts = [t or "" for t in table.column("text").to_pylist()] labels = np.asarray(table.column("label").to_pylist(), dtype=np.int64) if not set(np.unique(labels).tolist()) <= {0, 1}: raise ValueError(f"{path.name}: labels must be 0/1, got {np.unique(labels).tolist()}") if force_source is not None: sources = [force_source] * len(texts) else: sources = [str(s) for s in table.column("source").to_pylist()] if "source" in names else ["?"] * len(texts) if not three: return Labelled(texts, labels, labels.copy(), allowed_mask(labels, [], 2), sources) label3 = np.asarray(table.column("label3").to_pylist(), dtype=np.int64) label_set = [str(s) for s in table.column("label_set").to_pylist()] bad_sets = sorted(set(label_set) - set(LABEL_SETS)) if bad_sets: raise ValueError(f"{path.name}: label_set values outside {sorted(LABEL_SETS)}: {bad_sets}") expected3 = np.fromiter((LABEL3_OF_SET[s] for s in label_set), dtype=np.int64, count=len(label_set)) if not np.array_equal(expected3, label3): n_bad = int((expected3 != label3).sum()) raise ValueError(f"{path.name}: {n_bad} row(s) where label3 disagrees with label_set " f"(contract: legit->0, phish->1, spam->2, unwanted->-1)") expected_bin = (label3 != 0).astype(np.int64) if not np.array_equal(expected_bin, labels): raise ValueError(f"{path.name}: {int((expected_bin != labels).sum())} row(s) where the binary label disagrees " f"with label3 (label must be 1 for phish/spam/unwanted, 0 for legit)") return Labelled(texts, labels, label3, allowed_mask(label3, label_set, 3), sources) def concat_labelled(a: Labelled, b: Labelled) -> Labelled: return Labelled(a.texts + b.texts, np.concatenate([a.labels, b.labels]), np.concatenate([a.label3, b.label3]), np.concatenate([a.allowed, b.allowed]), a.sources + b.sources) def label_counts(label3: np.ndarray, num_labels: int) -> dict[str, int]: """Human-readable class counts for logs and records (keys: display names + 'unwanted').""" if num_labels == 2: return {LABEL_NAMES[0]: int((label3 == 0).sum()), UNWANTED_NAME: int((label3 == 1).sum())} out = {LABEL_NAMES[c]: int((label3 == c).sum()) for c in range(3)} out["unwanted (unsplit)"] = int((label3 == -1).sum()) return out def encode_split(name: str, tok, data: Labelled, max_len: int, multiple: int) -> EncodedSplit: """Encode once with the shared ``encode_batch`` contract, then store rows trimmed of padding.""" t0 = common.timer() texts, labels = data.texts, data.labels rows: list[np.ndarray] = [] for i in range(0, len(texts), ENCODE_CHUNK): ids, mask = encode_batch(tok, texts[i:i + ENCODE_CHUNK], max_len, pad_to_multiple_of=multiple) n = mask.sum(dim=1).numpy() ids_np = ids.numpy() rows.extend(ids_np[r, :n[r]].astype(np.int32) for r in range(ids_np.shape[0])) lengths = np.fromiter((len(r) for r in rows), dtype=np.int64, count=len(rows)) offsets = np.zeros(len(rows) + 1, dtype=np.int64) np.cumsum(lengths, out=offsets[1:]) flat = np.concatenate(rows) if rows else np.zeros(0, dtype=np.int32) split = EncodedSplit(name, flat, offsets, labels, data.label3, data.allowed, data.sources, texts) LOG.info("%s: %d rows encoded in %.1fs; tokens mean %.0f / median %.0f / max %d; %.1f%% truncated at %d; " "binary positives %d (%.1f%%); classes %s", name, len(rows), common.timer() - t0, lengths.mean(), np.median(lengths), lengths.max(), 100.0 * float((lengths >= max_len).mean()), max_len, int(labels.sum()), 100.0 * float(labels.mean()), label_counts(data.label3, split.num_labels)) return split def epoch_batches(lengths: np.ndarray, batch_size: int, seed: int, epoch: int) -> list[np.ndarray]: """Length-bucketed batch order for one epoch; a pure function of (seed, epoch) for resume.""" rng = np.random.default_rng([seed, epoch]) perm = rng.permutation(len(lengths)) mega = MEGA_BATCH_FACTOR * batch_size batches: list[np.ndarray] = [] for start in range(0, len(perm), mega): chunk = perm[start:start + mega] chunk = chunk[np.argsort(-lengths[chunk], kind="stable")] batches.extend(chunk[j:j + batch_size] for j in range(0, len(chunk), batch_size)) order = rng.permutation(len(batches)) return [batches[k] for k in order] def eval_batches(lengths: np.ndarray, batch_size: int) -> list[np.ndarray]: order = np.argsort(-lengths, kind="stable") return [order[i:i + batch_size] for i in range(0, len(order), batch_size)] class RowIndexDataset(Dataset): """``torch.utils.data.Dataset`` over an ``EncodedSplit``: an item is a row index. Membership and order come from the batch sampler (``epoch_batches`` / ``eval_batches``: seeded shuffling + length bucketing), padding from the collate function, so every batch the trainer sees goes through ``torch.utils.data.DataLoader`` exactly as the course reference does, while the resume-exact batch order is kept. """ def __init__(self, split: EncodedSplit) -> None: self.split = split def __len__(self) -> int: return len(self.split) def __getitem__(self, i: int) -> int: return int(i) def make_loader(split: EncodedSplit, batches: list[np.ndarray], multiple: int, micro_batch: int | None = None) -> DataLoader: """DataLoader whose batch sampler is the precomputed batch list. Yields ``(idx, ids, mask, labels, allowed)`` per batch, or with ``micro_batch`` set ``(idx, [(ids, mask, labels, allowed), ...])`` where each micro-batch is padded to ITS OWN longest row (the gradient-accumulation geometry the smoke test measured). Tensors are produced on the CPU; the training loop moves them to the device. ``num_workers=0`` keeps the order a pure function of (seed, epoch) and avoids forking next to a live CUDA context. """ cpu = torch.device("cpu") def collate(rows: list[int]): idx = np.asarray(rows, dtype=np.int64) if micro_batch is None: return (idx, *split.collate(idx, multiple, cpu)) return idx, [split.collate(idx[c:c + micro_batch], multiple, cpu) for c in range(0, len(idx), micro_batch)] # An explicit generator keeps the DataLoader from drawing its worker base seed from the # global torch RNG at iteration start, which would shift the dropout stream and break # bit-exact resume on the CPU dev path; the loader itself uses no randomness (num_workers=0, # explicit batch_sampler), so the seed value is irrelevant. return DataLoader(RowIndexDataset(split), batch_sampler=[[int(i) for i in b] for b in batches], collate_fn=collate, num_workers=0, generator=torch.Generator().manual_seed(0)) def load_data(cfg: dict, args: TrainArgs, tok) -> tuple[EncodedSplit, EncodedSplit, int]: processed = Path(cfg["paths"]["data_processed"]) holdout_source = str(cfg["data"]["holdout_test_source"]) max_len = int(cfg["model"]["max_seq_len"]) multiple = int(cfg["finetune"]["dynamic_padding_multiple"]) data = read_split(processed / TRAIN_FILE, holdout_source, args.limit_rows, args.seed, args.num_labels) n_synth = 0 if args.synthetic == "on": synth = processed / SYNTHETIC_FILE if synth.exists(): # Synthetic rows carry their own label3/label_set (category "1:*" -> phishing, "0:*" -> legit). sd = read_split(synth, holdout_source, int(cfg["data"]["synthetic"].get("max_rows", 2000)), args.seed, args.num_labels, force_source="synthetic") n_synth = len(sd) data = concat_labelled(data, sd) LOG.info("synthetic supplement: %d rows appended to TRAIN only; classes %s", n_synth, label_counts(sd.label3, args.num_labels)) else: LOG.warning("--synthetic on but %s is missing; training without synthetic rows", synth) train = encode_split("train", tok, data, max_len, multiple) vd = read_split(processed / VAL_FILE, holdout_source, args.limit_rows, args.seed, args.num_labels) if args.limit_rows is not None: LOG.warning("--limit-rows %d: validation is a seeded SUBSAMPLE too (dev only; metrics are not comparable)", args.limit_rows) val = encode_split("val", tok, vd, max_len, multiple) assert "synthetic" not in set(val.sources), "synthetic rows must never reach validation" if args.num_labels == 3 and args.selection_metric == "val_macro_f1_3way" and not (val.label3 >= 0).any(): raise ValueError("selection metric val_macro_f1_3way but finetune_val has no row with label3 in {0,1,2}") return train, val, n_synth def check_collate_matches_encode_batch(split: EncodedSplit, tok, idx: np.ndarray, max_len: int, multiple: int) -> None: """The trained-on tensors must equal what ``encode_batch`` yields for the same texts.""" ids, mask, _, _ = split.collate(idx, multiple, torch.device("cpu")) ref_ids, ref_mask = encode_batch(tok, [split.texts[int(i)] for i in idx], max_len, pad_to_multiple_of=multiple) if ids.shape != ref_ids.shape or not torch.equal(ids, ref_ids) or not torch.equal(mask, ref_mask): raise AssertionError("collated batch differs from tokenizer.encode_batch output") LOG.info("collate == encode_batch on the first batch (%s): OK", tuple(ids.shape)) # ---------------------------------------------------------------------------------- model def newest_pretrain_final(cfg: dict) -> Path | None: root = Path(cfg["paths"]["checkpoints"]) / "pretrain" finals = [p for p in root.glob("*/final") if (p / "model.safetensors").exists()] if root.exists() else [] return max(finals, key=lambda p: p.stat().st_mtime) if finals else None def read_pretrained(path: Path) -> tuple[dict[str, Tensor], dict[str, Any]]: """state_dict + EncoderConfig dict from a Phase 1 final/ or step-NNNNNNN/ directory.""" weights = path / "model.safetensors" if not weights.exists(): raise FileNotFoundError(f"{weights} not found") cfg_json = path / "config.json" if cfg_json.exists(): meta = common.read_json(cfg_json) else: with safe_open(str(weights), "pt") as fh: raw = (fh.metadata() or {}).get("config") if raw is None: raise RuntimeError(f"{path}: no config.json and no 'config' safetensors metadata") meta = json.loads(raw) model_cfg = meta["model"] if isinstance(meta, dict) and "model" in meta else meta return load_file(str(weights), device="cpu"), dict(model_cfg) def slice_layers(state: dict[str, Tensor], n_layers: int) -> dict[str, Tensor]: """Keep encoder.layers.0..n_layers-1 (the instructor's 'number of layers' dial on a pretrained stack).""" out: dict[str, Tensor] = {} dropped = set() for k, v in state.items(): if k.startswith("encoder.layers."): i = int(k.split(".")[2]) if i >= n_layers: dropped.add(i) continue out[k] = v if dropped: LOG.info("--n-layers %d: dropped pretrained layers %d..%d", n_layers, min(dropped), max(dropped)) return out def build_model(cfg: dict, args: TrainArgs, device: torch.device) -> tuple[EncoderForClassification, str]: """Returns (model, init_descriptor) where init_descriptor is 'scratch' or 'pretrained:'.""" # The head size comes from finetune.num_labels, never from config model: (the pretraining snapshot). mcfg_dict = {**cfg["model"], "num_labels": args.num_labels} if args.init == "scratch": mcfg_dict["n_layers"] = args.n_layers ecfg = EncoderConfig.from_dict(mcfg_dict) common.set_seed(args.seed) with torch.device(device): model = EncoderForClassification(ecfg) return model, "scratch" path = Path(args.init.split(":", 1)[1]) if ":" in args.init else newest_pretrain_final(cfg) if path is None: raise FileNotFoundError("no checkpoints/pretrain/*/final found; pass --init pretrained: or --init scratch") if not path.is_absolute(): path = common.project_root() / path state, pre_cfg = read_pretrained(path) if int(pre_cfg["n_layers"]) < args.n_layers: raise ValueError(f"--n-layers {args.n_layers} > pretrained n_layers {pre_cfg['n_layers']}") # Architecture from the checkpoint; regularisation from config.yaml; depth from the dial. merged = {**pre_cfg, "n_layers": args.n_layers, "num_labels": args.num_labels, "dropout": mcfg_dict["dropout"], "attention_dropout": mcfg_dict["attention_dropout"], "grad_checkpointing": mcfg_dict.get("grad_checkpointing", False), "tie_mlm_head": pre_cfg.get("tie_mlm_head", True)} ecfg = EncoderConfig.from_dict(merged) common.set_seed(args.seed) with torch.device(device): model = EncoderForClassification.from_pretrained_encoder(slice_layers(state, args.n_layers), ecfg) LOG.info("loaded pretrained encoder from %s (%d of %d layers)", path, args.n_layers, pre_cfg["n_layers"]) return model, f"pretrained:{path}" def freeze(model: EncoderForClassification, k: int) -> None: if k <= 0: return for p in model.encoder.embeddings.parameters(): p.requires_grad_(False) for layer in list(model.encoder.layers)[:k]: for p in layer.parameters(): p.requires_grad_(False) LOG.info("frozen: embeddings + encoder layers 0..%d", k - 1) # ----------------------------------------------------------------------------------- loss def class_weights(label3: np.ndarray, num_labels: int, device: torch.device) -> Tensor: """Inverse-frequency class weights ``w_c = N / (C * n_c)`` over the FULLY labelled rows (label3 >= 0) of the train split actually used. Unwanted rows have no class of their own; they receive mean(w_phish, w_spam) through ``row_weights`` at loss time.""" full = label3[label3 >= 0] counts = np.bincount(full, minlength=num_labels).astype(np.float64) w = full.shape[0] / (float(num_labels) * np.maximum(counts, 1)) return torch.tensor(w, dtype=torch.float32, device=device) def partial_label_nll(logits: Tensor, allowed: Tensor) -> Tensor: """Per-row ``-log sum_{c in allowed} softmax(logits)_c`` = logsumexp(all) - logsumexp(allowed). One allowed class -> exactly ``F.cross_entropy(reduction="none")``; {phish, spam} -> ``-log(p1 + p2)``. Computed in fp32 from bf16-autocast logits. """ logits = logits.float() if allowed.shape != logits.shape: raise ValueError(f"allowed mask {tuple(allowed.shape)} does not match logits {tuple(logits.shape)}") lse_all = torch.logsumexp(logits, dim=-1) lse_allowed = torch.logsumexp(logits.masked_fill(~allowed, float("-inf")), dim=-1) return lse_all - lse_allowed def row_weights(allowed: Tensor, weights: Tensor) -> Tensor: """Mean of the class weights over each row's allowed classes: w[c] for a full label, mean(w1, w2) for an "unwanted" row.""" a = allowed.to(weights.dtype) return (a @ weights) / a.sum(dim=-1) def compute_loss(logits: Tensor, allowed: Tensor, kind: str, weights: Tensor | None, gamma: float) -> Tensor: """The three loss dials, all defined over the partial-label NLL (see module docstring). cross_entropy mean(nll) weighted_cross_entropy sum(w_r nll_r) / sum(w_r) (F.cross_entropy(weight=)'s normalisation) focal mean((1 - exp(-nll))^gamma * nll) (= -(1 - p_t)^gamma log p_t; gamma 0 = plain) """ nll = partial_label_nll(logits, allowed) if kind == "cross_entropy": return nll.mean() if kind == "weighted_cross_entropy": if weights is None: raise ValueError("weighted_cross_entropy needs class weights") w = row_weights(allowed, weights) return (w * nll).sum() / w.sum() if kind == "focal": return ((1.0 - torch.exp(-nll)).pow(gamma) * nll).mean() raise ValueError(f"unknown loss {kind!r}") # -------------------------------------------------------------------------------- metrics def classification_metrics(labels: np.ndarray, probs: np.ndarray, threshold: float = DECISION_THRESHOLD) -> dict[str, Any]: from sklearn.metrics import (accuracy_score, average_precision_score, confusion_matrix, precision_recall_curve, precision_recall_fscore_support, roc_auc_score) pred = (probs >= threshold).astype(np.int64) p, r, f1, _ = precision_recall_fscore_support(labels, pred, average="binary", pos_label=1, zero_division=0) cm = confusion_matrix(labels, pred, labels=[0, 1]) try: roc = float(roc_auc_score(labels, probs)) pr_auc = float(average_precision_score(labels, probs)) pc, rc, th = precision_recall_curve(labels, probs) f1s = 2 * pc[:-1] * rc[:-1] / np.maximum(pc[:-1] + rc[:-1], 1e-12) j = int(np.argmax(f1s)) opt_t, opt_f1 = float(th[j]), float(f1s[j]) except ValueError: # a single class present (tiny dev subsets) roc = pr_auc = opt_t = opt_f1 = float("nan") return { "accuracy": float(accuracy_score(labels, pred)), "precision": float(p), "recall": float(r), "f1": float(f1), "roc_auc": roc, "pr_auc": pr_auc, "tn": int(cm[0, 0]), "fp": int(cm[0, 1]), "fn": int(cm[1, 0]), "tp": int(cm[1, 1]), "threshold": threshold, "f1_optimal_threshold": opt_t, "f1_at_optimal_threshold": opt_f1, } def three_way_metrics(label3: np.ndarray, probs: np.ndarray) -> dict[str, Any]: """Metric set (b) on the rows with a full three-way label (label3 in {0, 1, 2}). ``probs`` is ``[n, 3]``. Prediction = argmax over the three classes. Keys (flat, so they drop straight into epochs.csv): n_3way, accuracy_3way, macro_f1_3way (macro over ALL three classes, a class absent from the slice scores 0, so tiny dev subsets read low), {precision,recall,f1}_{legitimate,phishing,spam}, cm3_ with l/p/s initials (cm3_ps = true phishing predicted spam), and the phishing-vs-spam confusion restricted to true unwanted rows where the prediction is the argmax over {phishing, spam} only ("given it is unwanted, which kind?"): n_unwanted_3way, pvs_phish_as_phish, pvs_phish_as_spam, pvs_spam_as_phish, pvs_spam_as_spam. Returns None/0 placeholders (same keys) when the slice is empty or the model is binary: None, not NaN, so best/config.json and experiments.jsonl of a binary run stay strict JSON. """ from sklearn.metrics import confusion_matrix, precision_recall_fscore_support out: dict[str, Any] = {k: (0 if k.startswith(("n_", "cm3_", "pvs_")) else None) for k in THREE_WAY_KEYS} sel = label3 >= 0 if probs.ndim != 2 or probs.shape[1] != 3 or not sel.any(): return out y, p = label3[sel], probs[sel] pred = p.argmax(axis=1) cm = confusion_matrix(y, pred, labels=[0, 1, 2]) prec, rec, f1, _ = precision_recall_fscore_support(y, pred, labels=[0, 1, 2], zero_division=0) out.update(n_3way=int(sel.sum()), accuracy_3way=float((pred == y).mean()), macro_f1_3way=float(f1.mean())) for c, slug in enumerate(CLASS_SLUGS): out[f"precision_{slug}"], out[f"recall_{slug}"], out[f"f1_{slug}"] = float(prec[c]), float(rec[c]), float(f1[c]) for i in range(3): for j in range(3): out[f"cm3_{CLASS_SLUGS[i][0]}{CLASS_SLUGS[j][0]}"] = int(cm[i, j]) unw = y >= 1 if unw.any(): kind = 1 + p[unw][:, 1:3].argmax(axis=1) cm2 = confusion_matrix(y[unw], kind, labels=[1, 2]) out.update(n_unwanted_3way=int(unw.sum()), pvs_phish_as_phish=int(cm2[0, 0]), pvs_phish_as_spam=int(cm2[0, 1]), pvs_spam_as_phish=int(cm2[1, 0]), pvs_spam_as_spam=int(cm2[1, 1])) return out def metric_sets(res: dict[str, Any]) -> dict[str, Any]: """Structured view of a flat evaluation result for best/config.json and experiments.jsonl: ``binary_unwanted`` (metric set a) and ``three_way`` (metric set b, None for a binary model or an empty slice).""" binary = {k: res.get(k) for k in (*METRIC_KEYS, "tn", "fp", "fn", "tp", "threshold", "f1_optimal_threshold", "f1_at_optimal_threshold", "n")} binary["positive_class"] = UNWANTED_NAME three = None if res.get("n_3way"): three = { "n": res["n_3way"], "accuracy": res["accuracy_3way"], "macro_f1": res["macro_f1_3way"], "per_class": {LABEL_NAMES[c]: {m: res[f"{m}_{slug}"] for m in ("precision", "recall", "f1")} for c, slug in enumerate(CLASS_SLUGS)}, "confusion": {"order": [LABEL_NAMES[c] for c in range(3)], "rows_true_cols_pred": [[res[f"cm3_{a[0]}{b[0]}"] for b in CLASS_SLUGS] for a in CLASS_SLUGS]}, "phish_vs_spam_on_true_unwanted": { "n": res["n_unwanted_3way"], "order": [LABEL_NAMES[1], LABEL_NAMES[2]], "rows_true_cols_pred": [[res["pvs_phish_as_phish"], res["pvs_phish_as_spam"]], [res["pvs_spam_as_phish"], res["pvs_spam_as_spam"]]]}, } return {"binary_unwanted": binary, "three_way": three} def autocast_ctx(device: torch.device): """bf16 autocast on CUDA; a no-op on the CPU dev path (fp32 end to end).""" if device.type == "cuda": return torch.autocast(device_type="cuda", dtype=torch.bfloat16) return contextlib.nullcontext() def unwanted_prob(probs: Tensor) -> Tensor: """P(unwanted) = p_phish + p_spam (3 classes) or p_1 (binary).""" return probs[:, 1:].sum(dim=-1) @torch.inference_mode() def evaluate(model: EncoderForClassification, split: EncodedSplit, args: TrainArgs, multiple: int, device: torch.device) -> dict[str, Any]: """Both metric sets on ``split``. ``val_loss`` is the unweighted partial-label NLL (the plain cross_entropy dial) so it is comparable across loss variants.""" model.eval() t0 = common.timer() probs, labels, label3s, loss_sum = [], [], [], 0.0 loader = make_loader(split, eval_batches(split.lengths, args.eval_batch_size), multiple) for idx, ids, mask, y, allowed in loader: ids, mask, y, allowed = (t.to(device, non_blocking=True) for t in (ids, mask, y, allowed)) with autocast_ctx(device): logits = model(ids, mask) logits = logits.float() loss_sum += float(partial_label_nll(logits, allowed).sum()) probs.append(logits.softmax(-1).cpu().numpy()) labels.append(y.cpu().numpy()) # eval_batches() visits rows in length order, so EVERY target must be gathered with the same idx # as the predictions; split.label3 in storage order would score a permutation of the model. label3s.append(split.label3[idx]) model.train() probs_np, labels_np, label3_np = np.concatenate(probs), np.concatenate(labels), np.concatenate(label3s) out = classification_metrics(labels_np, unwanted_prob(torch.from_numpy(probs_np)).numpy()) out["f1_unwanted"] = out["f1"] out.update(three_way_metrics(label3_np, probs_np)) out["val_loss"] = loss_sum / max(1, len(labels_np)) out["n"] = int(len(labels_np)) out["eval_seconds"] = round(common.timer() - t0, 1) return out def selection_value(res: dict[str, Any], metric: str) -> float: v = res.get(SELECTION_KEY[metric]) return float("-inf") if v is None or (isinstance(v, float) and math.isnan(v)) else float(v) # ------------------------------------------------------------------------------- run state @dataclass class RunState: step: int = 0 epoch: int = 0 batch_in_epoch: int = 0 best_score: float = -1.0 # value of the selection metric at best/ (named best_f1 before entry 6) best_step: int = 0 best_epoch: int = 0 evals_without_improvement: int = 0 train_seconds: float = 0.0 eval_seconds: float = 0.0 peak_reserved_gb: float = 0.0 peak_allocated_gb: float = 0.0 best_val: dict[str, Any] | None = None last_val: dict[str, Any] | None = None n_train: int = 0 n_val: int = 0 n_synthetic: int = 0 init: str = "" stopped_by: str = "" resumed_from_step: int | None = None metrics_tail: list[dict[str, Any]] = field(default_factory=list) class StopFlag: """SIGINT/SIGTERM -> finish the current micro-step, save last/, exit 0. Second signal -> default.""" def __init__(self) -> None: self.requested: str | None = None self._old = {s: signal.signal(s, self._handle) for s in (signal.SIGINT, signal.SIGTERM)} def _handle(self, signum: int, frame: Any) -> None: if self.requested: signal.signal(signum, self._old[signum]) raise KeyboardInterrupt self.requested = signal.Signals(signum).name LOG.warning("%s received: finishing the current step, saving last/ and exiting", self.requested) def rng_states(device: torch.device) -> dict[str, Any]: """RNG snapshot for resume. The CUDA generator is read only when training ON CUDA: ``torch.cuda.get_rng_state_all`` goes through ``_lazy_init`` and would create a CUDA context on a GPU the --cpu dev path must leave alone (the production pretraining run lives there).""" return {"python": random.getstate(), "numpy": np.random.get_state(), "torch_cpu": torch.get_rng_state(), "torch_cuda": torch.cuda.get_rng_state_all() if device.type == "cuda" else []} def restore_rng(states: dict[str, Any], device: torch.device) -> None: random.setstate(states["python"]) np.random.set_state(states["numpy"]) torch.set_rng_state(states["torch_cpu"]) if device.type == "cuda" and states.get("torch_cuda"): torch.cuda.set_rng_state_all(states["torch_cuda"]) def json_clean(obj: Any) -> Any: """Recursively replace non-finite floats (NaN/inf from a degenerate metric) with None so every JSON artifact (best/config.json, last/config.json, experiments.jsonl) is strict JSON; the stdlib encoder would otherwise emit the bare token ``NaN`` that non-Python parsers reject.""" if isinstance(obj, float): return obj if math.isfinite(obj) else None if isinstance(obj, np.floating): return float(obj) if np.isfinite(obj) else None if isinstance(obj, dict): return {k: json_clean(v) for k, v in obj.items()} if isinstance(obj, (list, tuple)): return [json_clean(v) for v in obj] return obj def save_model_dir(model: EncoderForClassification, out: Path, step: int, extra: dict[str, Any]) -> None: out.mkdir(parents=True, exist_ok=True) sd = {k: v.detach().to("cpu").contiguous() for k, v in model.state_dict().items()} cfg_dict = model.config.to_dict() save_file(sd, str(out / "model.safetensors"), metadata={"format": "pt", "step": str(step), "config": json.dumps(cfg_dict)}) common.write_json(json_clean({"model": cfg_dict, "tokenizer": "tokenizer.json", **extra}), out / "config.json") def save_last(run_dir: Path, model: EncoderForClassification, opt: torch.optim.Optimizer, sched: torch.optim.lr_scheduler.LRScheduler, state: RunState, args: TrainArgs, schedule: dict[str, int], device: torch.device) -> None: t0 = common.timer() last = run_dir / "last" save_model_dir(model, last, state.step, {"finetune": args.to_record(), "state": _state_record(state)}) torch.save({ "optimizer": opt.state_dict(), "scheduler": sched.state_dict(), "state": asdict(state), "args": args.to_record(), "schedule": schedule, "rng": rng_states(device), }, last / "trainer_state.pt") LOG.info("saved last/ at step %d (%.1fs)", state.step, common.timer() - t0) def _state_record(state: RunState) -> dict[str, Any]: d = asdict(state) d.pop("metrics_tail", None) return d def write_csv_row(path: Path, row: dict[str, Any], header: list[str]) -> None: new = not path.exists() with path.open("a", newline="", encoding="utf-8") as fh: w = csv.DictWriter(fh, fieldnames=header, extrasaction="ignore") if new: w.writeheader() w.writerow({k: ("" if row.get(k) is None else row.get(k)) for k in header}) def trim_csv(path: Path, max_step: int) -> None: """Drop rows with step > max_step: a crash between checkpoints leaves rows past the last/ being resumed, and the resumed steps would otherwise be appended a second time.""" if not path.exists(): return with path.open(newline="", encoding="utf-8") as fh: reader = csv.DictReader(fh) header, rows = list(reader.fieldnames or []), list(reader) keep = [r for r in rows if r.get("step") and int(float(r["step"])) <= max_step] if len(keep) == len(rows): return tmp = path.with_suffix(path.suffix + ".tmp") with tmp.open("w", newline="", encoding="utf-8") as fh: w = csv.DictWriter(fh, fieldnames=header) w.writeheader() w.writerows(keep) os.replace(tmp, path) LOG.info("%s: dropped %d row(s) logged after step %d (crash between checkpoints)", path.name, len(rows) - len(keep), max_step) METRICS_HEADER = ["step", "epoch", "loss", "lr", "examples_per_s", "max_alloc_gb", "max_reserved_gb", "elapsed_s", "eta_s", "val_loss", "val_f1", "val_roc_auc", "val_macro_f1_3way"] # Binary roll-up columns first (unchanged names, so old epochs.csv readers keep working), then the three-way set. EPOCHS_HEADER = ["epoch", "step", "is_epoch_end", "n", "val_loss", "accuracy", "precision", "recall", "f1", "roc_auc", "pr_auc", "tn", "fp", "fn", "tp", "threshold", "f1_optimal_threshold", "f1_at_optimal_threshold", "f1_unwanted", *THREE_WAY_KEYS, "eval_seconds", "is_best"] # ----------------------------------------------------------------------------- experiments def experiment_record(args: TrainArgs, state: RunState, params: dict[str, Any], steps_per_epoch: int) -> dict[str, Any]: val = state.best_val or {} return { "run_name": args.run_name, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), "init": state.init, "dials": { "n_layers": args.n_layers, "frozen_layers": args.freeze_layers, "epochs_planned": args.epochs, "epochs_run": round(state.step / max(1, steps_per_epoch), 3), "loss": args.loss, "focal_gamma": args.focal_gamma if args.loss == "focal" else None, "lr": args.lr, "batch_size": args.batch_size, "micro_batch": args.micro_batch, "class_weights": "inverse_frequency" if args.loss == "weighted_cross_entropy" else None, "synthetic": args.synthetic, "synthetic_rows": state.n_synthetic, "warmup_fraction": args.warmup_fraction, "max_steps": args.max_steps, "limit_rows": args.limit_rows, "seed": args.seed, "schedule": args.schedule, "selection_metric": args.selection_metric, "pooling": args.pooling, "num_labels": args.num_labels, "device": "cpu" if args.cpu else "cuda", }, "params": params, "steps_run": state.step, "n_train": state.n_train, "n_val": state.n_val, "train_seconds": round(state.train_seconds, 1), "peak_vram_gb": round(state.peak_reserved_gb, 3), "best_epoch": state.best_epoch, "best_step": state.best_step, "best_score": state.best_score if val else None, # Old binary keys kept verbatim (finetune.py list / evaluate.py read them); the three-way set is added. "val": ({**{k: val.get(k) for k in METRIC_KEYS}, "f1_unwanted": val.get("f1"), "macro_f1_3way": val.get("macro_f1_3way"), "per_class": {LABEL_NAMES[c]: {m: val.get(f"{m}_{slug}") for m in ("precision", "recall", "f1")} for c, slug in enumerate(CLASS_SLUGS)} if val.get("n_3way") else None, "three_way": metric_sets(val)["three_way"]} if val else None), "val_threshold": DECISION_THRESHOLD, "label_names": [LABEL_NAMES[c] for c in range(args.num_labels)] if args.num_labels == 3 else [LABEL_NAMES[0], UNWANTED_NAME], "stopped_by": state.stopped_by, "resumed_from_step": state.resumed_from_step, "test": None, } def upsert_experiment(path: Path, rec: dict[str, Any]) -> str: """ONE train line per run_name: a resumed invocation replaces its earlier train line instead of appending a duplicate; evaluate.py's lines (test != null) are never touched.""" if rec.get("test") is not None: raise ValueError("upsert_experiment writes train lines only (test must be None); a test record would " "overwrite the run's train line") lines = path.read_text(encoding="utf-8").splitlines() if path.exists() else [] out, replaced = [], False rec = json_clean(rec) # strict JSON: no bare NaN tokens for binary runs' three-way placeholders for line in lines: if not line.strip(): continue try: old = json.loads(line) except json.JSONDecodeError: out.append(line) continue if not replaced and old.get("run_name") == rec["run_name"] and old.get("test") is None: out.append(json.dumps(rec, allow_nan=False)) replaced = True else: out.append(line) if not replaced: out.append(json.dumps(rec, allow_nan=False)) path.parent.mkdir(parents=True, exist_ok=True) tmp = path.with_suffix(".jsonl.tmp") tmp.write_text("\n".join(out) + "\n", encoding="utf-8") os.replace(tmp, path) return "replaced" if replaced else "appended" def run_log_block(args: TrainArgs, state: RunState, run_dir: Path, params: dict[str, Any]) -> str: v = state.best_val or {} nan = float("nan") three = "" if v.get("n_3way"): three = (f" Three-way on {v['n_3way']:,} fully-labelled val rows: macro-F1 {v['macro_f1_3way']:.4f}, accuracy " f"{v['accuracy_3way']:.4f}; per-class F1 " + ", ".join(f"{LABEL_NAMES[c]} {v[f'f1_{s}']:.4f}" for c, s in enumerate(CLASS_SLUGS)) + f"; phishing-vs-spam on true unwanted (n={v['n_unwanted_3way']:,}): phish->phish {v['pvs_phish_as_phish']}, " f"phish->spam {v['pvs_phish_as_spam']}, spam->phish {v['pvs_spam_as_phish']}, spam->spam {v['pvs_spam_as_spam']}.") return "\n".join([ f"## Fine-tune run `{args.run_name}` — {time.strftime('%Y-%m-%d %H:%M')}", f"**Dials.** init {state.init}; {args.num_labels} classes; {args.n_layers} layers, {args.freeze_layers} frozen; loss {args.loss}" f"{f' (gamma {args.focal_gamma})' if args.loss == 'focal' else ''}; lr {args.lr:g}; batch {args.batch_size} " f"(micro {args.micro_batch}); epochs {args.epochs} planned; synthetic {args.synthetic} ({state.n_synthetic} rows); " f"selection {args.selection_metric}; train rows {state.n_train:,}, val rows {state.n_val:,}; params {params['total']:,} " f"({params['trainable']:,} trainable).", f"**Result.** {state.step} steps, stopped by {state.stopped_by}; training time {state.train_seconds / 60:.1f} min " f"(eval {state.eval_seconds / 60:.1f} min of it); peak reserved {state.peak_reserved_gb:.2f} GiB. Best val at step " f"{state.best_step} (epoch {state.best_epoch}), {UNWANTED_NAME} roll-up: F1 {v.get('f1', nan):.4f}, precision " f"{v.get('precision', nan):.4f}, recall {v.get('recall', nan):.4f}, accuracy {v.get('accuracy', nan):.4f}, ROC-AUC " f"{v.get('roc_auc', nan):.4f}, PR-AUC {v.get('pr_auc', nan):.4f}; F1-optimal threshold {v.get('f1_optimal_threshold', nan):.3f}." f"{three} Checkpoints: {run_dir}/best, {run_dir}/last.", ]) # ------------------------------------------------------------------------------- training CPU_VRAM = {"allocated_gb": 0.0, "reserved_gb": 0.0, "max_allocated_gb": 0.0, "max_reserved_gb": 0.0} def backfill_legacy_args(saved: dict[str, Any], run_dir: Path) -> None: """A trainer_state.pt written before RUN_LOG entry 6 has no ``num_labels`` / ``selection_metric`` in its args. Fill them from what that run actually was (the head size recorded in last/config.json model.num_labels; the only selection metric that existed, val_f1_phishing) so the saved-overrides-config rule resolves the run as the binary run it is, instead of the current config's 3-class head crashing in load_state_dict with a size mismatch.""" args = saved.get("args") or {} if "num_labels" not in args: try: n = int(common.read_json(run_dir / "last" / "config.json")["model"]["num_labels"]) except (OSError, KeyError, TypeError, ValueError, json.JSONDecodeError): n = 2 args["num_labels"] = n LOG.info("--resume: legacy trainer_state.pt without num_labels; using %d from last/config.json", n) if "selection_metric" not in args: args["selection_metric"] = "val_f1_phishing" saved["args"] = args def lr_lambda_factory(total_steps: int, warmup_steps: int): def lr_lambda(completed: int) -> float: if completed < warmup_steps: return (completed + 1) / warmup_steps return max(0.0, (total_steps - completed) / max(1, total_steps - warmup_steps)) return lr_lambda def train(ns: argparse.Namespace, cfg: dict) -> int: t_invocation = common.timer() # Decide the device FIRST and touch torch.cuda only on the CUDA branch. Under --cpu the GPU may be # (and today is) occupied by the production pretraining run: torch.cuda.current_device() / # mem_get_info() / get_rng_state_all() all run _lazy_init(), which creates a CUDA context there. # Hiding the devices from this process is the belt to that braces: nothing has called CUDA yet. if ns.cpu: # Hide the devices before the CUDA runtime is first queried (torch import does not query it): then # torch.cuda.is_available() is False and torch's own AdamW.step() health check, which otherwise calls # torch.accelerator.current_stream() -> _lazy_init() on any visible GPU, is skipped. If a GPU is still # visible the contract cannot be honoured from inside this process, so refuse rather than risk it. os.environ["CUDA_VISIBLE_DEVICES"] = "" if torch.cuda.is_initialized() or torch.cuda.is_available(): LOG.error("--cpu: a CUDA device is still visible to this process (CUDA was queried before the devices could " "be hidden); run with CUDA_VISIBLE_DEVICES=\"\" in the environment so the occupied GPU is never touched") return 2 device = torch.device("cpu") device_info = {"device": "cpu", "torch": torch.__version__} LOG.warning("--cpu: dev path on the CPU with the config as given (no autocast, no VRAM guard, CUDA never " "initialised; CUDA_VISIBLE_DEVICES cleared for this process); not for real runs") elif not torch.cuda.is_available(): LOG.error("CUDA is required for fine-tuning (--cpu is the dev-only escape hatch)") return 2 else: device = torch.device("cuda") device_info = common.print_device_info() LOG.info("PYTORCH_CUDA_ALLOC_CONF=%s", os.environ.get("PYTORCH_CUDA_ALLOC_CONF")) run_dir = Path(cfg["paths"]["checkpoints"]) / "finetune" / ns.run_name last_state_path = run_dir / "last" / "trainer_state.pt" saved: dict[str, Any] | None = None if ns.resume: if not last_state_path.exists(): LOG.error("--resume: %s not found", last_state_path) return 2 # Own file (optimizer tensors + Python/NumPy RNG tuples), hence weights_only=False. saved = torch.load(last_state_path, map_location="cpu", weights_only=False) backfill_legacy_args(saved, run_dir) elif run_dir.exists() and any(run_dir.iterdir()): if not ns.overwrite: LOG.error("%s exists; use --resume to continue it, --overwrite to discard it, or a new --run-name", run_dir) return 2 import shutil shutil.rmtree(run_dir) LOG.warning("--overwrite: removed %s", run_dir) try: args = TrainArgs.resolve(ns, cfg, saved["args"] if saved else None) except ValueError as e: LOG.error("dial error: %s", e) return 2 args.run_name = ns.run_name args.cpu = bool(ns.cpu) # the device is this invocation's choice, never restored from last/ if saved: changed = {k: (saved["args"].get(k), getattr(ns, k)) for k in RESUME_FIXED_DIALS if getattr(ns, k, None) is not None and getattr(ns, k) != saved["args"].get(k)} if changed: LOG.error("--resume: %s cannot change on resume (saved -> requested); the restored optimizer state and " "batch order would not match. Start a new --run-name instead", changed) return 2 LOG.info("resolved dials: %s", args.to_record()) common.set_seed(args.seed) vr = cfg["vram"] max_len = int(cfg["model"]["max_seq_len"]) multiple = int(cfg["finetune"]["dynamic_padding_multiple"]) tok = load_tokenizer(cfg["paths"]["tokenizer"]) train_split, val_split, n_synth = load_data(cfg, args, tok) check_collate_matches_encode_batch(train_split, tok, epoch_batches(train_split.lengths, args.batch_size, args.seed, 0)[0], max_len, multiple) if saved and (saved["state"]["n_train"] != len(train_split) or saved["state"]["n_val"] != len(val_split)): LOG.error("--resume: data changed (train %d -> %d, val %d -> %d); the batch order would not line up", saved["state"]["n_train"], len(train_split), saved["state"]["n_val"], len(val_split)) return 2 try: model, init_desc = build_model(cfg, args, device) except (FileNotFoundError, ValueError, RuntimeError) as e: LOG.error("model init failed: %s", e) return 2 if saved: sd = load_file(str(run_dir / "last" / "model.safetensors"), device="cpu") try: model.load_state_dict(sd, strict=True) except RuntimeError as e: # head / depth mismatch that slipped past the dial checks LOG.error("--resume: last/model.safetensors does not fit the resolved dials (num_labels %d, n_layers %d): %s", args.num_labels, args.n_layers, str(e).splitlines()[0]) return 2 freeze(model, args.freeze_layers) model.train() params = count_parameters(model) LOG.info("EncoderForClassification params: %s", params) if tok.get_vocab_size() > model.config.vocab_size: raise AssertionError("tokenizer vocab larger than model vocab") steps_per_epoch = math.ceil(len(train_split) / args.batch_size) total_steps = steps_per_epoch * args.epochs if args.max_steps is not None: total_steps = min(total_steps, args.max_steps) warmup_steps = max(1, math.ceil(args.warmup_fraction * total_steps)) schedule = {"steps_per_epoch": steps_per_epoch, "total_steps": total_steps, "warmup_steps": warmup_steps} LOG.info("schedule: %d steps/epoch x %d epochs -> %d optimizer steps (max_steps=%s), warmup %d, " "micro-batch %d x %d accumulation", steps_per_epoch, args.epochs, total_steps, args.max_steps, warmup_steps, args.micro_batch, args.batch_size // args.micro_batch) opt_cfg = {"weight_decay": args.weight_decay, "betas": list(args.betas), "eps": args.eps} opt = build_optimizer(model, args.lr, opt_cfg) lr_lambda = lr_lambda_factory(total_steps, warmup_steps) sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda) weights = class_weights(train_split.label3, args.num_labels, device) if args.loss == "weighted_cross_entropy" else None if weights is not None: names = [LABEL_NAMES[c] for c in range(3)] if args.num_labels == 3 else [LABEL_NAMES[0], UNWANTED_NAME] LOG.info("inverse-frequency class weights %s from fully-labelled counts %s%s", dict(zip(names, [round(float(w), 4) for w in weights])), np.bincount(train_split.label3[train_split.label3 >= 0], minlength=args.num_labels).tolist(), f"; {int((train_split.label3 < 0).sum())} unwanted rows weighted mean(w_phish, w_spam)={float(weights[1:].mean()):.4f}" if args.num_labels == 3 else "") if args.num_labels == 3: LOG.info("label contract: classes %s; binary roll-up '%s' = P(%s) + P(%s); selection metric %s", LABEL_NAMES, UNWANTED_NAME, LABEL_NAMES[1], LABEL_NAMES[2], args.selection_metric) state = RunState(n_train=len(train_split), n_val=len(val_split), n_synthetic=n_synth, init=init_desc) if saved: st = dict(saved["state"]) if "best_f1" in st: # trainer_state.pt written before entry 6 st["best_score"] = st.pop("best_f1") known = {f.name for f in fields(RunState)} state = RunState(**{**{k: v for k, v in st.items() if k in known}, "stopped_by": "", "init": st["init"]}) state.resumed_from_step = state.step opt.load_state_dict(saved["optimizer"]) sched.load_state_dict(saved["scheduler"]) if args.lr != float(saved["args"]["lr"]): # load_state_dict restored the saved base lr; re-apply the override so the next optimizer step # (which runs before the next sched.step()) and every later one use the lr the record claims. scale = lr_lambda(sched.last_epoch) for g in opt.param_groups: g["initial_lr"] = args.lr g["lr"] = args.lr * scale sched.base_lrs = [args.lr] * len(opt.param_groups) LOG.warning("--lr %g overrides the saved %g from step %d onward (steps 1..%d used the old lr); current lr %.3e", args.lr, float(saved["args"]["lr"]), state.step, state.step, args.lr * scale) if saved["schedule"] != schedule: LOG.warning("schedule changed on resume %s -> %s; the LR lambda now follows the new totals from step %d", saved["schedule"], schedule, state.step) restore_rng(saved["rng"], device) LOG.info("resumed %s at step %d (epoch %d, batch %d/%d); best %s %.4f at step %d; %.1f s trained so far", ns.run_name, state.step, state.epoch, state.batch_in_epoch, steps_per_epoch, args.selection_metric, state.best_score, state.best_step, state.train_seconds) if args.max_steps is not None and state.step >= args.max_steps: LOG.warning("nothing to resume: step %d already reached --max-steps %d; raise --max-steps or --epochs to continue", state.step, args.max_steps) return 0 if state.epoch >= args.epochs: LOG.warning("nothing to resume: epoch %d already reached --epochs %d; raise --epochs to continue", state.epoch, args.epochs) return 0 if state.evals_without_improvement >= args.patience: LOG.warning("resuming after early stopping (%d/%d evaluations without improvement): the run stops again at the " "next evaluation unless %s improves", state.evals_without_improvement, args.patience, args.selection_metric) run_dir.mkdir(parents=True, exist_ok=True) metrics_csv, epochs_csv = run_dir / "metrics.csv", run_dir / "epochs.csv" if saved: trim_csv(metrics_csv, state.step) trim_csv(epochs_csv, state.step) stop = StopFlag() if device.type == "cuda": torch.cuda.reset_peak_memory_stats() band_warned = False label_names = [LABEL_NAMES[c] for c in range(3)] if args.num_labels == 3 else [LABEL_NAMES[0], UNWANTED_NAME] def record_eval(is_epoch_end: bool) -> None: nonlocal state res = evaluate(model, val_split, args, multiple, device) state.eval_seconds += res["eval_seconds"] state.last_val = res score = selection_value(res, args.selection_metric) improved = score > state.best_score if improved: state.best_score, state.best_step, state.best_epoch, state.best_val = score, state.step, state.epoch, res state.evals_without_improvement = 0 save_model_dir(model, run_dir / "best", state.step, { "finetune": args.to_record(), "best_epoch": state.epoch, "best_step": state.step, "num_labels": args.num_labels, "label_names": label_names, "binary_rollup_name": UNWANTED_NAME, "selection_metric": args.selection_metric, "selection_value": score, "val_metrics": res, "val_metric_sets": metric_sets(res), "thresholds": {"decision": DECISION_THRESHOLD, "decision_applies_to": "P(unwanted) = sum of non-legitimate probabilities", "val_f1_optimal": res["f1_optimal_threshold"], "val_f1_at_optimal": res["f1_at_optimal_threshold"]}, "init": state.init, "params": params, }) else: state.evals_without_improvement += 1 three = "" if res.get("n_3way"): three = (f" | 3-way n={res['n_3way']} macro-F1 {res['macro_f1_3way']:.4f} acc {res['accuracy_3way']:.4f} F1[L P S]=[%s] " f"cm3[rows true L,P,S; cols pred]=[%s] phish/spam on unwanted [pp ps sp ss]=[%d %d %d %d]" % ( " ".join(f"{res[f'f1_{s}']:.3f}" for s in CLASS_SLUGS), " ".join(f"{res[k]}" for k in CM3_KEYS), *(res[k] for k in PVS_KEYS))) LOG.info("eval step %d epoch %d%s: n=%d loss %.4f | unwanted roll-up acc %.4f P %.4f R %.4f F1 %.4f ROC-AUC %.4f " "PR-AUC %.4f cm[tn fp fn tp]=[%d %d %d %d] F1-opt thr %.3f (F1 %.4f)%s | %s=%.4f %s (%.1fs)", state.step, state.epoch, " [epoch end]" if is_epoch_end else "", res["n"], res["val_loss"], res["accuracy"], res["precision"], res["recall"], res["f1"], res["roc_auc"], res["pr_auc"], res["tn"], res["fp"], res["fn"], res["tp"], res["f1_optimal_threshold"], res["f1_at_optimal_threshold"], three, args.selection_metric, score, "NEW BEST -> best/" if improved else f"no improvement ({state.evals_without_improvement}/{args.patience})", res["eval_seconds"]) write_csv_row(epochs_csv, {"epoch": state.epoch, "step": state.step, "is_epoch_end": int(is_epoch_end), "is_best": int(improved), **res}, EPOCHS_HEADER) t_loop = t_inv_start = common.timer() steps_this_invocation = 0 overhead_inv = 0.0 # eval + checkpoint seconds this invocation, excluded from the ETA's per-step cost window_loss, window_rows, window_steps, t_window = 0.0, 0, 0, common.timer() last_eval_step = state.step if saved else -1 exit_code = 0 def flush_time() -> None: nonlocal t_loop now = common.timer() state.train_seconds += now - t_loop t_loop = now try: while state.epoch < args.epochs and not state.stopped_by: batches = epoch_batches(train_split.lengths, args.batch_size, args.seed, state.epoch) loader = make_loader(train_split, batches[state.batch_in_epoch:], multiple, args.micro_batch) for b, (idx, micro_batches) in enumerate(loader, start=state.batch_in_epoch): n_rows = len(idx) loss_acc = 0.0 for ids, mask, _y_bin, allowed in micro_batches: ids, mask, allowed = (t.to(device, non_blocking=True) for t in (ids, mask, allowed)) n_sub = int(ids.shape[0]) with autocast_ctx(device): logits = model(ids, mask) loss = compute_loss(logits, allowed, args.loss, weights, args.focal_gamma) (loss * (n_sub / n_rows)).backward() loss_acc += float(loss.detach()) * n_sub / n_rows torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) opt.step() opt.zero_grad(set_to_none=True) sched.step() state.step += 1 state.batch_in_epoch = b + 1 steps_this_invocation += 1 window_loss += loss_acc window_rows += n_rows window_steps += 1 if not math.isfinite(loss_acc): LOG.error("non-finite loss at step %d; stopping", state.step) state.stopped_by = "non_finite_loss" exit_code = 3 # Read the allocator only on CUDA; the CPU dev path reports zeros without touching torch.cuda. vram = common.vram_stats() if device.type == "cuda" else CPU_VRAM state.peak_reserved_gb = max(state.peak_reserved_gb, vram["max_reserved_gb"]) state.peak_allocated_gb = max(state.peak_allocated_gb, vram["max_allocated_gb"]) if device.type == "cuda" and vram["max_reserved_gb"] > float(vr["hard_cap_gb"]): LOG.error("VRAM hard cap breached: %.2f GiB reserved > %s GiB; lower --micro-batch. Saving last/ and aborting", vram["max_reserved_gb"], vr["hard_cap_gb"]) state.stopped_by = "vram_cap" exit_code = 3 elif not band_warned and vram["max_reserved_gb"] > float(vr["target_gb_high"]): LOG.warning("reserved peak %.2f GiB is above the target band [%s, %s] GiB", vram["max_reserved_gb"], vr["target_gb_low"], vr["target_gb_high"]) band_warned = True if stop.requested: state.stopped_by = f"signal:{stop.requested}" is_eval = args.eval_every > 0 and state.step % args.eval_every == 0 epoch_end = b == len(batches) - 1 hit_max = args.max_steps is not None and state.step >= args.max_steps if hit_max and not state.stopped_by: state.stopped_by = "max_steps" is_log = state.step % args.log_every == 0 or is_eval or epoch_end or bool(state.stopped_by) row: dict[str, Any] | None = None if is_log: now = common.timer() elapsed = state.train_seconds + (now - t_loop) per_step = (now - t_inv_start - overhead_inv) / max(1, steps_this_invocation) eta = per_step * max(0, total_steps - state.step) row = {"step": state.step, "epoch": state.epoch, "loss": round(window_loss / max(1, window_steps), 5), "lr": sched.get_last_lr()[0], "examples_per_s": round(window_rows / max(1e-9, now - t_window), 1), "max_alloc_gb": round(vram["max_allocated_gb"], 3), "max_reserved_gb": round(vram["max_reserved_gb"], 3), "elapsed_s": round(elapsed, 1), "eta_s": round(eta)} LOG.info("step %d/%d epoch %d loss %.4f lr %.2e %.1f ex/s alloc %.2f GiB reserved %.2f GiB elapsed %s ETA %s", state.step, total_steps, state.epoch, row["loss"], row["lr"], row["examples_per_s"], row["max_alloc_gb"], row["max_reserved_gb"], _hms(elapsed), _hms(eta)) window_loss, window_rows, window_steps, t_window = 0.0, 0, 0, now # Evaluate at the cadence, at every epoch end and when --max-steps ends the run; never # after a signal (fast exit), a VRAM breach or a NaN model (it must not become best/). do_eval = (is_eval or epoch_end or hit_max) and state.step != last_eval_step \ and state.stopped_by in ("", "max_steps") if do_eval: t_ev = common.timer() record_eval(epoch_end) overhead_inv += common.timer() - t_ev last_eval_step = state.step if row is not None: m3 = state.last_val.get("macro_f1_3way") row.update(val_loss=round(state.last_val["val_loss"], 5), val_f1=round(state.last_val["f1"], 5), val_roc_auc=round(state.last_val["roc_auc"], 5), val_macro_f1_3way=None if m3 is None or math.isnan(m3) else round(m3, 5)) if state.evals_without_improvement >= args.patience and not state.stopped_by: state.stopped_by = "early_stopping" LOG.info("early stopping: %d evaluations without %s improvement (best %.4f at step %d)", state.evals_without_improvement, args.selection_metric, state.best_score, state.best_step) if row is not None: write_csv_row(metrics_csv, row, METRICS_HEADER) state.metrics_tail = (state.metrics_tail + [row])[-20:] if state.stopped_by: break if do_eval: flush_time() t_sv = common.timer() save_last(run_dir, model, opt, sched, state, args, schedule, device) overhead_inv += common.timer() - t_sv t_window = common.timer() # examples/s must not absorb eval + checkpoint time if not state.stopped_by: state.batch_in_epoch = 0 state.epoch += 1 if state.epoch >= args.epochs: state.stopped_by = "epochs_done" except KeyboardInterrupt: LOG.error("second interrupt: hard stop without a checkpoint") return 130 flush_time() save_last(run_dir, model, opt, sched, state, args, schedule, device) if args.cpu and torch.cuda.is_initialized(): # The --cpu contract is "the GPU is never touched"; a regression here would OOM the pretraining run. LOG.error("--cpu run initialised CUDA; a code path touched torch.cuda on the CPU dev path (bug)") exit_code = exit_code or 3 rec = experiment_record(args, state, params, steps_per_epoch) exp_path = Path(cfg["paths"]["reports"]) / "experiments.jsonl" how = upsert_experiment(exp_path, rec) LOG.info("%s experiment record for %s in %s", how, args.run_name, exp_path) if state.step >= RUN_LOG_MIN_STEPS: common.append_run_log(run_log_block(args, state, run_dir, params)) LOG.info("appended a RUN_LOG.md entry") v = state.best_val or {} print(f"\nTotal training time: {state.train_seconds:.1f} s ({_hms(state.train_seconds)}) over {state.step} optimizer steps" f" [{state.stopped_by}]; this invocation {common.timer() - t_invocation:.1f} s wall") print(f"Peak VRAM reserved {state.peak_reserved_gb:.2f} GiB (allocated {state.peak_allocated_gb:.2f} GiB) on {device_info['device']}") if v: print(f"Best val by {args.selection_metric} (step {state.best_step}, epoch {state.best_epoch}), {UNWANTED_NAME} roll-up: " + " ".join(f"{k}={v[k]:.4f}" for k in METRIC_KEYS) + f" -> {run_dir / 'best'}") if v.get("n_3way"): print(f"Three-way (n={v['n_3way']}): macro-F1={v['macro_f1_3way']:.4f} accuracy={v['accuracy_3way']:.4f} " + " ".join(f"F1[{LABEL_NAMES[c]}]={v[f'f1_{s}']:.4f}" for c, s in enumerate(CLASS_SLUGS)) + f"; phishing-vs-spam on true unwanted [pp ps sp ss]=[{v['pvs_phish_as_phish']} {v['pvs_phish_as_spam']} " f"{v['pvs_spam_as_phish']} {v['pvs_spam_as_spam']}]") return exit_code def _hms(seconds: float) -> str: s = int(round(seconds)) return f"{s // 3600:d}:{(s % 3600) // 60:02d}:{s % 60:02d}" # ------------------------------------------------------------------------------------ plot def plot(ns: argparse.Namespace, cfg: dict) -> int: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt run_dir = Path(cfg["paths"]["checkpoints"]) / "finetune" / ns.run_name metrics_csv, epochs_csv = run_dir / "metrics.csv", run_dir / "epochs.csv" if not metrics_csv.exists(): LOG.error("%s not found", metrics_csv) return 2 m = _read_csv(metrics_csv) e = _read_csv(epochs_csv) if epochs_csv.exists() else [] out = Path(ns.out) if ns.out else Path(cfg["paths"]["reports"]) / "finetune_curves.png" out.parent.mkdir(parents=True, exist_ok=True) # Two panels, one y-axis each (loss and scores live on different scales; no dual axis). fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4.2), dpi=150) steps = [float(r["step"]) for r in m] ax1.plot(steps, [float(r["loss"]) for r in m], color=PALETTE[0], lw=2, label="train loss") if e: ax1.plot([float(r["step"]) for r in e], [float(r["val_loss"]) for r in e], color=PALETTE[1], lw=2, marker="o", ms=4, label="val loss (CE)") ax1.set_xlabel("optimizer step") ax1.set_ylabel("loss") ax1.set_title(f"{ns.run_name}: loss") ax1.legend(frameon=False) if e: es = [float(r["step"]) for r in e] series = [("f1", "val F1 (unwanted roll-up)", PALETTE[0]), ("roc_auc", "val ROC-AUC", PALETTE[1]), ("pr_auc", "val PR-AUC", PALETTE[2])] if any(_num(r.get("macro_f1_3way")) is not None for r in e): series.append(("macro_f1_3way", "val macro-F1 (3-way)", PALETTE[3])) lows = [] for key, label, color in series: pts = [(s, _num(r.get(key))) for s, r in zip(es, e)] pts = [(s, y) for s, y in pts if y is not None] if not pts: continue ax2.plot([p[0] for p in pts], [p[1] for p in pts], color=color, lw=2, marker="o", ms=4, label=label) lows.append(min(p[1] for p in pts)) best = max(e, key=lambda r: float(r["f1"])) ax2.annotate(f"best F1 {float(best['f1']):.4f}\nstep {int(float(best['step']))}", (float(best["step"]), float(best["f1"])), textcoords="offset points", xytext=(8, -28), fontsize=8, color="#444444") ax2.set_ylim(max(0.0, min(0.5, min(lows) - 0.05)) if lows else 0.0, 1.0) ax2.legend(frameon=False, loc="lower right") else: ax2.text(0.5, 0.5, "no evaluations yet", ha="center", va="center", transform=ax2.transAxes) ax2.set_xlabel("optimizer step") ax2.set_ylabel("score") ax2.set_title(f"{ns.run_name}: validation") for ax in (ax1, ax2): ax.grid(True, color="#e6e6e6", lw=0.6) for side in ("top", "right"): ax.spines[side].set_visible(False) fig.tight_layout() fig.savefig(out) LOG.info("wrote %s (%d train rows, %d eval rows)", out, len(m), len(e)) print(out) return 0 def _read_csv(path: Path) -> list[dict[str, str]]: with path.open(newline="", encoding="utf-8") as fh: return [r for r in csv.DictReader(fh) if r.get("step")] def _num(s: str | None) -> float | None: """CSV cell -> float, None for blank / NaN (a binary run has no three-way columns).""" if s is None or s == "": return None try: v = float(s) except ValueError: return None return None if math.isnan(v) else v # ------------------------------------------------------------------------------------ list def list_experiments(ns: argparse.Namespace, cfg: dict) -> int: path = Path(cfg["paths"]["reports"]) / "experiments.jsonl" if not path.exists(): print(f"no experiments yet ({path})") return 0 runs: dict[str, dict[str, Any]] = {} for line in path.read_text(encoding="utf-8").splitlines(): if not line.strip(): continue rec = json.loads(line) slot = runs.setdefault(rec["run_name"], {}) slot["test" if rec.get("test") is not None else "train"] = rec def f(x: Any, nd: int = 4) -> str: return "-" if x is None or (isinstance(x, float) and math.isnan(x)) else (f"{x:.{nd}f}" if isinstance(x, float) else str(x)) # "val acc/P/R/F1" are the Unwanted (phishing or spam) roll-up at 0.5, identical to the binary runs' numbers; # the three-way columns are blank for binary (num_labels 2) runs. head = ["run", "init", "classes", "layers", "frozen", "epochs run/planned", "loss", "lr", "batch", "synthetic", "params (M)", "train min", "peak GiB", "best ep", "val acc", "val P", "val R", "val F1", "val ROC-AUC", "val PR-AUC", "val macro-F1 3w", "val P phish", "val R phish", "test F1", "test ROC-AUC", "stopped by"] rows = [] for name, slot in runs.items(): tr, te = slot.get("train"), slot.get("test") base = tr or te d = base.get("dials", {}) v = (tr or {}).get("val") or {} t = (te or {}).get("test") or {} phish = (v.get("per_class") or {}).get(LABEL_NAMES[1]) or {} init = base.get("init", "-") init = "scratch" if init == "scratch" else "pretrained" if init.startswith("pretrained") else init rows.append([name, init, f(d.get("num_labels", 2)), f(d.get("n_layers")), f(d.get("frozen_layers")), f"{f(d.get('epochs_run'), 2)}/{f(d.get('epochs_planned'))}", f(d.get("loss")), f"{d['lr']:.0e}" if d.get("lr") else "-", f(d.get("batch_size")), f"{d.get('synthetic', '-')}" + (f" ({d['synthetic_rows']})" if d.get("synthetic_rows") else ""), f((base.get("params") or {}).get("total", 0) / 1e6, 1) if base.get("params") else "-", f((tr or {}).get("train_seconds", 0) / 60.0, 1) if tr else "-", f((tr or {}).get("peak_vram_gb"), 2), f((tr or {}).get("best_epoch")), f(v.get("accuracy")), f(v.get("precision")), f(v.get("recall")), f(v.get("f1")), f(v.get("roc_auc")), f(v.get("pr_auc")), f(v.get("macro_f1_3way")), f(phish.get("precision")), f(phish.get("recall")), f(t.get("f1")), f(t.get("roc_auc")), f((tr or {}).get("stopped_by"))]) print("| " + " | ".join(head) + " |") print("|" + "|".join("---" for _ in head) + "|") for r in rows: print("| " + " | ".join(r) + " |") print(f"\n{len(rows)} run(s) from {path}") return 0 # ------------------------------------------------------------------------------------- CLI def parse_args(argv: list[str] | None = None) -> argparse.Namespace: ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) ap.add_argument("--config", default="config.yaml") sub = ap.add_subparsers(dest="cmd", required=True) tr = sub.add_parser("train", help="fine-tune the classifier (dials default to config.yaml; --resume restores the saved ones)") tr.add_argument("--run-name", default="finetune") tr.add_argument("--init", default=None, help="pretrained | pretrained: | scratch (default: newest checkpoints/pretrain/*/final)") tr.add_argument("--n-layers", type=int, default=None, help="use only the first N encoder layers") tr.add_argument("--freeze-layers", type=int, default=None, help="freeze embeddings + first K layers") tr.add_argument("--loss", choices=LOSSES, default=None) tr.add_argument("--focal-gamma", type=float, default=None) tr.add_argument("--lr", type=float, default=None) tr.add_argument("--batch-size", type=int, default=None, help="logical batch (rows per optimizer step)") tr.add_argument("--micro-batch", type=int, default=None, help="rows per fwd/bwd; batch-size must be a multiple (default 16)") tr.add_argument("--epochs", type=int, default=None) tr.add_argument("--max-steps", type=int, default=None, help="dev: stop after this many optimizer steps") tr.add_argument("--eval-every", type=int, default=None, help="optimizer steps between full val evaluations") tr.add_argument("--log-every", type=int, default=None) tr.add_argument("--eval-batch-size", type=int, default=None) tr.add_argument("--synthetic", choices=("on", "off"), default=None, help="append data/processed/synthetic.parquet to TRAIN") tr.add_argument("--limit-rows", type=int, default=None, help="dev: seeded subsample of train AND val") tr.add_argument("--seed", type=int, default=None) tr.add_argument("--num-labels", type=int, choices=NUM_LABELS_ALLOWED, default=None, help="override finetune.num_labels: 3 = legitimate/phishing/spam (partial-label loss), 2 = binary baseline") tr.add_argument("--cpu", action="store_true", help="dev only: run on the CPU with the config as given (tiny encoder + fixtures)") tr.add_argument("--resume", action="store_true", help="continue /last (dials restored; CLI flags override, except --init/--n-layers/--freeze-layers/" "--batch-size/--seed/--num-labels which are refused; --lr is re-applied to the restored optimizer)") tr.add_argument("--overwrite", action="store_true", help="discard an existing run dir of the same name") pl = sub.add_parser("plot", help="train loss and val F1/AUC vs step -> reports/finetune_curves.png") pl.add_argument("--run-name", required=True) pl.add_argument("--out", default=None) sub.add_parser("list", help="reports/experiments.jsonl as a markdown table") return ap.parse_args(argv) def main(argv: list[str] | None = None) -> int: ns = parse_args(argv) cfg = common.load_config(ns.config) global LOG LOG = common.setup_logging("finetune", cfg["paths"]["logs"]) if ns.cmd == "train": return train(ns, cfg) if ns.cmd == "plot": return plot(ns, cfg) return list_experiments(ns, cfg) if __name__ == "__main__": sys.exit(main())