Blame for code/graph/training/trainer.py

"""Training loop with real batching, LR scheduling, AMP and distributed support.

Defects this replaces
---------------------
``batch_size`` was validated by the config and then never read: one example was
one optimizer step, so the 3B recipe of ``max_steps: 100000`` described batch-1
training.  ``warmup_ratio`` was likewise validated and ignored -- no scheduler
existed.  ``wrap_distributed`` was defined but never called by any trainer or
script, and the trainer defaulted to CPU with no way to ask for a GPU, so the
documented 3B command could not have used an accelerator.

Because graphs have different node and edge counts, a batch here is a gradient
accumulation group: ``batch_size`` examples are forwarded, their losses averaged,
and one optimizer step taken.  That is genuine batching for variable-size graph
inputs and it makes the configured batch size meaningful.
"""

from __future__ import annotations

import hashlib
import json
import math
import os
import platform
import random
import subprocess
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
from contextlib import nullcontext
from dataclasses import dataclass
from itertools import islice
from pathlib import Path
from typing import Any

from dataclasses import replace
import torch
from torch import Tensor
from torch.nn import functional

from graph.config import GraphConfig
from graph.export import load_checkpoint_file, load_state_dict_compat
from graph.graph import GraphStore
from graph.model import STATUS_INDEX, GraphModel
from graph.synthetic.generator import GENERATOR_VERSION, ReasoningExample
from graph.training.losses import COMPONENTS, LossInputs, MultiTaskLoss

#: Logit margin by which a missing-edge example's support must sit below its supported twin.
PAIR_MARGIN = 1.0


def rank_device(device: torch.device, local_rank: int) -> torch.device:
    """The device one data-parallel process trains on.

    Under torchrun every process would otherwise resolve to ``cuda:0``, and NCCL
    refuses two ranks on one device.  Each CUDA rank takes its own GPU.
    """
    if device.type != "cuda":
        return device
    return torch.device("cuda", local_rank)


def resolve_device(preference: str | torch.device | None = None) -> torch.device:
    """Pick the best available device unless one is named explicitly."""
    if preference is not None and str(preference) != "auto":
        return torch.device(preference)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


@dataclass(frozen=True)
class TrainStepOutput:
    step: int
    loss: float
    components: dict[str, float]
    grad_norm: float
    learning_rate: float
    examples: int


@dataclass(frozen=True)
class EvaluationEvent:
    step: int
    metrics: dict[str, float]


#: Bytes of training state per parameter: fp32 weights, fp32 gradients and the two
#: fp32 Adam moments.  Activations come on top, so the check keeps headroom.
TRAINING_BYTES_PER_PARAMETER = 16
DEVICE_MEMORY_HEADROOM = 0.85


def require_device_memory(
    parameter_count: int, device: torch.device, total_bytes: int | None = None
) -> None:
    """Refuse a CUDA run whose parameter and optimizer state cannot fit on the device.

    Without sharding, every data-parallel rank holds the full state, so the device
    must hold it outright.  Failing here names the problem; the alternative is an
    out-of-memory error partway through the first step.
    """
    if device.type != "cuda":
        return
    if total_bytes is None:
        total_bytes = torch.cuda.get_device_properties(device).total_memory
    need = parameter_count * TRAINING_BYTES_PER_PARAMETER
    budget = total_bytes * DEVICE_MEMORY_HEADROOM
    if need > budget:
        raise RuntimeError(
            f"model needs about {need / 1e9:.1f} GB of parameter and optimizer state "
            f"per device; {device} has {total_bytes / 1e9:.1f} GB "
            f"({budget / 1e9:.1f} GB usable). Use a larger GPU, or implement sharded "
            "training (FSDP), which is not available."
        )


class GraphTrainer:
    """Own the optimizer, schedule, precision policy and checkpoint format."""

    def __init__(
        self,
        model: GraphModel,
        config: GraphConfig,
        device: torch.device | str | None = "auto",
        *,
        total_steps: int | None = None,
        distributed: bool = False,
    ) -> None:
        self.config = config
        self.skip_grad_norm = 0.0
        self.skipped_updates = 0
        self._recent_grad_norms: list[float] = []
        self.device = resolve_device(device)
        if distributed:
            local_rank = int(os.environ.get("LOCAL_RANK", "0"))
            self.device = rank_device(self.device, local_rank)
            if self.device.type == "cuda":
                torch.cuda.set_device(self.device)
        self.model = model.to(self.device)
        require_device_memory(
            sum(parameter.numel() for parameter in model.parameters()), self.device
        )
        self.distributed = distributed
        self.rank, self.world_size = 0, 1
        if distributed:
            if config.training.use_fsdp:
                raise NotImplementedError(
                    "FSDP is not implemented; set training.use_fsdp: false to use "
                    "data parallelism"
                )
            if not torch.distributed.is_initialized():
                raise RuntimeError(
                    "distributed=True requires torch.distributed to be initialised "
                    "(launch with torchrun)"
                )
            self.rank = torch.distributed.get_rank()
            self.world_size = torch.distributed.get_world_size()
            # Data parallelism starts every rank from rank 0's weights.
            for tensor in [*self.model.parameters(), *self.model.buffers()]:
                torch.distributed.broadcast(tensor.data, src=0)
        self.loss_fn = MultiTaskLoss(config.loss)
        if config.training.freeze_query_encoder:
            # No gradient flows into the encoder at all, so it can neither
            # update nor dominate the clipped gradient norm.
            encoder = getattr(self.model, "module", self.model).query_encoder
            for parameter in encoder.parameters():
                parameter.requires_grad_(False)
        #: Learned link head (``link_head: mlp``); registered on the model so the
        #: optimizer and checkpoints include it.  None keeps the cosine form.
        self.link_mlp: torch.nn.Module | None = None
        if config.training.link_head == "mlp":
            hidden = config.hidden_dim
            self.link_mlp = torch.nn.Sequential(
                torch.nn.Linear(2 * hidden, hidden),
                torch.nn.GELU(),
                torch.nn.Linear(hidden, 1),
            ).to(self.device)
            self.model.add_module("link_mlp", self.link_mlp)
        self.optimizer = torch.optim.AdamW(
            self._optimizer_groups(),
            lr=config.training.learning_rate,
            betas=(0.9, 0.95),
        )
        self.total_steps = total_steps or config.training.max_steps
        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer, self._schedule
        )
        self.scaler = torch.amp.GradScaler(
            self.device.type,
            enabled=config.training.precision == "fp16" and self.device.type == "cuda",
        )
        self.step = 0
        self.best_score = -math.inf
        self.best_step: int | None = None
        self.stopped_early = False
        #: Set by a caller (a step or evaluation callback) to end training after
        #: the current step; recorded as the reason in the summary.
        self.stop_reason: str | None = None
        self._evaluations_without_improvement = 0
        self.batch_size = max(1, config.training.batch_size)
        self.accumulation = max(1, config.training.gradient_accumulation)
        #: Temperature for cosine link logits; keeps the link term on the same
        #: scale as the other objectives so loss weights mean what they say.
        self.link_logit_scale = 2.0
        #: Examples whose activations are held simultaneously for one backward.
        self.backward_chunk_size = max(1, config.training.backward_chunk_size)
        if self.device.type == "cuda":
            # TF32 matmuls measured ~1.8x faster here at no observable cost to
            # convergence; these are small matmuls where precision is not the
            # limiting factor.
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True

    def _schedule(self, step: int) -> float:
        """Linear warmup into cosine decay, honouring ``warmup_ratio``."""
        warmup = max(1, int(self.total_steps * self.config.training.warmup_ratio))
        if step < warmup:
            return (step + 1) / warmup
        progress = (step - warmup) / max(1, self.total_steps - warmup)
        return 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress))) * 0.9 + 0.1

    # -- loops --------------------------------------------------------------

    def train(
        self,
        examples: Iterable[ReasoningExample],
        *,
        max_steps: int | None = None,
        validation: Sequence[ReasoningExample] = (),
        evaluate: Callable[
            [Sequence[ReasoningExample], GraphModel], Mapping[str, float]
        ]
        | None = None,
        evaluation_every: int | None = None,
        checkpoint_directory: str | Path | None = None,
        checkpoint_every: int | None = None,
        dataset_hash: str = "unknown",
        on_step: Callable[[TrainStepOutput], None] | None = None,
        selection_score: Callable[[Mapping[str, float]], float] | None = None,
        best_directory: str | Path | None = None,
        patience: int = 0,
    ) -> tuple[list[TrainStepOutput], list[EvaluationEvent]]:
        """Train for ``max_steps`` optimizer steps over batched examples.

        With ``selection_score``, every evaluation is scored (higher is better)
        and the best checkpoint is saved to ``best_directory``.  With
        ``patience > 0`` training stops after that many evaluations without
        improvement.  This model degrades with over-training on
        generalisation-dependent metrics (model_accuracy.md section 8: 2500
        steps beat 6000), so a fixed step budget is not a safe stopping rule.
        """
        limit = max_steps or self.config.training.max_steps
        interval = (
            evaluation_every
            if evaluation_every is not None
            else self.config.training.evaluation_every
        )
        saving = (
            checkpoint_every
            if checkpoint_every is not None
            else self.config.training.checkpoint_every
        )
        if saving and checkpoint_directory is None:
            raise ValueError(
                "checkpoint_directory is required when checkpoint_every is set"
            )
        outputs: list[TrainStepOutput] = []
        events: list[EvaluationEvent] = []
        self.model.train()
        for batch in self._batches(self._shard(examples), limit):
            output = self.train_step(batch)
            outputs.append(output)
            if on_step is not None:
                on_step(output)
            if (
                interval
                and self.step % interval == 0
                and evaluate is not None
                and validation
            ):
                event = EvaluationEvent(self.step, self._evaluate(validation, evaluate))
                events.append(event)
                if selection_score is not None:
                    score = float(selection_score(event.metrics))
                    if not math.isnan(score) and score > self.best_score:
                        self.best_score = score
                        self.best_step = self.step
                        self._evaluations_without_improvement = 0
                        if best_directory is not None:
                            self.save_checkpoint(
                                best_directory, dataset_hash=dataset_hash
                            )
                    else:
                        self._evaluations_without_improvement += 1
                        if (
                            patience
                            and self._evaluations_without_improvement >= patience
                        ):
                            self.stopped_early = True
            if saving and self.step % saving == 0 and checkpoint_directory is not None:
                self.save_checkpoint(
                    Path(checkpoint_directory) / f"step-{self.step}",
                    dataset_hash=dataset_hash,
                )
            # One rank deciding to stop must stop them all, or the next collective
            # waits forever.
            if self._any_rank(self.stop_reason is not None or self.stopped_early):
                if self.stop_reason is None and not self.stopped_early:
                    self.stop_reason = "stopped on another rank"
                break
        return outputs, events

    @property
    def is_main(self) -> bool:
        return self.rank == 0

    def _shard(
        self, examples: Iterable[ReasoningExample]
    ) -> Iterable[ReasoningExample]:
        """This rank's share of the examples: every ``world_size``-th one."""
        if not self.distributed or self.world_size == 1:
            return examples
        return islice(examples, self.rank, None, self.world_size)

    def _all_ranks(self, flag: bool) -> bool:
        """True only if ``flag`` holds on every rank (a collective when distributed)."""
        return self._reduce_flag(
            flag, torch.distributed.ReduceOp.MIN if self.distributed else None
        )

    def _any_rank(self, flag: bool) -> bool:
        """True if ``flag`` holds on any rank (a collective when distributed)."""
        return self._reduce_flag(
            flag, torch.distributed.ReduceOp.MAX if self.distributed else None
        )

    def _reduce_flag(self, flag: bool, op) -> bool:
        if op is None:
            return flag
        value = torch.tensor([1 if flag else 0], device=self.device)
        torch.distributed.all_reduce(value, op=op)
        return bool(value.item())

    def _average_gradients(self) -> None:
        """Sum each parameter's gradient across ranks and divide by the world size.

        Every rank visits the same parameters in the same order, so the collectives
        match.  A parameter with no gradient on this rank contributes zeros.
        """
        for parameter in self.model.parameters():
            if not parameter.requires_grad:
                continue
            if parameter.grad is None:
                parameter.grad = torch.zeros_like(parameter)
            torch.distributed.all_reduce(parameter.grad)
            parameter.grad.div_(self.world_size)

    def _batches(
        self, examples: Iterable[ReasoningExample], limit: int
    ) -> Iterator[list[ReasoningExample]]:
        stream = iter(examples)
        group = self.batch_size * self.accumulation
        while self.step < limit:
            batch = list(islice(stream, group))
            # Ranks can run out at different points, so stop only when all have.
            if not self._all_ranks(bool(batch)):
                return
            yield batch

    def _evaluate(
        self,
        validation: Sequence[ReasoningExample],
        evaluate: Callable[
            [Sequence[ReasoningExample], GraphModel], Mapping[str, float]
        ],
    ) -> dict[str, float]:
        was_training = self.model.training
        self.model.eval()
        with torch.no_grad():
            metrics = evaluate(validation, self.model)
        if was_training:
            self.model.train()
        return {str(key): float(value) for key, value in metrics.items()}

    # -- one optimizer step -------------------------------------------------

    def train_step(self, batch: Sequence[ReasoningExample]) -> TrainStepOutput:
        """Forward every example in ``batch``, average, and take one step.

        Losses are summed over a chunk and backpropagated once rather than per
        example.  Graphs have different shapes so the forward passes cannot be
        fused, but one autograd invocation over the combined graph measured about
        1.4x faster than one per example: these are small matmuls and the cost is
        dominated by kernel launches and engine overhead, not arithmetic.
        ``backward_chunk_size`` bounds how many examples' activations are held at
        once so the saving does not turn into an out-of-memory failure.
        """
        if not batch:
            raise ValueError("train_step requires a non-empty batch")
        self.optimizer.zero_grad(set_to_none=True)
        # Detached per-example values, read back once per step: converting each
        # with float() forced a GPU sync for every component of every example.
        logged: list[Tensor] = []
        counted = 0
        divisor = float(len(batch))
        for start in range(0, len(batch), self.backward_chunk_size):
            chunk = batch[start : start + self.backward_chunk_size]
            accumulated: Tensor | None = None
            with self._autocast():
                losses = self._chunk_losses(chunk)
            for loss_output in losses:
                scaled = loss_output.total / divisor
                accumulated = scaled if accumulated is None else accumulated + scaled
                logged.append(
                    torch.stack(
                        [
                            loss_output.total.detach().float().reshape(()),
                            *(
                                loss_output.components[name]
                                .detach()
                                .float()
                                .reshape(())
                                for name in COMPONENTS
                            ),
                        ]
                    )
                )
                counted += 1
            if accumulated is None:
                continue
            if self.scaler.is_enabled():
                self.scaler.scale(accumulated).backward()
            else:
                accumulated.backward()
        if self.distributed:
            self._average_gradients()
        if self.scaler.is_enabled():
            self.scaler.unscale_(self.optimizer)
        grad_norm = float(
            torch.nn.utils.clip_grad_norm_(
                self.model.parameters(), self.config.training.grad_clip
            ).item()
        )
        # Read the rate that was actually applied to this step, before the
        # scheduler advances to the next one.
        applied_learning_rate = float(self.optimizer.param_groups[0]["lr"])
        # Skip an update whose pre-clip gradient norm is a spike relative to the
        # recent median of applied updates. The step still advances, so the
        # schedule and the curriculum are unaffected.
        recent = self._recent_grad_norms
        spike = (
            self.skip_grad_norm > 0
            and len(recent) >= 50
            and grad_norm > self.skip_grad_norm * sorted(recent)[len(recent) // 2]
        )
        if spike:
            self.optimizer.zero_grad(set_to_none=True)
            self.skipped_updates += 1
        else:
            recent.append(grad_norm)
            del recent[:-200]
            if self.scaler.is_enabled():
                self.scaler.step(self.optimizer)
                self.scaler.update()
            else:
                self.optimizer.step()
        self.scheduler.step()
        self.step += 1
        means = (
            torch.stack(logged).cpu().sum(0) / max(1, counted)
            if logged
            else torch.zeros(len(COMPONENTS) + 1)
        ).tolist()
        return TrainStepOutput(
            step=self.step,
            loss=means[0],
            components=dict(zip(COMPONENTS, means[1:], strict=True)),
            grad_norm=grad_norm,
            learning_rate=applied_learning_rate,
            examples=counted,
        )

    def _chunk_losses(self, chunk: Sequence[ReasoningExample]) -> list:
        """Forward a chunk (batched when enabled) and score each example.

        A missing-edge example with a supported twin (``pair``) is scored next to
        its twin, and a margin is added so the missing example's support logit sits
        at least ``PAIR_MARGIN`` below its twin's.
        """
        expanded: list[ReasoningExample] = []
        for example in chunk:
            expanded.append(example)
            if example.pair is not None:
                expanded.append(example.pair)
        if self.config.training.batched_forward and not self.config.on_policy_selection:
            forwards = self.model.forward_batch(
                [GraphStore.from_example(example.graph) for example in expanded],
                [example.query for example in expanded],
                [example.supporting_paths for example in expanded],
                [example.competing_paths for example in expanded],
            )
        else:
            forwards = [self._forward_one(example) for example in expanded]
        outputs = [
            self._loss_from_forward(example, forward)
            for example, forward in zip(expanded, forwards, strict=True)
        ]
        for index, example in enumerate(expanded):
            if example.pair is None:
                continue
            twin_index = index + 1
            gap = (
                forwards[twin_index].support_logit.reshape(-1)[0]
                - forwards[index].support_logit.reshape(-1)[0]
            )
            margin = torch.relu(PAIR_MARGIN - gap)
            components = dict(outputs[index].components)
            components["verification"] = components["verification"] + margin
            outputs[index] = replace(
                outputs[index],
                total=outputs[index].total + margin,
                components=components,
            )
        return outputs

    def _forward_one(self, example: ReasoningExample):
        store = GraphStore.from_example(example.graph)
        return self.model.forward_example(
            store,
            example.query,
            supporting_paths=example.supporting_paths,
            competing_paths=example.competing_paths,
            answers=example.answers,
            status=example.status,
        )

    def _example_loss(self, example: ReasoningExample):
        """Build every supervision tensor for one example and score it."""
        return self._loss_from_forward(example, self._forward_one(example))

    def _loss_from_forward(self, example: ReasoningExample, forward):
        """Score one example's forward outputs against its labels."""
        device = self.device
        order = {node_id: index for index, node_id in enumerate(forward.entity_order)}
        answer_labels = torch.zeros(len(forward.entity_order), device=device)
        for answer in example.answers:
            index = order.get(answer)
            if index is not None:
                answer_labels[index] = 1.0

        graph_tensors = example.graph.graph
        link_logits, link_labels, negative_pairs = self._link_supervision(
            forward.node_states, graph_tensors
        )
        supported = float(example.status.value == "supported")
        status_label = STATUS_INDEX[example.status.value]
        count_label = min(len(example.answers), self.config.max_count)
        count_logits = (
            forward.count_logits
            if example.query.answer_mode in {"count", "aggregate"}
            else None
        )
        relation_labels = (
            graph_tensors.edge_type_ids.remainder(self.config.num_relations)[:1].to(
                device
            )
            if graph_tensors.num_edges
            else None
        )
        return self.loss_fn(
            LossInputs(
                node_embeddings=forward.node_states,
                node_positive_pairs=graph_tensors.edge_index.t().to(device)
                if graph_tensors.num_edges
                else None,
                node_negative_pairs=negative_pairs,
                link_logits=link_logits,
                link_labels=link_labels,
                next_hop_logits=forward.next_hop_logits,
                next_hop_labels=forward.next_hop_labels,
                relation_logits=forward.relation_logits,
                relation_labels=relation_labels,
                path_scores=forward.selection_logits,
                path_labels=forward.selection_labels,
                selection_threshold=forward.selection_threshold,
                pat_positive=forward.pat_positive,
                pat_negative=forward.pat_negative,
                answer_logits=forward.entity_logits,
                answer_labels=answer_labels,
                verification_logits=forward.support_logit,
                verification_labels=torch.tensor([supported], device=device),
                contradiction_logits=forward.contradiction_logits,
                contradiction_labels=torch.tensor([status_label], device=device),
                count_logits=count_logits,
                count_labels=torch.tensor([count_label], device=device)
                if count_logits is not None
                else None,
                confidence=forward.confidence_logit.reshape(1),
                # Confidence must predict whether *this model's own selection* is
                # right, not merely whether the example is answerable.  The old
                # label was ``status == supported``, which is nearly constant
                # across the corpus, so the head learned "always confident": at a
                # reported confidence of 0.85 the measured accuracy was 0.130.
                correctness=self._selection_correctness(forward, supported),
            )
        )

    @staticmethod
    def _selection_correctness(forward, supported: float) -> Tensor:
        """Achieved set-F1 of the model's current selection, as a 0-1 target.

        Detached: this is a *label* for the confidence head, so it must not
        backpropagate into the selection scores it is measuring.
        """
        logits, labels = forward.selection_logits, forward.selection_labels
        threshold = forward.selection_threshold
        if logits.numel() == 0 or threshold is None:
            return logits.new_tensor([supported])
        keep = (logits.detach() > threshold.detach()).float()
        gold = labels.detach().float()
        true_positive = (keep * gold).sum()
        denominator = (
            2.0 * true_positive
            + (keep * (1.0 - gold)).sum()
            + ((1.0 - keep) * gold).sum()
        )
        if float(denominator) <= 0.0:
            return logits.new_tensor([supported])
        return ((2.0 * true_positive) / denominator).reshape(1)

    def _link_score(self, left: Tensor, right: Tensor) -> Tensor:
        """Link logit for node-state pairs: scaled cosine, or a learned pair head.

        The learned head sees the elementwise product and absolute difference of
        the two normalised states, so it can score relations the cosine cannot.
        """
        if self.link_mlp is None:
            return (left * right).sum(-1) * self.link_logit_scale
        features = torch.cat((left * right, (left - right).abs()), dim=-1)
        return self.link_mlp(features).squeeze(-1)

    def _link_supervision(
        self, node_states: Tensor, graph_tensors
    ) -> tuple[Tensor | None, Tensor | None, Tensor | None]:
        """Sample negative node pairs by rejection instead of enumerating N^2.

        The previous implementation built the full non-edge set with a nested
        comprehension over all node pairs, which is quadratic and dominated step
        time on graphs of a few hundred nodes.
        """
        if not graph_tensors.num_edges:
            return None, None, None
        # Rejection sampling reads the CPU copy; reading the device copy back
        # would stall the GPU once per example.
        cpu_source, cpu_target = graph_tensors.edge_index.cpu()
        source, target = graph_tensors.edge_index.to(node_states.device)
        # Cosine-scaled logits.  Raw dot products of unnormalised node states
        # reached magnitudes in the tens, so the link term alone contributed a
        # loss of 30-50 and dominated the weighted total, leaving every other
        # objective to be optimised only through the gradient-clipped remainder.
        # The scale multiplies the cosine once.  It used to multiply both unit
        # vectors, making the logits 64 * cosine: the link loss saturated near 30
        # and, over a 12M-example run, node states collapsed to a mean pairwise
        # cosine of 0.98 (artifacts/retrain-20260928-full-epoch/freeze.json).
        node_states = functional.normalize(node_states, dim=-1)
        positive = self._link_score(node_states[source], node_states[target])
        count = int(source.shape[0])
        nodes = node_states.shape[0]
        if nodes < 2:
            return positive, torch.ones_like(positive), None
        existing = set(zip(cpu_source.tolist(), cpu_target.tolist(), strict=True))
        generator = random.Random(count * 31 + nodes)
        pairs: list[tuple[int, int]] = []
        for _ in range(count * 4):
            if len(pairs) >= count:
                break
            a, b = generator.randrange(nodes), generator.randrange(nodes)
            if a != b and (a, b) not in existing:
                pairs.append((a, b))
        if not pairs:
            return positive, torch.ones_like(positive), None
        index = torch.tensor(pairs, device=node_states.device, dtype=torch.long)
        negative = self._link_score(
            node_states[index[:, 0]], node_states[index[:, 1]]
        )
        return (
            torch.cat((positive, negative)),
            torch.cat((torch.ones_like(positive), torch.zeros_like(negative))),
            index,
        )

    def _autocast(self):
        """Mixed precision on CUDA only.

        ``precision`` describes the accelerator's compute dtype.  CPU bf16
        autocast is emulated on the platforms this runs on and measured roughly
        35x slower in the backward pass than plain fp32, with no memory saving
        that matters at these sizes, so it is deliberately not used.  MPS
        autocast coverage is likewise incomplete for these ops.  Both degrade to
        fp32, which is correct and still accelerated on MPS.
        """
        precision = self.config.training.precision
        if precision == "fp32" or self.device.type != "cuda":
            return nullcontext()
        dtype = torch.float16 if precision == "fp16" else torch.bfloat16
        return torch.autocast(device_type="cuda", dtype=dtype)

    @property
    def effective_precision(self) -> str:
        """The dtype actually used, which may differ from the configured one."""
        if self.device.type == "cuda":
            return self.config.training.precision
        return "fp32"

    # -- checkpoints --------------------------------------------------------

    def save_checkpoint(
        self, directory: str | Path, *, dataset_hash: str = "unknown"
    ) -> Path:
        destination = Path(directory)
        checkpoint = destination / "model.pt"
        if not self.is_main:
            # Every rank holds the same weights; only rank 0 writes them.
            return checkpoint
        destination.mkdir(parents=True, exist_ok=True)
        torch.save(
            {
                "model": self.model.state_dict(),
                "optimizer": self.optimizer.state_dict(),
                "scheduler": self.scheduler.state_dict(),
                "scaler": self.scaler.state_dict(),
                "step": self.step,
                "total_steps": self.total_steps,
                "config": self.config.to_dict(),
                "rng_state": torch.get_rng_state(),
            },
            checkpoint,
        )
        manifest = build_run_manifest(
            self.config,
            dataset_hash=dataset_hash,
            model=self.model,
            step=self.step,
            device=self.device,
        )
        manifest["checkpoint_sha256"] = _sha256_file(checkpoint)
        # The link objective changed form once (scale on each unit vector, giving
        # scale**2 * cosine, then scale * cosine).  Recording it lets a checkpoint
        # say which objective trained it.
        manifest["link_objective"] = {
            "form": "link_logit_scale * cosine",
            "link_logit_scale": self.link_logit_scale,
        }
        (destination / "run_manifest.json").write_text(
            json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8"
        )
        return checkpoint

    def _encoder_parameters(self) -> set[int]:
        encoder = getattr(self.model, "module", self.model).query_encoder
        return {id(parameter) for parameter in encoder.parameters()}

    def _layout(self, *, split: bool, with_encoder: bool) -> list[list[Tensor]]:
        """Parameters per optimizer group, in group order.

        Weight-decay and no-decay groups, each optionally split into the rest of
        the model followed by the query encoder; ``with_encoder=False`` leaves
        the encoder out (it is frozen).  Covers every grouping a checkpoint can
        have been saved with, so one can be matched by its group sizes.
        """
        encoder = self._encoder_parameters()
        decay: list[Tensor] = []
        no_decay: list[Tensor] = []
        for name, parameter in self.model.named_parameters():
            if not with_encoder and id(parameter) in encoder:
                continue
            if parameter.ndim <= 1 or name.endswith((".threshold", ".scale")):
                no_decay.append(parameter)
            else:
                decay.append(parameter)
        if not split:
            return [decay, no_decay]
        return [
            part
            for params in (decay, no_decay)
            for part in (
                [p for p in params if id(p) not in encoder],
                [p for p in params if id(p) in encoder],
            )
        ]

    def _optimizer_groups(self) -> list[dict[str, Any]]:
        """Weight-decay groups; the query encoder split out, scaled or frozen.

        With ``query_encoder_lr_scale == 1`` and the encoder trainable these are
        the original two groups, so existing checkpoints load unchanged.
        """
        training = self.config.training
        frozen = training.freeze_query_encoder
        split = training.query_encoder_lr_scale != 1.0 and not frozen
        encoder = self._encoder_parameters()
        groups: list[dict[str, Any]] = []
        for index, params in enumerate(
            self._layout(split=split, with_encoder=not frozen)
        ):
            if not params:
                continue
            weight_decay = training.weight_decay if index < (2 if split else 1) else 0.0
            group: dict[str, Any] = {"params": params, "weight_decay": weight_decay}
            if split and id(params[0]) in encoder:
                group["lr"] = training.learning_rate * training.query_encoder_lr_scale
            groups.append(group)
        return groups

    def _load_regrouped_optimizer(self, saved: Mapping[str, Any]) -> None:
        """Load Adam state saved under a different grouping, parameter by parameter.

        Optimizer state is keyed by position across the groups.  The layout the
        checkpoint was written with is identified by its group sizes, and each
        parameter's state moves to that parameter's position now; parameters no
        longer optimised (a frozen encoder) drop theirs.  Learning rates are set
        from the schedule at the restored step, since the saved scheduler state
        describes the old groups.
        """
        sizes = [len(group["params"]) for group in saved["param_groups"]]
        old_order: list[Tensor] | None = None
        for split in (False, True):
            for with_encoder in (True, False):
                layout = [
                    group
                    for group in self._layout(split=split, with_encoder=with_encoder)
                    if group
                ]
                if [len(group) for group in layout] == sizes:
                    old_order = [p for group in layout for p in group]
                    break
            if old_order is not None:
                break
        if old_order is None:
            raise RuntimeError(f"cannot match saved optimizer groups of sizes {sizes}")
        new_index = {
            id(parameter): index
            for index, parameter in enumerate(
                p for group in self.optimizer.param_groups for p in group["params"]
            )
        }
        current = self.optimizer.state_dict()
        current["state"] = {
            new_index[id(old_order[int(index)])]: value
            for index, value in saved["state"].items()
            if id(old_order[int(index)]) in new_index
        }
        self.optimizer.load_state_dict(current)
        factor = self._schedule(self.step)
        for group, base in zip(
            self.optimizer.param_groups, self.scheduler.base_lrs, strict=True
        ):
            group["lr"] = base * factor
        self.scheduler.last_epoch = self.step

    def load_checkpoint(self, checkpoint: str | Path) -> None:
        state = load_checkpoint_file(checkpoint, map_location=self.device)
        load_state_dict_compat(self.model, state["model"])
        self.step = int(state["step"])
        self.total_steps = int(state.get("total_steps", self.total_steps))
        # Same group count is not same grouping: a frozen encoder and an unfrozen
        # one both have two groups, with different parameters in them.
        saved_sizes = [
            len(group["params"]) for group in state["optimizer"]["param_groups"]
        ]
        current_sizes = [len(group["params"]) for group in self.optimizer.param_groups]
        if saved_sizes == current_sizes:
            self.optimizer.load_state_dict(state["optimizer"])
            if "scheduler" in state:
                self.scheduler.load_state_dict(state["scheduler"])
        else:
            self._load_regrouped_optimizer(state["optimizer"])
        if "scaler" in state:
            self.scaler.load_state_dict(state["scaler"])
        if "rng_state" in state:
            torch.set_rng_state(state["rng_state"].to(torch.uint8).cpu())


#: Bumped when the architecture changes in a way that invalidates old checkpoints.
#: v3 adds the verification evidence signals; the "graph-" prefix replaces "grm-".
ARCHITECTURE_VERSION = "graph-architecture-v3"


def build_run_manifest(
    config: GraphConfig,
    *,
    dataset_hash: str,
    model: GraphModel,
    step: int,
    device: torch.device | None = None,
) -> dict[str, object]:
    """Reproducibility record required by PLAN section 62."""
    config_hash = hashlib.sha256(
        json.dumps(config.to_dict(), sort_keys=True).encode()
    ).hexdigest()
    pat_hash = hashlib.sha256(
        json.dumps(config.pat.__dict__, sort_keys=True).encode()
    ).hexdigest()
    buckets = dict(model.parameter_buckets())
    return {
        "model_version": config.model_version,
        "config_hash": config_hash,
        "architecture_version": ARCHITECTURE_VERSION,
        "dataset_hash": dataset_hash,
        "generator_version": GENERATOR_VERSION,
        "seed": config.training.seed,
        "step": step,
        "python": platform.python_version(),
        "platform": platform.platform(),
        "torch": torch.__version__,
        "cuda": torch.version.cuda,
        "device": str(device or next(model.parameters()).device),
        "pat_config": config.pat.__dict__,
        "pat_config_hash": pat_hash,
        "generator_hash": hashlib.sha256(GENERATOR_VERSION.encode()).hexdigest(),
        "model_hash": _model_hash(model),
        "git_commit": _git_commit(),
        "git_dirty": _git_dirty(),
        "parameter_count": sum(parameter.numel() for parameter in model.parameters()),
        "parameter_buckets": buckets,
        "deterministic_requested": config.training.deterministic,
        "cublas_workspace_config": os.environ.get(CUBLAS_WORKSPACE_ENV),
        "environment": {
            key: os.environ[key]
            for key in ("CUDA_VISIBLE_DEVICES", "RANK", "WORLD_SIZE", "LOCAL_RANK")
            if key in os.environ
        },
    }


def _model_hash(model: GraphModel) -> str:
    digest = hashlib.sha256()
    for name, parameter in model.state_dict().items():
        digest.update(name.encode())
        digest.update(
            parameter.detach().to(torch.float32).cpu().contiguous().numpy().tobytes()
        )
    return digest.hexdigest()


def _sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _git_commit() -> str | None:
    try:
        result = subprocess.run(
            ("git", "rev-parse", "HEAD"), check=True, capture_output=True, text=True
        )
    except (OSError, subprocess.CalledProcessError):
        return None
    return result.stdout.strip() or None


def _git_dirty() -> bool | None:
    try:
        result = subprocess.run(
            ("git", "status", "--porcelain"), check=True, capture_output=True, text=True
        )
    except (OSError, subprocess.CalledProcessError):
        return None
    return bool(result.stdout.strip())


CUBLAS_WORKSPACE_ENV = "CUBLAS_WORKSPACE_CONFIG"


def seed_everything(seed: int, *, deterministic: bool = True) -> bool:
    """Seed every RNG and request deterministic kernels.

    Returns whether full determinism is actually achievable in this process.
    cuBLAS GEMMs are non-deterministic on CUDA >= 10.2 unless
    ``CUBLAS_WORKSPACE_CONFIG`` is set *before* the CUDA context is created, so
    setting it here only helps if CUDA has not been initialised yet.  The return
    value is recorded in the run manifest rather than assumed, because PLAN
    section 62 asks for reproducibility to be recorded, not asserted.
    """
    random.seed(seed)
    torch.manual_seed(seed)
    reproducible = True
    if deterministic:
        if os.environ.get(CUBLAS_WORKSPACE_ENV) is None:
            if torch.cuda.is_available() and torch.cuda.is_initialized():
                # Too late: the context already exists and cuBLAS will pick a
                # non-deterministic workspace.  Report it instead of pretending.
                reproducible = False
            else:
                os.environ[CUBLAS_WORKSPACE_ENV] = ":4096:8"
        torch.use_deterministic_algorithms(True, warn_only=True)
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    return reproducible


def setup_distributed(backend: str = "nccl") -> None:
    if backend == "nccl" and not torch.cuda.is_available():
        backend = "gloo"
    if (
        torch.distributed.is_available()
        and int(os.environ.get("WORLD_SIZE", "1")) > 1
        and not torch.distributed.is_initialized()
    ):
        torch.distributed.init_process_group(backend=backend)