"""Train the shared decision head over frozen native-rate signal embeddings.""" import argparse import json import os from pathlib import Path import random from time import perf_counter import torch from torch.nn import functional as F import yaml from .checkpoint import save_checkpoint, load_checkpoint from .data import ChoiceDataset, assert_disjoint, make_batch from .evaluate import collect, metrics from .features import open_feature_cache from .model import PayaModel from .tracking import open_run, run_identity, log_best from .atomic import atomic_output def choice_loss(logits, targets, mask, smoothing=0.0, reduction="mean"): if targets.dtype != torch.long or targets.shape != (len(logits),): raise ValueError("Targets must be integer candidate indices [B]") if ((targets < 0) | (targets >= logits.shape[1])).any(): raise ValueError("Target out of range") if not mask.gather(1, targets[:, None]).all(): raise ValueError("Cannot train on a padded candidate") if not 0 <= smoothing < 1 or reduction not in ("mean","none"): raise ValueError("Invalid candidate loss smoothing or reduction") if smoothing==0: return F.cross_entropy(logits, targets,reduction=reduction) logp=logits.log_softmax(-1) nll=-logp.gather(1,targets[:,None]).squeeze(1) uniform=-logp.masked_fill(~mask,0).sum(-1)/mask.sum(-1) values=(1-smoothing)*nll+smoothing*uniform return values.mean() if reduction=="mean" else values def train(config_path, run_dir, device="cpu", resume=None): config = yaml.safe_load(Path(config_path).read_text(encoding="utf-8")) destination = Path(run_dir) if destination.exists() and any(destination.iterdir()) and resume is None: raise FileExistsError("Use a fresh run directory or supply --resume") destination.mkdir(parents=True, exist_ok=True) seed = config.get("seed", 42) random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) torch.set_num_threads(config.get("threads", 4)) datasets = {name: ChoiceDataset(path) for name, path in config["data"].items() if name in ("train", "selection", "calibration", "test") and path} if not {"train", "selection"} <= datasets.keys(): raise ValueError("Training and independent selection datasets are required") assert_disjoint(*datasets.values()) model = PayaModel.from_assets(config["assets"], config.get("preparation"), config.get("head"), device) cache = open_feature_cache(config) opt = config["optim"] epochs = opt["epochs"] if isinstance(epochs, bool) or not isinstance(epochs, int) or epochs < 1: raise ValueError("Need at least one training epoch") optimizer = torch.optim.AdamW(model.trainable_parameters(), lr=opt["lr"], weight_decay=opt.get("weight_decay", 0.01)) interval = opt.get("checkpoint_every_steps", 100) if isinstance(interval, bool) or not isinstance(interval, int) or interval < 1: raise ValueError("checkpoint_every_steps must be a positive integer") bindings = {name: data.sha256 for name, data in datasets.items()} forbidden = {"datasets": [{"path": str(datasets[name].path), "sha256": datasets[name].sha256} for name in ("train", "selection")]} start, best, step, previous_run = 0, float("inf"), 0, None next_batch, epoch_total, epoch_count, restored_state = 0, 0.0, 0, None # Load lazy text assets before restoring RNG, so initialization cannot perturb resumed dropout. make_batch(model, [datasets["train"].rows[0]]) if resume: payload = load_checkpoint(resume, model) state = payload["training"] if not state or ({k: v for k, v in state["config"].items() if k != "wandb"} != {k: v for k, v in config.items() if k != "wandb"}): raise ValueError("Resume requires the same model, data and optimizer configuration") optimizer.load_state_dict(state["optimizer"]) if state.get("dataset_sha256") != bindings: raise ValueError("Resume dataset content differs from checkpoint") start, best = state["epoch"], state["best_nll"] step, previous_run = state["global_step"], state["wandb"] next_batch = state.get("next_batch", 0) epoch_total, epoch_count = state.get("epoch_loss_sum", 0.0), state.get("epoch_count", 0) restored_state = state history = [] completed_result = restored_state.get("completed_epoch_result") if restored_state else None last_batch_metrics = restored_state.get("last_batch_metrics") if restored_state else None history_path = destination/"history.jsonl" saved_history = [] if history_path.exists(): with history_path.open(encoding="utf-8") as source: saved_history = [json.loads(line) for line in source if line.strip()] saved_history = [row for row in saved_history if row["epoch"] <= start] with atomic_output(history_path,"w") as output: for row in saved_history: output.write(json.dumps(row)+"\n") steps_path = destination/"steps.jsonl" if restored_state and "metrics_bytes" in restored_state: committed_bytes = restored_state["metrics_bytes"] if not steps_path.exists() or steps_path.stat().st_size < committed_bytes: raise ValueError("Resume requires the committed steps.jsonl beside the checkpoint") with steps_path.open("r+b") as output: output.truncate(committed_bytes) output.flush() os.fsync(output.fileno()) if start >= epochs: return model, history tracking = config.get("wandb", {}) with open_run(config, destination, model, datasets, previous_run) as run, steps_path.open("a+b") as step_log: if restored_state: torch.set_rng_state(restored_state["torch_rng"]) random.setstate(restored_state["python_rng"]) if torch.cuda.is_available() and restored_state["cuda_rng"]: torch.cuda.set_rng_state_all(restored_state["cuda_rng"]) if run: run.log({"recovery/checkpoint_step": step, "recovery/next_batch": next_batch}) if restored_state.get("metrics_bytes"): with steps_path.open(encoding="utf-8") as source: for line in source: record=json.loads(line) if record["step"] % tracking.get("log_every",10)==0 or record["step"]==step: run.log({**record,"recovery/replayed":True}) elif last_batch_metrics: run.log({**last_batch_metrics,"recovery/replayed":True}) for completed in saved_history: run.log({"step":completed["step"],"epoch":completed["epoch"],"recovery/replayed":True, **{f"selection/{key}":value for key,value in completed["selection"].items() if key!='reliability'}}) def progress(epoch, cursor, total, count): step_log.flush() os.fsync(step_log.fileno()) return {"config": config, "epoch": epoch, "next_batch": cursor, "metrics_bytes":step_log.tell(), "epoch_loss_sum": total, "epoch_count": count, "completed_epoch_result": completed_result, "last_batch_metrics": last_batch_metrics, "best_nll": best, "global_step": step, "dataset_sha256": bindings, "wandb": run_identity(run, tracking.get("mode", "disabled")), "optimizer": optimizer.state_dict(), "torch_rng": torch.get_rng_state(), "python_rng": random.getstate(), "cuda_rng": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else [], "forbidden_calibration": forbidden} if not resume: save_checkpoint(destination/"last.pt", model, progress(0, 0, 0.0, 0)) for epoch in range(start, epochs): started = perf_counter() for data in datasets.values(): data.assert_unchanged() total, count = (epoch_total, epoch_count) if epoch == start else (0.0, 0) cursor = next_batch if epoch == start else 0 for batch_index, rows in enumerate(datasets["train"].batches( config.get("batch_size", 8), seed+epoch, start_batch=cursor), start=cursor): batch = make_batch(model, rows) targets = torch.tensor([row["target"] for row in rows], device=model.device) optimizer.zero_grad(set_to_none=True) model.train() result = model(batch, cache=cache) loss = choice_loss(result["logits"], targets, batch.criteria_mask,opt.get("label_smoothing",0.0)) if not torch.isfinite(loss): raise FloatingPointError("Nonfinite training loss") loss.backward() gradient_norm = torch.nn.utils.clip_grad_norm_(list(model.trainable_parameters()), 1.0) if not torch.isfinite(gradient_norm): raise FloatingPointError("Nonfinite training gradients") optimizer.step() step += 1 total += loss.item()*len(rows); count += len(rows) cursor = batch_index+1 last_batch_metrics = {"step": step, "epoch": epoch+1, "train/loss": loss.item(), "train/answer_loss": loss.item(), "train/gradient_norm": float(gradient_norm), "train/learning_rate": optimizer.param_groups[0]["lr"]} step_log.write((json.dumps(last_batch_metrics,allow_nan=False)+"\n").encode("utf-8")) if run and step % tracking.get("log_every", 10) == 0: run.log(last_batch_metrics) if step % interval == 0: save_checkpoint(destination/"last.pt", model, progress(epoch, cursor, total, count)) records = collect(model, datasets["selection"], config.get("batch_size", 8), cache, output_path=destination/"selection-predictions.jsonl") selection = metrics(records) improved = selection["nll"] < best if improved: best = selection["nll"] row = {"epoch": epoch+1, "phase": "decision_head", "step": step, "train_loss": total/max(1, count), "selection": selection} completed_result = row state = progress(epoch+1, 0, 0.0, 0) if improved: save_checkpoint(destination/"best.pt", model, state) saved_history = [previous for previous in saved_history if previous["epoch"] < epoch+1]+[row] with atomic_output(history_path,"w") as log: for previous in saved_history: log.write(json.dumps(previous)+"\n") save_checkpoint(destination/"last.pt", model, state) history.append(row) if run: logged = {"epoch": epoch+1, "step": step, "phase": row["phase"], "train/epoch_loss": row["train_loss"], "train/epoch_seconds": perf_counter()-started, **{f"selection/{k}": selection[k] for k in ("accuracy", "nll", "brier", "ece_15", "count")}} run.log(logged) if improved: run.summary.update({"best/epoch":epoch+1,"best/selection_nll":best}) if tracking.get("log_models",True): log_best(run, model, destination, epoch+1, best) print(json.dumps(row), flush=True) return model, history if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", type=Path, default=Path("config/paya.yaml")) parser.add_argument("--run-dir", type=Path, required=True) parser.add_argument("--device", default="cpu") parser.add_argument("--resume", type=Path) args = parser.parse_args() train(args.config, args.run_dir, args.device, args.resume)