Efim Beshmenev
1830 lines
67 KiB
Python
1830 lines
67 KiB
Python
#!/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, "<f4", FEATURE_FLOATS
|
|
)
|
|
anchor = strided_array(
|
|
blob, count, record_size, ANCHOR_FEATURE_OFFSET, "<f4", FEATURE_FLOATS
|
|
)
|
|
topology = strided_array(
|
|
blob, count, record_size, TOPOLOGY_OFFSET, "<u4"
|
|
).astype(np.int64)
|
|
gains = strided_array(blob, count, record_size, GAIN_OFFSET, "<i4", 3)
|
|
outcomes = strided_array(blob, count, record_size, OUTCOME_OFFSET, "u1", 4)
|
|
actions = strided_array(blob, count, record_size, ACTION_OFFSET, "<i4", 3)
|
|
propensities = strided_array(
|
|
blob, count, record_size, PROPENSITY_OFFSET, "<f4", 3
|
|
)
|
|
replicas = strided_array(
|
|
blob, count, record_size, REPLICA_COUNT_OFFSET, "<u4"
|
|
).astype(np.float32)
|
|
iterations = strided_array(
|
|
blob, count, record_size, ROLLOUT_ITERATIONS_OFFSET, "<u8"
|
|
).astype(np.float64)
|
|
|
|
if not np.isfinite(proposal).all() or not np.isfinite(anchor).all():
|
|
raise ValueError("exporter emitted non-finite transformer features")
|
|
if not np.isfinite(propensities).all():
|
|
raise ValueError("exporter emitted non-finite action propensities")
|
|
if (topology < 0).any() or (topology >= 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("<f4", copy=False)
|
|
if not np.isfinite(array).all():
|
|
raise ValueError("model contains NaN or infinity")
|
|
return array.tobytes(order="C")
|
|
|
|
|
|
def member_tensors(member: RankerMember) -> 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(
|
|
"<IIII",
|
|
header,
|
|
8,
|
|
MODEL_FORMAT_VERSION,
|
|
FEATURE_FORMAT_VERSION,
|
|
OBJECTIVE_VERSION,
|
|
1, # approved
|
|
)
|
|
struct.pack_into(
|
|
"<9I",
|
|
header,
|
|
24,
|
|
FACE_COUNT,
|
|
FACE_FEATURES,
|
|
GLOBAL_FEATURES,
|
|
MODEL_WIDTH,
|
|
ATTENTION_HEADS,
|
|
LAYER_COUNT,
|
|
FEED_FORWARD_WIDTH,
|
|
ENSEMBLE_SIZE,
|
|
TOPOLOGY_COUNT,
|
|
)
|
|
struct.pack_into("<I", header, 60, 0)
|
|
struct.pack_into("<Q", header, 64, training_seed & 0xFFFFFFFFFFFFFFFF)
|
|
header[72:104] = corpus_digest
|
|
struct.pack_into(
|
|
"<4f",
|
|
header,
|
|
104,
|
|
float(validation["average_precision"]),
|
|
float(validation["brier_score"]),
|
|
float(validation["expected_calibration_error"]),
|
|
float(validation["top_k_gain"]),
|
|
)
|
|
struct.pack_into("<Q", header, 120, len(payload))
|
|
struct.pack_into("<I", header, 128, payload_crc)
|
|
struct.pack_into("<I", header, 132, 0)
|
|
struct.pack_into("<Q", header, 136, float_count)
|
|
struct.pack_into("<QQ", header, 144, 0, 0)
|
|
header_crc = zlib.crc32(header) & 0xFFFFFFFF
|
|
struct.pack_into("<I", header, 132, header_crc)
|
|
return bytes(header) + payload
|
|
|
|
|
|
def publish_model(
|
|
output_directory: Path,
|
|
model_bytes: bytes,
|
|
report: dict[str, object],
|
|
) -> 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)
|