feat: add ROOT-GNN tasks and training lifecycle

src/gnn4colliders/models/__init__.py

@@ -0,0 +1 @@
+"""Architecture-specific machine-learning models."""

src/gnn4colliders/models/root_gnn/__init__.py

@@ -0,0 +1,6 @@
+"""ROOT-GNN model components."""
+
+from .edge_network import EdgeNetwork
+from .transfer import FineTunedEdgeNetwork, load_legacy_edge_network_state_dict
+
+__all__ = ["EdgeNetwork", "FineTunedEdgeNetwork", "load_legacy_edge_network_state_dict"]

src/gnn4colliders/models/root_gnn/blocks.py

@@ -0,0 +1,39 @@
+"""Small neural-network blocks used by the active ROOT-GNN model."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from torch import nn
+
+
+def make_mlp(
+    in_size: int,
+    hidden_size: int,
+    out_size: int,
+    n_layers: int,
+    *,
+    dropout: float = 0.0,
+    activation: Callable[[], nn.Module] = nn.ReLU,
+) -> nn.Sequential:
+    """Build the legacy ``Make_MLP`` block.
+
+    Every linear layer is followed by activation and dropout, including the
+    final linear layer; LayerNorm is then applied to the output.  This order
+    is part of the ROOT-GNN forward contract.
+    """
+    if n_layers < 1:
+        raise ValueError("n_layers must be positive")
+    layers: list[nn.Module] = []
+    if n_layers == 1:
+        sizes = [(in_size, out_size)]
+    else:
+        sizes = [(in_size, hidden_size)]
+        sizes.extend((hidden_size, hidden_size) for _ in range(n_layers - 2))
+        sizes.append((hidden_size, out_size))
+    for input_size, output_size in sizes:
+        layers.extend(
+            (nn.Linear(input_size, output_size), activation(), nn.Dropout(dropout))
+        )
+    layers.append(nn.LayerNorm(out_size))
+    return nn.Sequential(*layers)

src/gnn4colliders/models/root_gnn/edge_network.py

@@ -0,0 +1,188 @@
+"""The active edge-message-passing ROOT-GNN architecture."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch import nn
+
+from .blocks import make_mlp
+
+
+def _broadcast(values: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
+    return torch.repeat_interleave(values, counts, dim=0)
+
+
+def _copy_destination(edges: Any) -> dict[str, torch.Tensor]:
+    """DGL 2.x has ``copy_u`` but no built-in destination copy primitive."""
+    return {"m_v": edges.dst["h"]}
+
+
+class EdgeNetwork(nn.Module):
+    """ROOT-GNN ``Edge_Network`` with an explicit feature/classifier split."""
+
+    def __init__(
+        self,
+        sample_graph: Any,
+        sample_global: torch.Tensor | None,
+        hid_size: int,
+        out_size: int,
+        n_layers: int,
+        n_proc_steps: int,
+        dropout: float = 0.0,
+        **_: Any,
+    ) -> None:
+        super().__init__()
+        if n_proc_steps < 0:
+            raise ValueError("n_proc_steps must be non-negative")
+        node_features = sample_graph.ndata["features"]
+        edge_features = sample_graph.edata["features"]
+        global_width = 0 if sample_global is None else sample_global.shape[1]
+        self.has_global = global_width != 0
+        if not self.has_global:
+            global_width = 1
+        self.hid_size = hid_size
+        self.n_layers = n_layers
+        self.n_proc_steps = n_proc_steps
+        self.node_feature_size = int(node_features.shape[1])
+        self.edge_feature_size = int(edge_features.shape[1])
+        self.global_feature_size = (
+            int(sample_global.shape[1]) if sample_global is not None else 0
+        )
+        self.dropout = float(dropout)
+
+        self.node_encoder = make_mlp(
+            node_features.shape[1], hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.edge_encoder = make_mlp(
+            edge_features.shape[1], hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.global_encoder = make_mlp(
+            global_width, hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.node_update = make_mlp(
+            3 * hid_size, hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.edge_update = make_mlp(
+            4 * hid_size, hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.global_update = make_mlp(
+            3 * hid_size, hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.global_decoder = make_mlp(
+            hid_size, hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.classifier = nn.Linear(hid_size, out_size)
+
+    def checkpoint_config(self) -> dict[str, Any]:
+        """Return constructor information without serializing a sample graph."""
+        return {
+            "family": "root_gnn",
+            "class": type(self).__name__,
+            "node_feature_size": self.node_feature_size,
+            "edge_feature_size": self.edge_feature_size,
+            "global_feature_size": self.global_feature_size,
+            "hid_size": self.hid_size,
+            "out_size": int(self.classifier.out_features),
+            "n_layers": self.n_layers,
+            "n_proc_steps": self.n_proc_steps,
+            "dropout": self.dropout,
+        }
+
+    # Compatibility with the legacy public attribute while keeping the new
+    # name in the model API.
+    @property
+    def classify(self) -> nn.Linear:
+        return self.classifier
+
+    def _features(
+        self, graph: Any, global_features: torch.Tensor | None
+    ) -> torch.Tensor:
+        if not self.has_global:
+            return graph.batch_num_nodes()[:, None].to(
+                device=graph.ndata["features"].device, dtype=torch.float32
+            )
+        if global_features is None:
+            raise ValueError("global_features are required for this model")
+        if global_features.ndim == 1:
+            if graph.batch_num_nodes().numel() != 1:
+                raise ValueError(
+                    "one-dimensional global_features are valid only for one graph"
+                )
+            global_features = global_features.unsqueeze(0)
+        return global_features
+
+    def forward_features(
+        self, graph: Any, global_features: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """Return the decoded graph representation before classification."""
+        try:
+            import dgl
+        except ImportError as error:  # pragma: no cover - optional dependency
+            raise ImportError("EdgeNetwork requires the 'root-gnn' extra") from error
+        if hasattr(graph, "graph") and hasattr(graph, "global_features"):
+            global_features = graph.global_features
+            graph = graph.graph
+        with graph.local_scope():
+            h = self.node_encoder(graph.ndata["features"])
+            e = self.edge_encoder(graph.edata["features"])
+            graph.ndata["h"] = h
+            graph.edata["e"] = e
+            h_global = self.global_encoder(self._features(graph, global_features))
+            counts = graph.batch_num_nodes()
+            for _ in range(self.n_proc_steps):
+                graph.apply_edges(dgl.function.copy_u("h", "m_u"))
+                graph.apply_edges(_copy_destination)
+                graph.edata["e"] = self.edge_update(
+                    torch.cat(
+                        (
+                            graph.edata["e"],
+                            graph.edata["m_u"],
+                            graph.edata["m_v"],
+                            _broadcast(h_global, graph.batch_num_edges()),
+                        ),
+                        dim=1,
+                    )
+                )
+                graph.update_all(
+                    dgl.function.copy_e("e", "m"), dgl.function.sum("m", "h_e")
+                )
+                graph.ndata["h"] = self.node_update(
+                    torch.cat(
+                        (
+                            graph.ndata["h"],
+                            graph.ndata["h_e"],
+                            _broadcast(h_global, counts),
+                        ),
+                        dim=1,
+                    )
+                )
+                if "w" in graph.ndata:
+                    mask = torch.any(graph.ndata["features"] != 0, dim=1)
+                    valid_counts = []
+                    start = 0
+                    for count in counts.tolist():
+                        valid_counts.append(mask[start : start + count].sum())
+                        start += count
+                    denominator = (
+                        torch.stack(valid_counts).to(h.dtype).clamp_min(1)[:, None]
+                    )
+                    mean_nodes = dgl.sum_nodes(graph, "h", "w") / denominator
+                else:
+                    mean_nodes = dgl.mean_nodes(graph, "h")
+                h_global = self.global_update(
+                    torch.cat((h_global, mean_nodes, dgl.mean_edges(graph, "e")), dim=1)
+                )
+            return self.global_decoder(h_global)
+
+    def forward(
+        self, graph: Any, global_features: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        return self.classifier(self.forward_features(graph, global_features))
+
+    def representation(
+        self, graph: Any, global_features: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """Legacy alias for the reusable decoded representation."""
+        return self.forward_features(graph, global_features)

src/gnn4colliders/models/root_gnn/transfer.py

@@ -0,0 +1,83 @@
+"""Explicit transfer-learning wrapper for ROOT-GNN."""
+
+from __future__ import annotations
+
+import copy
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+from torch import nn
+
+from gnn4colliders.compat.checkpoint import map_legacy_edge_network_state_dict
+
+from .edge_network import EdgeNetwork
+
+
+def load_legacy_edge_network_state_dict(
+    model: nn.Module, state: Mapping[str, Any]
+) -> None:
+    """Load a legacy checkpoint/state dict after historical prefix cleanup."""
+    state_dict = state.get("model_state_dict", state)
+    model.load_state_dict(map_legacy_edge_network_state_dict(state_dict), strict=True)
+
+
+class FineTunedEdgeNetwork(nn.Module):
+    """A pretrained ROOT-GNN backbone with a task-specific classifier."""
+
+    def __init__(
+        self, backbone: EdgeNetwork, out_size: int, *, freeze_backbone: bool = False
+    ) -> None:
+        super().__init__()
+        # Own an independent classifier-free copy: constructing a transfer
+        # model must not mutate the pretrained model supplied by the caller.
+        self.backbone = copy.deepcopy(backbone)
+        self.backbone.classifier = nn.Identity()
+        self.classifier = nn.Linear(backbone.hid_size, out_size)
+        self.freeze_backbone = freeze_backbone
+        self.set_backbone_trainable(not freeze_backbone)
+
+    def checkpoint_config(self) -> dict[str, Any]:
+        config = self.backbone.checkpoint_config()
+        config.update(
+            {
+                "class": type(self).__name__,
+                "out_size": int(self.classifier.out_features),
+                "freeze_backbone": self.freeze_backbone,
+            }
+        )
+        return config
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained: EdgeNetwork,
+        out_size: int,
+        *,
+        freeze_backbone: bool = False,
+        state_dict: Mapping[str, Any] | None = None,
+    ) -> "FineTunedEdgeNetwork":
+        if state_dict is not None:
+            load_legacy_edge_network_state_dict(pretrained, state_dict)
+        return cls(pretrained, out_size, freeze_backbone=freeze_backbone)
+
+    def set_backbone_trainable(self, trainable: bool) -> None:
+        for parameter in self.backbone.parameters():
+            parameter.requires_grad = trainable
+        for parameter in self.classifier.parameters():
+            parameter.requires_grad = True
+
+    def forward_features(
+        self, graph: Any, global_features: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        return self.backbone.forward_features(graph, global_features)
+
+    def forward(
+        self, graph: Any, global_features: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        return self.classifier(self.forward_features(graph, global_features))
+
+    def representation(
+        self, graph: Any, global_features: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        return self.forward_features(graph, global_features)

src/gnn4colliders/tasks/__init__.py

@@ -0,0 +1,6 @@
+"""Learning-task semantics for interpreting model outputs."""
+
+from .binary_classification import BinaryClassificationTask
+from .multiclass_classification import MulticlassClassificationTask
+
+__all__ = ["BinaryClassificationTask", "MulticlassClassificationTask"]

src/gnn4colliders/tasks/_common.py

@@ -0,0 +1,78 @@
+"""Small shared helpers for classification tasks."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+def batch_labels(batch: Any) -> torch.Tensor:
+    """Return labels from a GraphBatch-like object."""
+
+    try:
+        return batch.labels
+    except AttributeError as error:
+        raise TypeError("batch must provide a named 'labels' field") from error
+
+
+def batch_weights(
+    batch: Any, *, device: torch.device, dtype: torch.dtype
+) -> torch.Tensor:
+    """Return event weights through named metadata, never positional tracking."""
+
+    try:
+        weights = batch.metadata.weight
+    except AttributeError as error:
+        raise TypeError("batch must provide named metadata.weight") from error
+    return torch.as_tensor(weights, device=device, dtype=dtype).reshape(-1)
+
+
+def per_label_weighted_mean(
+    elementwise_loss: torch.Tensor,
+    labels: torch.Tensor,
+    weights: torch.Tensor,
+) -> torch.Tensor:
+    """Reproduce the legacy weighted loss reduction.
+
+    Each label receives its own weighted mean, and those means are then
+    averaged equally.  In particular, this is not ordinary weighted BCE/CE.
+    """
+
+    result = elementwise_loss.new_zeros(())
+    unique_labels = torch.unique(labels)
+    for label in unique_labels:
+        mask = labels == label
+        result = (
+            result
+            + (weights[mask] * elementwise_loss[mask]).sum() / weights[mask].sum()
+        )
+    return result / len(unique_labels)
+
+
+def positive_weight_mask(weights: torch.Tensor) -> torch.Tensor:
+    """The legacy metric path excludes non-positive original weights."""
+
+    return weights > 0
+
+
+def weighted_binary_auc(
+    labels: torch.Tensor, scores: torch.Tensor, weights: torch.Tensor
+) -> float:
+    """Compute weighted binary ROC AUC, including half-credit ties."""
+
+    labels = labels.detach().cpu().to(dtype=torch.bool)
+    scores = scores.detach().cpu().to(dtype=torch.float64)
+    weights = weights.detach().cpu().to(dtype=torch.float64)
+    positive = labels
+    negative = ~labels
+    positive_weight = weights[positive].sum()
+    negative_weight = weights[negative].sum()
+    if positive_weight <= 0 or negative_weight <= 0:
+        raise ValueError("ROC AUC requires positive and negative samples")
+    pair_scores = scores[positive, None] - scores[None, negative]
+    pair_value = (pair_scores > 0).to(torch.float64) + 0.5 * (pair_scores == 0)
+    pair_weights = weights[positive, None] * weights[None, negative]
+    return float(
+        (pair_value * pair_weights).sum() / (positive_weight * negative_weight)
+    )

src/gnn4colliders/tasks/binary_classification.py

@@ -0,0 +1,78 @@
+"""Binary classification task semantics."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from ._common import (
+    batch_labels,
+    batch_weights,
+    per_label_weighted_mean,
+    positive_weight_mask,
+    weighted_binary_auc,
+)
+
+
+class BinaryClassificationTask:
+    """Interpret one-logit model outputs as a weighted binary task."""
+
+    def __init__(self, *, absolute_weights: bool = False, threshold: float = 0.5):
+        self.absolute_weights = absolute_weights
+        self.threshold = threshold
+
+    @staticmethod
+    def _logits(logits: torch.Tensor) -> torch.Tensor:
+        if logits.ndim == 2 and logits.shape[1] == 1:
+            return logits[:, 0]
+        if logits.ndim == 1:
+            return logits
+        raise ValueError("binary logits must have shape [B] or [B, 1]")
+
+    @staticmethod
+    def _targets(batch: Any, *, device: torch.device) -> torch.Tensor:
+        labels = torch.as_tensor(batch_labels(batch), device=device)
+        if labels.ndim == 2 and labels.shape[1] == 1:
+            labels = labels[:, 0]
+        elif labels.ndim != 1:
+            raise ValueError("binary labels must have shape [B] or [B, 1]")
+        return labels.to(dtype=torch.float32)
+
+    def loss(self, logits: torch.Tensor, batch: Any) -> torch.Tensor:
+        raw_logits = self._logits(logits)
+        labels = self._targets(batch, device=raw_logits.device)
+        if labels.shape != raw_logits.shape:
+            raise ValueError("binary logits and labels must have the same batch size")
+        weights = batch_weights(batch, device=raw_logits.device, dtype=raw_logits.dtype)
+        if weights.shape != raw_logits.shape:
+            raise ValueError("metadata.weight must contain one value per event")
+        if self.absolute_weights:
+            weights = weights.abs()
+        elementwise = torch.nn.functional.binary_cross_entropy_with_logits(
+            raw_logits, labels.to(dtype=raw_logits.dtype), reduction="none"
+        )
+        return per_label_weighted_mean(elementwise, labels, weights)
+
+    def predictions(self, logits: torch.Tensor) -> dict[str, torch.Tensor]:
+        scores = torch.sigmoid(self._logits(logits))
+        return {"scores": scores, "predictions": scores >= self.threshold}
+
+    predict = predictions
+
+    def metrics(self, logits: torch.Tensor, batch: Any) -> dict[str, float]:
+        scores = self.predictions(logits)["scores"]
+        labels = self._targets(batch, device=scores.device)
+        weights = batch_weights(batch, device=scores.device, dtype=scores.dtype)
+        metric_weights = weights.abs() if self.absolute_weights else weights
+        mask = positive_weight_mask(weights)
+        try:
+            auc = weighted_binary_auc(
+                labels[mask] == 1, scores[mask], metric_weights[mask]
+            )
+        except ValueError:
+            auc = float("nan")
+        accuracy = (
+            (self.predictions(logits)["predictions"] == (labels == 1)).float().mean()
+        )
+        return {"accuracy": float(accuracy), "roc_auc": float(auc), "auc": float(auc)}

src/gnn4colliders/tasks/multiclass_classification.py

@@ -0,0 +1,80 @@
+"""Multiclass classification task semantics."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from ._common import (
+    batch_labels,
+    batch_weights,
+    per_label_weighted_mean,
+    positive_weight_mask,
+    weighted_binary_auc,
+)
+
+
+class MulticlassClassificationTask:
+    """Interpret ``[B, C]`` logits with legacy per-class weighting."""
+
+    def __init__(self, *, absolute_weights: bool = False):
+        self.absolute_weights = absolute_weights
+
+    @staticmethod
+    def _targets(batch: Any, *, device: torch.device) -> torch.Tensor:
+        labels = torch.as_tensor(batch_labels(batch), device=device)
+        if labels.ndim == 2 and labels.shape[1] == 1:
+            labels = labels[:, 0]
+        if labels.ndim != 1:
+            raise ValueError("multiclass labels must have shape [B] or [B, 1]")
+        return labels.to(dtype=torch.long)
+
+    def loss(self, logits: torch.Tensor, batch: Any) -> torch.Tensor:
+        if logits.ndim != 2 or logits.shape[1] < 2:
+            raise ValueError("multiclass logits must have shape [B, C], with C >= 2")
+        labels = self._targets(batch, device=logits.device)
+        if labels.shape[0] != logits.shape[0]:
+            raise ValueError(
+                "multiclass logits and labels must have the same batch size"
+            )
+        weights = batch_weights(batch, device=logits.device, dtype=logits.dtype)
+        if weights.shape != labels.shape:
+            raise ValueError("metadata.weight must contain one value per event")
+        if self.absolute_weights:
+            weights = weights.abs()
+        elementwise = torch.nn.functional.cross_entropy(
+            logits, labels, reduction="none"
+        )
+        return per_label_weighted_mean(elementwise, labels, weights)
+
+    @staticmethod
+    def predictions(logits: torch.Tensor) -> dict[str, torch.Tensor]:
+        if logits.ndim != 2 or logits.shape[1] < 2:
+            raise ValueError("multiclass logits must have shape [B, C], with C >= 2")
+        scores = torch.softmax(logits, dim=1)
+        return {"scores": scores, "predictions": scores.argmax(dim=1)}
+
+    predict = predictions
+
+    def metrics(self, logits: torch.Tensor, batch: Any) -> dict[str, float]:
+        output = self.predictions(logits)
+        labels = self._targets(batch, device=logits.device)
+        weights = batch_weights(batch, device=logits.device, dtype=logits.dtype)
+        metric_weights = weights.abs() if self.absolute_weights else weights
+        mask = positive_weight_mask(weights)
+        accuracy = (output["predictions"] == labels).float().mean()
+        try:
+            class_aucs = []
+            for class_index in range(output["scores"].shape[1]):
+                class_aucs.append(
+                    weighted_binary_auc(
+                        labels[mask] == class_index,
+                        output["scores"][mask, class_index],
+                        metric_weights[mask],
+                    )
+                )
+            auc = sum(class_aucs) / len(class_aucs)
+        except ValueError:
+            auc = float("nan")
+        return {"accuracy": float(accuracy), "roc_auc": float(auc), "auc": float(auc)}

src/gnn4colliders/training/__init__.py

@@ -0,0 +1,37 @@
+"""Architecture-independent training infrastructure."""
+
+from .checkpoint import (
+    CHECKPOINT_SCHEMA_VERSION,
+    CheckpointManager,
+    build_model_from_checkpoint,
+    load_legacy_checkpoint,
+    load_model_weights,
+    load_pretrained_edge_network,
+    restore_training_state,
+)
+from .early_stopping import EarlyStopping
+from .optim import build_optimizer, trainable_parameters
+from .reproducibility import seed_everything
+from .schedulers import build_scheduler
+from .state import BatchResult, EpochResult, TrainerState, TrainingHistory
+from .trainer import Trainer
+
+__all__ = [
+    "BatchResult",
+    "EarlyStopping",
+    "EpochResult",
+    "Trainer",
+    "TrainerState",
+    "TrainingHistory",
+    "build_optimizer",
+    "build_scheduler",
+    "seed_everything",
+    "trainable_parameters",
+    "CHECKPOINT_SCHEMA_VERSION",
+    "CheckpointManager",
+    "build_model_from_checkpoint",
+    "load_legacy_checkpoint",
+    "load_model_weights",
+    "load_pretrained_edge_network",
+    "restore_training_state",
+]

src/gnn4colliders/training/checkpoint.py

@@ -0,0 +1,328 @@
+"""Small, explicit checkpoint persistence and compatibility adapters.
+
+Checkpoint files are trusted PyTorch artifacts.  The payload itself contains
+only state dictionaries and primitive configuration/metadata, which keeps the
+new format independent of live trainer or task objects.
+"""
+
+from __future__ import annotations
+
+import os
+import random
+import re
+import socket
+import subprocess
+import tempfile
+from collections.abc import Mapping
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Callable
+
+import numpy as np
+import torch
+from torch import nn
+
+from gnn4colliders.compat.checkpoint import (
+    load_legacy_checkpoint,
+    map_legacy_edge_network_state_dict,
+    normalize_legacy_state_dict_keys,
+)
+from gnn4colliders.data.metadata import FEATURE_SCHEMA_VERSION, GRAPH_SCHEMA_VERSION
+
+from .state import TrainerState
+
+CHECKPOINT_SCHEMA_VERSION = 1
+_EPOCH_RE = re.compile(r"(?:epoch|model_epoch)[_-](\d+)")
+
+
+def _primitive(value: Any) -> Any:
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    if isinstance(value, Mapping):
+        return {str(key): _primitive(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_primitive(item) for item in value]
+    if isinstance(value, np.generic):
+        return value.item()
+    raise TypeError(f"checkpoint configuration must be primitive, got {type(value)!r}")
+
+
+def _state_dict(model: nn.Module) -> dict[str, Any]:
+    # Compiled and DDP models expose the useful state on the wrapped module.
+    target = getattr(model, "_orig_mod", model)
+    state = target.state_dict()
+    return {
+        key.removeprefix("module."): value.detach().cpu()
+        for key, value in state.items()
+    }
+
+
+def _rng_state() -> dict[str, Any]:
+    state: dict[str, Any] = {
+        "python": random.getstate(),
+        "numpy": np.random.get_state(),
+        "torch": torch.get_rng_state(),
+    }
+    if torch.cuda.is_available():
+        state["torch_cuda"] = torch.cuda.get_rng_state_all()
+    return state
+
+
+def _restore_rng(state: Mapping[str, Any]) -> None:
+    if "python" in state:
+        random.setstate(state["python"])
+    if "numpy" in state:
+        np.random.set_state(state["numpy"])
+    if "torch" in state:
+        torch.set_rng_state(state["torch"])
+    if "torch_cuda" in state and torch.cuda.is_available():
+        torch.cuda.set_rng_state_all(state["torch_cuda"])
+
+
+def _git_commit() -> str | None:
+    try:
+        result = subprocess.run(
+            ["git", "rev-parse", "HEAD"], capture_output=True, text=True, check=True
+        )
+    except (OSError, subprocess.CalledProcessError):
+        return None
+    return result.stdout.strip() or None
+
+
+def _model_config(model: nn.Module) -> dict[str, Any]:
+    config_method = getattr(model, "checkpoint_config", None)
+    if callable(config_method):
+        return _primitive(config_method())
+    return {"class": f"{type(model).__module__}.{type(model).__qualname__}"}
+
+
+def _task_config(task: Any) -> dict[str, Any]:
+    values = {"type": type(task).__name__}
+    for name in ("absolute_weights", "threshold", "num_classes"):
+        if hasattr(task, name):
+            values[name] = _primitive(getattr(task, name))
+    return values
+
+
+def load_model_weights(model: nn.Module, checkpoint: Mapping[str, Any]) -> None:
+    """Load only model weights; optimizer and lifecycle state are untouched."""
+    state = checkpoint.get("model_state_dict", checkpoint)
+    if not isinstance(state, Mapping):
+        raise ValueError("checkpoint model_state_dict must be a mapping")
+    normalized = normalize_legacy_state_dict_keys(state)
+    # Task 8 owns the historical classify -> classifier rename.  Use its one
+    # compatibility path when applicable, while preserving generic models.
+    if type(model).__module__.startswith("gnn4colliders.models.root_gnn"):
+        normalized = map_legacy_edge_network_state_dict(normalized)
+    try:
+        model.load_state_dict(normalized, strict=True)
+    except RuntimeError as error:
+        raise ValueError(
+            f"checkpoint weights are incompatible with {type(model).__name__}"
+        ) from error
+
+
+class CheckpointManager:
+    """Save and discover versioned checkpoints in a caller-selected directory."""
+
+    def __init__(self, directory: str | os.PathLike[str]) -> None:
+        self.directory = Path(directory)
+        self.directory.mkdir(parents=True, exist_ok=True)
+
+    def save(
+        self,
+        *,
+        model: nn.Module,
+        trainer_state: TrainerState,
+        optimizer: torch.optim.Optimizer | None = None,
+        scheduler: Any | None = None,
+        early_stopping: Any | None = None,
+        model_config: Mapping[str, Any] | None = None,
+        task_config: Mapping[str, Any] | None = None,
+        metadata: Mapping[str, Any] | None = None,
+        name: str | None = None,
+        monitor_name: str | None = None,
+        monitor_value: float | None = None,
+        include_rng: bool = True,
+    ) -> Path:
+        epoch = int(trainer_state.epoch)
+        payload: dict[str, Any] = {
+            "schema_version": CHECKPOINT_SCHEMA_VERSION,
+            "epoch": epoch,
+            "global_step": int(trainer_state.global_step),
+            "model_state_dict": _state_dict(model),
+            "optimizer_state_dict": optimizer.state_dict()
+            if optimizer is not None
+            else None,
+            "scheduler_state_dict": scheduler.state_dict()
+            if scheduler is not None
+            else None,
+            "early_stopping_state": early_stopping.state_dict()
+            if early_stopping is not None
+            else None,
+            "trainer_state": {
+                "epoch": epoch,
+                "global_step": int(trainer_state.global_step),
+            },
+            "model_config": _primitive(model_config)
+            if model_config is not None
+            else _model_config(model),
+            "task_config": _primitive(task_config) if task_config is not None else None,
+            "metadata": {
+                "created_at": datetime.now(timezone.utc).isoformat(),
+                "hostname": socket.gethostname(),
+                "git_commit": _git_commit(),
+                "feature_schema_version": FEATURE_SCHEMA_VERSION,
+                "graph_schema_version": GRAPH_SCHEMA_VERSION,
+                **(_primitive(metadata) if metadata is not None else {}),
+            },
+            "monitor_name": monitor_name,
+            "monitor_value": None if monitor_value is None else float(monitor_value),
+        }
+        if include_rng:
+            payload["rng_state"] = _rng_state()
+        filename = name or f"epoch_{epoch:04d}.pt"
+        target = self.directory / filename
+        with tempfile.NamedTemporaryFile(
+            dir=self.directory, prefix=f".{filename}.", suffix=".tmp", delete=False
+        ) as handle:
+            temporary = Path(handle.name)
+        try:
+            torch.save(payload, temporary)
+            os.replace(temporary, target)
+        finally:
+            temporary.unlink(missing_ok=True)
+        return target
+
+    @staticmethod
+    def _epoch(path: Path) -> int | None:
+        match = _EPOCH_RE.search(path.stem)
+        return int(match.group(1)) if match else None
+
+    def checkpoints(self) -> list[Path]:
+        return sorted(
+            (
+                path
+                for path in self.directory.glob("*.pt")
+                if self._epoch(path) is not None
+            ),
+            key=lambda path: self._epoch(path),
+        )
+
+    def latest(self) -> Path:
+        paths = self.checkpoints()
+        if not paths:
+            raise FileNotFoundError(f"no epoch checkpoints found in {self.directory}")
+        return paths[-1]
+
+    def best(self, *, mode: str = "min", monitor: str | None = None) -> Path:
+        if mode not in {"min", "max"}:
+            raise ValueError("mode must be 'min' or 'max'")
+        candidates: list[tuple[float, int, Path]] = []
+        for path in self.checkpoints():
+            payload = self.load(path)
+            if monitor is not None and payload.get("monitor_name") != monitor:
+                continue
+            value = payload.get("monitor_value")
+            if value is not None:
+                candidates.append((float(value), self._epoch(path) or -1, path))
+        if not candidates:
+            raise FileNotFoundError(
+                "no checkpoints contain the requested monitor value"
+            )
+        return (min if mode == "min" else max)(
+            candidates, key=lambda item: (item[0], item[1])
+        )[2]
+
+    @staticmethod
+    def load(
+        path: str | os.PathLike[str], *, map_location: Any = "cpu"
+    ) -> dict[str, Any]:
+        payload = torch.load(path, map_location=map_location, weights_only=False)
+        if not isinstance(payload, Mapping):
+            raise ValueError("checkpoint payload must be a mapping")
+        if "schema_version" not in payload:
+            return load_legacy_checkpoint(payload)
+        version = payload["schema_version"]
+        if version != CHECKPOINT_SCHEMA_VERSION:
+            raise ValueError(f"unsupported checkpoint schema version: {version!r}")
+        return dict(payload)
+
+
+def restore_training_state(
+    checkpoint: Mapping[str, Any],
+    *,
+    model: nn.Module,
+    trainer: Any | None = None,
+    optimizer: torch.optim.Optimizer | None = None,
+    scheduler: Any | None = None,
+    early_stopping: Any | None = None,
+    restore_rng: bool = True,
+) -> TrainerState:
+    """Restore explicitly selected resume components and return trainer state."""
+    load_model_weights(model, checkpoint)
+    if optimizer is not None and checkpoint.get("optimizer_state_dict") is not None:
+        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+    if scheduler is not None and checkpoint.get("scheduler_state_dict") is not None:
+        scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
+    if (
+        early_stopping is not None
+        and checkpoint.get("early_stopping_state") is not None
+    ):
+        early_stopping.load_state_dict(checkpoint["early_stopping_state"])
+    state_data = checkpoint.get("trainer_state", checkpoint)
+    state = TrainerState(
+        epoch=int(state_data.get("epoch", -1)),
+        global_step=int(state_data.get("global_step", 0)),
+    )
+    if trainer is not None:
+        trainer.state = state
+    if restore_rng and checkpoint.get("rng_state") is not None:
+        _restore_rng(checkpoint["rng_state"])
+    return state
+
+
+def build_model_from_checkpoint(
+    checkpoint: Mapping[str, Any], factory: Callable[[Mapping[str, Any]], nn.Module]
+) -> nn.Module:
+    """Build a model through an explicit caller-owned factory."""
+    config = checkpoint.get("model_config")
+    if not isinstance(config, Mapping):
+        raise ValueError("checkpoint does not contain reconstructable model_config")
+    model = factory(config)
+    load_model_weights(model, checkpoint)
+    return model
+
+
+def load_pretrained_edge_network(
+    checkpoint: Mapping[str, Any] | str | os.PathLike[str],
+    *,
+    sample_graph: Any,
+    sample_global: torch.Tensor | None = None,
+    map_location: Any = "cpu",
+    **overrides: Any,
+) -> nn.Module:
+    """Construct and load a new ROOT-GNN backbone for transfer learning."""
+    from gnn4colliders.models.root_gnn import EdgeNetwork
+
+    payload = (
+        CheckpointManager.load(checkpoint, map_location=map_location)
+        if not isinstance(checkpoint, Mapping)
+        else checkpoint
+    )
+    config = dict(payload.get("model_config") or {})
+    config.update(overrides)
+    config.pop("class", None)
+    required = {"hid_size", "out_size", "n_layers", "n_proc_steps"}
+    missing = required - config.keys()
+    if missing:
+        raise ValueError(
+            f"checkpoint lacks EdgeNetwork configuration: {sorted(missing)}"
+        )
+    model = EdgeNetwork(
+        sample_graph,
+        sample_global,
+        **{key: config[key] for key in required | {"dropout"} if key in config},
+    )
+    load_model_weights(model, payload)
+    return model

src/gnn4colliders/training/early_stopping.py

@@ -0,0 +1,72 @@
+"""In-memory early-stopping policy."""
+
+from __future__ import annotations
+
+import math
+
+
+class EarlyStopping:
+    """Stop after ``patience`` consecutive non-improvements.
+
+    Equality is a non-improvement, matching the legacy ``EarlyStop`` policy.
+    """
+
+    def __init__(
+        self,
+        *,
+        monitor: str = "loss",
+        mode: str = "min",
+        patience: int = 15,
+        min_delta: float = 1e-8,
+    ) -> None:
+        if mode not in {"min", "max"}:
+            raise ValueError("mode must be 'min' or 'max'")
+        if patience < 1:
+            raise ValueError("patience must be positive")
+        if min_delta < 0:
+            raise ValueError("min_delta must be non-negative")
+        self.monitor, self.mode, self.patience, self.min_delta = (
+            monitor,
+            mode,
+            patience,
+            min_delta,
+        )
+        self.best = math.inf if mode == "min" else -math.inf
+        self.num_bad_epochs = 0
+        self.should_stop = False
+
+    def update(self, value: float) -> bool:
+        value = float(value)
+        improved = (
+            value < self.best - self.min_delta
+            if self.mode == "min"
+            else value > self.best + self.min_delta
+        )
+        if improved:
+            self.best = value
+            self.num_bad_epochs = 0
+        else:
+            self.num_bad_epochs += 1
+        self.should_stop = self.num_bad_epochs >= self.patience
+        return self.should_stop
+
+    def state_dict(self) -> dict[str, object]:
+        """Return the primitive state needed to resume this policy."""
+        return {
+            "monitor": self.monitor,
+            "mode": self.mode,
+            "patience": self.patience,
+            "min_delta": self.min_delta,
+            "best": self.best,
+            "num_bad_epochs": self.num_bad_epochs,
+            "should_stop": self.should_stop,
+        }
+
+    def load_state_dict(self, state: dict[str, object]) -> None:
+        """Restore policy state, rejecting incompatible policy settings."""
+        for name in ("monitor", "mode", "patience", "min_delta"):
+            if name in state and getattr(self, name) != state[name]:
+                raise ValueError(f"early-stopping {name} does not match checkpoint")
+        self.best = float(state["best"])
+        self.num_bad_epochs = int(state["num_bad_epochs"])
+        self.should_stop = bool(state["should_stop"])

src/gnn4colliders/training/optim.py

@@ -0,0 +1,38 @@
+"""Explicit construction helpers for the active optimizer workflow."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+
+
+def trainable_parameters(model: nn.Module) -> Iterable[nn.Parameter]:
+    """Yield only parameters enabled for optimization."""
+
+    return (parameter for parameter in model.parameters() if parameter.requires_grad)
+
+
+def build_optimizer(
+    parameters: Iterable[nn.Parameter] | nn.Module,
+    *,
+    name: str = "adam",
+    learning_rate: float = 1e-4,
+    weight_decay: float = 0.0,
+    **kwargs: object,
+) -> torch.optim.Optimizer:
+    """Build an optimizer used by standard ROOT-GNN configurations."""
+
+    if isinstance(parameters, nn.Module):
+        parameters = trainable_parameters(parameters)
+    normalized = name.lower().replace("_", "")
+    if normalized == "adam":
+        return torch.optim.Adam(
+            parameters, lr=learning_rate, weight_decay=weight_decay, **kwargs
+        )
+    if normalized == "adamw":
+        return torch.optim.AdamW(
+            parameters, lr=learning_rate, weight_decay=weight_decay, **kwargs
+        )
+    raise ValueError(f"unsupported optimizer {name!r}; use 'adam' or 'adamw'")

src/gnn4colliders/training/reproducibility.py

@@ -0,0 +1,20 @@
+"""Explicit process-level reproducibility configuration."""
+
+from __future__ import annotations
+
+import random
+
+import numpy as np
+import torch
+
+
+def seed_everything(seed: int, *, deterministic: bool = False) -> None:
+    """Seed Python, NumPy, and Torch; optionally request strict Torch behavior."""
+
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    if deterministic:
+        torch.use_deterministic_algorithms(True)

src/gnn4colliders/training/schedulers.py

@@ -0,0 +1,22 @@
+"""Scheduler construction with explicit epoch/metric stepping."""
+
+from __future__ import annotations
+
+import torch
+
+
+def build_scheduler(
+    optimizer: torch.optim.Optimizer,
+    *,
+    name: str = "exponential",
+    gamma: float = 1.0,
+    **kwargs: object,
+) -> torch.optim.lr_scheduler.LRScheduler | torch.optim.lr_scheduler.ReduceLROnPlateau:
+    """Build the active epoch scheduler or a validation-driven scheduler."""
+
+    normalized = name.lower().replace("_", "")
+    if normalized in {"exponential", "exponentiallr"}:
+        return torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma)
+    if normalized in {"reducelronplateau", "plateau"}:
+        return torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, **kwargs)
+    raise ValueError(f"unsupported scheduler {name!r}")

src/gnn4colliders/training/state.py

@@ -0,0 +1,31 @@
+"""Small in-memory types used by the training lifecycle."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Mapping
+
+
+@dataclass(frozen=True)
+class BatchResult:
+    loss: float
+    logits: object
+
+
+@dataclass(frozen=True)
+class EpochResult:
+    loss: float
+    metrics: Mapping[str, float] = field(default_factory=dict)
+    num_samples: int = 0
+
+
+@dataclass
+class TrainingHistory:
+    train: list[EpochResult] = field(default_factory=list)
+    validation: list[EpochResult] = field(default_factory=list)
+
+
+@dataclass
+class TrainerState:
+    epoch: int = -1
+    global_step: int = 0

src/gnn4colliders/training/trainer.py

@@ -0,0 +1,272 @@
+"""Single-process, architecture-independent model training lifecycle."""
+
+from __future__ import annotations
+
+import inspect
+import logging
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+from torch import nn
+
+from gnn4colliders.data.graph_dataset import GraphBatch
+from gnn4colliders.distributed import (
+    DistributedContext,
+    broadcast_bool,
+    gather_objects,
+    gather_tensor,
+    prepare_model,
+)
+
+from .early_stopping import EarlyStopping
+from .state import BatchResult, EpochResult, TrainerState, TrainingHistory
+
+logger = logging.getLogger(__name__)
+
+
+def _model_forward(model: nn.Module, batch: GraphBatch) -> torch.Tensor:
+    """Call both supported model boundaries without masking model errors."""
+
+    try:
+        signature = inspect.signature(model.forward)
+        positional = [
+            parameter
+            for parameter in signature.parameters.values()
+            if parameter.kind
+            in (parameter.POSITIONAL_ONLY, parameter.POSITIONAL_OR_KEYWORD)
+        ]
+    except (TypeError, ValueError):
+        positional = []
+    if len(positional) <= 1:
+        return model(batch)
+    return model(batch.graph, batch.global_features)
+
+
+def _combined_batch(batches: list[GraphBatch], logits: torch.Tensor) -> Any:
+    labels = torch.cat([batch.labels.reshape(-1) for batch in batches])
+    weights = torch.cat([batch.metadata.weight.reshape(-1) for batch in batches])
+    sample_ids = tuple(
+        sample_id for batch in batches for sample_id in batch.metadata.sample_id
+    )
+    return SimpleNamespace(
+        labels=labels,
+        metadata=SimpleNamespace(weight=weights, sample_id=sample_ids),
+        logits=logits,
+    )
+
+
+class Trainer:
+    """Orchestrate model, task, optimizer, and optional scheduler."""
+
+    def __init__(
+        self,
+        model: nn.Module,
+        task: Any,
+        optimizer: torch.optim.Optimizer,
+        scheduler: Any | None = None,
+        *,
+        device: torch.device | str = "cpu",
+        early_stopping: EarlyStopping | None = None,
+        scheduler_step: str = "epoch",
+        distributed_context: DistributedContext | None = None,
+    ) -> None:
+        if scheduler_step not in {"epoch", "validation"}:
+            raise ValueError("scheduler_step must be 'epoch' or 'validation'")
+        self.distributed_context = (
+            distributed_context or DistributedContext.single_process(device)
+        )
+        self.device = self.distributed_context.device
+        self.model = prepare_model(model, self.distributed_context)
+        self.task = task
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.early_stopping = early_stopping
+        self.scheduler_step = scheduler_step
+        self.state = TrainerState()
+
+    def _train_moved_batch(self, batch: GraphBatch) -> BatchResult:
+        self.optimizer.zero_grad(set_to_none=True)
+        logits = _model_forward(self.model, batch)
+        loss = self.task.loss(logits, batch)
+        if not torch.isfinite(loss):
+            raise FloatingPointError("non-finite training loss")
+        loss.backward()
+        self.optimizer.step()
+        self.state.global_step += 1
+        return BatchResult(loss=float(loss.detach().cpu()), logits=logits.detach())
+
+    def train_batch(self, batch: GraphBatch) -> BatchResult:
+        """Run one optimizer step, moving the supplied batch to the device."""
+
+        return self._train_moved_batch(batch.to(self.device))
+
+    def train_epoch(self, loader: Any) -> EpochResult:
+        self.model.train()
+        losses: list[float] = []
+        batches: list[GraphBatch] = []
+        outputs: list[torch.Tensor] = []
+        num_samples = 0
+        for batch in loader:
+            moved = batch.to(self.device)
+            result = self._train_moved_batch(moved)
+            losses.append(result.loss)
+            batches.append(moved)
+            outputs.append(result.logits)
+            num_samples += int(moved.labels.shape[0])
+        if not losses:
+            raise ValueError("cannot train on an empty loader")
+        logits = torch.cat(outputs, dim=0)
+        combined = _combined_batch(batches, logits)
+        if self.distributed_context.enabled:
+            combined = self._gather_combined(combined)
+            # Compute the reported loss over the global event set.  Gradients
+            # still came from local losses and were synchronized by DDP.
+            global_loss = float(self.task.loss(combined.logits, combined).cpu())
+        else:
+            global_loss = sum(losses) / len(losses)
+        metrics = {
+            key: float(value)
+            for key, value in self.task.metrics(logits, combined).items()
+        }
+        # Legacy training reports the arithmetic mean of batch losses, while
+        # metrics are evaluated over every event in the complete epoch.
+        return EpochResult(
+            loss=global_loss, metrics=metrics, num_samples=int(combined.labels.shape[0])
+        )
+
+    def evaluate(
+        self, loader: Any, *, return_outputs: bool = False
+    ) -> EpochResult | Any:
+        self.model.eval()
+        batches: list[GraphBatch] = []
+        outputs: list[torch.Tensor] = []
+        with torch.inference_mode():
+            for batch in loader:
+                moved = batch.to(self.device)
+                outputs.append(_model_forward(self.model, moved).detach())
+                batches.append(moved)
+        if not outputs:
+            raise ValueError("cannot evaluate an empty loader")
+        logits = torch.cat(outputs, dim=0)
+        combined = _combined_batch(batches, logits)
+        if self.distributed_context.enabled:
+            combined = self._gather_combined(combined)
+        loss = self.task.loss(logits, combined)
+        metrics = {
+            key: float(value)
+            for key, value in self.task.metrics(logits, combined).items()
+        }
+        result = EpochResult(
+            loss=float(loss.cpu()),
+            metrics=metrics,
+            num_samples=int(logits.shape[0]),
+        )
+        if return_outputs:
+            return (
+                result,
+                logits.cpu(),
+                combined.labels.cpu(),
+                combined.metadata.weight.cpu(),
+                combined.metadata.sample_id,
+            )
+        return result
+
+    def _gather_combined(self, combined: Any) -> Any:
+        logits = gather_tensor(combined.logits, self.distributed_context)
+        labels = gather_tensor(combined.labels, self.distributed_context)
+        weights = gather_tensor(combined.metadata.weight, self.distributed_context)
+        ids = tuple(
+            item
+            for group in gather_objects(
+                combined.metadata.sample_id, self.distributed_context
+            )
+            for item in group
+        )
+        return SimpleNamespace(
+            labels=labels,
+            metadata=SimpleNamespace(weight=weights, sample_id=ids),
+            logits=logits,
+        )
+
+    def fit(
+        self,
+        train_loader: Any,
+        validation_loader: Any | None = None,
+        *,
+        epochs: int,
+    ) -> TrainingHistory:
+        """Train using validation data for every-epoch model selection.
+
+        The held-out test split is intentionally absent from this method and
+        should be passed to :meth:`evaluate` only after model selection is
+        complete.  This avoids the legacy train/test/val naming inversion.
+        """
+
+        if epochs < 1:
+            raise ValueError("epochs must be positive")
+        if self.early_stopping is not None and validation_loader is None:
+            raise ValueError("early stopping requires a validation loader")
+        history = TrainingHistory()
+        # ``state.epoch`` is the last completed epoch.  This makes a saved
+        # state unambiguous: a resumed fit starts at ``epoch + 1``.
+        start_epoch = self.state.epoch + 1
+        for epoch in range(start_epoch, start_epoch + epochs):
+            if hasattr(train_loader, "set_epoch"):
+                train_loader.set_epoch(epoch)
+            if validation_loader is not None and hasattr(
+                validation_loader, "set_epoch"
+            ):
+                validation_loader.set_epoch(epoch)
+            train_result = self.train_epoch(train_loader)
+            history.train.append(train_result)
+            validation_result = None
+            if validation_loader is not None:
+                validation_result = self.evaluate(validation_loader)
+                history.validation.append(validation_result)
+            if self.distributed_context.is_main_process:
+                logger.info("epoch=%d train_loss=%s", epoch, train_result.loss)
+            if (
+                validation_result is not None
+                and self.distributed_context.is_main_process
+            ):
+                logger.info(
+                    "epoch=%d validation_loss=%s metrics=%s",
+                    epoch,
+                    validation_result.loss,
+                    validation_result.metrics,
+                )
+            # Record the completed epoch before policy logic can terminate
+            # the loop, so a checkpoint after early stopping is resumable.
+            self.state.epoch = epoch
+            if self.early_stopping is not None:
+                if self.early_stopping.monitor == "loss":
+                    monitored = validation_result.loss
+                else:
+                    if validation_result is None:
+                        raise ValueError("early stopping monitor requires validation")
+                    try:
+                        monitored = validation_result.metrics[
+                            self.early_stopping.monitor
+                        ]
+                    except KeyError as error:
+                        raise KeyError(
+                            f"metric {self.early_stopping.monitor!r} was not produced"
+                        ) from error
+                stop = (
+                    self.early_stopping.update(monitored)
+                    if self.distributed_context.is_main_process
+                    else False
+                )
+                if broadcast_bool(stop, self.distributed_context):
+                    break
+            if self.scheduler is not None:
+                if self.scheduler_step == "validation":
+                    if validation_result is None:
+                        raise ValueError(
+                            "validation scheduler requires a validation loader"
+                        )
+                    self.scheduler.step(validation_result.loss)
+                else:
+                    self.scheduler.step()
+        return history

tests/integration/test_real_root_model.py

@@ -0,0 +1,109 @@
+"""ROOT-GNN model coverage over every entry in the 64-event ROOT fixture."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import awkward as ak
+import pytest
+import torch
+
+from gnn4colliders.data import EventMetadata, GraphSample, batch_graph_samples
+from gnn4colliders.features import build_node_features
+from gnn4colliders.graphs import build_dgl_graph
+from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork
+
+pytest.importorskip("dgl")
+uproot = pytest.importorskip("uproot")
+
+FEATURE_BRANCHES = [
+    ["jet_pt", "ele_pt", "mu_pt", "ph_pt", "MET_met"],
+    ["jet_eta", "ele_eta", "mu_eta", "ph_eta", 0],
+    ["jet_phi", "ele_phi", "mu_phi", "ph_phi", "MET_phi"],
+    "CALC_E",
+    ["jet_btag", 0, 0, 0, 0],
+    [0, "ele_charge", "mu_charge", 0, 0],
+    "NODE_TYPE",
+]
+OBJECT_TYPES = ["vector", "vector", "vector", "vector", "single"]
+FEATURE_SCALES = [0.1, 1, 1, 0.1, 1, 1, 1]
+
+
+def _fixture_path() -> Path:
+    configured = os.environ.get("GNN4COLLIDERS_ROOT_FIXTURE")
+    return Path(configured) if configured else Path("data/processed/ttH_NLO_64.root")
+
+
+@pytest.fixture(scope="module")
+def root_graph_samples():
+    path = _fixture_path()
+    if not path.exists():
+        pytest.skip(f"ROOT sample fixture is absent: {path}")
+    with uproot.open(path) as root_file:
+        arrays = root_file["output"].arrays(entry_start=0, entry_stop=64, library="ak")
+    samples = []
+    for index, event in enumerate(ak.Array(arrays)):
+        features, _ = build_node_features(
+            event, FEATURE_BRANCHES, OBJECT_TYPES, FEATURE_SCALES
+        )
+        samples.append(
+            GraphSample(
+                graph=build_dgl_graph(features),
+                label=torch.tensor(index % 2, dtype=torch.long),
+                global_features=torch.tensor(
+                    [float(event["Number"]), float(event["weight"])],
+                    dtype=torch.float32,
+                ),
+                metadata=EventMetadata(
+                    fold=index % 4,
+                    weight=float(event["weight"]),
+                    sample_id=f"fixture:{index}",
+                ),
+            )
+        )
+    return samples
+
+
+def test_all_64_real_root_graphs_run_as_one_graph_batch(root_graph_samples):
+    batch = batch_graph_samples(root_graph_samples)
+    model = EdgeNetwork(
+        root_graph_samples[0].graph,
+        torch.zeros(1, 2),
+        hid_size=16,
+        out_size=12,
+        n_layers=2,
+        n_proc_steps=2,
+    ).eval()
+    transferred = FineTunedEdgeNetwork.from_pretrained(
+        model, out_size=1, freeze_backbone=True
+    ).eval()
+    with torch.no_grad():
+        logits = model(batch)
+        transfer_logits = transferred(batch)
+    assert batch.graph.batch_num_nodes().shape == (64,)
+    assert logits.shape == (64, 12)
+    assert transfer_logits.shape == (64, 1)
+    assert torch.isfinite(logits).all()
+    assert torch.isfinite(transfer_logits).all()
+
+
+def test_real_root_graph_batch_matches_individual_forward(root_graph_samples):
+    model = EdgeNetwork(
+        root_graph_samples[0].graph,
+        torch.zeros(1, 2),
+        hid_size=8,
+        out_size=1,
+        n_layers=2,
+        n_proc_steps=1,
+    ).eval()
+    batch = batch_graph_samples(root_graph_samples)
+    with torch.no_grad():
+        batched = model(batch)
+        individual = torch.cat(
+            [
+                model(sample.graph, sample.global_features)
+                for sample in root_graph_samples
+            ]
+        )
+    assert torch.allclose(batched, individual, atol=1e-5, rtol=1e-5)

tests/integration/test_real_root_sample.py

@@ -0,0 +1,74 @@
+"""Smoke tests against a small fixture derived from the real ttH ROOT sample."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import awkward as ak
+import pytest
+import torch
+import uproot
+
+from gnn4colliders.features import build_node_features
+from gnn4colliders.graphs import build_dgl_graph
+
+pytest.importorskip("dgl")
+
+FEATURE_BRANCHES = [
+    ["jet_pt", "ele_pt", "mu_pt", "ph_pt", "MET_met"],
+    ["jet_eta", "ele_eta", "mu_eta", "ph_eta", 0],
+    ["jet_phi", "ele_phi", "mu_phi", "ph_phi", "MET_phi"],
+    "CALC_E",
+    ["jet_btag", 0, 0, 0, 0],
+    [0, "ele_charge", "mu_charge", 0, 0],
+    "NODE_TYPE",
+]
+OBJECT_TYPES = ["vector", "vector", "vector", "vector", "single"]
+FEATURE_SCALES = [0.1, 1, 1, 0.1, 1, 1, 1]
+
+
+def _fixture_path() -> Path:
+    configured = os.environ.get("GNN4COLLIDERS_ROOT_FIXTURE")
+    if configured:
+        return Path(configured)
+    return Path("data/processed/ttH_NLO_64.root")
+
+
+@pytest.fixture(scope="module")
+def root_tree():
+    path = _fixture_path()
+    if not path.exists():
+        pytest.skip(
+            f"ROOT sample fixture is absent: {path}; download ttH_NLO.root and "
+            "create the reduced fixture first"
+        )
+    root_file = uproot.open(path)
+    tree = root_file["output"]
+    try:
+        yield tree
+    finally:
+        root_file.close()
+
+
+def test_real_root_sample_has_expected_fixture_shape(root_tree):
+    assert root_tree.num_entries == 64
+    assert {"jet_pt", "MET_met", "weight", "Number"}.issubset(root_tree.keys())
+
+
+def test_real_root_events_build_features_and_graphs(root_tree):
+    arrays = root_tree.arrays(entry_start=0, entry_stop=8, library="ak")
+    for event in ak.Array(arrays):
+        event_features, lengths = build_node_features(
+            event, FEATURE_BRANCHES, OBJECT_TYPES, FEATURE_SCALES
+        )
+        graph = build_dgl_graph(event_features)
+
+        assert event_features.dtype == torch.float32
+        assert event_features.shape == (sum(lengths), 7)
+        assert graph.number_of_nodes() == event_features.shape[0]
+        expected_edges = event_features.shape[0] * max(event_features.shape[0] - 1, 1)
+        assert graph.number_of_edges() == expected_edges
+        assert graph.edata["features"].shape == (expected_edges, 3)
+        assert torch.isfinite(event_features).all()
+        assert torch.isfinite(graph.edata["features"]).all()

tests/integration/test_root_gnn_model.py

@@ -0,0 +1,23 @@
+import pytest
+import torch
+
+from gnn4colliders.data import EventMetadata, GraphSample, batch_graph_samples
+from gnn4colliders.graphs import build_dgl_graph
+from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork
+
+pytest.importorskip("dgl")
+
+
+def test_graph_batch_runs_through_pretraining_and_transfer_models():
+    graph = build_dgl_graph(torch.tensor([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]]))
+    sample = GraphSample(
+        graph=graph,
+        label=torch.tensor(1),
+        global_features=torch.tensor([1.0, 2.0]),
+        metadata=EventMetadata(fold=0, weight=1.0, sample_id="fixture:0"),
+    )
+    batch = batch_graph_samples([sample])
+    model = EdgeNetwork(graph, torch.zeros(1, 2), 8, 12, 2, 1)
+    transferred = FineTunedEdgeNetwork.from_pretrained(model, 1)
+    assert model(batch).shape == (1, 12)
+    assert transferred(batch).shape == (1, 1)

tests/unit/models/root_gnn/test_blocks.py

@@ -0,0 +1,22 @@
+import pytest
+from torch import nn
+
+from gnn4colliders.models.root_gnn.blocks import make_mlp
+
+
+def test_make_mlp_matches_legacy_layer_order():
+    block = make_mlp(3, 4, 5, 2, dropout=0.1)
+    assert [type(layer) for layer in block] == [
+        nn.Linear,
+        nn.ReLU,
+        nn.Dropout,
+        nn.Linear,
+        nn.ReLU,
+        nn.Dropout,
+        nn.LayerNorm,
+    ]
+
+
+def test_make_mlp_rejects_empty_network():
+    with pytest.raises(ValueError, match="positive"):
+        make_mlp(3, 4, 5, 0)

tests/unit/models/root_gnn/test_edge_cases.py

@@ -0,0 +1,89 @@
+import pytest
+import torch
+
+from gnn4colliders.graphs import build_dgl_graph
+from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork
+
+pytest.importorskip("dgl")
+
+
+def _graph(node_features):
+    graph = build_dgl_graph(torch.tensor(node_features, dtype=torch.float32))
+    return graph, torch.tensor([[1.0]], dtype=torch.float32)
+
+
+def test_no_globals_uses_graph_node_count_as_legacy_fallback():
+    graph, _ = _graph([[10.0, 0.2, 0.1]])
+    model = EdgeNetwork(graph, torch.empty((1, 0)), 4, 1, 1, 0).eval()
+    with torch.no_grad():
+        logits = model(graph)
+    assert logits.shape == (1, 1)
+    assert torch.isfinite(logits).all()
+
+
+def test_missing_or_mismatched_global_shapes_fail_clearly():
+    graph, globals_ = _graph([[10.0, 0.2, 0.1]])
+    model = EdgeNetwork(graph, globals_, 4, 1, 1, 0)
+    with pytest.raises(ValueError, match="global_features are required"):
+        model(graph)
+    batched_graph = build_dgl_graph(
+        torch.tensor([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]], dtype=torch.float32)
+    )
+    batched_graph = __import__("dgl").batch([graph, batched_graph])
+    with pytest.raises(ValueError, match="one-dimensional"):
+        model(batched_graph, torch.tensor([1.0]))
+
+
+def test_model_preserves_float64_dtype():
+    graph = build_dgl_graph(
+        torch.tensor([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]], dtype=torch.float64)
+    )
+    globals_ = torch.tensor([[1.0]], dtype=torch.float64)
+    model = EdgeNetwork(graph, globals_, 4, 1, 1, 1).double().eval()
+    assert model(graph, globals_).dtype == torch.float64
+
+
+def test_negative_processing_steps_are_rejected():
+    graph, globals_ = _graph([[10.0, 0.2, 0.1]])
+    with pytest.raises(ValueError, match="non-negative"):
+        EdgeNetwork(graph, globals_, 4, 1, 1, -1)
+
+
+def test_single_node_graph_and_zero_processing_step_are_supported():
+    graph, globals_ = _graph([[10.0, 0.2, 0.1]])
+    model = EdgeNetwork(graph, globals_, 4, 1, 1, 0).eval()
+    assert graph.number_of_edges() == 1
+    assert model(graph, globals_).shape == (1, 1)
+
+
+def test_weighted_padded_nodes_produce_finite_logits():
+    graph, globals_ = _graph([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2], [0.0, 0.0, 0.0]])
+    graph.ndata["w"] = torch.tensor([[1.0], [1.0], [0.0]])
+    model = EdgeNetwork(graph, globals_, 4, 1, 1, 1).eval()
+    logits = model(graph, globals_)
+    assert logits.shape == (1, 1)
+    assert torch.isfinite(logits).all()
+
+
+def test_unfrozen_transfer_backbone_receives_gradients():
+    graph, globals_ = _graph([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]])
+    source = EdgeNetwork(graph, globals_, 4, 3, 1, 1)
+    transfer = FineTunedEdgeNetwork.from_pretrained(source, 1, freeze_backbone=False)
+    transfer(graph, globals_).sum().backward()
+    assert any(
+        parameter.grad is not None for parameter in transfer.backbone.parameters()
+    )
+    assert all(
+        parameter.grad is not None for parameter in transfer.classifier.parameters()
+    )
+
+
+def test_frozen_transfer_backbone_has_no_gradients_and_source_is_unchanged():
+    graph, globals_ = _graph([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]])
+    source = EdgeNetwork(graph, globals_, 4, 3, 1, 1)
+    source_classifier_type = type(source.classifier)
+    transfer = FineTunedEdgeNetwork.from_pretrained(source, 1, freeze_backbone=True)
+    transfer(graph, globals_).sum().backward()
+    assert isinstance(source.classifier, source_classifier_type)
+    assert isinstance(transfer.backbone.classifier, torch.nn.Identity)
+    assert all(parameter.grad is None for parameter in transfer.backbone.parameters())

tests/unit/models/root_gnn/test_edge_network.py

@@ -0,0 +1,41 @@
+import pytest
+import torch
+
+from gnn4colliders.graphs import build_dgl_graph
+from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork
+
+dgl = pytest.importorskip("dgl")
+
+
+def _model(out_size=1):
+    graph = build_dgl_graph(torch.tensor([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]]))
+    return EdgeNetwork(graph, torch.zeros(1, 2), 8, out_size, 2, 2), graph
+
+
+def test_edge_network_returns_raw_binary_logits_without_graph_leakage():
+    model, graph = _model()
+    globals_ = torch.tensor([[1.0, 2.0]])
+    before = set(graph.ndata.keys()) | set(graph.edata.keys())
+    logits = model(graph, globals_)
+    again = model(graph, globals_)
+    assert logits.shape == (1, 1)
+    assert logits.dtype == torch.float32
+    assert torch.allclose(logits, again)
+    assert torch.equal(
+        model.representation(graph, globals_), model.forward_features(graph, globals_)
+    )
+    assert model.classify is model.classifier
+    assert before == (set(graph.ndata.keys()) | set(graph.edata.keys()))
+
+
+def test_edge_network_supports_multiclass_and_transfer_freezing():
+    model, graph = _model(out_size=12)
+    transferred = FineTunedEdgeNetwork.from_pretrained(model, 1, freeze_backbone=True)
+    assert transferred(graph, torch.zeros(1, 2)).shape == (1, 1)
+    assert not any(
+        parameter.requires_grad for parameter in transferred.backbone.parameters()
+    )
+    assert all(
+        parameter.requires_grad for parameter in transferred.classifier.parameters()
+    )
+    assert model(graph, torch.zeros(1, 2)).shape == (1, 12)

tests/unit/tasks/test_classification.py

@@ -0,0 +1,66 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from gnn4colliders.tasks import BinaryClassificationTask, MulticlassClassificationTask
+
+
+def _batch(labels, weights):
+    return SimpleNamespace(
+        labels=torch.tensor(labels),
+        metadata=SimpleNamespace(weight=torch.tensor(weights)),
+    )
+
+
+def test_binary_normalizes_column_or_vector_targets_and_uses_named_weight():
+    task = BinaryClassificationTask()
+    batch = _batch([[0], [0], [1], [1]], [1.0, 2.0, 1.0, 3.0])
+    logits = torch.tensor([[-1.0], [0.0], [0.5], [1.0]])
+    expected_elementwise = torch.nn.functional.binary_cross_entropy_with_logits(
+        logits[:, 0], batch.labels[:, 0].float(), reduction="none"
+    )
+    expected = (
+        (expected_elementwise[:2] * torch.tensor([1.0, 2.0])).sum() / 3
+        + (expected_elementwise[2:] * torch.tensor([1.0, 3.0])).sum() / 4
+    ) / 2
+    assert torch.allclose(task.loss(logits, batch), expected)
+
+
+def test_binary_negative_weights_are_not_absolute_by_default():
+    batch = _batch([0, 0, 1, 1], [1.0, -2.0, 1.0, 1.0])
+    logits = torch.tensor([[0.0], [2.0], [0.0], [2.0]])
+    task = BinaryClassificationTask()
+    elementwise = torch.nn.functional.binary_cross_entropy_with_logits(
+        logits[:, 0], batch.labels.float(), reduction="none"
+    )
+    expected = (
+        (elementwise[:2] * torch.tensor([1.0, -2.0])).sum() / -1
+        + elementwise[2:].mean()
+    ) / 2
+    assert task.loss(logits, batch).item() == pytest.approx(expected.item())
+    assert BinaryClassificationTask(absolute_weights=True).loss(
+        logits, batch
+    ).item() != pytest.approx(expected.item())
+
+
+def test_binary_predictions_and_undefined_auc():
+    task = BinaryClassificationTask()
+    batch = _batch([1, 1], [1.0, 2.0])
+    output = task.predictions(torch.tensor([[-2.0], [2.0]]))
+    assert output["scores"].shape == (2,)
+    assert output["predictions"].tolist() == [False, True]
+    assert (
+        task.metrics(torch.tensor([[-2.0], [2.0]]), batch)["roc_auc"]
+        != task.metrics(torch.tensor([[-2.0], [2.0]]), batch)["roc_auc"]
+    )
+
+
+def test_multiclass_loss_and_metrics():
+    task = MulticlassClassificationTask()
+    batch = _batch([0, 1, 2], [1.0, 1.0, 1.0])
+    logits = torch.eye(3)
+    assert task.loss(logits, batch).item() > 0
+    output = task.predict(logits)
+    assert output["predictions"].tolist() == [0, 1, 2]
+    assert task.metrics(logits, batch)["accuracy"] == 1.0

tests/unit/training/test_checkpoint.py

@@ -0,0 +1,107 @@
+import random
+
+import numpy as np
+import torch
+from torch import nn
+
+from gnn4colliders.training import (
+    CheckpointManager,
+    EarlyStopping,
+    TrainerState,
+    load_legacy_checkpoint,
+    load_model_weights,
+    restore_training_state,
+)
+
+
+def test_checkpoint_round_trip_and_selection(tmp_path):
+    model = nn.Linear(2, 1)
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
+    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.5)
+    model(torch.ones(2)).sum().backward()
+    optimizer.step()
+    scheduler.step()
+    early = EarlyStopping(patience=2)
+    early.update(3.0)
+    state = TrainerState(epoch=9, global_step=17)
+    manager = CheckpointManager(tmp_path)
+    manager.save(
+        model=model,
+        optimizer=optimizer,
+        scheduler=scheduler,
+        early_stopping=early,
+        trainer_state=state,
+        task_config={"type": "binary_classification", "threshold": 0.5},
+        monitor_name="loss",
+        monitor_value=0.4,
+    )
+    state.epoch = 10
+    manager.save(
+        model=model,
+        optimizer=optimizer,
+        scheduler=scheduler,
+        early_stopping=early,
+        trainer_state=state,
+        monitor_name="loss",
+        monitor_value=0.8,
+    )
+
+    assert manager.latest().name == "epoch_0010.pt"
+    assert manager.best(monitor="loss").name == "epoch_0009.pt"
+    payload = manager.load(manager.latest())
+    assert payload["schema_version"] == 1
+    assert payload["task_config"] is None  # omitted explicitly for the second save
+
+    restored_model = nn.Linear(2, 1)
+    restored_optimizer = torch.optim.Adam(restored_model.parameters(), lr=0.1)
+    restored_scheduler = torch.optim.lr_scheduler.ExponentialLR(
+        restored_optimizer, gamma=0.5
+    )
+    restored_early = EarlyStopping(patience=2)
+    restored = restore_training_state(
+        payload,
+        model=restored_model,
+        optimizer=restored_optimizer,
+        scheduler=restored_scheduler,
+        early_stopping=restored_early,
+    )
+    assert restored == TrainerState(epoch=10, global_step=17)
+    assert restored_scheduler.last_epoch == scheduler.last_epoch
+    assert restored_early.best == early.best
+    assert all(
+        torch.equal(left, right)
+        for left, right in zip(model.parameters(), restored_model.parameters())
+    )
+
+
+def test_weight_only_load_does_not_need_optimizer(tmp_path):
+    model = nn.Linear(1, 1)
+    manager = CheckpointManager(tmp_path)
+    path = manager.save(model=model, trainer_state=TrainerState(epoch=0))
+    fresh = nn.Linear(1, 1)
+    load_model_weights(fresh, manager.load(path))
+    assert torch.equal(model.weight, fresh.weight)
+    assert torch.equal(model.bias, fresh.bias)
+
+
+def test_legacy_checkpoint_normalizes_supported_prefixes():
+    state = {"module._orig_mod.linear.weight": torch.ones(1, 1)}
+    normalized = load_legacy_checkpoint(
+        {"epoch": 4, "model_state_dict": state, "early_stop": {"count": 2}}
+    )
+    assert list(normalized["model_state_dict"]) == ["linear.weight"]
+    assert normalized["epoch"] == 4
+    assert normalized["early_stopping_state"]["num_bad_epochs"] == 2
+
+
+def test_rng_state_is_restored(tmp_path):
+    manager = CheckpointManager(tmp_path)
+    model = nn.Linear(1, 1)
+    path = manager.save(model=model, trainer_state=TrainerState(epoch=0))
+    expected = (random.random(), np.random.random(), torch.rand(1))
+    payload = manager.load(path)
+    restore_training_state(payload, model=nn.Linear(1, 1))
+    actual = (random.random(), np.random.random(), torch.rand(1))
+    assert expected[0] == actual[0]
+    assert expected[1] == actual[1]
+    assert torch.equal(expected[2], actual[2])

tests/unit/training/test_trainer.py

@@ -0,0 +1,81 @@
+import torch
+from torch import nn
+
+from gnn4colliders.data import BatchMetadata, EventMetadata, GraphBatch
+from gnn4colliders.tasks import BinaryClassificationTask
+from gnn4colliders.training import (
+    EarlyStopping,
+    Trainer,
+    build_optimizer,
+    build_scheduler,
+)
+
+
+class TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(1, 1)
+
+    def forward(self, batch):
+        return self.linear(batch.graph)
+
+
+def _batch(values, labels, weights=None):
+    values = torch.tensor(values, dtype=torch.float32).reshape(-1, 1)
+    if weights is None:
+        weights = [1.0] * len(labels)
+    return GraphBatch(
+        graph=values,
+        labels=torch.tensor(labels),
+        global_features=None,
+        metadata=BatchMetadata.from_events(
+            [
+                EventMetadata(fold=0, weight=weight, sample_id=f"event:{index}")
+                for index, weight in enumerate(weights)
+            ]
+        ),
+    )
+
+
+def test_graph_batch_to_moves_named_tensors_without_losing_ids():
+    batch = _batch([1, 2], [0, 1])
+    moved = batch.to("cpu")
+    assert moved.labels.device.type == "cpu"
+    assert moved.metadata.weight.device.type == "cpu"
+    assert moved.metadata.sample_id == ("event:0", "event:1")
+
+
+def test_trainer_aggregates_metrics_over_the_complete_epoch():
+    torch.manual_seed(4)
+    model = TinyModel()
+    optimizer = build_optimizer(model, learning_rate=0.01)
+    scheduler = build_scheduler(optimizer, gamma=0.5)
+    trainer = Trainer(model, BinaryClassificationTask(), optimizer, scheduler)
+    loader = [
+        _batch([1, 2], [0, 1]),
+        _batch([3, 4], [0, 1]),
+    ]
+    before = model.linear.weight.detach().clone()
+    result = trainer.train_epoch(loader)
+    assert result.num_samples == 4
+    assert set(result.metrics) == {"accuracy", "roc_auc", "auc"}
+    assert torch.isfinite(torch.tensor(result.loss))
+    assert not torch.equal(before, model.linear.weight.detach())
+    trainer.scheduler.step()
+    assert optimizer.param_groups[0]["lr"] == 0.005
+
+
+def test_fit_uses_validation_for_early_stopping_and_records_history():
+    model = TinyModel()
+    optimizer = build_optimizer(model, learning_rate=0.01)
+    trainer = Trainer(
+        model,
+        BinaryClassificationTask(),
+        optimizer,
+        early_stopping=EarlyStopping(patience=1),
+    )
+    loader = [_batch([1, 2], [0, 1])]
+    history = trainer.fit(loader, loader, epochs=4)
+    assert len(history.train) == len(history.validation)
+    assert 1 <= len(history.train) <= 4
+    assert all(isinstance(item.loss, float) for item in history.train)