| """Full-weight cross-entropy training with periodic, fixed-fold evaluation.""" |
|
|
| import argparse |
| import copy |
| import hashlib |
| import json |
| import math |
| import os |
| from pathlib import Path |
| import random |
| import subprocess |
| import shutil |
| import time |
| from collections.abc import Callable, Sequence |
| from typing import cast |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from autojev.evaluate import ( |
| Metrics, Prediction, calibration_ok, evaluate_logits, fit_temperature, |
| hard_label, label_index, metrics, options, read_predictions, read_rows, |
| selection_key, validate_coverage, write_json, |
| ) |
| from autojev.events import record |
| from autojev.model import DecisionModel |
| from autojev.optim import CPUOffloadAdamW |
| from autojev.sft_pipeline import Receipt, ids_sha256, load_schedule, path_at |
| from autojev.types import Example, JSONValue |
|
|
|
|
| class Arguments(argparse.Namespace): |
| train: str |
| schedule: str | None |
| development: str |
| temperature: str |
| reference: str |
| public: str | None |
| guard_data: str | None |
| guard_reference: str | None |
| quick_development: str | None |
| quick_guard_data: str | None |
| quick_temperature: str | None |
| quick_eval_every: int |
| run: str |
| output: str |
| base_model: str |
| revision: str |
| cache_dir: str | None |
| epochs: int |
| batch_size: int |
| effective_batch_size: int |
| token_budget: int |
| max_length: int |
| lr: float |
| weight_decay: float |
| seed: int |
| eval_every: int |
| public_eval_every: int |
| resume_every: int |
| stop_after: int | None |
| resume: str | None |
| initial_checkpoint: str | None |
| cpu_threads: int |
| allow_shared_families: bool |
| keep_only_selected: bool |
| no_resume: bool |
| resume_unlink_first: bool |
|
|
|
|
| KEEP_ONLY_SELECTED = False |
|
|
|
|
| def digest(path: str | Path) -> str: |
| with Path(path).open("rb") as stream: |
| return hashlib.file_digest(stream, "sha256").hexdigest() |
|
|
|
|
| def initial_identity(path: str, base_model: str, revision: str) -> dict[str, JSONValue]: |
| """Bind a decision artifact used to start a new stage, not an optimizer resume.""" |
| directory = Path(path).resolve() |
| saved = cast(dict[str, JSONValue], json.loads((directory / "decision_config.json").read_text())) |
| if saved.get("format_version") != 1 or saved.get("base_model") != base_model or saved.get("revision") != revision: |
| raise ValueError("Initial checkpoint format/base model/revision differs from the requested model") |
| temperature = saved.get("temperature") |
| if not isinstance(temperature, (int, float)) or isinstance(temperature, bool) or not math.isfinite(temperature) or temperature <= 0: |
| raise ValueError("Initial checkpoint temperature must be positive and finite") |
| index = cast(dict[str, JSONValue], json.loads((directory / "model.safetensors.index.json").read_text())) |
| weights = index.get("weight_map") |
| if not isinstance(weights, dict) or not weights: |
| raise ValueError("Initial checkpoint must have an indexed set of backbone weights") |
| names = {"config.json", "decision_config.json", "model.safetensors.index.json", "readout.safetensors", |
| "tokenizer.json", "tokenizer_config.json", "processor_config.json", "chat_template.jinja"} |
| for name in weights.values(): |
| if not isinstance(name, str) or Path(name).name != name or not name.endswith(".safetensors"): |
| raise ValueError("Checkpoint shard names must be local safetensors filenames") |
| names.add(name) |
| |
| hashes: dict[str, JSONValue] = {name: digest(directory / name) for name in sorted(names) if name != "processor_config.json" or (directory / name).exists()} |
| result: dict[str, JSONValue] = { |
| "path": str(directory), "base_model": base_model, "revision": revision, |
| "temperature": temperature, "decision_config": saved, "files_sha256": hashes, |
| "artifact_sha256": hashlib.sha256(json.dumps(hashes, sort_keys=True, separators=(",", ":")).encode()).hexdigest(), |
| } |
| release_manifest = directory / "release-manifest.json" |
| if release_manifest.exists(): |
| result["release_manifest_sha256"] = digest(release_manifest) |
| return result |
|
|
|
|
| def append(path: Path, value: object) -> None: |
| with path.open("a") as stream: |
| stream.write(json.dumps(value, ensure_ascii=False) + "\n") |
|
|
|
|
| class ScheduledRows: |
| """A frozen ID order whose audited text may arrive in immutable tranches.""" |
|
|
| def __init__(self, path: Path, calibration: Path, protected: Sequence[Example], run: str) -> None: |
| self.path = path.absolute() |
| self.schedule = load_schedule(self.path) |
| self.sha256 = digest(self.path) |
| self.calibration = calibration |
| self.calibration_sha256 = digest(calibration) |
| self.run = run |
| self.ids = self.schedule["ordered_ids"] |
| self.protected_ids = {row["id"] for row in protected} |
| self.protected_families = {(str(row["source"].get("dataset", row["suite"])), row["family"]) for row in protected} |
| if set(self.ids) & self.protected_ids: |
| raise ValueError("Scheduled training IDs overlap a held-out fold") |
| for source, checksum in ((self.schedule["plan_path"], self.schedule["plan_sha256"]), |
| (self.schedule["original_train_path"], self.schedule["original_train_sha256"])): |
| if digest(path_at(self.path, source)) != checksum: |
| raise ValueError(f"Scheduled source changed: {source}") |
| if path_at(self.path, self.schedule["calibration_path"]).resolve() != calibration.resolve(): |
| raise ValueError("--temperature must be the calibration fold bound by the schedule") |
| self.tranche_for_id = {identifier: part["index"] for part in self.schedule["tranches"] for identifier in part["ordered_ids"]} |
| self.receipts: list[dict[str, JSONValue]] = [] |
| self.cache: dict[int, dict[str, Example]] = {} |
|
|
| def _read(self, index: int) -> tuple[dict[str, JSONValue], list[Example]]: |
| if digest(self.path) != self.sha256 or digest(self.calibration) != self.calibration_sha256: |
| raise ValueError("The frozen schedule or calibration fold changed") |
| part = self.schedule["tranches"][index] |
| manifest = path_at(self.path, part["manifest_path"]) |
| receipt = cast(Receipt, json.loads(manifest.read_text())) |
| expected = {"schema_version": 1, "complete": True, "schedule_sha256": self.sha256, |
| "tranche": index, "rows": len(part["ordered_ids"]), |
| "ordered_ids_sha256": ids_sha256(part["ordered_ids"]), |
| "calibration_sha256": self.calibration_sha256, |
| "previous_receipts_sha256": [value["manifest_sha256"] for value in self.receipts[:index]]} |
| if any(receipt.get(key) != value for key, value in expected.items()): |
| raise ValueError(f"Tranche {index} readiness receipt differs from the frozen schedule") |
| data_path = path_at(self.path, part["data_path"]) |
| audit_path = path_at(self.path, receipt["audit_path"]) |
| for path, checksum in ((data_path, receipt["data_sha256"]), (audit_path, receipt["audit_sha256"]), |
| (path_at(self.path, part["synthetic_path"]), receipt["synthetic_sha256"]), |
| (path_at(self.path, self.schedule["controls_audit_path"]), receipt["controls_audit_sha256"])): |
| if digest(path) != checksum: |
| raise ValueError(f"Tranche {index} bound file changed: {path}") |
| audit = cast(dict[str, JSONValue], json.loads(audit_path.read_text())) |
| if (audit.get("complete") is not True or audit.get("phase") != "tranche" |
| or audit.get("schedule_sha256") != self.sha256 or audit.get("tranche") != index |
| or audit.get("accepted_rows") != len(part["synthetic_ids"])): |
| raise ValueError(f"Tranche {index} has no complete matching synthetic audit") |
| audited_train = cast(dict[str, JSONValue], cast(dict[str, JSONValue], audit["outputs"])["train"]) |
| if audited_train.get("sha256") != receipt["synthetic_sha256"]: |
| raise ValueError(f"Tranche {index} synthetic rows differ from their audit") |
| controls = cast(dict[str, JSONValue], json.loads(path_at(self.path, self.schedule["controls_audit_path"]).read_text())) |
| if controls.get("complete") is not True or controls.get("phase") != "controls": |
| raise ValueError("The fixed synthetic controls must pass their complete audit first") |
| rows = read_rows(data_path) |
| if [row["id"] for row in rows] != part["ordered_ids"]: |
| raise ValueError(f"Tranche {index} text IDs/order differ from the frozen schedule") |
| families = {(str(row["source"].get("dataset", row["suite"])), row["family"]) for row in rows} |
| if families & self.protected_families: |
| raise ValueError(f"Tranche {index} has a family overlapping a held-out fold") |
| provenance = cast(dict[str, JSONValue], dict(receipt)) |
| provenance.update(manifest_path=str(manifest), manifest_sha256=digest(manifest), data_path=str(data_path)) |
| return provenance, rows |
|
|
| def batch(self, step: int, batch_size: int, before_wait: Callable[[], None] | None = None) -> list[Example]: |
| selected_ids = self.ids[step * batch_size:(step + 1) * batch_size] |
| if not selected_ids: |
| raise ValueError("The training cursor is outside the frozen schedule") |
| needed = {self.tranche_for_id[identifier] for identifier in selected_ids} |
| for index in sorted(needed): |
| if index in self.cache: |
| continue |
| manifest = path_at(self.path, self.schedule["tranches"][index]["manifest_path"]) |
| waiting_started = time.monotonic() |
| waiting = False |
| while not manifest.exists(): |
| if not waiting and before_wait is not None: |
| before_wait() |
| waiting = True |
| record("training_waiting_for_data", run=self.run, step=step, next_step=step + 1, |
| tranche=index, schedule_sha256=self.sha256, manifest=str(manifest), |
| waiting_seconds=time.monotonic() - waiting_started) |
| time.sleep(20) |
| provenance, rows = self._read(index) |
| if index == len(self.receipts): |
| self.receipts.append(provenance) |
| record("training_tranche_loaded", run=self.run, step=step, tranche=index, |
| receipt=provenance, waiting_seconds=time.monotonic() - waiting_started if waiting else 0) |
| elif index >= len(self.receipts) or self.receipts[index] != provenance: |
| raise ValueError("A previously validated tranche receipt changed, or a tranche was skipped") |
| self.cache[index] = {row["id"]: row for row in rows} |
| result = [self.cache[self.tranche_for_id[identifier]][identifier] for identifier in selected_ids] |
| self.cache = {index: rows for index, rows in self.cache.items() if index in needed} |
| return sorted(result, key=length_estimate) |
|
|
| def snapshot(self, examples_seen: int) -> dict[str, JSONValue]: |
| if not 0 <= examples_seen <= len(self.ids): |
| raise ValueError("Invalid consumed-row cursor") |
| return {"schedule_path": str(self.path), "schedule_sha256": self.sha256, |
| "scheduled_rows": len(self.ids), "examples_seen": examples_seen, |
| "consumed_ids_sha256": ids_sha256(self.ids[:examples_seen]), |
| "validated_tranches": cast(JSONValue, copy.deepcopy(self.receipts))} |
|
|
| def restore(self, saved: dict[str, JSONValue], examples_seen: int) -> None: |
| receipts = saved.get("validated_tranches") |
| if not isinstance(receipts, list) or len(receipts) > len(self.schedule["tranches"]): |
| raise ValueError("Invalid saved tranche receipt prefix") |
| if examples_seen > sum(len(part["ordered_ids"]) for part in self.schedule["tranches"][:len(receipts)]): |
| raise ValueError("The saved cursor includes rows without validated tranche receipts") |
| for index, previous in enumerate(receipts): |
| provenance, _ = self._read(index) |
| if provenance != previous: |
| raise ValueError(f"Tranche {index} changed since the saved optimizer cursor") |
| self.receipts.append(provenance) |
| if self.snapshot(examples_seen) != saved: |
| raise ValueError("The saved training schedule/cursor differs from this run") |
|
|
|
|
| def length_estimate(row: Example) -> int: |
| measured = row["source"].get("input_tokens") |
| if isinstance(measured, int) and not isinstance(measured, bool): |
| return measured + 16 |
| return len(json.dumps([row["state"], row["question"]], ensure_ascii=False)) // 3 + 192 + 512 * len(row.get("images", [])) |
|
|
|
|
| def fixed_subset(path: str | None, full: Sequence[Example]) -> list[Example]: |
| """Keep diagnostic panels identical to their frozen parent rows, including option order.""" |
| if path is None: |
| return [] |
| rows = read_rows(Path(path)) |
| originals = {row["id"]: json.dumps(row, ensure_ascii=False) for row in full} |
| if (not rows or len({row["id"] for row in rows}) != len(rows) |
| or any(originals.get(row["id"]) != json.dumps(row, ensure_ascii=False) for row in rows)): |
| raise ValueError(f"Quick evaluation file is not an exact, unique subset of its full fold: {path}") |
| return rows |
|
|
|
|
| def microbatches(rows: Sequence[Example], batch_size: int, token_budget: int) -> list[list[Example]]: |
| result: list[list[Example]] = [] |
| pending: list[Example] = [] |
| longest = 0 |
| for row in rows: |
| length = length_estimate(row) |
| if pending and (len(pending) == batch_size or max(longest, length) * (len(pending) + 1) > token_budget): |
| result.append(pending) |
| pending, longest = [], 0 |
| pending.append(row) |
| longest = max(longest, length) |
| if pending: |
| result.append(pending) |
| return result |
|
|
|
|
| def targets(rows: Sequence[Example], device: torch.device) -> torch.Tensor: |
| values = torch.zeros((len(rows), 255), dtype=torch.float32, device=device) |
| for index, row in enumerate(rows): |
| labels, target = options(row["question"]), row["target"] |
| if isinstance(target, list): |
| distribution = target |
| elif row["question"]["type"] == "noul": |
| positive = float(cast(float, target)) |
| distribution = [1.0 - positive, positive] |
| else: |
| distribution = [float(label == target) for label in labels] |
| if len(distribution) != len(labels) or any(not math.isfinite(p) or p < 0 for p in distribution) or abs(sum(distribution) - 1) > 1e-6: |
| raise ValueError(f"Invalid training target: {row['id']}") |
| values[index, :len(distribution)] = torch.tensor(distribution, device=device) |
| return values |
|
|
|
|
| def augment(rows: Sequence[Example], rng: random.Random) -> list[Example]: |
| result = copy.deepcopy(list(rows)) |
| for row in result: |
| if row["question"]["type"] == "choice": |
| criteria = row["question"]["criteria"] |
| target = row["target"] |
| weights = dict(zip(criteria, target, strict=True)) if isinstance(target, list) else None |
| items = list(criteria.items()) |
| rng.shuffle(items) |
| row["question"]["criteria"] = dict(items) |
| if weights is not None: |
| row["target"] = [weights[key] for key, _ in items] |
| return result |
|
|
|
|
| @torch.inference_mode() |
| def infer(model: DecisionModel, rows: Sequence[Example], args: Arguments) -> list[list[float]]: |
| was_training = model.training |
| model.eval() |
| result: list[list[float]] = [] |
| for batch in microbatches(rows, args.batch_size, args.token_budget): |
| logits = model(model.prepare(batch, max_length=args.max_length)).detach().cpu() |
| for row, values in zip(batch, logits, strict=True): |
| result.append(cast(list[float], values[:len(options(row["question"]))].tolist())) |
| model.train(was_training) |
| return result |
|
|
|
|
| def save_predictions(path: Path, predictions: Sequence[Prediction]) -> None: |
| with path.open("w") as stream: |
| for prediction in predictions: |
| stream.write(json.dumps(prediction, ensure_ascii=False) + "\n") |
|
|
|
|
| def sync_directory(path: Path) -> None: |
| descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) |
| try: |
| os.fsync(descriptor) |
| finally: |
| os.close(descriptor) |
|
|
|
|
| def select_checkpoint(root: Path, step: int) -> None: |
| pending = root / "selected.pending" |
| pending.unlink(missing_ok=True) |
| pending.symlink_to(f"step-{step:05d}", target_is_directory=True) |
| os.replace(pending, root / "selected") |
| sync_directory(root) |
|
|
|
|
| def save_selected(model: DecisionModel, root: Path, temperature: float, step: int, provenance: dict[str, str], |
| initial_artifact: dict[str, JSONValue] | None = None, |
| training_schedule: dict[str, JSONValue] | None = None) -> None: |
| destination = root / f"step-{step:05d}" |
| metadata: dict[str, JSONValue] = {"step": step, "provenance": cast(JSONValue, provenance)} |
| if initial_artifact is not None: |
| metadata["initial_artifact"] = initial_artifact |
| if training_schedule is not None: |
| metadata["training_schedule"] = training_schedule |
| previous = (root / "selected").resolve() if (root / "selected").is_symlink() else None |
| model.save(destination, temperature=temperature, **metadata) |
| for file in destination.rglob("*"): |
| if file.is_file(): |
| with file.open("rb") as stream: |
| os.fsync(stream.fileno()) |
| sync_directory(destination) |
| select_checkpoint(root, step) |
| if KEEP_ONLY_SELECTED and previous is not None and previous != destination.resolve() and previous.parent == root.resolve() and previous.name.startswith("step-") and previous.is_dir(): |
| shutil.rmtree(previous) |
| record("checkpoint_pruned", path=str(previous), selected_step=step) |
|
|
|
|
| def archive_interrupted_tail(run: Path, output: Path, step: int) -> None: |
| """Keep abandoned work as history while the active trajectory resumes its saved cursor.""" |
| suffix = str(time.time_ns()) |
| archive = run / "interrupted" / suffix |
| for name in ("training.jsonl", "evaluations.jsonl", "public-evaluations.jsonl", "quick-evaluations.jsonl"): |
| path = run / name |
| if not path.exists(): |
| continue |
| lines = path.read_text().splitlines() |
| tail = [line for line in lines if json.loads(line)["step"] > step] |
| if tail: |
| archive.mkdir(parents=True, exist_ok=True) |
| (archive / name).write_text("\n".join(tail) + "\n") |
| path.write_text("\n".join(line for line in lines if json.loads(line)["step"] <= step) + "\n") |
| for prefix in ("development", "temperature", "public", "guard", "quick-development", "quick-temperature", "quick-guard"): |
| for path in run.glob(f"{prefix}-*.jsonl"): |
| suffix_step = path.stem.rsplit("-", 1)[1] |
| if suffix_step.isdigit() and int(suffix_step) > step: |
| archive.mkdir(parents=True, exist_ok=True) |
| path.rename(archive / path.name) |
| for path in output.glob("step-*"): |
| if path.is_dir() and path.name.rsplit("-", 1)[1].isdigit() and int(path.name.rsplit("-", 1)[1]) > step: |
| destination = output / "interrupted" / suffix |
| destination.mkdir(parents=True, exist_ok=True) |
| path.rename(destination / path.name) |
| record("training_resumed", run=run.name, resume_step=step, archived_tail=archive.exists(), attempt=suffix) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| for name, default in [("train", "data/train.jsonl"), ("development", "data/dev.jsonl"), |
| ("temperature", "data/temperature.jsonl"), ("reference", "progress/baselines/jev-development.predictions.jsonl")]: |
| parser.add_argument(f"--{name}", default=default) |
| parser.add_argument("--schedule", help="Frozen global ID schedule with immutable audited tranches; replaces --train") |
| parser.add_argument("--public") |
| parser.add_argument("--guard-data", help="Separate development cohort whose calibration also gates selection") |
| parser.add_argument("--guard-reference", help="Complete fixed Jev predictions for --guard-data") |
| parser.add_argument("--quick-development", help="Frozen subset of development for diagnostic curves only") |
| parser.add_argument("--quick-guard-data", help="Frozen subset of the guard cohort for diagnostic curves only") |
| parser.add_argument("--quick-temperature", help="Frozen subset of temperature data, used only to calibrate quick curves") |
| parser.add_argument("--quick-eval-every", type=int, default=20) |
| parser.add_argument("--run", required=True) |
| parser.add_argument("--output", required=True) |
| parser.add_argument("--base-model", default="Qwen/Qwen3.8-27B") |
| parser.add_argument("--revision", default="1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0") |
| parser.add_argument("--cache-dir") |
| parser.add_argument("--epochs", type=int, default=1) |
| parser.add_argument("--batch-size", type=int, default=32) |
| parser.add_argument("--effective-batch-size", type=int, default=256) |
| parser.add_argument("--token-budget", type=int, default=8192) |
| parser.add_argument("--max-length", type=int, default=8192) |
| parser.add_argument("--lr", type=float, default=2e-6) |
| parser.add_argument("--weight-decay", type=float, default=0.01) |
| parser.add_argument("--seed", type=int, default=20260920) |
| parser.add_argument("--eval-every", type=int, default=20) |
| parser.add_argument("--public-eval-every", type=int, default=100) |
| parser.add_argument("--resume-every", type=int, default=100) |
| parser.add_argument("--stop-after", type=int) |
| parser.add_argument("--resume") |
| parser.add_argument("--initial-checkpoint", help="Start a new full-weight SFT stage from this artifact with a fresh optimizer") |
| parser.add_argument("--cpu-threads", type=int, default=32) |
| |
| parser.add_argument("--allow-shared-families", action="store_true", |
| help="Held-out folds share (dataset, family) with train; IDs and source.group_id must still be disjoint") |
| parser.add_argument("--keep-only-selected", action="store_true", |
| help="Delete the previously selected step-* artifact after a newer one is selected (disk budget)") |
| parser.add_argument("--no-resume", action="store_true", help="Never write resume.pt (optimizer state is ~12 bytes/param)") |
| parser.add_argument("--resume-unlink-first", action="store_true", |
| help="Remove the old resume.pt before writing the new one (halves peak disk; not atomic)") |
| args = parser.parse_args(namespace=Arguments()) |
| global KEEP_ONLY_SELECTED |
| KEEP_ONLY_SELECTED = args.keep_only_selected |
| if min(args.epochs, args.batch_size, args.effective_batch_size, args.token_budget, args.eval_every, args.public_eval_every, args.resume_every, args.quick_eval_every) < 1: |
| raise ValueError("Batch, epoch and interval settings must be positive") |
| if args.stop_after is not None and args.stop_after < 1: |
| raise ValueError("A pilot must include at least one training step") |
| if bool(args.guard_data) != bool(args.guard_reference): |
| raise ValueError("--guard-data and --guard-reference must be supplied together") |
| quick_paths = (args.quick_development, args.quick_guard_data, args.quick_temperature) |
| if any(quick_paths) and (not all(quick_paths) or not args.guard_data): |
| raise ValueError("Quick evaluation needs all three frozen subsets and a full guard cohort/reference") |
| if args.initial_checkpoint and not os.environ.get("AUTOJEV_EVENTS"): |
| raise ValueError("Set AUTOJEV_EVENTS to the new stage's event log before continuing from a checkpoint") |
| if args.schedule and (args.epochs != 1 or not os.environ.get("AUTOJEV_EVENTS")): |
| raise ValueError("A streamed schedule requires one epoch and an explicit AUTOJEV_EVENTS path") |
| run, output = Path(args.run), Path(args.output) |
| if args.initial_checkpoint and any(path.resolve().is_relative_to(Path(args.initial_checkpoint).resolve()) for path in (run, output)): |
| raise ValueError("New run/output directories must be outside the initial checkpoint artifact") |
| run.mkdir(parents=True, exist_ok=True) |
| output.mkdir(parents=True, exist_ok=True) |
| if (run / "config.json").exists() and not args.resume: |
| raise ValueError("Run already exists; resume explicitly or choose a new run ID") |
| train = [] if args.schedule else read_rows(Path(args.train)) |
| development, temperature_rows = (read_rows(Path(path)) for path in (args.development, args.temperature)) |
| if (not train and not args.schedule) or not development or not temperature_rows: |
| raise ValueError("Training, development and temperature folds must be nonempty") |
| reference_predictions = read_predictions(Path(args.reference)) |
| validate_coverage(development, reference_predictions) |
| reference = metrics(reference_predictions) |
| public_rows = read_rows(Path(args.public)) if args.public else [] |
| guard_rows = read_rows(Path(args.guard_data)) if args.guard_data else [] |
| guard_reference: Metrics | None = None |
| guard_predictions: list[Prediction] = [] |
| if args.guard_data: |
| if not guard_rows: |
| raise ValueError("The guard cohort must be nonempty") |
| guard_predictions = read_predictions(Path(cast(str, args.guard_reference))) |
| validate_coverage(guard_rows, guard_predictions) |
| guard_reference = metrics(guard_predictions) |
| quick_development = fixed_subset(args.quick_development, development) |
| quick_guard = fixed_subset(args.quick_guard_data, guard_rows) |
| quick_temperature = fixed_subset(args.quick_temperature, temperature_rows) |
| quick_ids = {row["id"] for row in quick_development} |
| quick_guard_ids = {row["id"] for row in quick_guard} |
| quick_reference = metrics([value for value in reference_predictions if value["id"] in quick_ids]) if quick_development else None |
| quick_guard_reference = metrics([value for value in guard_predictions if value["id"] in quick_guard_ids]) if quick_guard else None |
| partitions = {"train": train, "development": development, "temperature": temperature_rows, |
| "public": public_rows, "guard": guard_rows} |
| families: dict[str, set[tuple[str, str]]] = {} |
| identifiers: dict[str, set[str]] = {} |
| source_groups: dict[str, set[tuple[str, str]]] = {} |
| for name, rows in partitions.items(): |
| identifiers[name] = {row["id"] for row in rows} |
| if len(identifiers[name]) != len(rows): |
| raise ValueError(f"Duplicate IDs in {name}") |
| families[name] = {(str(row["source"].get("dataset", row["suite"])), row["family"]) for row in rows} |
| source_groups[name] = {(row["family"], str(row["source"].get("group_id", row["id"]))) for row in rows} |
| for other in identifiers: |
| if name == other: |
| continue |
| if identifiers[name] & identifiers[other]: |
| raise ValueError(f"Partitions overlap: {name}/{other}") |
| if args.allow_shared_families: |
| if source_groups[name] & source_groups[other]: |
| raise ValueError(f"Partitions share source groups: {name}/{other}") |
| elif families[name] & families[other]: |
| raise ValueError(f"Partitions overlap: {name}/{other}") |
| scheduled = ScheduledRows(Path(args.schedule), Path(args.temperature), development + temperature_rows + public_rows + guard_rows, run.name) if args.schedule else None |
| if scheduled and scheduled.schedule["seed"] != args.seed: |
| raise ValueError("The training seed must match the frozen schedule seed") |
| train_rows = len(scheduled.ids) if scheduled else len(train) |
| hashes = {name: digest(path) for name, path in [("development", args.development), ("temperature", args.temperature), ("reference", args.reference)]} |
| hashes["train_schedule" if scheduled else "train"] = scheduled.sha256 if scheduled else digest(args.train) |
| if args.public: |
| hashes["public"] = digest(args.public) |
| if args.guard_data: |
| hashes["guard_data"] = digest(args.guard_data) |
| hashes["guard_reference"] = digest(cast(str, args.guard_reference)) |
| for key, path in (("quick_development", args.quick_development), ("quick_guard_data", args.quick_guard_data), ("quick_temperature", args.quick_temperature)): |
| if path is not None: |
| hashes[key] = digest(path) |
| initial_artifact = initial_identity(args.initial_checkpoint, args.base_model, args.revision) if args.initial_checkpoint else None |
| revision = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip() |
| package = Path(__file__).resolve().parent |
| code_hashes = {name: digest(package / name) for name in ("train.py", "model.py", "optim.py", "evaluate.py", "types.py", "events.py")} |
| if scheduled: |
| code_hashes["sft_pipeline.py"] = digest(package / "sft_pipeline.py") |
| code_hashes["uv.lock"] = digest(package.parents[1] / "uv.lock") |
| config = {**vars(args), "data_sha256": hashes, "code_sha256": code_hashes, "git_commit": revision, |
| "initial_artifact": initial_artifact, "events_path": os.getenv("AUTOJEV_EVENTS", "progress/events.jsonl"), |
| "train_rows": train_rows, "development_rows": len(development), "temperature_rows": len(temperature_rows), |
| "guard_rows": len(guard_rows), "quick_development_rows": len(quick_development), |
| "quick_guard_rows": len(quick_guard), "quick_temperature_rows": len(quick_temperature)} |
| if not args.resume: |
| write_json(run / "config.json", config) |
| if scheduled: |
| scheduled.batch(0, args.effective_batch_size) |
| random.seed(args.seed) |
| torch.manual_seed(args.seed) |
| torch.cuda.manual_seed_all(args.seed) |
| started = time.monotonic() |
| model = DecisionModel(checkpoint=args.initial_checkpoint, train=True, base_model=args.base_model, |
| revision=args.revision, cache_dir=args.cache_dir, gradient_checkpointing=True, |
| cpu_threads=args.cpu_threads) |
| optimizer = CPUOffloadAdamW(model.named_parameters(), lr=args.lr, weight_decay=args.weight_decay) |
| parameters = [parameter for parameter in model.parameters() if parameter.requires_grad] |
| groups: list[list[Example]] = [] |
| for epoch in range(args.epochs): |
| ordered = list(train) |
| random.Random(args.seed + epoch).shuffle(ordered) |
| for offset in range(0, len(ordered), args.effective_batch_size): |
| groups.append(sorted(ordered[offset:offset + args.effective_batch_size], key=length_estimate)) |
| total_steps = math.ceil(train_rows / args.effective_batch_size) if scheduled else len(groups) |
| step, examples_seen = 0, 0 |
| last_resume_step: int | None = None |
| best: Metrics | None = None |
| best_guard: Metrics | None = None |
| best_step: int | None = None |
| selected_temperature: float | None = None |
| if args.resume: |
| state = torch.load(args.resume, map_location="cpu", weights_only=True) |
| if state["data_sha256"] != hashes or state["total_steps"] != total_steps: |
| raise ValueError("Resume data/schedule differs from the saved run") |
| if state["config"]["code_sha256"] != code_hashes: |
| raise ValueError("Training implementation differs from the saved run") |
| if state["config"].get("initial_artifact") != initial_artifact: |
| raise ValueError("Initial checkpoint content or provenance differs from the saved stage") |
| if state["config"].get("events_path") != config["events_path"]: |
| raise ValueError("Resume event log differs from the saved stage") |
| for key in vars(args): |
| if key in {"resume", "stop_after"}: |
| continue |
| if state["config"][key] != vars(args)[key]: |
| raise ValueError(f"Resume configuration differs: {key}") |
| step, examples_seen = state["step"], state["examples_seen"] |
| if scheduled: |
| if examples_seen != min(step * args.effective_batch_size, train_rows): |
| raise ValueError("The saved consumed-row cursor differs from its global update") |
| scheduled.restore(state["training_schedule"], examples_seen) |
| optimizer.load_state_dict(state["optimizer"]) |
| last_resume_step = step |
| if args.stop_after is not None and args.stop_after <= step: |
| raise ValueError("The requested stopping step must follow the saved step") |
| best, best_step, selected_temperature = state["best"], state["best_step"], state["selected_temperature"] |
| best_guard = state.get("best_guard") |
| random.setstate(state["python_rng"]) |
| torch.set_rng_state(state["torch_rng"]) |
| torch.cuda.set_rng_state_all(state["cuda_rng"]) |
| del state |
| archive_interrupted_tail(run, output, step) |
| if best_step is not None: |
| selected = output / f"step-{best_step:05d}" |
| if not selected.is_dir(): |
| raise ValueError("The checkpoint selected by the resume state is missing") |
| select_checkpoint(output, best_step) |
| else: |
| (output / "selected").unlink(missing_ok=True) |
| record("training_started", run=run.name, git_commit=revision, data_sha256=hashes, |
| initial_artifact=initial_artifact, initialization="exact_resume" if args.resume else "checkpoint_fresh_optimizer" if initial_artifact else "base_fresh_optimizer", |
| step=step, total_steps=total_steps, training_schedule=scheduled.snapshot(examples_seen) if scheduled else None, |
| trainable_parameters=sum(p.numel() for p in parameters), base_model=model.base_model, |
| base_revision=model.revision, text_only=model.text_only) |
|
|
| def evaluate(include_public: bool = False, diagnostic: bool = False) -> None: |
| nonlocal best, best_guard, best_step, selected_temperature |
| began = time.monotonic() |
| temperature_logits = infer(model, temperature_rows, args) |
| fitted = fit_temperature(temperature_logits, [label_index(options(row["question"]), hard_label(row)) for row in temperature_rows]) |
| logits = infer(model, development, args) |
| raw_predictions, fitted_predictions = evaluate_logits(development, logits), evaluate_logits(development, logits, fitted) |
| raw, calibrated = metrics(raw_predictions), metrics(fitted_predictions) |
| guard_raw: Metrics | None = None |
| guard_calibrated: Metrics | None = None |
| guard_ok = True |
| if guard_rows: |
| guard_logits = infer(model, guard_rows, args) |
| guard_raw = metrics(evaluate_logits(guard_rows, guard_logits)) |
| guard_predictions = evaluate_logits(guard_rows, guard_logits, fitted) |
| guard_calibrated = metrics(guard_predictions) |
| guard_ok = calibration_ok(guard_calibrated, cast(Metrics, guard_reference)) |
| save_predictions(run / f"guard-{step:05d}.jsonl", guard_predictions) |
| public_ok = calibration_ok(calibrated, reference) |
| eligible = public_ok and guard_ok |
| improved = not diagnostic and eligible and (best is None or selection_key(calibrated, step) < selection_key(best, cast(int, best_step))) |
| if improved: |
| save_selected(model, output, fitted, step, {"run": run.name, "git_commit": revision, **hashes, **code_hashes}, |
| initial_artifact, scheduled.snapshot(examples_seen) if scheduled else None) |
| best, best_guard, best_step, selected_temperature = calibrated, guard_calibrated, step, fitted |
| save_predictions(run / f"development-{step:05d}.jsonl", fitted_predictions) |
| save_predictions(run / f"temperature-{step:05d}.jsonl", evaluate_logits(temperature_rows, temperature_logits, fitted)) |
| guard_target = guard_calibrated is not None and guard_reference is not None and guard_ok and guard_calibrated["accuracy"] > guard_reference["accuracy"] |
| value = record("evaluation", run=run.name, step=step, examples_seen=examples_seen, raw=raw, fitted=calibrated, |
| temperature=fitted, jev=reference, eligible=eligible, public_calibration_ok=public_ok, |
| guard_raw=guard_raw, guard_fitted=guard_calibrated, guard_jev=guard_reference, |
| guard_calibration_ok=guard_ok, guard_empirical_target_met=guard_target, |
| development_target_met=eligible and calibrated["accuracy"] > reference["accuracy"], |
| diagnostic=diagnostic, selected_step=best_step, elapsed_seconds=time.monotonic() - started, |
| evaluation_seconds=time.monotonic() - began) |
| append(run / "evaluations.jsonl", value) |
| print(json.dumps(value), flush=True) |
| if include_public and public_rows: |
| predictions = evaluate_logits(public_rows, infer(model, public_rows, args), fitted) |
| save_predictions(run / f"public-{step:05d}.jsonl", predictions) |
| value = record("public_evaluation", run=run.name, step=step, temperature=fitted, metrics=metrics(predictions), status="known_suite_development_check") |
| append(run / "public-evaluations.jsonl", value) |
|
|
| def save_resume() -> None: |
| nonlocal last_resume_step |
| began = time.monotonic() |
| state = {"optimizer": optimizer.state_dict(), "step": step, "examples_seen": examples_seen, "total_steps": total_steps, "data_sha256": hashes, "config": config, |
| "best": best, "best_guard": best_guard, "best_step": best_step, "selected_temperature": selected_temperature, |
| "python_rng": random.getstate(), "torch_rng": torch.get_rng_state(), "cuda_rng": torch.cuda.get_rng_state_all()} |
| if scheduled: |
| state["training_schedule"] = scheduled.snapshot(examples_seen) |
| pending = output / "resume.pt.pending" |
| if args.resume_unlink_first: |
| (output / "resume.pt").unlink(missing_ok=True) |
| torch.save(state, pending) |
| with pending.open("rb") as stream: |
| os.fsync(stream.fileno()) |
| os.replace(pending, output / "resume.pt") |
| sync_directory(output) |
| last_resume_step = step |
| record("resume_saved", run=run.name, step=step, bytes=(output / "resume.pt").stat().st_size, seconds=time.monotonic() - began) |
|
|
| def evaluate_quick() -> None: |
| began = time.monotonic() |
| temperature_logits = infer(model, quick_temperature, args) |
| fitted = fit_temperature(temperature_logits, [label_index(options(row["question"]), hard_label(row)) for row in quick_temperature]) |
| logits = infer(model, quick_development, args) |
| predictions = evaluate_logits(quick_development, logits, fitted) |
| guard_logits = infer(model, quick_guard, args) |
| guard_predictions = evaluate_logits(quick_guard, guard_logits, fitted) |
| save_predictions(run / f"quick-development-{step:05d}.jsonl", predictions) |
| save_predictions(run / f"quick-guard-{step:05d}.jsonl", guard_predictions) |
| save_predictions(run / f"quick-temperature-{step:05d}.jsonl", evaluate_logits(quick_temperature, temperature_logits, fitted)) |
| value = record("quick_evaluation", run=run.name, step=step, examples_seen=examples_seen, |
| raw=metrics(evaluate_logits(quick_development, logits)), fitted=metrics(predictions), |
| guard_raw=metrics(evaluate_logits(quick_guard, guard_logits)), guard_fitted=metrics(guard_predictions), |
| temperature=fitted, jev=quick_reference, guard_jev=quick_guard_reference, |
| diagnostic=True, selection_candidate=False, selected_step=best_step, |
| temperature_rows=len(quick_temperature), development_rows=len(quick_development), guard_rows=len(quick_guard), |
| elapsed_seconds=time.monotonic() - started, evaluation_seconds=time.monotonic() - began) |
| append(run / "quick-evaluations.jsonl", value) |
| print(json.dumps(value), flush=True) |
|
|
| def save_before_wait() -> None: |
| if step > 0 and last_resume_step != step: |
| save_resume() |
|
|
| if step == 0: |
| evaluate(include_public=bool(public_rows)) |
| stop = min(total_steps, args.stop_after) if args.stop_after is not None else total_steps |
| while step < stop: |
| rows = scheduled.batch(step, args.effective_batch_size, save_before_wait) if scheduled else groups[step] |
| began = time.monotonic() |
| model.train() |
| group = augment(rows, random.Random(args.seed + 100003 * (step + 1))) |
| optimizer.zero_grad() |
| total_loss = 0.0 |
| input_tokens = 0 |
| for batch in microbatches(group, args.batch_size, args.token_budget): |
| prepared = model.prepare(batch, max_length=args.max_length) |
| logits = model(prepared) |
| target = targets(batch, logits.device) |
| loss = -(target * F.log_softmax(logits, dim=-1)).sum(-1).mean() |
| if not torch.isfinite(loss): |
| raise RuntimeError(f"Nonfinite loss at step {step + 1}") |
| (loss * len(batch) / len(group)).backward() |
| total_loss += float(loss.detach()) * len(batch) |
| input_tokens += prepared.input_tokens |
| del logits, loss, target, prepared |
| gradient_norm = float(torch.nn.utils.clip_grad_norm_(parameters, 1.0)) |
| if not math.isfinite(gradient_norm): |
| raise RuntimeError(f"Nonfinite gradient at step {step + 1}") |
| step += 1 |
| warmup = max(1, int(0.05 * total_steps)) |
| factor = step / warmup if step <= warmup else 0.1 + 0.45 * (1 + math.cos(math.pi * (step - warmup) / max(1, total_steps - warmup))) |
| for group_parameters in optimizer.param_groups: |
| group_parameters["lr"] = args.lr * factor |
| optimizer_started = time.monotonic() |
| optimizer.step() |
| torch.cuda.synchronize() |
| examples_seen += len(group) |
| value = record("training_step", run=run.name, step=step, loss=total_loss / len(group), learning_rate=args.lr * factor, |
| examples_seen=examples_seen, group_examples=len(group), input_tokens=input_tokens, gradient_norm=gradient_norm, |
| step_seconds=time.monotonic() - began, optimizer_seconds=time.monotonic() - optimizer_started, |
| elapsed_seconds=time.monotonic() - started, gpu_peak_gb=torch.cuda.max_memory_allocated() / 1e9) |
| append(run / "training.jsonl", value) |
| print(json.dumps(value), flush=True) |
| if step % args.eval_every == 0 or step == stop: |
| evaluate(include_public=step % args.public_eval_every == 0 or step == total_steps, |
| diagnostic=step % args.eval_every != 0 and step != total_steps) |
| elif quick_development and step % args.quick_eval_every == 0: |
| evaluate_quick() |
| if not args.no_resume and (step % args.resume_every == 0 or step == stop): |
| save_resume() |
| summary = {"run": run.name, "steps": step, "planned_steps": total_steps, "complete": step == total_steps, |
| "examples_seen": examples_seen, "best_step": best_step, "selected_temperature": selected_temperature, |
| "selected_metrics": best, "selected_guard_metrics": best_guard, "jev": reference, "guard_jev": guard_reference, |
| "initial_artifact": initial_artifact, "data_sha256": hashes, "git_commit": revision, |
| "elapsed_seconds": time.monotonic() - started, "checkpoint": str(output / "selected") if best else None} |
| if scheduled: |
| summary["training_schedule"] = scheduled.snapshot(examples_seen) |
| write_json(run / "summary.json", summary) |
| record("training_finished" if step == total_steps else "training_paused", **summary) |
| print(json.dumps(summary, indent=2), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
