"""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)