run.py

# 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("<eos>", "").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()