"""Single-GPU training with exact question weighting and durable resume state.""" import importlib.metadata import json import math import random import shutil import signal import threading import time import uuid from pathlib import Path import numpy as np import torch from .budget import Budget from .checkpoint import ( checkpoint_path, load_training, restore_rng, rng_state, save_checkpoint, select_best, ) from .common import atomic_json, digest, source_hash from .data import encoded_row, load_rows, prepared_directory, tokenizer_for, verify_prepared from .evaluation import evaluate, fit_temperatures, grouped_metrics, permutation_sensitivity from .model import collate, load_model, soft_ce_sum from .tracking import Tracker def seed_all(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def epoch_order(length, seed, epoch): order = list(range(length)) random.Random(seed + epoch).shuffle(order) return order def make_scheduler(optimizer, cfg, total_steps): warmup = max(1, int(total_steps * cfg["train"]["warmup_fraction"])) floor = cfg["train"]["final_learning_rate"] / cfg["train"]["learning_rate"] def factor(step): if step < warmup: return (step + 1) / warmup progress = min(1.0, (step - warmup) / max(1, total_steps - warmup)) return floor + (1 - floor) * 0.5 * (1 + math.cos(math.pi * progress)) return torch.optim.lr_scheduler.LambdaLR(optimizer, factor) def optimize_window(model, codec, encoded, optimizer, microbatch, clip, device): """Sum losses over decisions, then divide once by this whole accumulation window.""" total = sum(len(row["positions"]) for row in encoded) start_rng = rng_state() while True: optimizer.zero_grad(set_to_none=True) losses = 0.0 try: for start in range(0, len(encoded), microbatch): batch = collate(encoded[start : start + microbatch], codec.pad, device) with torch.autocast( device.type, dtype=torch.bfloat16, enabled=device.type == "cuda" ): logits = model(batch) loss = soft_ce_sum(logits, batch["target"]) / total if not torch.isfinite(loss): raise FloatingPointError("Non-finite training loss") loss.backward() losses += float(loss.detach()) del logits, loss, batch grad = torch.nn.utils.clip_grad_norm_( [p for p in model.parameters() if p.requires_grad], clip, error_if_nonfinite=True ) optimizer.step() optimizer.zero_grad(set_to_none=True) return losses, float(grad), microbatch except torch.cuda.OutOfMemoryError: optimizer.zero_grad(set_to_none=True) # Do not retain a failed graph in local variables or replay partial gradients. logits = loss = batch = None torch.cuda.empty_cache() restore_rng(start_rng) if microbatch != 8: raise microbatch = 4 def train( cfg, run_dir, resume=False, directory=None, device="cuda", commit=None, preparation_cost=0.0, tracking_disabled=False, continuation=None, ): run_dir = Path(run_dir) directory = Path(directory) if directory else prepared_directory(cfg) manifest = verify_prepared(cfg, directory) config_hash = digest(cfg) identity_file = run_dir / "run.json" if identity_file.exists(): if not resume: raise FileExistsError( "Run exists; use explicit --resume to continue the same experiment" ) identity = json.loads(identity_file.read_text()) if ( identity["config_hash"] != config_hash or (identity["code_hash"] != source_hash() and continuation is None) or identity["prepared_manifest_hash"] != manifest["manifest_hash"] ): raise ValueError("Resume requires the same code, configuration and prepared data") else: if resume: raise FileNotFoundError("Cannot resume a run that has not started") identity = { "run_id": uuid.uuid4().hex[:12], "config_hash": config_hash, "code_hash": source_hash(), "prepared_manifest_hash": manifest["manifest_hash"], } atomic_json(identity_file, identity) if continuation is not None: if not resume: raise ValueError("Three-epoch continuation requires explicit resume") from .continuation import archive_previous_export, validate_continuation validate_continuation(cfg, run_dir, continuation) archive_previous_export(run_dir) atomic_json(run_dir / "continuation.json", continuation) budget = Budget( cfg, run_dir, preparation_cost, extension=continuation["limits"] if continuation is not None else None, ) previous_handler = None if budget.stop(training=True): budget.persist(clean=True) if commit: commit() raise RuntimeError("This experiment has exhausted its remaining training budget") if threading.current_thread() is threading.main_thread(): previous_handler = signal.getsignal(signal.SIGTERM) signal.signal(signal.SIGTERM, lambda *_: setattr(budget, "shutdown_requested", True)) tracker = None try: seed_all(cfg["seed"]) codec = tokenizer_for(cfg, directory / "tokenizer") saved = checkpoint_path(run_dir) if resume else None model = load_model(cfg, codec, device=device, adapter=saved / "adapter" if saved else None) device = next(model.parameters()).device optimizer = torch.optim.AdamW( [p for p in model.parameters() if p.requires_grad], lr=cfg["train"]["learning_rate"], weight_decay=cfg["train"]["weight_decay"], ) rows = load_rows(directory, "train") validation = load_rows(directory, "validation") window_size = cfg["train"]["microbatch"] * cfg["train"]["accumulation"] total_steps = math.ceil(len(rows) / window_size) * cfg["train"]["epochs"] scheduler = make_scheduler(optimizer, cfg, total_steps) state = { "step": 0, "epoch": 0, "position": 0, "best_loss": float("inf"), "bad_evaluations": 0, "decisions_processed": 0, "last_eval_step": -1, "microbatch": cfg["train"]["microbatch"], } if resume: state = load_training(saved, optimizer, scheduler) tracker = Tracker( cfg, run_dir, identity["run_id"], resume=resume, disabled=tracking_disabled, resume_step=state["step"] + 1 if resume else None, ) tracker.workspace() versions = { name: importlib.metadata.version(name) for name in ("torch", "transformers", "peft", "wandb", "tokenizers", "huggingface-hub") } provenance = { "identity": identity, "config": cfg, "versions": versions, "model_revision": cfg["model"]["revision"], "dataset_revision": cfg["data"]["revision"], "prepared": manifest, "device": str(device), "gpu": torch.cuda.get_device_name() if device.type == "cuda" else None, "trainable_parameters": sum(p.numel() for p in model.parameters() if p.requires_grad), "executing_code_hash": source_hash(), "continuation": continuation, } atomic_json(run_dir / "provenance.json", provenance) def persist_checkpoint(): budget.persist() return save_checkpoint(run_dir, model, optimizer, scheduler, state, cfg, codec, commit) def validate(): predicted, completed = evaluate( model, codec, validation, cfg["evaluation"]["batch_size"], budget ) if not completed: return False metrics = grouped_metrics(predicted) tracker.evaluation(metrics, state["step"]) state["last_eval_step"] = state["step"] loss = metrics["overall"]["soft_ce"] improvement = state["best_loss"] - loss state["bad_evaluations"] = ( 0 if improvement >= cfg["train"]["minimum_improvement"] else state["bad_evaluations"] + 1 ) if loss < state["best_loss"]: state["best_loss"] = loss checkpoint = persist_checkpoint() select_best(run_dir, checkpoint, loss, commit) return True if not resume: initial, completed = evaluate( model, codec, validation, cfg["evaluation"]["batch_size"], budget ) if not completed: raise RuntimeError("Budget exhausted during initial validation") tracker.evaluation(grouped_metrics(initial), 0, section="initial_validation") persist_checkpoint() model.train() started_step = state["step"] measured_seconds = 0.0 measured_decisions = measured_tokens = 0 last_commit = time.monotonic() stopping = None while state["epoch"] < cfg["train"]["epochs"] and not stopping: order = epoch_order(len(rows), cfg["seed"], state["epoch"]) while state["position"] < len(order): stopping = budget.stop(training=True) if stopping: break indices = order[state["position"] : state["position"] + window_size] encoded = [ encoded_row( rows[index], codec, seed=int(digest([cfg["seed"], state["epoch"], rows[index]["id"]])[:16], 16), ) for index in indices ] tick = time.monotonic() loss, grad, state["microbatch"] = optimize_window( model, codec, encoded, optimizer, state["microbatch"], cfg["train"]["gradient_clip"], device, ) scheduler.step() if device.type == "cuda": torch.cuda.synchronize() elapsed = time.monotonic() - tick count = sum(len(row["positions"]) for row in encoded) tokens = sum(len(row["input_ids"]) for row in encoded) state["position"] += len(indices) state["step"] += 1 state["decisions_processed"] += count measured_seconds += elapsed measured_decisions += count measured_tokens += tokens metrics = { "train/soft_ce": loss, "train/gradient_norm": grad, "train/learning_rate": optimizer.param_groups[0]["lr"], "train/epoch": state["epoch"] + state["position"] / len(rows), "train/decisions_processed": state["decisions_processed"], "train/microbatch": state["microbatch"], "performance/step_seconds": elapsed, "performance/decisions_second": measured_decisions / measured_seconds, "performance/tokens_second": measured_tokens / measured_seconds, "performance/peak_memory_gib": torch.cuda.max_memory_allocated() / 2**30 if device.type == "cuda" else 0, **budget.metrics(), } completed_steps = state["step"] - started_step if completed_steps >= 50: metrics["performance/projected_remaining_seconds"] = ( (total_steps - state["step"]) * measured_seconds / completed_steps ) if state["step"] % cfg["train"]["log_steps"] == 0 or completed_steps == 50: tracker.log(metrics, state["step"]) print( json.dumps( { "step": state["step"], "loss": loss, "epoch": metrics["train/epoch"], "estimated_usd": budget.cost, } ), flush=True, ) budget.persist() if commit and time.monotonic() - last_commit >= 60: commit() last_commit = time.monotonic() if state["step"] % cfg["train"]["eval_steps"] == 0: if not validate(): stopping = "evaluation_budget" break if ( continuation is None and state["epoch"] >= 1 and state["bad_evaluations"] >= cfg["train"]["patience"] ): stopping = "early_stopping" break if state["step"] % cfg["train"]["save_steps"] == 0: persist_checkpoint() if state["position"] == len(order): state["epoch"] += 1 state["position"] = 0 if state["last_eval_step"] != state["step"] and not validate(): stopping = "evaluation_budget" persist_checkpoint() if continuation is None and state["bad_evaluations"] >= cfg["train"]["patience"]: stopping = stopping or "early_stopping" stopping = stopping or "epochs_complete" if budget.shutdown_requested: persist_checkpoint() summary = { "status": "interrupted", "stopping_reason": "shutdown", "steps": state["step"], "run_id": identity["run_id"], "estimated_usd": budget.cost, } atomic_json(run_dir / "summary.json", summary) budget.persist(clean=True) if commit: commit() return summary if state["last_eval_step"] != state["step"] and not budget.stop(): validate() persist_checkpoint() if not (run_dir / "best.json").exists(): raise RuntimeError("No fully validated trained checkpoint fits within the budget") best = checkpoint_path(run_dir, "best") best_step = json.loads((best / "complete.json").read_text())["step"] model = optimizer = scheduler = None if device.type == "cuda": torch.cuda.empty_cache() model = load_model( cfg, codec, device=str(device), adapter=best / "adapter", trainable=False ) model.eval() calibration_rows, calibration_complete = evaluate( model, codec, load_rows(directory, "calibration"), cfg["evaluation"]["batch_size"], budget, ) temperatures = ( fit_temperatures(calibration_rows) if calibration_complete else {k: 1.0 for k in ("choice", "noul", "score")} ) atomic_json( run_dir / "calibration.json", { "temperatures": temperatures, "completed": calibration_complete, "decision_count": len(calibration_rows), }, ) test_rows = load_rows(directory, "test") test, test_complete = evaluate( model, codec, test_rows, cfg["evaluation"]["batch_size"], budget ) raw, calibrated = grouped_metrics(test), grouped_metrics(test, temperatures) permutation = permutation_sensitivity( model, codec, test_rows, cfg["evaluation"]["permutation_examples"], cfg["seed"], budget ) evaluation = { "test_completed": test_complete, "calibration_completed": calibration_complete, "raw": raw, "calibrated": calibrated, "permutation": permutation, } atomic_json(run_dir / "evaluation.json", evaluation) with (run_dir / "test-predictions.jsonl").open("w") as stream: for row in test: stream.write(json.dumps(row) + "\n") if test: tracker.evaluation(raw, state["step"], "test_raw") tracker.evaluation(calibrated, state["step"], "test_calibrated") tracker.final_plots( raw, calibrated, test, temperatures, cfg["evaluation"]["error_examples"] ) tracker.log( {"calibration/temperature_" + k: v for k, v in temperatures.items()}, state["step"] ) tracker.log( { "final/permutation_" + k: v for k, v in permutation.items() if isinstance(v, (int, float)) }, state["step"], ) export = run_dir / "export" if export.exists(): shutil.rmtree(export) shutil.copytree(best / "adapter", export) shutil.copyfile(run_dir / "calibration.json", export / "calibration.json") summary = { "status": "completed" if test_complete and calibration_complete else "partial_evaluation", "stopping_reason": stopping, "steps": state["step"], "selected_step": best_step, "completed_epochs": state["epoch"], "fractional_epochs": state["epoch"] + state["position"] / len(rows), "best_validation_soft_ce": state["best_loss"], "estimated_usd": budget.cost, "run_id": identity["run_id"], "test_completed": test_complete, "calibration_completed": calibration_complete, } tracker.finish(export, summary) budget.persist(clean=True) if commit: commit() return summary except BaseException as exc: atomic_json( run_dir / "failure.json", {"type": type(exc).__name__, "message": str(exc)[:500]} ) budget.persist(clean=True) if tracker: tracker.finish( None, { "status": "failed", "error_type": type(exc).__name__, "estimated_usd": budget.cost, }, ) if commit: commit() raise finally: if previous_handler is not None: signal.signal(signal.SIGTERM, previous_handler)