"""Backtracking (stage-replay) training loop for the recurrent latent model. Motivation ---------- Standard curriculum training runs EVERY batch at `num_passes = curriculum.stage` (the current frontier depth). When the curriculum advances to stage i, the earlier, shallower readouts (stage 1..i-1, each of which should be produced by applying the tied recurrent step exactly t times) drift / are forgotten. Backtracking = interleaved replay with the recurrence depth MATCHED to the replayed stage: stage t <-> difficulty level t+2 <-> num_passes = t <-> t active slots With cand_slot_mode="depth" and passes_per_stage=pps the depth becomes pps*t and the active-slot count follows num_passes rather than the puzzle level. Replay is still safe because that mode maps slot j -> stage j: slots [0, pps*t) hold the right targets no matter which frontier stage the sampler built the batch for. * frontier step : level (i+2) puzzles, num_passes = i -> trains stage-i readout * backtrack step: level (t+2) puzzles, num_passes = t -> re-derives stage-t for a randomly chosen earlier t in {1..i-1} Each optimisation step is frontier with prob (1 - backtrack_prob) and a backtrack step with prob backtrack_prob (t uniform over the earlier stages). Because the recurrent step is weight-tied, replaying "apply f exactly t times = stage t" keeps f reusable at every depth instead of specialising to the frontier depth. Everything else (model, data, per-slot BCE + LM CE loss, eval, promotion) is reused unchanged from the standard pipeline. """ import functools import math import os from absl import logging from clu import metric_writers from flax import jax_utils from flax.training import checkpoints from flax.training import common_utils import jax import numpy as np import tensorflow as tf from train import data from train import evaluater from train import model from train import trainer class StageBatchSampler: """Draws minibatches of a *specific* difficulty level from the train set. Reuses SudokuDataset (which loads the puzzles + staged candidate masks and builds a level->indices map) and assembles numpy batches in the exact tuple layout the train step expects: (seq, puzzle, start_index, level, cand_targets). """ def __init__(self, config, seed=0): self.ds = data.SudokuDataset(config, train=True) self.rng = np.random.RandomState(seed) # Only levels that actually have puzzles. self.available = {lvl: idx for lvl, idx in self.ds.level_index.items() if len(idx) > 0} def has_level(self, level): return level in self.available def _gather(self, idxs): seqs, puzzles, starts, levels, cands = [], [], [], [], [] for idx in idxs: seq, puzzle, start_index, lvl, cand = self.ds.__getitem__(int(idx)) seqs.append(seq) puzzles.append(puzzle) starts.append(start_index) levels.append(lvl) cands.append(cand) return ( np.stack(seqs).astype(np.int32), np.stack(puzzles).astype(np.int32), np.stack(starts).astype(np.int32), np.stack(levels).astype(np.int32), np.stack(cands).astype(np.int32), ) def sample_all(self, bs): """Batch drawn uniformly from the whole corpus, matching the standard loop's sampler. Difficulty is not a curriculum axis, so every step -- frontier or repair -- uses this same distribution and differs only in the recurrence depth it trains at.""" return self._gather(self.rng.randint(len(self.ds), size=bs)) def sample(self, level, bs): idx_pool = self.available[level] idxs = idx_pool[self.rng.randint(len(idx_pool), size=bs)] return self._gather(idxs) def sample_mixed(self, levels, bs): """Batch drawn uniformly over `levels` (level first, then a puzzle of that level) — the exact distribution the standard curriculum sampler uses. Frontier steps must use this, NOT single-level batches: training exclusively on the frontier level starves every easier level and is a different (worse) recipe than the standard loop, not "standard + BT".""" levels = [l for l in levels if l in self.available] idxs = [] for _ in range(bs): lvl = levels[self.rng.randint(len(levels))] pool = self.available[lvl] idxs.append(int(pool[self.rng.randint(len(pool))])) return self._gather(idxs) def _run_train_step(step_fn, state, batch_tuple, dropout_rngs): """Shard a single numpy batch and run one (pmapped) train step.""" inputs, _, start_index, levels, cand_targets = batch_tuple inputs = common_utils.shard(jax.tree_util.tree_map(np.asarray, inputs)) start_index = common_utils.shard(jax.tree_util.tree_map(np.asarray, start_index)) levels = common_utils.shard(jax.tree_util.tree_map(np.asarray, levels)) cand_targets = np.asarray(cand_targets) _nd = jax.local_device_count() cand_targets = cand_targets.reshape( (_nd, cand_targets.shape[0] // _nd) + cand_targets.shape[1:]) state, metrics, _ = step_fn( state, inputs, start_index, levels, cand_targets, dropout_rng=dropout_rngs) return state, metrics def _prepare_backtrack(config, workdir): """Build the model, state, pmapped steps, sampler, and writers shared by the probabilistic and adaptive backtracking loops. Returns a dict of everything both loops need. `state` is already replicated. """ K = int(config.num_latent_slots) assert K > 0, "Backtracking is for the recurrent latent model (K>0)." logging.info("Creating datasets (backtracking loop)") curriculum = data.CurriculumState( stage=int(getattr(config, "curriculum_start_stage", 1)), max_stage=int(getattr(config, "curriculum_max_stage", 6))) # Per-level train sampler (replaces the uniform-over-unlocked sampler so we # can draw a batch of a specific replay stage's level on demand). sampler = StageBatchSampler(config, seed=int(config.seed)) eval_data_iter = data.create_iter(config, config.minibatch_size, train=False) model_config = model.TransformerConfig( dtype=config.dtype, vocab_size=config.vocab_size, seq_len=config.seq_len, num_heads=config.num_heads, num_layers=config.num_layers, emb_dim=config.emb_dim, qkv_dim=config.qkv_dim, mlp_dim=config.mlp_dim, dropout_rate=config.dropout_rate, attention_dropout_rate=config.attention_dropout_rate, deterministic=False, num_latent_slots=K, inject_latents=bool(int(getattr(config, "recurrent_latent", 1))), ) print(str(model_config.__dict__), flush=True) rng = jax.random.PRNGKey(config.seed) rng, init_rng, dropout_rng = jax.random.split(rng, 3) net = model.TransformerLMHeadModel(model_config) dummy_latents = jax.numpy.zeros( (config.minibatch_size, K, config.emb_dim), model_config.dtype) dummy_positions = jax.numpy.zeros((config.minibatch_size, K), jax.numpy.int32) dummy_active = jax.numpy.zeros((config.minibatch_size, K), bool) _, initial_variables = jax.jit(net.init_with_output)( {"params": init_rng, "dropout": dropout_rng}, jax.numpy.ones((config.minibatch_size, config.seq_len), jax.numpy.int32), dummy_latents, dummy_positions, dummy_active) state, lr_scheduler_fn = trainer.get_state(config, net, initial_variables) start_step = 0 if config.resume_training: state = checkpoints.restore_checkpoint(config.ckpt_loc, state) start_step = int(state.step) print("----------Restored model from", config.ckpt_loc, f"at step {start_step}-----------") writer = metric_writers.create_default_writer( workdir, asynchronous=False, just_logging=(jax.process_index() > 0)) tf_summary_writer = tf.summary.create_file_writer(workdir) state = jax_utils.replicate(state) dropout_rngs = jax.random.split(rng, jax.local_device_count()) def make_p_train_step(num_passes, backtrack): return jax.pmap( functools.partial( trainer.train_step, config=model_config, hyperparams=config, learning_rate_fn=lr_scheduler_fn, num_passes=num_passes, backtrack=backtrack), axis_name="batch", donate_argnums=(0,)) # Precompile, for every recurrence depth 1..K, both a FRONTIER step (normal # loss: LM CE + all active-slot BCE) and a strict BACKTRACK step (only the # last active slot's readout, no CE). Frontier is used when t == curriculum # stage i; the strict backtrack step is used by the probabilistic loop for # replayed earlier stages t < i. # Stage -> recurrence depth. With passes_per_stage>1 a stage advances the # chain by more than one slot (pps=2, K=12: stage 6 -> all 12 slots), so the # replay depth for stage t is pps*t, not t. Keyed by stage either way. pps = int(getattr(config, "passes_per_stage", 1)) _depth = lambda t: min(pps * t, K) p_frontier = {t: make_p_train_step(_depth(t), False) for t in range(1, K + 1)} p_backtrack = {t: make_p_train_step(_depth(t), True) for t in range(1, K + 1)} p_eval_step = jax.pmap(functools.partial( evaluater.eval_step, config=model_config.replace(deterministic=True)), axis_name="batch") hooks, report_progress, _ = trainer.get_metrics_report_progress( config, workdir, writer) return { "K": K, "start_step": start_step, "curriculum": curriculum, "sampler": sampler, "eval_data_iter": eval_data_iter, "model_config": model_config, "state": state, "dropout_rngs": dropout_rngs, "p_frontier": p_frontier, "p_backtrack": p_backtrack, "p_eval_step": p_eval_step, "hooks": hooks, "report_progress": report_progress, "writer": writer, "tf_summary_writer": tf_summary_writer, "promote_threshold": float(getattr(config, "promote_acc_threshold", 0.85)), "promote_patience": int(getattr(config, "promote_patience_steps", 8000)), "min_stage_steps": int(getattr(config, "min_stage_steps", 2000)), "ckpt_keep": int(getattr(config, "ckpt_keep", 100)), "stage_ckpt_dir": os.path.join(workdir, "stage_ckpts"), } def train_and_evaluate_backtrack(config, workdir): """Backtracking curriculum training loop (stage replay with matched depth).""" workdir = os.path.abspath(workdir) backtrack_prob = float(getattr(config, "backtrack_prob", 0.5)) frontier_mix = bool(int(getattr(config, "backtrack_frontier_mix", 1))) ctx = _prepare_backtrack(config, workdir) K = ctx["K"] curriculum = ctx["curriculum"] sampler = ctx["sampler"] eval_data_iter = ctx["eval_data_iter"] state = ctx["state"] dropout_rngs = ctx["dropout_rngs"] p_frontier = ctx["p_frontier"] p_backtrack = ctx["p_backtrack"] p_eval_step = ctx["p_eval_step"] hooks = ctx["hooks"] writer = ctx["writer"] tf_summary_writer = ctx["tf_summary_writer"] promote_threshold = ctx["promote_threshold"] promote_patience = ctx["promote_patience"] min_stage_steps = ctx["min_stage_steps"] ckpt_keep = ctx["ckpt_keep"] stage_ckpt_dir = ctx["stage_ckpt_dir"] stage_started_at = 0 sched_rng = np.random.RandomState(int(config.seed) + 1) # Bookkeeping: how often each replay depth is actually trained. stage_step_counts = {t: 0 for t in range(1, K + 1)} def sample_target_stage(i): """Frontier stage i with prob (1-p); else a backtrack stage 1..i-1.""" if i <= 1 or sched_rng.rand() >= backtrack_prob: return i return int(sched_rng.randint(1, i)) # uniform in {1..i-1} with metric_writers.ensure_flushes(writer): for step in range(0, config.max_steps): if step % 10000 == 0: print("Step:", step, flush=True) i = curriculum.stage t = sample_target_stage(i) # Frontier step (t == i): normal loss. Backtrack step (t < i): strict # replay of the depth-t readout only. Both draw from the whole # corpus: depth, not difficulty, is the replayed axis. step_fn = p_frontier[t] if t == i else p_backtrack[t] batch = sampler.sample_all(config.minibatch_size) state, metrics = _run_train_step(step_fn, state, batch, dropout_rngs) stage_step_counts[t] += 1 for h in hooks: h(step) if math.isnan(metrics["loss"][0]): print("Loss became nan; stopping.", flush=True) break if step % config.eval_every_steps == 0: eval_metrics = evaluater.get_eval_metrics( state, eval_data_iter, p_eval_step, config) per_level = eval_metrics.pop("per_level_acc") per_depth = eval_metrics.pop("per_slot_acc_changed", {}) per_depth_all = eval_metrics.pop("per_slot_acc", {}) if not any(v >= 0 for v in per_depth.values()): per_depth = per_depth_all def _m(key): v = eval_metrics.get(key, []) return round(float(np.mean(v)), 4) if len(v) else -1.0 print(step, "stage", curriculum.stage, "target_t", t, "loss", round(float(metrics["loss"].mean()), 4), "ce", round(float(metrics["ce_loss"].mean()), 4), "aux_bce", round(float(metrics["aux_loss"].mean()), 4), "| val_acc", _m("acc"), "loc_acc", _m("loc_acc"), "val|loc", _m("val_given_loc_acc"), "| cand_bit_acc", _m("cand_bit_acc"), "cand_set_acc", _m("cand_set_acc"), "cand_set_chg", _m("cand_set_acc_changed"), "| replay_counts", dict(stage_step_counts), flush=True) with tf_summary_writer.as_default(): tf.summary.scalar("loss", metrics["loss"].mean(), step=step) tf.summary.scalar("ce_loss", metrics["ce_loss"].mean(), step=step) tf.summary.scalar("aux_bce_loss", metrics["aux_loss"].mean(), step=step) tf.summary.scalar("curriculum_stage", curriculum.stage, step=step) for key in eval_metrics.keys(): tf.summary.scalar("eval_" + key, np.array(eval_metrics[key]).mean(), step=step) for lvl, v in per_level.items(): if v >= 0: tf.summary.scalar(f"eval_acc_level_{lvl}", v, step=step) # ---- Curriculum promotion (same rule as the standard loop) ---- if curriculum.stage < curriculum.max_stage: frontier_acc = per_depth.get(curriculum.stage, -1.0) steps_in_stage = step - stage_started_at hit_threshold = frontier_acc >= promote_threshold patience_over = steps_in_stage >= promote_patience if steps_in_stage >= min_stage_steps and (hit_threshold or patience_over): reason = "threshold" if hit_threshold else "patience" curriculum.stage += 1 stage_started_at = step print(f"[curriculum] step {step}: promote to stage " f"{curriculum.stage} ({reason}; graduated depth " f"{curriculum.stage - 1} cand-set " f"acc={frontier_acc:.3f}); " f"backtrack pool now depths 1..{curriculum.stage-1}", flush=True) if config.save_checkpoint: unrep = jax_utils.unreplicate(state) checkpoints.save_checkpoint_multiprocess( workdir, unrep, step, keep=ckpt_keep, overwrite=True) checkpoints.save_checkpoint_multiprocess( stage_ckpt_dir, unrep, step, keep=100, overwrite=True, prefix=f"stage{curriculum.stage}_") if config.save_checkpoint and step > 0 and step % config.save_every_steps == 0: checkpoints.save_checkpoint_multiprocess( workdir, jax_utils.unreplicate(state), step, keep=ckpt_keep, overwrite=True) if config.save_checkpoint: checkpoints.save_checkpoint_multiprocess( workdir, jax_utils.unreplicate(state), config.max_steps, keep=ckpt_keep, overwrite=True) def train_and_evaluate_backtrack_adaptive(config, workdir): """Deficit-driven ("adaptive") backtracking loop. Difference from the probabilistic loop -------------------------------------- The probabilistic loop replays a *uniformly random* earlier stage every step with fixed probability p, regardless of whether that stage needs help. This loop instead *watches* each earlier stage's held-out accuracy and only goes back to repair a stage when it has actually regressed: * Reference: when the curriculum promotes past stage t, we record that stage's frontier-level (t+2) val_acc as its "graduation" accuracy. * Trigger (relative drop): while at frontier stage i, at every eval we scan earlier stages 1..i-1; a stage t is *in deficit* if its current level-(t+2) val_acc has fallen below (graduation_acc[t] - margin). * Selection (most-deficient first): if any stage is in deficit we switch the training target to the single most-regressed stage and train there. * Repair step: a FULL frontier-style step at the matched depth (level t+2, num_passes=t, LM CE + all-active-slot BCE) -- NOT the strict readout-only backtrack step. The metric we are trying to restore is placement accuracy (driven by the LM CE), so the repair objective must include it. * Exit (recover-or-cap): stay on the stage until its val_acc climbs back above (graduation_acc[t] - margin), or until a max-repair-steps cap fires (so a stuck stage cannot stall the frontier forever). On exit we re-scan and either move to the next most-deficient stage or return to the frontier. Frontier promotion is paused while repairing, and the repair time is credited back to the frontier stage's patience clock on return. Depth-matched eval note: the evaluator scores a level-L puzzle with exactly k=L-2 active latent slots, so per_level_acc[t+2] is a faithful measurement of the stage-t readout -- a clean trigger signal with no extra instrumentation. """ workdir = os.path.abspath(workdir) margin = float(getattr(config, "backtrack_margin", 0.03)) max_repair_steps = int(getattr(config, "backtrack_max_repair_steps", 6000)) min_frontier_steps = int(getattr(config, "backtrack_min_frontier_steps", 0)) frontier_target_acc = float( getattr(config, "backtrack_frontier_target_acc", 0.0)) grad_acc_seed_raw = str(getattr(config, "backtrack_grad_acc_seed", "") or "") max_repair_fraction = float( getattr(config, "backtrack_max_repair_fraction", 0.0)) grad_decay = float(getattr(config, "backtrack_grad_decay", 0.0)) freeze_after_step = int(getattr(config, "backtrack_freeze_after_step", 0)) frontier_mix = bool(int(getattr(config, "backtrack_frontier_mix", 1))) ctx = _prepare_backtrack(config, workdir) K = ctx["K"] start_step = ctx["start_step"] curriculum = ctx["curriculum"] sampler = ctx["sampler"] eval_data_iter = ctx["eval_data_iter"] state = ctx["state"] dropout_rngs = ctx["dropout_rngs"] p_frontier = ctx["p_frontier"] p_eval_step = ctx["p_eval_step"] hooks = ctx["hooks"] writer = ctx["writer"] tf_summary_writer = ctx["tf_summary_writer"] promote_threshold = ctx["promote_threshold"] promote_loc_threshold = float(getattr(config, "promote_loc_threshold", 0.70)) instance_mode = bool(getattr(config, "instance_dir", None)) promote_patience = ctx["promote_patience"] min_stage_steps = ctx["min_stage_steps"] ckpt_keep = ctx["ckpt_keep"] stage_ckpt_dir = ctx["stage_ckpt_dir"] # Controller state. grad_acc = {} # stage t -> level-(t+2) val_acc at graduation plateau_steps = int(getattr(config, "plateau_steps", 0)) plateau_delta = float(getattr(config, "plateau_delta", 0.005)) stage_best_acc = -1.0 stage_best_step = start_step mode = "frontier" # "frontier" | "repair" repair_stage = None # stage currently being repaired repair_started_at = 0 # step the current repair stage began (cap) repair_episode_start = 0 # step we first left the frontier (clock credit) stage_started_at = start_step # step the current frontier stage began last_repair_return_step = start_step # for min-frontier cooldown step_counts = {t: 0 for t in range(1, K + 1)} # steps trained at each depth repair_steps_total = 0 # for duty-cycle cap print(f"[repair] knobs: margin={margin} max_repair_steps={max_repair_steps} " f"min_frontier={min_frontier_steps} frontier_target={frontier_target_acc} " f"max_repair_frac={max_repair_fraction} grad_decay={grad_decay} " f"freeze_after={freeze_after_step} frontier_mix={frontier_mix}", flush=True) # When seeding from a mid-curriculum checkpoint (start stage > 1), the # earlier stages graduated in the *source* run so we have no reference for # them. Prefer an explicit SUDOKU_GRAD_ACC_SEED (true graduation refs); # otherwise seed from the first eval (old behavior — can over-trigger). seed_start_stage = int(getattr(config, "curriculum_start_stage", 1)) grad_acc_seeded = seed_start_stage <= 1 if (not grad_acc_seeded) and grad_acc_seed_raw.strip(): try: vals = [float(x) for x in grad_acc_seed_raw.split(",") if x.strip()] for s, v in enumerate(vals, start=1): if s < seed_start_stage: grad_acc[s] = v if grad_acc: grad_acc_seeded = True print(f"[repair] seeded graduation refs from env: " f"{dict((k, round(v, 3)) for k, v in grad_acc.items())}", flush=True) except ValueError: print(f"[repair] WARNING: bad SUDOKU_GRAD_ACC_SEED=" f"{grad_acc_seed_raw!r}; falling back to first-eval seeding", flush=True) def compute_deficits(frontier_stage, per_depth): """{depth t: graduation_acc[t] - current_acc} over earlier graduated depths whose current candidate-set accuracy is measured. Keyed on reasoning depth, not difficulty level: depth t's accuracy is how well wave snapshot t is predicted, which is exactly what stage t taught. A drop there means that propagation block has been forgotten.""" d = {} for t in range(1, frontier_stage): if t in grad_acc and per_depth.get(t, -1.0) >= 0: d[t] = grad_acc[t] - per_depth.get(t, -1.0) return d def effective_margin(per_depth, frontier_stage): """Widen the repair trigger while the frontier itself is still weak.""" m = margin if frontier_target_acc > 0: f_acc = per_depth.get(frontier_stage, -1.0) if 0.0 <= f_acc < frontier_target_acc: # Only repair clearer regressions until the frontier is good. m = max(m, margin + (frontier_target_acc - f_acc)) return m def most_deficient(frontier_stage, per_depth, use_margin=None): """Most-regressed depth whose drop exceeds the margin, else None.""" d = compute_deficits(frontier_stage, per_depth) if not d: return None m = margin if use_margin is None else use_margin t = max(d, key=d.get) return t if d[t] > m else None with metric_writers.ensure_flushes(writer): for step in range(start_step, config.max_steps): if step % 10000 == 0: print("Step:", step, flush=True) i = curriculum.stage # Target depth: the frontier when training normally, else the depth # we are repairing. Difficulty is never gated, so a repair differs # from a frontier step only in the recurrence depth it trains at: # replaying depth t re-supervises wave snapshots 1..t on the same # full-corpus batch distribution. t = repair_stage if mode == "repair" else i batch = sampler.sample_all(config.minibatch_size) state, metrics = _run_train_step( p_frontier[t], state, batch, dropout_rngs) step_counts[t] += 1 if mode == "repair": repair_steps_total += 1 for h in hooks: h(step) if math.isnan(metrics["loss"][0]): print("Loss became nan; stopping.", flush=True) break if step % config.eval_every_steps == 0: eval_metrics = evaluater.get_eval_metrics( state, eval_data_iter, p_eval_step, config) per_level = eval_metrics.pop("per_level_acc") # Depth accuracy drives the controller. Changed-cells-only, so # a slot that merely copies its predecessor scores zero credit. per_depth = eval_metrics.pop("per_slot_acc_changed", {}) per_depth_all = eval_metrics.pop("per_slot_acc", {}) per_stage_inset = eval_metrics.pop("per_stage_inset_acc", {}) if instance_mode: # Repair if a stage's in-set rate falls behind its # graduation value. The candidate head is off. per_depth = per_stage_inset elif not any(v >= 0 for v in per_depth.values()): per_depth = per_depth_all def _m(key): v = eval_metrics.get(key, []) return round(float(np.mean(v)), 4) if len(v) else -1.0 # Seed graduation refs for stages inherited from a checkpoint. if not grad_acc_seeded: for s in range(1, seed_start_stage): acc_s = per_depth.get(s, -1.0) if acc_s >= 0: grad_acc[s] = float(acc_s) grad_acc_seeded = True print(f"[repair] step {step}: seeded graduation refs from " f"resume: " f"{dict((k, round(v, 3)) for k, v in grad_acc.items())}", flush=True) # Soft graduation refs: forgive chronic mild regression. if grad_decay > 0 and grad_acc: for s in list(grad_acc.keys()): cur_s = per_depth.get(s, -1.0) if cur_s >= 0 and cur_s < grad_acc[s]: old = grad_acc[s] grad_acc[s] = ( (1.0 - grad_decay) * grad_acc[s] + grad_decay * float(cur_s)) if step % (config.eval_every_steps * 5) == 0: print(f"[repair] soft-grad depth {s}: " f"{old:.3f}->{grad_acc[s]:.3f} " f"(cur={cur_s:.3f})", flush=True) frontier_acc = per_depth.get(curriculum.stage, -1.0) eff_margin = effective_margin(per_depth, i) deficits = compute_deficits(i, per_depth) elapsed = max(1, step - start_step + 1) repair_frac = repair_steps_total / float(elapsed) bt_frozen = ( freeze_after_step > 0 and step >= freeze_after_step) print(step, "stage", curriculum.stage, "mode", mode, "repair_stage", repair_stage, "target_t", t, "loss", round(float(metrics["loss"].mean()), 4), "ce", round(float(metrics["ce_loss"].mean()), 4), "aux_bce", round(float(metrics["aux_loss"].mean()), 4), "| val_acc", _m("acc"), "loc_acc", _m("loc_acc"), "val|loc", _m("val_given_loc_acc"), "| frontier_acc", round(float(frontier_acc), 4), "eff_margin", round(float(eff_margin), 4), "| cand_bit_acc", _m("cand_bit_acc"), "cand_set_acc", _m("cand_set_acc"), "cand_set_chg", _m("cand_set_acc_changed"), "| grad_acc", {k: round(v, 3) for k, v in grad_acc.items()}, "deficits", {k: round(v, 3) for k, v in deficits.items()}, "step_counts", dict(step_counts), "repair_frac", round(repair_frac, 3), "bt_frozen", bt_frozen, flush=True) with tf_summary_writer.as_default(): tf.summary.scalar("loss", metrics["loss"].mean(), step=step) tf.summary.scalar("ce_loss", metrics["ce_loss"].mean(), step=step) tf.summary.scalar("aux_bce_loss", metrics["aux_loss"].mean(), step=step) tf.summary.scalar("curriculum_stage", curriculum.stage, step=step) tf.summary.scalar("repair_mode", 1 if mode == "repair" else 0, step=step) tf.summary.scalar("repair_stage", repair_stage or 0, step=step) if frontier_acc >= 0: tf.summary.scalar("frontier_depth_acc", frontier_acc, step=step) tf.summary.scalar("eff_repair_margin", eff_margin, step=step) for key in eval_metrics.keys(): tf.summary.scalar("eval_" + key, np.array(eval_metrics[key]).mean(), step=step) for lvl, v in per_level.items(): if v >= 0: tf.summary.scalar(f"eval_acc_level_{lvl}", v, step=step) for s, v in per_depth_all.items(): if v >= 0: tf.summary.scalar(f"eval_cand_depth_{s}", v, step=step) def _save_stage_ckpt(tag): if config.save_checkpoint: unrep = jax_utils.unreplicate(state) checkpoints.save_checkpoint_multiprocess( workdir, unrep, step, keep=ckpt_keep, overwrite=True) checkpoints.save_checkpoint_multiprocess( stage_ckpt_dir, unrep, step, keep=100, overwrite=True, prefix=f"{tag}_") # ---------------- Deficit-driven controller ---------------- # Freeze / duty-cycle: force frontier-only when budget exhausted. duty_ok = ( max_repair_fraction <= 0 or repair_frac < max_repair_fraction) if bt_frozen and mode == "repair": print(f"[repair] step {step}: freeze_after=" f"{freeze_after_step}; leaving repair", flush=True) mode = "frontier" repair_stage = None last_repair_return_step = step if (not duty_ok) and mode == "repair": print(f"[repair] step {step}: duty-cycle cap " f"(repair_frac={repair_frac:.3f}>=" f"{max_repair_fraction}); return to frontier", flush=True) mode = "frontier" repair_stage = None last_repair_return_step = step if mode == "frontier": cooldown_ok = ( min_frontier_steps <= 0 or (step - last_repair_return_step) >= min_frontier_steps) can_repair = (not bt_frozen) and duty_ok and cooldown_ok worst = (most_deficient(i, per_depth, use_margin=eff_margin) if can_repair else None) if worst is not None: mode = "repair" repair_stage = worst repair_started_at = step repair_episode_start = step print(f"[repair] step {step}: enter repair of depth " f"{worst} (snapshot {worst} acc=" f"{per_depth.get(worst, -1.0):.3f} < grad " f"{grad_acc.get(worst, -1.0):.3f} - " f"eff_margin {eff_margin:.3f}; " f"frontier_acc={frontier_acc:.3f})", flush=True) elif bt_frozen and most_deficient( i, per_depth, use_margin=eff_margin) is not None: if step % (config.eval_every_steps * 5) == 0: print(f"[repair] step {step}: deficit present but " f"BT frozen after {freeze_after_step}", flush=True) elif (not duty_ok) and most_deficient( i, per_depth, use_margin=eff_margin) is not None: if step % (config.eval_every_steps * 5) == 0: print(f"[repair] step {step}: deficit present but " f"duty-cycle cap " f"(frac={repair_frac:.3f})", flush=True) elif (not cooldown_ok) and most_deficient( i, per_depth, use_margin=eff_margin) is not None: print(f"[repair] step {step}: deficit present but " f"frontier cooldown " f"({step - last_repair_return_step}/" f"{min_frontier_steps}); staying on stage {i}", flush=True) elif curriculum.stage < curriculum.max_stage: # Normal promotion (same rule as the standard loop). steps_in_stage = step - stage_started_at # A negative accuracy means "not measured this eval": it # must not reset the plateau tracker nor satisfy the # threshold/plateau rule. Only patience fires unmeasured. measured = frontier_acc >= 0 if measured and frontier_acc > stage_best_acc + plateau_delta: stage_best_acc = frontier_acc stage_best_step = step loc_now = _m("loc_acc") loc_ready = (not instance_mode) or loc_now >= promote_loc_threshold hit_threshold = (measured and loc_ready and frontier_acc >= promote_threshold) stalled = (measured and loc_ready and plateau_steps > 0 and (step - stage_best_step) >= plateau_steps) patience_over = steps_in_stage >= promote_patience if steps_in_stage >= min_stage_steps and ( hit_threshold or stalled or patience_over): reason = ("threshold" if hit_threshold else "plateau" if stalled else "patience") # Record this stage's graduation accuracy BEFORE moving on. grad_acc[curriculum.stage] = float(frontier_acc) curriculum.stage += 1 stage_started_at = step stage_best_acc = -1.0 stage_best_step = step print(f"[curriculum] step {step}: promote to stage " f"{curriculum.stage} ({reason}; graduated " f"depth {curriculum.stage - 1} cand-set " f"acc={frontier_acc:.3f})", flush=True) _save_stage_ckpt(f"stage{curriculum.stage}") elif (frontier_target_acc > 0 and 0.0 <= frontier_acc < frontier_target_acc and step % (config.eval_every_steps * 5) == 0): print(f"[frontier] step {step}: depth {i} " f"acc={frontier_acc:.3f} " f"< target {frontier_target_acc:.3f}; " f"keeping frontier priority", flush=True) else: # mode == "repair" r = repair_stage cur = per_depth.get(r, -1.0) ref = grad_acc.get(r, -1.0) recovered = cur >= (ref - margin) capped = (step - repair_started_at) >= max_repair_steps if recovered or capped: why = "recovered" if recovered else "cap" print(f"[repair] step {step}: depth {r} done ({why}; " f"snapshot acc={cur:.3f} vs grad {ref:.3f})", flush=True) _save_stage_ckpt(f"repair{r}") # Re-scan: chain to the next most-deficient stage, or # return to the frontier and credit repair time back to # the frontier stage's patience clock. nxt = None if (not bt_frozen) and duty_ok: nxt = most_deficient( i, per_depth, use_margin=eff_margin) if nxt is not None: repair_stage = nxt repair_started_at = step print(f"[repair] step {step}: chain to stage {nxt}", flush=True) else: mode = "frontier" repair_stage = None last_repair_return_step = step stage_started_at += (step - repair_episode_start) if config.save_checkpoint and step > 0 and step % config.save_every_steps == 0: checkpoints.save_checkpoint_multiprocess( workdir, jax_utils.unreplicate(state), step, keep=ckpt_keep, overwrite=True) if config.save_checkpoint: checkpoints.save_checkpoint_multiprocess( workdir, jax_utils.unreplicate(state), config.max_steps, keep=ckpt_keep, overwrite=True)