#!/usr/bin/env python3 """Train and publish the Szilassi proposal transformer. The native TransformerTrainingExport executable is the only reader of the versioned .sztd archive. This script snapshots its fixed-width stream once, then performs bounded-memory PyTorch training from that immutable snapshot. Search may continue sealing new shards while the snapshot is trained; those new shards are deliberately left for the next retraining run. """ from __future__ import annotations import argparse import contextlib import ctypes import dataclasses import hashlib import json import math import os import random import struct import subprocess import sys import time import uuid import zlib from collections import defaultdict from pathlib import Path from typing import BinaryIO, Iterable, Iterator, Sequence # cuBLAS requires this setting for deterministic CUDA GEMMs. It has to be set # before torch initializes CUDA. os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") try: import numpy as np except ImportError as exc: # pragma: no cover - dependency failure path raise SystemExit("NumPy is required to train the transformer.") from exc try: import torch import torch.nn as nn import torch.nn.functional as F except ImportError as exc: # pragma: no cover - dependency failure path raise SystemExit( "PyTorch with CUDA support is required. Install a CUDA PyTorch build " "for this Python interpreter." ) from exc # Whole-run split fixed for the first 6.5 GB corpus. Never split individual # records: neighboring CUDA trajectories are strongly correlated. REPORT_TRAIN_RUN = "889b3dfb-118c-48ca-ace4-853da7900f9c" REPORT_VALIDATION_RUN = "b84f957c-0e29-412c-ad04-f3530092f02e" REPORT_TEST_RUN = "39968628-bdef-4c6f-9617-c1fcd722ec3a" REPORT_STRESS_RUN = "59995ced-8ccb-4ce4-9549-612f3b8fd0d1" SPLIT_RUNS = { "train": REPORT_TRAIN_RUN, "validation": REPORT_VALIDATION_RUN, "test": REPORT_TEST_RUN, "stress": REPORT_STRESS_RUN, } RUN_BYTES_TO_SPLIT = { uuid.UUID(run_id).bytes: split for split, run_id in SPLIT_RUNS.items() } # TransformerTrainingExport stream v1. The exporter writes packed # little-endian records with no native padding. EXPORT_MAGIC = b"SZTXP001" EXPORT_VERSION = 1 EXPORT_HEADER = struct.Struct("<8sII") FACE_COUNT = 12 FACE_FEATURES = 16 GLOBAL_FEATURES = 24 FEATURE_FLOATS = FACE_COUNT * FACE_FEATURES + GLOBAL_FEATURES RUN_ID_OFFSET = 0 SEQUENCE_OFFSET = 16 TOPOLOGY_OFFSET = 24 PROPOSAL_FEATURE_OFFSET = 28 ANCHOR_FEATURE_OFFSET = PROPOSAL_FEATURE_OFFSET + FEATURE_FLOATS * 4 GAIN_OFFSET = ANCHOR_FEATURE_OFFSET + FEATURE_FLOATS * 4 OUTCOME_OFFSET = GAIN_OFFSET + 3 * 4 ACTION_OFFSET = OUTCOME_OFFSET + 4 PROPENSITY_OFFSET = ACTION_OFFSET + 3 * 4 REPLICA_COUNT_OFFSET = PROPENSITY_OFFSET + 3 * 4 ROLLOUT_ITERATIONS_OFFSET = REPLICA_COUNT_OFFSET + 4 TRAJECTORY_FLAGS_OFFSET = ROLLOUT_ITERATIONS_OFFSET + 8 ANCHOR_METRIC_FLAGS_OFFSET = TRAJECTORY_FLAGS_OFFSET + 4 INJECTED_METRIC_FLAGS_OFFSET = ANCHOR_METRIC_FLAGS_OFFSET + 4 RESULT_METRIC_FLAGS_OFFSET = INJECTED_METRIC_FLAGS_OFFSET + 4 EXPECTED_RECORD_SIZE = RESULT_METRIC_FLAGS_OFFSET + 4 # SZTRNK01 model format v1, byte-exact with TransformerRanker.h. MODEL_MAGIC = b"SZTRNK01" MODEL_FORMAT_VERSION = 1 FEATURE_FORMAT_VERSION = 1 OBJECTIVE_VERSION = 5 MODEL_HEADER_BYTES = 160 MODEL_WIDTH = 64 ATTENTION_HEADS = 4 LAYER_COUNT = 3 FEED_FORWARD_WIDTH = 256 ENSEMBLE_SIZE = 3 TOPOLOGY_COUNT = 59 MOVE_COUNT = 5 SCALE_COUNT = 3 def eprint(*values: object) -> None: print(*values, file=sys.stderr, flush=True) def set_below_normal_priority() -> None: """Keep CPU-side decoding/training orchestration below interactive work.""" if os.name != "nt": try: os.nice(5) except OSError: pass return BELOW_NORMAL_PRIORITY_CLASS = 0x00004000 kernel32 = ctypes.windll.kernel32 # type: ignore[attr-defined] get_current_process = kernel32.GetCurrentProcess get_current_process.argtypes = [] get_current_process.restype = ctypes.c_void_p set_priority_class = kernel32.SetPriorityClass set_priority_class.argtypes = [ctypes.c_void_p, ctypes.c_uint32] set_priority_class.restype = ctypes.c_int handle = get_current_process() if not set_priority_class(handle, BELOW_NORMAL_PRIORITY_CLASS): eprint("Warning: could not lower trainer process priority.") def atomic_write(path: Path, payload: bytes) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name( f"{path.name}.tmp.{os.getpid()}.{uuid.uuid4().hex}" ) try: with temporary.open("wb") as output: output.write(payload) output.flush() os.fsync(output.fileno()) os.replace(temporary, path) finally: with contextlib.suppress(FileNotFoundError): temporary.unlink() def atomic_write_json(path: Path, value: object) -> None: payload = ( json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n" ).encode("utf-8") atomic_write(path, payload) def read_exact(stream: BinaryIO, size: int) -> bytes: parts: list[bytes] = [] remaining = size while remaining: part = stream.read(remaining) if not part: raise EOFError(f"unexpected EOF: wanted {size} bytes") parts.append(part) remaining -= len(part) return b"".join(parts) def find_exporter(repository: Path, explicit: str | None) -> Path: if explicit: candidate = Path(explicit).expanduser().resolve() if candidate.is_file(): return candidate raise FileNotFoundError(f"training exporter not found: {candidate}") names = [ repository / "build/msbuild/bin/x64/Release/TransformerTrainingExport.exe", repository / "build/vs2026/Release/TransformerTrainingExport.exe", repository / "build/Release/TransformerTrainingExport.exe", repository / "build/bin/Release/TransformerTrainingExport.exe", ] for candidate in names: if candidate.is_file(): return candidate.resolve() matches = sorted( repository.glob("build/**/TransformerTrainingExport.exe"), key=lambda item: ("Release" not in item.parts, len(item.parts)), ) if matches: return matches[0].resolve() raise FileNotFoundError( "TransformerTrainingExport.exe was not found. Build its Release x64 " "project before running the trainer, or pass --exporter." ) @dataclasses.dataclass(frozen=True) class SnapshotInfo: path: Path corpus_digest: bytes record_count: int split_counts: dict[str, int] record_size: int created_by_trainer: bool includes_additional_runs: bool # Contiguous record-index ranges. The native exporter orders whole runs, # so training can globally shuffle bounded chunks without a RAM-sized index. split_segments: dict[str, tuple[tuple[int, int], ...]] def validate_stream_header(header: bytes) -> int: if len(header) != EXPORT_HEADER.size: raise ValueError("training stream header is truncated") magic, version, record_size = EXPORT_HEADER.unpack(header) if magic != EXPORT_MAGIC: raise ValueError(f"unexpected training stream magic: {magic!r}") if version != EXPORT_VERSION: raise ValueError(f"unsupported training stream version: {version}") if record_size != EXPECTED_RECORD_SIZE: raise ValueError( f"training record size mismatch: exporter={record_size}, " f"trainer={EXPECTED_RECORD_SIZE}" ) return record_size def classify_run(run_id: bytes, includes_additional_runs: bool) -> str | None: split = RUN_BYTES_TO_SPLIT.get(run_id) if split is not None: return split return "deployment_extra" if includes_additional_runs else None def inspect_snapshot( path: Path, compute_digest: bool = True, includes_additional_runs: bool = False, ) -> SnapshotInfo: digest = hashlib.sha256() counts: dict[str, int] = defaultdict(int) segments: dict[str, list[tuple[int, int]]] = defaultdict(list) total = 0 active_split: str | None = None active_start = 0 with path.open("rb") as stream: header = read_exact(stream, EXPORT_HEADER.size) record_size = validate_stream_header(header) while True: record = stream.read(record_size) if not record: break if len(record) != record_size: raise ValueError("training snapshot ends with a partial record") split = classify_run(record[:16], includes_additional_runs) if split is None: raise ValueError( "exporter emitted a run outside the requested four-run corpus" ) if split != active_split: if active_split is not None: segments[active_split].append((active_start, total - active_start)) active_split = split active_start = total counts[split] += 1 total += 1 if compute_digest: digest.update(record) if active_split is not None: segments[active_split].append((active_start, total - active_start)) return SnapshotInfo( path=path, corpus_digest=digest.digest(), record_count=total, split_counts=dict(counts), record_size=record_size, created_by_trainer=False, includes_additional_runs=includes_additional_runs, split_segments={key: tuple(value) for key, value in segments.items()}, ) def create_training_snapshot( repository: Path, archive_root: Path, exporter: Path, snapshot_directory: Path, includes_additional_runs: bool, ) -> SnapshotInfo: snapshot_directory.mkdir(parents=True, exist_ok=True) base = snapshot_directory / f"snapshot_{os.getpid()}_{time.time_ns()}" partial_path = base.with_suffix(".sztx.partial") final_path = base.with_suffix(".sztx") command = [str(exporter), "--root", str(archive_root)] if not includes_additional_runs: for run_id in SPLIT_RUNS.values(): command.extend(["--include-run", run_id]) eprint("Exporting an immutable four-run training snapshot...") eprint(" " + subprocess.list2cmdline(command)) creationflags = 0 if os.name == "nt": creationflags = getattr(subprocess, "BELOW_NORMAL_PRIORITY_CLASS", 0) process = subprocess.Popen( command, cwd=repository, stdout=subprocess.PIPE, stderr=None, creationflags=creationflags, ) assert process.stdout is not None digest = hashlib.sha256() counts: dict[str, int] = defaultdict(int) segments: dict[str, list[tuple[int, int]]] = defaultdict(list) record_count = 0 active_split: str | None = None active_start = 0 try: header = read_exact(process.stdout, EXPORT_HEADER.size) record_size = validate_stream_header(header) with partial_path.open("wb") as output: output.write(header) buffered = bytearray() while True: chunk = process.stdout.read(8 * 1024 * 1024) if not chunk: break output.write(chunk) buffered.extend(chunk) complete_bytes = len(buffered) - len(buffered) % record_size offset = 0 while offset < complete_bytes: view = memoryview(buffered)[offset : offset + record_size] split = classify_run( bytes(view[:16]), includes_additional_runs ) if split is None: raise ValueError( "exporter emitted a run outside the requested corpus" ) if split != active_split: if active_split is not None: segments[active_split].append( (active_start, record_count - active_start) ) active_split = split active_start = record_count digest.update(view) counts[split] += 1 record_count += 1 offset += record_size if complete_bytes: buffered = bytearray(buffered[complete_bytes:]) if buffered: raise ValueError("exporter ended with a partial training record") output.flush() os.fsync(output.fileno()) return_code = process.wait() if return_code != 0: raise RuntimeError( f"TransformerTrainingExport failed with exit code {return_code}; " "its partial output is invalid" ) missing = sorted(set(SPLIT_RUNS) - set(counts)) if missing: raise ValueError( "training export contains no records for split(s): " + ", ".join(missing) ) if active_split is not None: segments[active_split].append( (active_start, record_count - active_start) ) os.replace(partial_path, final_path) return SnapshotInfo( path=final_path, corpus_digest=digest.digest(), record_count=record_count, split_counts=dict(counts), record_size=record_size, created_by_trainer=True, includes_additional_runs=includes_additional_runs, split_segments={key: tuple(value) for key, value in segments.items()}, ) except BaseException: if process.poll() is None: process.terminate() with contextlib.suppress(subprocess.TimeoutExpired): process.wait(timeout=5) if process.poll() is None: process.kill() raise finally: process.stdout.close() with contextlib.suppress(FileNotFoundError): partial_path.unlink() def iter_record_batches( snapshot: SnapshotInfo, splits: set[str], batch_size: int, shuffle_records: int, seed: int, max_records_per_split: int, ) -> Iterator[list[bytes]]: if batch_size <= 0: raise ValueError("batch_size must be positive") rng = random.Random(seed) # Export order is whole-run order. Shuffle a tiny list of disk chunk # descriptors first, then shuffle records inside a bounded RAM buffer. This # interleaves train/validation/test/stress during deployment retraining # without retaining a 570k-record index or loading the 1 GB snapshot. io_chunk_records = max(batch_size, min(4096, max(batch_size, shuffle_records))) descriptors: list[tuple[int, int]] = [] for split in splits: remaining = ( max_records_per_split if max_records_per_split else snapshot.split_counts.get(split, 0) ) for segment_start, segment_count in snapshot.split_segments.get(split, ()): take = min(segment_count, remaining) for local_start in range(0, take, io_chunk_records): descriptors.append( ( segment_start + local_start, min(io_chunk_records, take - local_start), ) ) remaining -= take if remaining <= 0: break if shuffle_records > 1: rng.shuffle(descriptors) else: descriptors.sort(key=lambda value: value[0]) capacity = max(batch_size, shuffle_records) buffer: list[bytes] = [] with snapshot.path.open("rb") as stream: record_size = validate_stream_header(read_exact(stream, EXPORT_HEADER.size)) for record_index, record_count in descriptors: stream.seek(EXPORT_HEADER.size + record_index * record_size) block = read_exact(stream, record_count * record_size) buffer.extend( block[offset : offset + record_size] for offset in range(0, len(block), record_size) ) if len(buffer) < capacity: continue if shuffle_records > 1: rng.shuffle(buffer) complete = len(buffer) - len(buffer) % batch_size for start in range(0, complete, batch_size): yield buffer[start : start + batch_size] buffer = buffer[complete:] if buffer: if shuffle_records > 1: rng.shuffle(buffer) yield buffer def strided_array( blob: bytes, batch_size: int, record_size: int, offset: int, dtype: str, width: int = 1, ) -> np.ndarray: item_size = np.dtype(dtype).itemsize if width == 1: shape = (batch_size,) strides = (record_size,) else: shape = (batch_size, width) strides = (record_size, item_size) return np.ndarray( shape=shape, dtype=np.dtype(dtype), buffer=blob, offset=offset, strides=strides, ).copy() @dataclasses.dataclass class TrainingBatch: proposal_face: torch.Tensor proposal_global: torch.Tensor anchor_face: torch.Tensor anchor_global: torch.Tensor topology: torch.Tensor improved: torch.Tensor better: torch.Tensor signed_gain: torch.Tensor plane_action: torch.Tensor move_action: torch.Tensor scale_action: torch.Tensor plane_propensity: torch.Tensor move_propensity: torch.Tensor scale_propensity: torch.Tensor sample_weight: torch.Tensor @property def size(self) -> int: return int(self.topology.shape[0]) def to(self, device: torch.device) -> "TrainingBatch": return TrainingBatch(**{ field.name: getattr(self, field.name).to(device, non_blocking=False) for field in dataclasses.fields(self) }) def decode_batch(records: Sequence[bytes]) -> TrainingBatch: if not records: raise ValueError("cannot decode an empty batch") count = len(records) blob = b"".join(records) record_size = len(records[0]) if record_size != EXPECTED_RECORD_SIZE or any( len(record) != record_size for record in records ): raise ValueError("inconsistent fixed training record size") proposal = strided_array( blob, count, record_size, PROPOSAL_FEATURE_OFFSET, "= TOPOLOGY_COUNT).any(): raise ValueError("exporter emitted an invalid topology index") # A replica and a longer short rollout provide somewhat stronger evidence, # but max-of-many descendants must not dominate single-chain observations. replica_weight = 1.0 + 0.12 * np.log2(np.maximum(1.0, replicas)) effort = 0.80 + 0.40 * np.clip( np.log1p(iterations) / math.log(1_000_001.0), 0.0, 1.0 ) sample_weight = np.clip(replica_weight * effort, 0.60, 2.0).astype(np.float32) proposal_face = proposal[:, : FACE_COUNT * FACE_FEATURES].reshape( count, FACE_COUNT, FACE_FEATURES ) proposal_global = proposal[:, FACE_COUNT * FACE_FEATURES :] anchor_face = anchor[:, : FACE_COUNT * FACE_FEATURES].reshape( count, FACE_COUNT, FACE_FEATURES ) anchor_global = anchor[:, FACE_COUNT * FACE_FEATURES :] return TrainingBatch( proposal_face=torch.from_numpy(proposal_face), proposal_global=torch.from_numpy(proposal_global), anchor_face=torch.from_numpy(anchor_face), anchor_global=torch.from_numpy(anchor_global), topology=torch.from_numpy(topology), improved=torch.from_numpy(outcomes[:, 0].astype(np.float32)), better=torch.from_numpy(outcomes[:, 1].astype(np.bool_)), # Export field 2 is signed total defect gain: anchor C+I minus result C+I. signed_gain=torch.from_numpy(gains[:, 2].astype(np.float32)), plane_action=torch.from_numpy(actions[:, 0].astype(np.int64)), move_action=torch.from_numpy(actions[:, 1].astype(np.int64)), scale_action=torch.from_numpy(actions[:, 2].astype(np.int64)), plane_propensity=torch.from_numpy(propensities[:, 0]), move_propensity=torch.from_numpy(propensities[:, 1]), scale_propensity=torch.from_numpy(propensities[:, 2]), sample_weight=torch.from_numpy(sample_weight), ) class PreNormEncoderBlock(nn.Module): def __init__(self) -> None: super().__init__() self.norm1 = nn.LayerNorm(MODEL_WIDTH) self.self_attn = nn.MultiheadAttention( MODEL_WIDTH, ATTENTION_HEADS, dropout=0.0, batch_first=True, ) self.norm2 = nn.LayerNorm(MODEL_WIDTH) self.linear1 = nn.Linear(MODEL_WIDTH, FEED_FORWARD_WIDTH) self.linear2 = nn.Linear(FEED_FORWARD_WIDTH, MODEL_WIDTH) def forward(self, value: torch.Tensor) -> torch.Tensor: normalized = self.norm1(value) attended, _ = self.self_attn( normalized, normalized, normalized, need_weights=False, ) value = value + attended normalized = self.norm2(value) return value + self.linear2(F.silu(self.linear1(normalized))) class RankerMember(nn.Module): def __init__(self) -> None: super().__init__() self.face_proj = nn.Linear(FACE_FEATURES, MODEL_WIDTH) self.global_proj = nn.Linear(GLOBAL_FEATURES, MODEL_WIDTH) self.face_pos = nn.Parameter(torch.empty(FACE_COUNT, MODEL_WIDTH)) self.topology_embedding = nn.Embedding(TOPOLOGY_COUNT, MODEL_WIDTH) self.layers = nn.ModuleList( [PreNormEncoderBlock() for _ in range(LAYER_COUNT)] ) self.final_norm = nn.LayerNorm(MODEL_WIDTH) self.value_head = nn.Linear(MODEL_WIDTH, 2) self.plane_head = nn.Linear(MODEL_WIDTH, 1) self.move_head = nn.Linear(MODEL_WIDTH, MOVE_COUNT) self.scale_head = nn.Linear(MODEL_WIDTH, SCALE_COUNT) self.reset_parameters() def reset_parameters(self) -> None: nn.init.normal_(self.face_pos, mean=0.0, std=0.02) nn.init.normal_(self.topology_embedding.weight, mean=0.0, std=0.02) for module in self.modules(): if isinstance(module, nn.Linear): nn.init.xavier_uniform_(module.weight) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.LayerNorm): nn.init.ones_(module.weight) nn.init.zeros_(module.bias) # MultiheadAttention owns its packed projection outside nn.Linear. for layer in self.layers: nn.init.xavier_uniform_(layer.self_attn.in_proj_weight) if layer.self_attn.in_proj_bias is not None: nn.init.zeros_(layer.self_attn.in_proj_bias) def forward( self, face: torch.Tensor, global_features: torch.Tensor, topology: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: topology_token = self.topology_embedding(topology) cls = self.global_proj(global_features) + topology_token faces = ( self.face_proj(face) + self.face_pos.unsqueeze(0) + topology_token.unsqueeze(1) ) tokens = torch.cat((cls.unsqueeze(1), faces), dim=1) for layer in self.layers: tokens = layer(tokens) tokens = self.final_norm(tokens) cls = tokens[:, 0] face_tokens = tokens[:, 1:] return ( self.value_head(cls), self.plane_head(face_tokens).squeeze(-1), self.move_head(cls), self.scale_head(cls), ) def initialize_member(seed: int, device: torch.device) -> RankerMember: torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) return RankerMember().to(device) def weighted_mean(values: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: return (values * weights).sum() / weights.sum().clamp_min(1.0e-6) @dataclasses.dataclass class LossParts: total: torch.Tensor bce: torch.Tensor gain: torch.Tensor policy: torch.Tensor useful: int def member_loss( member: RankerMember, batch: TrainingBatch, bootstrap_weight: torch.Tensor, gain_clip: float, ips_clip: float, amp_enabled: bool, ) -> LossParts: count = batch.size combined_face = torch.cat((batch.proposal_face, batch.anchor_face), dim=0) combined_global = torch.cat( (batch.proposal_global, batch.anchor_global), dim=0 ) combined_topology = torch.cat((batch.topology, batch.topology), dim=0) with torch.autocast( device_type="cuda", dtype=torch.bfloat16, enabled=amp_enabled, ): value, plane_logits, move_logits, scale_logits = member( combined_face, combined_global, combined_topology ) proposal_value = value[:count] anchor_plane = plane_logits[count:] anchor_move = move_logits[count:] anchor_scale = scale_logits[count:] weights = batch.sample_weight * bootstrap_weight weights = weights / weights.mean().clamp_min(1e-6) bce_values = F.binary_cross_entropy_with_logits( proposal_value[:, 0].float(), batch.improved, reduction="none" ) bce = weighted_mean(bce_values, weights) gain_target = ( batch.signed_gain.clamp(-gain_clip, gain_clip) if gain_clip > 0.0 else batch.signed_gain ) gain_values = F.smooth_l1_loss( proposal_value[:, 1].float(), gain_target, reduction="none" ) gain = weighted_mean(gain_values, weights) useful = (batch.improved > 0.5) | batch.better valid = ( (batch.plane_action >= 0) & (batch.plane_action < FACE_COUNT) & (batch.move_action >= 0) & (batch.move_action < MOVE_COUNT) & (batch.scale_action >= 0) & (batch.scale_action < SCALE_COUNT) & (batch.plane_propensity > 0.0) & (batch.move_propensity > 0.0) & (batch.scale_propensity > 0.0) ) policy_mask = useful & valid useful_count = int(policy_mask.sum().item()) if useful_count: policy_weight = weights[policy_mask] plane_ips = ( 1.0 / (FACE_COUNT * batch.plane_propensity[policy_mask].clamp_min(1e-6)) ).clamp(0.25, ips_clip) move_ips = ( 1.0 / (MOVE_COUNT * batch.move_propensity[policy_mask].clamp_min(1e-6)) ).clamp(0.25, ips_clip) scale_ips = ( 1.0 / (SCALE_COUNT * batch.scale_propensity[policy_mask].clamp_min(1e-6)) ).clamp(0.25, ips_clip) plane_ce = F.cross_entropy( anchor_plane[policy_mask].float(), batch.plane_action[policy_mask], reduction="none", ) move_ce = F.cross_entropy( anchor_move[policy_mask].float(), batch.move_action[policy_mask], reduction="none", ) scale_ce = F.cross_entropy( anchor_scale[policy_mask].float(), batch.scale_action[policy_mask], reduction="none", ) policy = ( weighted_mean(plane_ce, policy_weight * plane_ips) + weighted_mean(move_ce, policy_weight * move_ips) + weighted_mean(scale_ce, policy_weight * scale_ips) ) / 3.0 else: policy = anchor_plane.sum().float() * 0.0 total = bce + 0.25 * gain + policy return LossParts(total, bce, gain, policy, useful_count) @dataclasses.dataclass class TrainEpochStats: records: int = 0 batches: int = 0 useful_policy_records: int = 0 loss: float = 0.0 bce: float = 0.0 gain: float = 0.0 policy: float = 0.0 def as_dict(self) -> dict[str, float | int]: denominator = max(1, self.batches * ENSEMBLE_SIZE) return { "records": self.records, "batches": self.batches, "useful_policy_records_member_sum": self.useful_policy_records, "loss": self.loss / denominator, "bce": self.bce / denominator, "gain_smooth_l1": self.gain / denominator, "policy_ce": self.policy / denominator, } def learning_rate_factor(step: int, total_steps: int) -> float: warmup = max(1, int(total_steps * 0.05)) if step < warmup: return max(0.05, float(step + 1) / float(warmup)) progress = (step - warmup) / max(1, total_steps - warmup) return 0.05 + 0.95 * 0.5 * (1.0 + math.cos(math.pi * progress)) def train_ensemble( snapshot: SnapshotInfo, members: list[RankerMember], splits: set[str], epochs: int, batch_size: int, shuffle_records: int, max_records_per_split: int, seed: int, learning_rate: float, weight_decay: float, gain_clip: float, ips_clip: float, gradient_clip: float, amp_enabled: bool, device: torch.device, validation_callback=None, ) -> tuple[list[dict[str, object]], int]: optimizers = [ torch.optim.AdamW( member.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(0.9, 0.999), ) for member in members ] selected_count = sum(snapshot.split_counts.get(split, 0) for split in splits) if max_records_per_split: selected_count = sum( min(snapshot.split_counts.get(split, 0), max_records_per_split) for split in splits ) steps_per_epoch = max(1, math.ceil(selected_count / batch_size)) total_steps = max(1, steps_per_epoch * epochs) schedulers = [ torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda=lambda step, total=total_steps: learning_rate_factor(step, total), ) for optimizer in optimizers ] history: list[dict[str, object]] = [] best_epoch = epochs best_score = -math.inf best_states: list[dict[str, torch.Tensor]] | None = None bootstrap_rngs = [ np.random.default_rng( (seed + (member_index + 1) * 0x9E3779B1) & 0xFFFFFFFFFFFFFFFF ) for member_index in range(ENSEMBLE_SIZE) ] for epoch in range(1, epochs + 1): for member in members: member.train() stats = TrainEpochStats() batches = iter_record_batches( snapshot, splits, batch_size, shuffle_records, seed + epoch * 0x9E3779B1, max_records_per_split, ) for raw_records in batches: batch = decode_batch(raw_records).to(device) for optimizer in optimizers: optimizer.zero_grad(set_to_none=True) bootstrap_weights: list[torch.Tensor] = [] for bootstrap_rng in bootstrap_rngs: bootstrap = np.minimum( bootstrap_rng.poisson(1.0, batch.size), 4 ).astype(np.float32) if not bootstrap.any(): bootstrap[0] = 1.0 bootstrap_weights.append( torch.from_numpy(bootstrap).to(device, non_blocking=False) ) losses = [ member_loss( member, batch, bootstrap_weight, gain_clip, ips_clip, amp_enabled, ) for member, bootstrap_weight in zip(members, bootstrap_weights) ] # Members have disjoint parameters. Summing preserves each member's # natural loss scale (averaging would divide every gradient by 3). sum(part.total for part in losses).backward() for member, optimizer, scheduler in zip( members, optimizers, schedulers ): torch.nn.utils.clip_grad_norm_(member.parameters(), gradient_clip) optimizer.step() scheduler.step() stats.records += batch.size stats.batches += 1 for part in losses: stats.loss += float(part.total.detach().cpu()) stats.bce += float(part.bce.detach().cpu()) stats.gain += float(part.gain.detach().cpu()) stats.policy += float(part.policy.detach().cpu()) stats.useful_policy_records += part.useful epoch_report: dict[str, object] = { "epoch": epoch, "learning_rate": optimizers[0].param_groups[0]["lr"], "training": stats.as_dict(), } if validation_callback is not None: validation = validation_callback(members, epoch) epoch_report["validation"] = validation score = ( float(validation["average_precision"]) - 0.10 * float(validation["brier_score"]) + 0.01 * float(validation["top_k_gain"]) ) if math.isfinite(score) and score > best_score: best_score = score best_epoch = epoch best_states = [ { key: value.detach().cpu().clone() for key, value in member.state_dict().items() } for member in members ] history.append(epoch_report) eprint( f"Epoch {epoch}/{epochs}: loss={stats.as_dict()['loss']:.6f}, " f"records={stats.records:,}" ) if best_states is not None: for member, state in zip(members, best_states): member.load_state_dict(state) return history, best_epoch @dataclasses.dataclass class MetricAccumulator: probability_bins: int = 4096 calibration_bins: int = 20 def __post_init__(self) -> None: self.count = 0 self.positive = 0.0 self.brier_sum = 0.0 self.bce_sum = 0.0 self.gain_abs_sum = 0.0 self.gain_sum = 0.0 self.score_count = np.zeros(self.probability_bins, dtype=np.float64) self.score_positive = np.zeros(self.probability_bins, dtype=np.float64) self.score_gain = np.zeros(self.probability_bins, dtype=np.float64) self.cal_count = np.zeros(self.calibration_bins, dtype=np.float64) self.cal_probability = np.zeros(self.calibration_bins, dtype=np.float64) self.cal_positive = np.zeros(self.calibration_bins, dtype=np.float64) self.policy_useful = 0 self.plane_correct = 0 self.move_correct = 0 self.scale_correct = 0 self.plane_ce_sum = 0.0 self.move_ce_sum = 0.0 self.scale_ce_sum = 0.0 def add( self, probability: np.ndarray, gain_prediction: np.ndarray, improved: np.ndarray, signed_gain: np.ndarray, useful: np.ndarray, plane_probability: np.ndarray, move_probability: np.ndarray, scale_probability: np.ndarray, plane_action: np.ndarray, move_action: np.ndarray, scale_action: np.ndarray, ) -> None: probability = np.clip(probability.astype(np.float64), 1e-7, 1.0 - 1e-7) improved = improved.astype(np.float64) signed_gain = signed_gain.astype(np.float64) gain_prediction = gain_prediction.astype(np.float64) count = probability.size self.count += count self.positive += float(improved.sum()) self.brier_sum += float(np.square(probability - improved).sum()) self.bce_sum += float( (-(improved * np.log(probability) + (1.0 - improved) * np.log1p(-probability))).sum() ) self.gain_abs_sum += float(np.abs(gain_prediction - signed_gain).sum()) self.gain_sum += float(signed_gain.sum()) score_index = np.minimum( (probability * self.probability_bins).astype(np.int64), self.probability_bins - 1, ) self.score_count += np.bincount( score_index, minlength=self.probability_bins ) self.score_positive += np.bincount( score_index, weights=improved, minlength=self.probability_bins ) self.score_gain += np.bincount( score_index, weights=signed_gain, minlength=self.probability_bins ) cal_index = np.minimum( (probability * self.calibration_bins).astype(np.int64), self.calibration_bins - 1, ) self.cal_count += np.bincount(cal_index, minlength=self.calibration_bins) self.cal_probability += np.bincount( cal_index, weights=probability, minlength=self.calibration_bins ) self.cal_positive += np.bincount( cal_index, weights=improved, minlength=self.calibration_bins ) useful = useful.astype(bool) self.policy_useful += int(useful.sum()) if useful.any(): rows = np.flatnonzero(useful) useful_plane_actions = plane_action[useful].astype(np.int64) useful_move_actions = move_action[useful].astype(np.int64) useful_scale_actions = scale_action[useful].astype(np.int64) self.plane_ce_sum += float( -np.log( np.clip( plane_probability[rows, useful_plane_actions], 1e-7, 1.0 ) ).sum() ) self.move_ce_sum += float( -np.log( np.clip( move_probability[rows, useful_move_actions], 1e-7, 1.0 ) ).sum() ) self.scale_ce_sum += float( -np.log( np.clip( scale_probability[rows, useful_scale_actions], 1e-7, 1.0 ) ).sum() ) plane_prediction = plane_probability.argmax(axis=-1) move_prediction = move_probability.argmax(axis=-1) scale_prediction = scale_probability.argmax(axis=-1) self.plane_correct += int( (plane_prediction[useful] == plane_action[useful]).sum() ) self.move_correct += int( (move_prediction[useful] == move_action[useful]).sum() ) self.scale_correct += int( (scale_prediction[useful] == scale_action[useful]).sum() ) def finalize(self) -> dict[str, float | int]: if self.count == 0: raise ValueError("cannot evaluate an empty split") prevalence = self.positive / self.count total_positive = self.positive true_positive = 0.0 false_positive = 0.0 previous_recall = 0.0 average_precision = 0.0 if 0.0 < total_positive < self.count: for index in range(self.probability_bins - 1, -1, -1): true_positive += self.score_positive[index] false_positive += ( self.score_count[index] - self.score_positive[index] ) recall = true_positive / total_positive precision = true_positive / max(1.0, true_positive + false_positive) average_precision += precision * (recall - previous_recall) previous_recall = recall ece = 0.0 for index in range(self.calibration_bins): count = self.cal_count[index] if count: mean_probability = self.cal_probability[index] / count mean_positive = self.cal_positive[index] / count ece += (count / self.count) * abs(mean_probability - mean_positive) top_target = max(1.0, math.ceil(self.count * 0.10)) top_seen = 0.0 top_gain = 0.0 for index in range(self.probability_bins - 1, -1, -1): count = self.score_count[index] if count <= 0.0: continue take = min(count, top_target - top_seen) top_gain += self.score_gain[index] * (take / count) top_seen += take if top_seen >= top_target: break overall_gain = self.gain_sum / self.count top_mean_gain = top_gain / max(1.0, top_seen) null_brier = prevalence * (1.0 - prevalence) null_bce = 0.0 if 0.0 < prevalence < 1.0: null_bce = -( prevalence * math.log(prevalence) + (1.0 - prevalence) * math.log1p(-prevalence) ) policy_denominator = max(1, self.policy_useful) plane_ce = self.plane_ce_sum / policy_denominator move_ce = self.move_ce_sum / policy_denominator scale_ce = self.scale_ce_sum / policy_denominator return { "records": self.count, "positives": int(round(self.positive)), "prevalence": prevalence, "average_precision": average_precision, "brier_score": self.brier_sum / self.count, "null_brier_score": null_brier, "binary_cross_entropy": self.bce_sum / self.count, "null_binary_cross_entropy": null_bce, "expected_calibration_error": ece, "gain_mae": self.gain_abs_sum / self.count, "mean_signed_gain": overall_gain, "top_decile_mean_signed_gain": top_mean_gain, "top_k_gain": top_mean_gain - overall_gain, "useful_policy_records": self.policy_useful, "plane_top1_accuracy": self.plane_correct / policy_denominator, "move_top1_accuracy": self.move_correct / policy_denominator, "scale_top1_accuracy": self.scale_correct / policy_denominator, "plane_cross_entropy": plane_ce, "move_cross_entropy": move_ce, "scale_cross_entropy": scale_ce, "policy_cross_entropy": (plane_ce + move_ce + scale_ce) / 3.0, "plane_uniform_cross_entropy": math.log(FACE_COUNT), "move_uniform_cross_entropy": math.log(MOVE_COUNT), "scale_uniform_cross_entropy": math.log(SCALE_COUNT), "policy_uniform_cross_entropy": ( math.log(FACE_COUNT) + math.log(MOVE_COUNT) + math.log(SCALE_COUNT) ) / 3.0, } @torch.no_grad() def evaluate_ensemble( snapshot: SnapshotInfo, members: list[RankerMember], split: str, batch_size: int, max_records_per_split: int, gain_clip: float, amp_enabled: bool, device: torch.device, ) -> dict[str, float | int]: for member in members: member.eval() metrics = MetricAccumulator() batches = iter_record_batches( snapshot, {split}, batch_size, 0, 0, max_records_per_split, ) for raw_records in batches: batch = decode_batch(raw_records).to(device) count = batch.size combined_face = torch.cat((batch.proposal_face, batch.anchor_face), dim=0) combined_global = torch.cat( (batch.proposal_global, batch.anchor_global), dim=0 ) combined_topology = torch.cat((batch.topology, batch.topology), dim=0) probabilities = [] gains = [] plane_probabilities = [] move_probabilities = [] scale_probabilities = [] with torch.autocast( device_type="cuda", dtype=torch.bfloat16, enabled=amp_enabled, ): for member in members: value, plane, move, scale = member( combined_face, combined_global, combined_topology ) probabilities.append(torch.sigmoid(value[:count, 0].float())) gains.append(value[:count, 1].float()) plane_probabilities.append(torch.softmax(plane[count:].float(), dim=-1)) move_probabilities.append(torch.softmax(move[count:].float(), dim=-1)) scale_probabilities.append(torch.softmax(scale[count:].float(), dim=-1)) probability = torch.stack(probabilities).mean(0) gain = torch.stack(gains).mean(0) plane_probability = torch.stack(plane_probabilities).mean(0) move_probability = torch.stack(move_probabilities).mean(0) scale_probability = torch.stack(scale_probabilities).mean(0) valid_actions = ( (batch.plane_action >= 0) & (batch.plane_action < FACE_COUNT) & (batch.move_action >= 0) & (batch.move_action < MOVE_COUNT) & (batch.scale_action >= 0) & (batch.scale_action < SCALE_COUNT) ) useful = ((batch.improved > 0.5) | batch.better) & valid_actions metrics.add( probability.cpu().numpy(), gain.cpu().numpy(), batch.improved.cpu().numpy(), ( batch.signed_gain.clamp(-gain_clip, gain_clip) if gain_clip > 0.0 else batch.signed_gain ).cpu().numpy(), useful.cpu().numpy(), plane_probability.cpu().numpy(), move_probability.cpu().numpy(), scale_probability.cpu().numpy(), batch.plane_action.cpu().numpy(), batch.move_action.cpu().numpy(), batch.scale_action.cpu().numpy(), ) return metrics.finalize() def report_gate( validation: dict[str, float | int], test: dict[str, float | int], stress: dict[str, float | int], minimum_records: int, ) -> tuple[bool, list[str]]: reasons: list[str] = [] for name, metrics in ( ("validation", validation), ("test", test), ("stress", stress), ): numeric = [float(value) for value in metrics.values()] if not all(math.isfinite(value) for value in numeric): reasons.append(f"{name}: non-finite metric") if int(metrics["records"]) < minimum_records: reasons.append(f"{name}: fewer than {minimum_records} records") prevalence = float(metrics["prevalence"]) if not 0.0 < prevalence < 1.0: reasons.append(f"{name}: split needs both positive and negative outcomes") for name, metrics in (("validation", validation), ("test", test)): prevalence = float(metrics["prevalence"]) if float(metrics["average_precision"]) <= prevalence: reasons.append(f"{name}: AP does not beat prevalence") if float(metrics["top_k_gain"]) <= 0.0: reasons.append(f"{name}: top-decile rank has no signed-gain uplift") if float(metrics["brier_score"]) > 1.05 * float(metrics["null_brier_score"]): reasons.append(f"{name}: Brier score is worse than the null model by >5%") if float(metrics["policy_cross_entropy"]) > 1.02 * float( metrics["policy_uniform_cross_entropy"] ): reasons.append(f"{name}: policy CE is worse than uniform by >2%") if float(stress["average_precision"]) < 0.95 * float(stress["prevalence"]): reasons.append("stress: AP falls materially below prevalence") if float(stress["top_k_gain"]) < -0.05: reasons.append("stress: top-decile signed-gain uplift is materially negative") return not reasons, reasons def tensor_bytes(tensor: torch.Tensor) -> bytes: array = tensor.detach().cpu().float().contiguous().numpy().astype(" Iterable[torch.Tensor]: yield member.face_proj.weight yield member.face_proj.bias yield member.global_proj.weight yield member.global_proj.bias yield member.face_pos yield member.topology_embedding.weight for layer in member.layers: yield layer.norm1.weight yield layer.norm1.bias yield layer.self_attn.in_proj_weight assert layer.self_attn.in_proj_bias is not None yield layer.self_attn.in_proj_bias yield layer.self_attn.out_proj.weight assert layer.self_attn.out_proj.bias is not None yield layer.self_attn.out_proj.bias yield layer.norm2.weight yield layer.norm2.bias yield layer.linear1.weight yield layer.linear1.bias yield layer.linear2.weight yield layer.linear2.bias yield member.final_norm.weight yield member.final_norm.bias yield member.value_head.weight yield member.value_head.bias yield member.plane_head.weight yield member.plane_head.bias yield member.move_head.weight yield member.move_head.bias yield member.scale_head.weight yield member.scale_head.bias def expected_payload_float_count() -> int: per_member = 0 per_member += MODEL_WIDTH * FACE_FEATURES + MODEL_WIDTH per_member += MODEL_WIDTH * GLOBAL_FEATURES + MODEL_WIDTH per_member += FACE_COUNT * MODEL_WIDTH per_member += TOPOLOGY_COUNT * MODEL_WIDTH per_layer = 0 per_layer += MODEL_WIDTH * 2 per_layer += (3 * MODEL_WIDTH) * MODEL_WIDTH + 3 * MODEL_WIDTH per_layer += MODEL_WIDTH * MODEL_WIDTH + MODEL_WIDTH per_layer += MODEL_WIDTH * 2 per_layer += FEED_FORWARD_WIDTH * MODEL_WIDTH + FEED_FORWARD_WIDTH per_layer += MODEL_WIDTH * FEED_FORWARD_WIDTH + MODEL_WIDTH per_member += LAYER_COUNT * per_layer per_member += MODEL_WIDTH * 2 per_member += 2 * MODEL_WIDTH + 2 per_member += MODEL_WIDTH + 1 per_member += MOVE_COUNT * MODEL_WIDTH + MOVE_COUNT per_member += SCALE_COUNT * MODEL_WIDTH + SCALE_COUNT return ENSEMBLE_SIZE * per_member def serialize_model( members: Sequence[RankerMember], training_seed: int, corpus_digest: bytes, validation: dict[str, float | int], ) -> bytes: if len(members) != ENSEMBLE_SIZE: raise ValueError(f"expected {ENSEMBLE_SIZE} ensemble members") if len(corpus_digest) != 32: raise ValueError("training corpus digest must be SHA-256") payload = b"".join( tensor_bytes(tensor) for member in members for tensor in member_tensors(member) ) if len(payload) % 4: raise AssertionError("float payload is not aligned") float_count = len(payload) // 4 expected_count = expected_payload_float_count() if float_count != expected_count: raise AssertionError( f"payload tensor order/shape mismatch: {float_count} != {expected_count}" ) payload_crc = zlib.crc32(payload) & 0xFFFFFFFF header = bytearray(MODEL_HEADER_BYTES) header[0:8] = MODEL_MAGIC struct.pack_into( " tuple[str, Path]: digest = hashlib.sha256(model_bytes).hexdigest() generation_path = output_directory / f"model_{digest}.sztf" if generation_path.exists(): existing_digest = hashlib.sha256(generation_path.read_bytes()).hexdigest() if existing_digest != digest: raise RuntimeError("immutable model generation digest collision") else: atomic_write(generation_path, model_bytes) # Search loads current.sztf. Replacing it never mutates immutable generations. atomic_write(output_directory / "current.sztf", model_bytes) report["model_digest_sha256"] = digest report["model_generation"] = generation_path.name report_payload = ( json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n" ).encode("utf-8") generation_report = output_directory / f"report_{digest}.json" if not generation_report.exists(): atomic_write(generation_report, report_payload) atomic_write(output_directory / "report.json", report_payload) return digest, generation_path def parse_args(argv: Sequence[str]) -> argparse.Namespace: parser = argparse.ArgumentParser( description="Train the CUDA PyTorch Szilassi proposal transformer." ) parser.add_argument( "--repository", default=str(Path(__file__).resolve().parents[1]), help="repository root", ) parser.add_argument( "--archive-root", default="results/search", help="search archive root, relative to repository by default", ) parser.add_argument("--exporter", help="TransformerTrainingExport.exe path") parser.add_argument( "--snapshot", help="reuse an existing SZTXP001 snapshot instead of launching exporter", ) parser.add_argument( "--snapshot-directory", default="runtime/transformer_training", help="temporary fixed-stream snapshot directory", ) parser.add_argument( "--keep-snapshot", action="store_true", help="retain generated snapshot" ) parser.add_argument( "--deployment-all-runs", action="store_true", help=( "include every sealed run in final deployment retraining while " "keeping the four fixed reporting splits; intended for the later " "full-corpus/200 GB retrain" ), ) parser.add_argument( "--output-directory", default="results/search/neural/transformer", ) parser.add_argument("--seed", type=int, default=0x5A17A551) parser.add_argument("--report-epochs", type=int, default=5) parser.add_argument( "--deployment-epochs", type=int, default=0, help="0 uses the best reporting epoch", ) parser.add_argument("--batch-size", type=int, default=1024) parser.add_argument("--shuffle-records", type=int, default=32768) parser.add_argument("--learning-rate", type=float, default=3.0e-4) parser.add_argument("--weight-decay", type=float, default=1.0e-4) parser.add_argument("--gradient-clip", type=float, default=1.0) parser.add_argument( "--gain-clip", type=float, default=0.0, help="absolute signed-gain target clip; 0 keeps the complete target", ) parser.add_argument("--ips-clip", type=float, default=4.0) parser.add_argument( "--max-records-per-split", type=int, default=0, help="diagnostic bound; 0 trains on every exported record", ) parser.add_argument("--minimum-gate-records", type=int, default=1000) parser.add_argument( "--no-bf16", action="store_true", help="disable BF16 autocast" ) parser.add_argument( "--force-deploy", action="store_true", help="publish despite a failed statistical gate (diagnostics only)", ) parser.add_argument( "--nondeterministic", action="store_true", help="allow nondeterministic CUDA kernels", ) return parser.parse_args(argv) def resolve_under_repository(repository: Path, value: str) -> Path: path = Path(value).expanduser() if not path.is_absolute(): path = repository / path return path.resolve() def validate_options(args: argparse.Namespace) -> None: positive_integer_names = ( "report_epochs", "batch_size", "shuffle_records", "minimum_gate_records", ) for name in positive_integer_names: if getattr(args, name) <= 0: raise ValueError(f"--{name.replace('_', '-')} must be positive") if args.deployment_epochs < 0 or args.max_records_per_split < 0: raise ValueError("epoch/record bounds cannot be negative") for name in ( "learning_rate", "weight_decay", "gradient_clip", "ips_clip", ): value = float(getattr(args, name)) if not math.isfinite(value) or value <= 0.0: raise ValueError(f"--{name.replace('_', '-')} must be positive and finite") if not math.isfinite(args.gain_clip) or args.gain_clip < 0.0: raise ValueError("--gain-clip must be finite and non-negative") def main(argv: Sequence[str] | None = None) -> int: args = parse_args(argv or sys.argv[1:]) validate_options(args) set_below_normal_priority() if not torch.cuda.is_available(): raise RuntimeError( "CUDA PyTorch is mandatory for transformer training; no CUDA device is available" ) device = torch.device("cuda:0") properties = torch.cuda.get_device_properties(device) amp_enabled = not args.no_bf16 and properties.major >= 8 repository = Path(args.repository).expanduser().resolve() archive_root = resolve_under_repository(repository, args.archive_root) snapshot_directory = resolve_under_repository( repository, args.snapshot_directory ) output_directory = resolve_under_repository( repository, args.output_directory ) output_directory.mkdir(parents=True, exist_ok=True) random.seed(args.seed) np.random.seed(args.seed & 0xFFFFFFFF) torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) if not args.nondeterministic: torch.use_deterministic_algorithms(True) torch.backends.cudnn.benchmark = False torch.backends.cuda.matmul.allow_tf32 = False if hasattr(torch.backends.cuda, "enable_flash_sdp"): torch.backends.cuda.enable_flash_sdp(False) torch.backends.cuda.enable_mem_efficient_sdp(False) torch.backends.cuda.enable_math_sdp(True) eprint(f"CUDA device: {properties.name}") eprint(f"PyTorch: {torch.__version__}; BF16 autocast: {amp_enabled}") snapshot: SnapshotInfo | None = None try: if args.snapshot: snapshot_path = resolve_under_repository(repository, args.snapshot) snapshot = inspect_snapshot( snapshot_path, includes_additional_runs=args.deployment_all_runs, ) else: exporter = find_exporter(repository, args.exporter) snapshot = create_training_snapshot( repository, archive_root, exporter, snapshot_directory, args.deployment_all_runs, ) eprint( f"Snapshot: {snapshot.record_count:,} records, " f"SHA-256 {snapshot.corpus_digest.hex()}" ) for split in SPLIT_RUNS: eprint( f" {split:10s} {snapshot.split_counts.get(split, 0):,} records " f"({SPLIT_RUNS[split]})" ) if snapshot.split_counts.get("deployment_extra", 0): eprint( " deployment_extra " f"{snapshot.split_counts['deployment_extra']:,} records" ) reporting_members = [ initialize_member(args.seed + member * 0x9E3779B1, device) for member in range(ENSEMBLE_SIZE) ] def validate_during_training( members: list[RankerMember], epoch: int ) -> dict[str, float | int]: del epoch return evaluate_ensemble( snapshot, members, "validation", args.batch_size, args.max_records_per_split, args.gain_clip, amp_enabled, device, ) report_history, best_epoch = train_ensemble( snapshot=snapshot, members=reporting_members, splits={"train"}, epochs=args.report_epochs, batch_size=args.batch_size, shuffle_records=args.shuffle_records, max_records_per_split=args.max_records_per_split, seed=args.seed, learning_rate=args.learning_rate, weight_decay=args.weight_decay, gain_clip=args.gain_clip, ips_clip=args.ips_clip, gradient_clip=args.gradient_clip, amp_enabled=amp_enabled, device=device, validation_callback=validate_during_training, ) reporting_metrics = { split: evaluate_ensemble( snapshot, reporting_members, split, args.batch_size, args.max_records_per_split, args.gain_clip, amp_enabled, device, ) for split in ("validation", "test", "stress") } approved, gate_reasons = report_gate( reporting_metrics["validation"], reporting_metrics["test"], reporting_metrics["stress"], args.minimum_gate_records, ) if args.force_deploy and not approved: eprint("WARNING: statistical deployment gate was overridden.") approved = True base_report: dict[str, object] = { "format": "szilassi-transformer-training-report-v1", "approved": approved, "gate_overridden": bool(args.force_deploy and gate_reasons), "gate_reasons": gate_reasons, "corpus_digest_sha256": snapshot.corpus_digest.hex(), "corpus_records": snapshot.record_count, "split_runs": SPLIT_RUNS, "split_records": snapshot.split_counts, "deployment_all_runs": snapshot.includes_additional_runs, "report_best_epoch": best_epoch, "report_training_history": report_history, "reporting_metrics": reporting_metrics, "architecture": { "ensemble": ENSEMBLE_SIZE, "tokens": "CLS+12 faces", "face_features": FACE_FEATURES, "global_features": GLOBAL_FEATURES, "d_model": MODEL_WIDTH, "layers": LAYER_COUNT, "heads": ATTENTION_HEADS, "feed_forward": FEED_FORWARD_WIDTH, "activation": "SiLU", "normalization": "pre-LayerNorm plus final LayerNorm", "dropout": 0.0, "ensemble_bootstrap": "independent capped Poisson(1)", }, "loss": { "value": "BCE(improved) + 0.25*SmoothL1(signed defect gain)", "policy": "mean clipped normalized-IPS CE on useful outcomes", "gain_clip": args.gain_clip if args.gain_clip > 0.0 else None, "ips_clip": args.ips_clip, }, "training": { "seed": args.seed, "batch_size": args.batch_size, "shuffle_records": args.shuffle_records, "learning_rate": args.learning_rate, "weight_decay": args.weight_decay, "gradient_clip": args.gradient_clip, "report_epochs": args.report_epochs, "bf16_autocast": amp_enabled, "cuda_device": properties.name, "pytorch_version": torch.__version__, "deterministic": not args.nondeterministic, }, "created_unix_time_ns": time.time_ns(), } if not approved: rejected_path = output_directory / ( f"rejected_report_{snapshot.corpus_digest.hex()}_{time.time_ns()}.json" ) atomic_write_json(rejected_path, base_report) eprint("Deployment gate failed; current.sztf was not replaced:") for reason in gate_reasons: eprint(f" - {reason}") return 4 del reporting_members torch.cuda.empty_cache() deployment_epochs = args.deployment_epochs or best_epoch deployment_members = [ initialize_member(args.seed + member * 0x9E3779B1, device) for member in range(ENSEMBLE_SIZE) ] deployment_splits = set(SPLIT_RUNS) if snapshot.includes_additional_runs: deployment_splits.add("deployment_extra") deployment_history, _ = train_ensemble( snapshot=snapshot, members=deployment_members, splits=deployment_splits, epochs=deployment_epochs, batch_size=args.batch_size, shuffle_records=args.shuffle_records, max_records_per_split=args.max_records_per_split, seed=args.seed ^ 0xD1B54A32, learning_rate=args.learning_rate, weight_decay=args.weight_decay, gain_clip=args.gain_clip, ips_clip=args.ips_clip, gradient_clip=args.gradient_clip, amp_enabled=amp_enabled, device=device, validation_callback=None, ) validation_metrics = reporting_metrics["validation"] model_bytes = serialize_model( deployment_members, args.seed, snapshot.corpus_digest, validation_metrics, ) base_report["deployment_epochs"] = deployment_epochs base_report["deployment_training_history"] = deployment_history digest, generation = publish_model( output_directory, model_bytes, base_report, ) eprint(f"Published immutable generation: {generation}") eprint(f"Current model digest: {digest}") return 0 finally: if ( snapshot is not None and snapshot.created_by_trainer and not args.keep_snapshot ): with contextlib.suppress(FileNotFoundError): snapshot.path.unlink() if __name__ == "__main__": try: raise SystemExit(main()) except KeyboardInterrupt: eprint("Training interrupted; current.sztf was not changed.") raise SystemExit(130) except Exception as exc: eprint(f"Transformer training failed: {exc}") raise SystemExit(2)