# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. import torch import torch.distributed import torch.optim as optim from transformers import AutoModelForCausalLM, AutoConfig from stokenizer import STokenizer from graph_metrics import perhop_categorize, category_log_dict, finalonly_categorize import wandb from torch.nn.parallel import DistributedDataParallel as DDP from torch.distributed.fsdp import FullyShardedDataParallel as FSDP import torch.distributed as dist from torch.utils.data.distributed import DistributedSampler from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy from transformers.models.llama.modeling_llama import LlamaDecoderLayer from transformers.models.gpt2.modeling_gpt2 import GPT2Block from coconut import Coconut from dataset import ( MyCollator, get_graph_latent_question_dataset, get_graph_no_latent_question_dataset, get_graph_latent_cot_dataset, get_graph_latent_cot_dataset_backtrack, get_graph_finalonly_dataset, get_graph_no_cot_dataset, get_graph_cot_dataset, ) from tqdm import tqdm import os, sys import time import yaml import json import gc import argparse import functools from utils import Config, set_seed def main(): parser = argparse.ArgumentParser(description="coconut") parser.add_argument("config_file") args = parser.parse_args() # init distributed environment dist.init_process_group("nccl") local_rank = int(os.environ["LOCAL_RANK"]) rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) torch.cuda.set_device(local_rank) # load the configuration file with open(args.config_file) as f: config_dict = yaml.safe_load(f) if rank == 0: print("Config:", config_dict) configs = Config(config_dict) set_seed(configs.seed) save_dir = os.path.join(configs.save_path, configs.name) if not os.path.exists(save_dir) and rank == 0: os.makedirs(save_dir) torch.distributed.barrier() cur_ckpts = os.listdir(save_dir) # check if the job is preempted and resumed. checkpoints = [f for f in cur_ckpts if f.startswith("checkpoint_")] if len(checkpoints) > 0 and not configs.only_eval: # if there are previous checkpoints, and only_eval is False # it means the previous run was preempted and the program is restarted. # need to find the latest checkpoint and resume from that. if rank == 0: print( f"Warning: found previous run and gonna resume from that. the inputted `resume` argument is ignored!" ) checkpoints.sort(key=lambda x: int(x.split("_")[1])) # Get the last item in the sorted list latest_checkpoint = checkpoints[-1] configs.resume = int(latest_checkpoint.split("_")[1]) load_dir = os.path.join(configs.save_path, configs.name, latest_checkpoint) configs.load_model_path = load_dir print(f"Loading from previous run epoch_{configs.resume}!") elif configs.resume != 0: # by setting `resume`, we can skip a few epoches at the beginning. if configs.load_model_path == "None": print( f"Warning: you want to skip the first {configs.resume} but you are not loading any existing checkpoint!" ) # not an intended use case at this point print( f"Loading from {configs.load_model_path} and skip the first {configs.resume} epochs" ) model = AutoModelForCausalLM.from_config( AutoConfig.from_pretrained(configs.model_id) ) print(model) tokenizer = STokenizer() latent_id = tokenizer.convert_tokens_to_ids("<|latent|>") start_id = tokenizer.convert_tokens_to_ids("<|start-latent|>") end_id = tokenizer.convert_tokens_to_ids("<|end-latent|>") loaded = False if configs.load_model_path != "None": saved_weights = torch.load( configs.load_model_path, map_location=torch.device(rank) ) if configs.coconut and not any( [k.startswith("base_causallm") for k in saved_weights.keys()] ): # we are loading a base model into coconut model # e.g., for GSM8k, we used a SFTed model to skip the stage 0 loaded = True print(model.load_state_dict(saved_weights, strict=False)) elif not configs.coconut and any( [k.startswith("base_causallm") for k in saved_weights.keys()] ): raise ValueError("Cannot load coconut model weights into a causallm model") elif configs.coconut and any( [k.startswith("base_causallm") for k in saved_weights.keys()] ): # loading from preempted run # will handle later pass else: # resume or evaluate sft model loaded = True print(model.load_state_dict(saved_weights, strict=False)) if configs.no_thoughts: configs.c_thought = 0 configs.coconut = False if configs.coconut: model = Coconut( model, latent_id, start_id, end_id, tokenizer.eos_token_id, backprop_depth=getattr(configs, "backprop_depth", None), ) if configs.load_model_path != "None" and not loaded: print(model.load_state_dict(saved_weights, strict=False)) print(f"Running FSDP on rank = {rank}, world size = {world_size}") model = model.to(rank) llama_auto_wrap_policy = functools.partial( transformer_auto_wrap_policy, transformer_layer_cls={ # GPT2Block, # for GPT2, we don't need to shard layers (it becomes DDP) LlamaDecoderLayer # only shard llama's layers. }, ) if configs.bf16: model.to(torch.bfloat16) # if only eval, use ddp (to avoid bugs in fsdp) if configs.only_eval: parallel_model = DDP(model, device_ids=[rank]) else: parallel_model = FSDP( model, auto_wrap_policy=llama_auto_wrap_policy, device_id=rank ) del model if rank == 0: print(parallel_model) answers_val = [ d["target"] for d in json.load(open(configs.val_path)) ] if "gsm" in configs.val_path: max_new_tokens = 64 else: max_new_tokens = 128 total_train_steps = 0 if not configs.debug and not configs.only_eval and rank == 0: # Persist a wandb run id in the run dir so a preempted + auto-resumed job # continues the SAME wandb run (one continuous x-axis) instead of opening a # fresh run whose step resets to 0. Wiping the run dir => fresh id => new run. run_dir = os.path.join(configs.save_path, configs.name) os.makedirs(run_dir, exist_ok=True) id_path = os.path.join(run_dir, "wandb_run_id.txt") if os.path.exists(id_path): with open(id_path) as f: wandb_id = f.read().strip() wandb_resume = "allow" else: wandb_id = wandb.util.generate_id() with open(id_path, "w") as f: f.write(wandb_id) wandb_resume = None wandb_run = wandb.init(project=configs.project, name=configs.name, id=wandb_id, resume=wandb_resume) wandb_run.config.update(configs, allow_val_change=True) # Plot epoch-keyed metrics against the (resume-monotonic) training epoch so # eval / train curves align and stitch cleanly across resumes. wandb_run.define_metric("train/step") wandb_run.define_metric("train/epoch") wandb_run.define_metric("train/loss", step_metric="train/step") wandb_run.define_metric("eval/*", step_metric="train/epoch") wandb_run.define_metric("revert/*", step_metric="train/epoch") text_table = wandb.Table(columns=["step", "text"]) else: wandb_run = None optimizer = optim.AdamW( parallel_model.parameters(), lr=configs.lr, weight_decay=configs.weight_decay, ) best_acc = 0 collator = MyCollator(tokenizer, latent_id=latent_id, label_pad_token_id=-100) revert_next_stage = 0 # ---- Backtracking state ------------------------------------------------- # Training is IDENTICAL to the no-backtrack arm except when a previously- # mastered stage regresses below `backtrack_detect_threshold` in the per-hop # eval: `bt_target_stage` is then set to the earliest regressed stage and # training is pointed back at it (same vanilla dataset builder) until it # recovers, after which bt_target_stage returns to None (frontier training). backtrack = getattr(configs, "backtrack", False) bt_detect_threshold = getattr(configs, "backtrack_detect_threshold", 0.9) bt_target_stage = None # None = no regression -> train at the frontier # ---- Accuracy-gated curriculum promotion (vs. fixed epochs-per-stage) ----- # When `accuracy_staging` is on, the latent frontier `cur_stage` only advances # once every stage 1..cur_stage has reached `promote_threshold` (frontier acc # for BFS). Promotion is thus driven by measured accuracy, not by the epoch # counter, and is held whenever an earlier stage regresses (backtracking then # rehearses the regressed stages until they recover). We also record how long # (epochs + wall-clock) each stage took to solve. acc_staging = getattr(configs, "accuracy_staging", False) promote_threshold = getattr(configs, "promote_threshold", bt_detect_threshold) # ---- Loss-gated curriculum promotion ------------------------------------ # When `loss_staging` is on, advance only when the current-stage eval CE loss # falls to <= `promote_loss_threshold`. Pinning is done with # max_latent_stage == init_stage (never promotes). Prefer this over fixed # epochs_per_stage when deeper graphs need longer stage-0 warmup. loss_staging = getattr(configs, "loss_staging", False) promote_loss_threshold = float(getattr(configs, "promote_loss_threshold", 1.5)) cur_stage = int(getattr(configs, "init_stage", 0 if (acc_staging or loss_staging) else 1)) run_start_time = time.time() stage_start_time = run_start_time stage_start_epoch = configs.resume # Two SEPARATE gates (do not conflate): # promote_metric + promote_threshold -> stage i -> i+1 # backtrack_metric + backtrack_detect_threshold -> retrain earlier stage # Legacy `staging_metric` sets BOTH when the new keys are omitted. _default_key = "frontier" if getattr(configs, "bfs_variant", False) else "optimal" _legacy = getattr(configs, "staging_metric", None) or _default_key promote_metric = getattr(configs, "promote_metric", None) or _legacy backtrack_metric = getattr(configs, "backtrack_metric", None) or _legacy # Optional soft deadline: if a stage has not cleared the promote gate after # this many epochs, force-promote anyway. None / <=0 disables (default). max_epochs_per_stage = int(getattr(configs, "max_epochs_per_stage", 0) or 0) if acc_staging and rank == 0: print(f"[acc-stage] accuracy-gated curriculum ON: init_stage={cur_stage} " f"promote=({promote_metric}>={promote_threshold}) " f"backtrack=({backtrack_metric}>={bt_detect_threshold} if BT else off) " f"max_latent_stage={configs.max_latent_stage}" + (f" max_epochs_per_stage={max_epochs_per_stage}" if max_epochs_per_stage > 0 else "") + (f" stage_matched_q={bool(getattr(configs, 'stage_matched_q', False))}" if getattr(configs, "stage_matched_q", False) else "")) if loss_staging and rank == 0: print(f"[loss-stage] loss-gated curriculum ON: init_stage={cur_stage} " f"promote_loss_threshold={promote_loss_threshold} " f"max_latent_stage={configs.max_latent_stage}") for epoch in range(configs.resume, configs.num_epochs): if configs.cot or configs.no_cot: scheduled_stage = 0 elif acc_staging or loss_staging: scheduled_stage = cur_stage elif getattr(configs, "revert_staging", False): scheduled_stage = revert_next_stage else: scheduled_stage = epoch // configs.epochs_per_stage # Gate cheap train/eval-loss prints (and the val-CE forward) to `log_every`. # Gate expensive generation/per-hop eval to `eval_every`. Default 1 = every epoch. log_every = int(getattr(configs, "log_every", 1)) eval_every = int(getattr(configs, "eval_every", 1)) do_log = ( configs.only_eval or ((epoch + 1) % log_every == 0) or (epoch + 1 == configs.num_epochs) or (epoch + 1 == configs.resume + 1) # always log first epoch after resume ) do_eval = ( configs.only_eval or ((epoch + 1) % eval_every == 0) or (epoch + 1 == configs.num_epochs) ) if rank == 0 and do_log: print("scheduled_stage", scheduled_stage) if True: if configs.cot or configs.no_cot: dataset_gen_val = get_graph_no_latent_question_dataset( configs.val_path, configs, tokenizer, ) else: dataset_gen_val = get_graph_latent_question_dataset( configs.val_path, scheduled_stage, configs, tokenizer, ) valid_gen_dataloader = torch.utils.data.DataLoader( dataset_gen_val, num_workers=1, pin_memory=True, batch_size=1, collate_fn=collator, sampler=DistributedSampler(dataset_gen_val, shuffle=False), ) if not configs.only_eval: if configs.cot: dataset_train = get_graph_cot_dataset( configs.train_path, configs, tokenizer, ) elif configs.no_cot: dataset_train = get_graph_no_cot_dataset( configs.train_path, configs, tokenizer, ) elif getattr(configs, "final_only", False): dataset_train = get_graph_finalonly_dataset( configs.train_path, scheduled_stage, configs, tokenizer, ) elif backtrack: # Backtracking = identical training to the no-backtrack arm, EXCEPT # when a previously-mastered stage has regressed (bt_target_stage # set from the per-hop eval): then train at that earlier stage until # it recovers, after which training returns to the frontier. Uses # the exact same vanilla dataset builder as the control arm. _train_stage = ( scheduled_stage if bt_target_stage is None else bt_target_stage ) dataset_train = get_graph_latent_cot_dataset( configs.train_path, _train_stage, configs, tokenizer, ) if rank == 0 and bt_target_stage is not None: print(f"[backtrack] RETRAIN stage {bt_target_stage} " f"(frontier={scheduled_stage}, thr={bt_detect_threshold})") else: dataset_train = get_graph_latent_cot_dataset( configs.train_path, scheduled_stage, configs, tokenizer, ) train_dataloader = torch.utils.data.DataLoader( dataset_train, num_workers=1, shuffle=False, pin_memory=True, batch_size=configs.batch_size_training, collate_fn=collator, sampler=DistributedSampler(dataset_train, shuffle=True), ) # the sampler is deterministic even if shuffle is set to True # so we have shuffled the dataset when it's constructed (at every epoch). if configs.cot: dataset_loss_val = get_graph_cot_dataset( configs.val_path, configs, tokenizer, ) elif configs.no_cot: dataset_loss_val = get_graph_no_cot_dataset( configs.val_path, configs, tokenizer, ) elif getattr(configs, "final_only", False): dataset_loss_val = get_graph_finalonly_dataset( configs.val_path, scheduled_stage, configs, tokenizer, ) else: dataset_loss_val = get_graph_latent_cot_dataset( configs.val_path, scheduled_stage, configs, tokenizer, ) valid_loss_dataloader = torch.utils.data.DataLoader( dataset_loss_val, num_workers=1, shuffle=False, pin_memory=True, batch_size=configs.batch_size_training, collate_fn=collator, sampler=DistributedSampler(dataset_loss_val, shuffle=False), ) if configs.reset_optimizer and scheduled_stage < configs.max_latent_stage: del optimizer optimizer = optim.AdamW( parallel_model.parameters(), lr=configs.lr, weight_decay=configs.weight_decay, ) parallel_model.module.train() # Epoch-level logging only (no per-batch tqdm / print — those blow up logs). epoch_loss_sum = 0.0 epoch_loss_n = 0 for step, batch in enumerate(train_dataloader): # NOTE: removed per-epoch "logging training data" dump. It was not # loading the dataset — only pretty-printing batch-0 tokens into a # wandb Table that was never logged (wandb_run.log commented out), # and it spammed the log every epoch. total_train_steps += 1 batch = { key: batch[key].to(rank) for key in batch.keys() if key != "idx" } outputs = parallel_model(**batch) loss = outputs.loss / configs.gradient_accumulation_steps loss.backward() epoch_loss_sum += float( (loss.detach() * configs.gradient_accumulation_steps).float().item() ) epoch_loss_n += 1 if (step + 1) % configs.gradient_accumulation_steps == 0 or step == len( train_dataloader ) - 1: # Linear LR warmup over the first `warmup_steps` optimizer steps # (stabilizes the start; L20 long sequences diverged without it). _warmup = getattr(configs, "warmup_steps", 0) if _warmup and total_train_steps <= _warmup: _scale = total_train_steps / max(1, _warmup) for _pg in optimizer.param_groups: _pg["lr"] = configs.lr * _scale # Gradient clipping to prevent the divergence seen at L20. # NOTE: under FSDP the params are sharded, so the plain # torch.nn.utils.clip_grad_norm_ computes the norm over only the # local shard and effectively does not clip. FSDP provides its own # clip_grad_norm_ that all-reduces the global norm across ranks. _clip = getattr(configs, "grad_clip", 0.0) if _clip and _clip > 0: if isinstance(parallel_model, FSDP): parallel_model.clip_grad_norm_(_clip) else: torch.nn.utils.clip_grad_norm_( parallel_model.parameters(), _clip ) optimizer.step() optimizer.zero_grad() # Train/eval-loss logging throttled by `log_every` (still train every epoch). _tl = torch.tensor( [epoch_loss_sum, float(epoch_loss_n)], device=rank, dtype=torch.float64 ) dist.all_reduce(_tl, op=dist.ReduceOp.SUM) avg_train_loss = (_tl[0] / _tl[1]).item() if _tl[1] > 0 else float("nan") if do_log and rank == 0: print( f"train epoch {epoch+1}/{configs.num_epochs} " f"stage={scheduled_stage} loss={avg_train_loss:.4f}" ) if wandb_run: wandb_run.log({ "train/epoch": epoch + 1, "train/loss": avg_train_loss, "train/scheduled_stage": scheduled_stage, }) dist.barrier() if ( not configs.save_only_improve and not configs.debug and not configs.only_eval ): # Optional cadence: save_every=N keeps every Nth epoch (+ always epoch 1). # Default 1 preserves previous "save every epoch" behaviour. _save_every = int(getattr(configs, "save_every", 1)) if _save_every <= 1 or (epoch + 1) == 1 or (epoch + 1) % _save_every == 0: states = parallel_model.state_dict() if rank == 0: torch.save( states, os.path.join(save_dir, f"checkpoint_{epoch + 1}") ) print("saving model.") dist.barrier() del states gc.collect() torch.cuda.empty_cache() # val loss (only on log epochs — skip the forward the rest of the time) if do_log: total_loss = 0 with torch.no_grad(): parallel_model.module.eval() for step, batch in enumerate(valid_loss_dataloader): batch = { key: batch[key].to(rank) for key in batch.keys() if key != "idx" } outputs = parallel_model(**batch) loss = outputs.loss dist.all_reduce(loss, op=dist.ReduceOp.SUM) total_loss += loss.item() / world_size avg_eval_loss = total_loss / len(valid_loss_dataloader) if rank == 0: print("eval loss", avg_eval_loss) if wandb_run: wandb_run.log({ "eval/loss": avg_eval_loss, "eval/scheduled_stage": scheduled_stage, "train/epoch": epoch + 1, }) # ---- Loss-gated promotion (on log epochs; uses cheap val CE) ---- if loss_staging: _now = time.time() if (avg_eval_loss <= promote_loss_threshold and cur_stage < configs.max_latent_stage): if rank == 0: print( f"[loss-stage] PROMOTE stage {cur_stage} -> {cur_stage + 1} " f"| eval_loss={avg_eval_loss:.4f} <= {promote_loss_threshold} " f"in {epoch + 1 - stage_start_epoch} epochs / " f"{_now - stage_start_time:.0f}s | total {_now - run_start_time:.0f}s" ) if wandb_run: wandb_run.log({ "loss_stage/solved_stage": cur_stage, "loss_stage/stage_epochs": epoch + 1 - stage_start_epoch, "loss_stage/stage_time_s": _now - stage_start_time, "loss_stage/cur_stage": cur_stage + 1, "train/epoch": epoch + 1, }) cur_stage += 1 stage_start_time = _now stage_start_epoch = epoch + 1 elif rank == 0: _reason = ( f"pinned (max_latent_stage={configs.max_latent_stage})" if cur_stage >= configs.max_latent_stage else f"eval_loss={avg_eval_loss:.4f} > {promote_loss_threshold}" ) print( f"[loss-stage] HOLD at stage {cur_stage} ({_reason}) " f"| {epoch + 1 - stage_start_epoch} epochs / " f"{_now - stage_start_time:.0f}s in stage" ) if wandb_run: wandb_run.log({ "loss_stage/cur_stage": cur_stage, "loss_stage/eval_loss": avg_eval_loss, "train/epoch": epoch + 1, }) # if scheduled_stage >= configs.max_latent_stage: if do_eval: # val generation accuracy total_length = len(valid_gen_dataloader) cor, cor_cot, total = ( torch.tensor(0, device=rank), torch.tensor(0, device=rank), torch.tensor(0, device=rank), ) with torch.no_grad(): parallel_model.module.eval() for idx, batch in enumerate(valid_gen_dataloader): test_idx = batch["idx"][0] batch = { k: v.to(rank) for k, v in batch.items() if v != None and k not in ["idx", "position_ids"] } # https://github.com/huggingface/transformers/issues/32492 assert len(batch["input_ids"]) == 1 answer = str(answers_val[test_idx.cpu().item()]) # answer_cot = cot_val[test_idx.cpu().item()] # question = question_val[test_idx.cpu().item()] total += 1 # synced_gpus=True in FSDP mode, as we need to keep # forward pass the same on each device if configs.cot: outputs = parallel_model.module.generate( **batch, max_new_tokens=64, synced_gpus=not configs.only_eval, eos_token_id=tokenizer.eos_token_id, ) elif configs.no_cot: outputs = parallel_model.module.generate( **batch, max_new_tokens=64, synced_gpus=not configs.only_eval, eos_token_id=tokenizer.eos_token_id, ) else: outputs = parallel_model.module.generate( **batch, max_new_tokens=1, synced_gpus=not configs.only_eval, eos_token_id=tokenizer.eos_token_id, ) text_output = tokenizer.decode(outputs[0], skip_special_tokens=True).replace("", "").strip() answer_output = text_output.split("[A]")[-1].replace(",", "").strip() cot_output = ( ("\n".join(text_output.split("\n")[1:])).split("#")[0].strip() ) if idx < 5 and rank == 0: # print some examples print( f"Question {test_idx}: Answer = '{answer}'" ) print(f"Full output: '{tokenizer.decode(outputs[0])}'") print(f"Extracted Output: '{answer_output}'") cor += answer_output == answer # cor_cot += cot_output == answer_cot if rank == 0: print(f"Device {rank}: Cor={cor}, Total={total}") dist.all_reduce(cor_cot, op=dist.ReduceOp.SUM) dist.all_reduce(cor, op=dist.ReduceOp.SUM) dist.all_reduce(total, op=dist.ReduceOp.SUM) # cor_cot = cor_cot.item() cor = cor.item() total = total.item() if rank == 0: print(f"Accuracy on validation set: {cor} / {total} = {cor/total}") # print(f"CoT match on validation set: {cor_cot} / {total} = {cor_cot/total}") sys.stdout.flush() if wandb_run: wandb_run.log({"eval/acc": cor / total, "train/epoch": epoch + 1}) if not configs.only_eval and not (configs.cot or configs.no_cot): if getattr(configs, "final_only", False): eval_cats = finalonly_categorize(parallel_model, configs.val_path, tokenizer, collator, rank) train_cats = finalonly_categorize(parallel_model, configs.train_path, tokenizer, collator, rank, max_samples=getattr(configs, "perhop_train_samples", 256)) if rank == 0: if wandb_run: log_cat = category_log_dict("eval", eval_cats, "acc") log_cat.update(category_log_dict("train", train_cats, "acc")) log_cat["eval/scheduled_stage"] = scheduled_stage log_cat["train/epoch"] = epoch + 1 wandb_run.log(log_cat) print("final-only per-depth:", {k: {m: round(v, 3) for m, v in c.items()} for k, c in eval_cats.items()}) if getattr(configs, "revert_staging", False): _accs = {k: c[getattr(configs, "revert_metric", "acc")] for k, c in eval_cats.items()} _thr = getattr(configs, "revert_threshold", 0.9) revert_next_stage = next((k - 1 for k in sorted(_accs) if _accs[k] < _thr), configs.max_latent_stage) if rank == 0: print(" -> revert: next scheduled_stage", revert_next_stage) if wandb_run: wandb_run.log({"revert/stage": revert_next_stage, "train/epoch": epoch + 1}) sys.stdout.flush() else: # promote_metric / backtrack_metric are independent. All of # frontier / optimal / superposition / ce_score are always # computed and logged; only the chosen keys gate decisions. _smq = bool(getattr(configs, "stage_matched_q", False)) eval_cats = perhop_categorize( parallel_model, configs.val_path, tokenizer, collator, rank, max_samples=getattr(configs, "perhop_val_samples", None), stage_matched_q=_smq, ) train_cats = perhop_categorize( parallel_model, configs.train_path, tokenizer, collator, rank, max_samples=getattr(configs, "perhop_train_samples", 256), stage_matched_q=_smq, ) if rank == 0: if wandb_run: log_cat = category_log_dict("eval", eval_cats, promote_metric) log_cat.update(category_log_dict("train", train_cats, promote_metric)) log_cat["eval/scheduled_stage"] = scheduled_stage _m2i = {"frontier": 0, "optimal": 1, "superposition": 2, "ce_score": 3} log_cat["eval/promote_metric"] = _m2i.get(promote_metric, -1) log_cat["eval/backtrack_metric"] = _m2i.get(backtrack_metric, -1) log_cat["train/epoch"] = epoch + 1 wandb_run.log(log_cat) # Compact one-liner by default (full dicts make morning logs # unreadable). Set eval_print_full: True to dump everything. _hops = sorted(eval_cats) _fr = " ".join(f"{k}:{eval_cats[k]['frontier']:.2f}" for k in _hops) _ce = " ".join(f"{k}:{eval_cats[k]['ce_score']:.2f}" for k in _hops) print(f"eval (prom={promote_metric}@{promote_threshold} " f"bt={backtrack_metric}@{bt_detect_threshold}) " f"frontier=[{_fr}] ce_score=[{_ce}]") if getattr(configs, "eval_print_full", False): print("eval per-hop full:", {k: {m: round(v, 3) for m, v in c.items()} for k, c in eval_cats.items()}) if getattr(configs, "revert_staging", False): _accs = {k: c[getattr(configs, "revert_metric", "frontier")] for k, c in eval_cats.items()} _thr = getattr(configs, "revert_threshold", 0.9) revert_next_stage = next((k - 1 for k in sorted(_accs) if _accs[k] < _thr), configs.max_latent_stage) if rank == 0: print(" -> revert: next scheduled_stage", revert_next_stage) if wandb_run: wandb_run.log({"revert/stage": revert_next_stage, "train/epoch": epoch + 1}) if backtrack: # BACKTRACK gate (independent of promote): check mastered # hops with backtrack_metric. Earliest hop below # bt_detect_threshold -> retrain stage (hop-1). _accs_bt = {k: c[backtrack_metric] for k, c in eval_cats.items()} _regressed_hop = next( (k for k in sorted(_accs_bt) if k <= cur_stage and _accs_bt[k] < bt_detect_threshold), None, ) bt_target_stage = ( None if _regressed_hop is None else _regressed_hop - 1 ) if rank == 0: _mastered = [_accs_bt[k] for k in sorted(_accs_bt) if k <= cur_stage] _minacc = min(_mastered) if _mastered else 1.0 print(f" -> backtrack[{backtrack_metric}>={bt_detect_threshold}]: " f"target_stage={bt_target_stage} " f"min_mastered={_minacc:.3f}") if wandb_run: wandb_run.log({ "backtrack/target_stage": -1 if bt_target_stage is None else bt_target_stage, "backtrack/min_mastered": _minacc, "train/epoch": epoch + 1, }) # ---- PROMOTE gate (independent of backtrack) -------------- # Advance stage i -> i+1 when promote_metric clears # promote_threshold. Default: require hops 1..cur_stage+1 # (retention on the PROMOTE metric). With # promote_on_current_only: only hop cur_stage+1. if acc_staging: _accs_p = {k: c[promote_metric] for k, c in eval_cats.items()} if getattr(configs, "promote_on_current_only", False): _within = [_accs_p.get(cur_stage + 1, 0.0)] else: _within = [_accs_p[k] for k in sorted(_accs_p) if 1 <= k <= cur_stage + 1] _min_within = min(_within) if _within else 0.0 _now = time.time() _stage_epochs = epoch + 1 - stage_start_epoch _acc_ok = _min_within >= promote_threshold _force_ok = ( max_epochs_per_stage > 0 and _stage_epochs >= max_epochs_per_stage ) if (_acc_ok or _force_ok) and cur_stage < configs.max_latent_stage: _how = ( f"promote[{promote_metric}] min hops 1..{cur_stage + 1} = " f"{_min_within:.3f} >= {promote_threshold}" if _acc_ok else f"FORCE after {_stage_epochs} epochs " f"(promote[{promote_metric}]={_min_within:.3f} " f"< {promote_threshold}, max_epochs_per_stage={max_epochs_per_stage})" ) if rank == 0: print(f"[acc-stage] PROMOTE stage {cur_stage} -> {cur_stage + 1} " f"| {_how} in " f"{_stage_epochs} epochs / " f"{_now - stage_start_time:.0f}s | total {_now - run_start_time:.0f}s") if wandb_run: wandb_run.log({ "acc_stage/solved_stage": cur_stage, "acc_stage/stage_epochs": _stage_epochs, "acc_stage/stage_time_s": _now - stage_start_time, "acc_stage/total_time_s": _now - run_start_time, "acc_stage/cur_stage": cur_stage + 1, "acc_stage/force_promote": int(not _acc_ok), "train/epoch": epoch + 1, }) cur_stage += 1 stage_start_time = _now stage_start_epoch = epoch + 1 # Always snapshot on promote — these are the warm-start # points for later W=1 (single-latent BPTT) transfer runs. if not configs.debug and not configs.only_eval: states = parallel_model.state_dict() if rank == 0: _p = os.path.join( save_dir, f"checkpoint_{epoch + 1}" ) torch.save(states, _p) print(f"saving model (promote -> stage {cur_stage}).") dist.barrier() del states gc.collect() torch.cuda.empty_cache() else: if rank == 0: print(f"[acc-stage] HOLD at stage {cur_stage} " f"(promote[{promote_metric}] min hops 1..{cur_stage + 1} = " f"{_min_within:.3f} < {promote_threshold}) " f"| {_stage_epochs} epochs / {_now - stage_start_time:.0f}s in stage") if wandb_run: wandb_run.log({"acc_stage/cur_stage": cur_stage, "acc_stage/min_within_acc": _min_within, "train/epoch": epoch + 1}) sys.stdout.flush() if configs.only_eval: break dist.barrier() if ( cor / total > best_acc and configs.save_only_improve and not configs.debug and not configs.only_eval ): states = parallel_model.state_dict() if rank == 0: torch.save(states, os.path.join(save_dir, f"checkpoint_{epoch + 1}")) print("saving model.") best_acc = cor / total dist.barrier() del states gc.collect() torch.cuda.empty_cache() if __name__ == "__main__": main()