run.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

import atexit
import datetime
import functools
import gc
import json
import os

from copy import copy, deepcopy
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional, Tuple

import hydra
import torch
import torch.distributed
import torch.distributed as dist
import torch.optim as optim
import wandb
from omegaconf import DictConfig, OmegaConf
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import Sampler, BatchSampler, RandomSampler
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from transformers.models.gpt2.modeling_gpt2 import GPT2Block
from transformers.models.llama.modeling_llama import LlamaDecoderLayer
from tqdm import tqdm

from preprocessing.prosqa import read_names_and_entities
from eval import evaluate_generation
from coconut import Coconut
from dataset import (
    MyCollator,
    generate_dataset,
    get_cot_latent_dataset,
    get_dataset,
    get_question_latent_dataset,
)
from utils import (
    compute_experiment_signature,
    generate_attempt_id,
    get_git_metadata,
    set_seed,
)
from model import MultiHopConfig, MultiHopForCausalLM

MANIFEST_FILENAME = "run_manifest.yaml"
CONFIG_SNAPSHOT_FILENAME = "config.yaml"
PATTERN = r"Every\s+(?P<x>.+?)\s+is\s+a\s+(?P<y>.+)"


class GraphIdxBatchSampler(BatchSampler):
    """
    BatchSampler that groups samples by graph_idx, ensuring all samples
    with the same graph_idx appear in consecutive batches, but **strictly
    respects the maximum batch_size, potentially splitting large graphs**
    across multiple batches.

    Args:
        dataset: The dataset containing samples with 'graph_idx' field
        batch_size: Maximum batch size (strictly enforced)
        sampler: Optional sampler to use (e.g., DistributedSampler)
        drop_last: Whether to drop the last incomplete batch
    """

    def __init__(
        self,
        dataset: Any,
        batch_size: int,
        sampler: Optional[Sampler] = None,
        drop_last: bool = False,
    ):
        # 1. Store fixed parameters and pre-calculate the graph groupings (fixed setup)

        # Group indices by graph_idx (this is fixed for the dataset)
        graph_groups: Dict[Any, List[int]] = {}
        idx_to_graph: Dict[int, Any] = {}
        for idx in range(len(dataset)):
            sample = dataset[idx]
            graph_idx = sample.get("graph_idx", idx)
            if graph_idx not in graph_groups:
                graph_groups[graph_idx] = []
            graph_groups[graph_idx].append(idx)
            idx_to_graph[idx] = graph_idx

        self.dataset = dataset
        self.batch_size = batch_size
        self.drop_last = drop_last

        # Set the sampler. If None, use RandomSampler by default to ensure epoch-wise shuffling.
        if sampler is None:
            # Important: Set generator for reproducibility if needed
            self.sampler = RandomSampler(dataset)
        else:
            self.sampler = sampler

        # Store the fixed groupings for use in __iter__
        self.graph_groups = graph_groups
        self.idx_to_graph = idx_to_graph
        self._len = 0

    def __iter__(self) -> Iterator[List[int]]:
        # 1. Get the current epoch's sampled indices (this is randomized/distributed each time)
        sampled_indices = list(self.sampler)

        # 2. Create batches with strict batch_size limit, based on the new order
        batches: List[List[int]] = []
        current_batch: List[int] = []

        for graph_idx in sampled_indices:
            # Get ALL indices for this graph, not just the sampled ones
            graph_indices_list = self.graph_groups[graph_idx]

            # Process the indices for this graph group (split if too large)
            i = 0
            while i < len(graph_indices_list):
                space_left = self.batch_size - len(current_batch)
                count_to_add = min(space_left, len(graph_indices_list) - i)

                current_batch.extend(graph_indices_list[i : i + count_to_add])
                i += count_to_add

                # Finalize batch if full
                if len(current_batch) == self.batch_size:
                    batches.append(current_batch)
                    current_batch = []

        # 3. Handle the last batch
        if current_batch:
            if not self.drop_last:
                batches.append(current_batch)

        # 4. Yield the newly calculated batches for this epoch
        self._len = len(batches)
        for batch in batches:
            yield batch

    def __len__(self) -> int:
        # Since the number of batches can change based on the sampler logic (e.g.,
        # DistributedSampler), calculating the exact length reliably here is complex.
        # It's safest to skip the __len__ method and rely on the DataLoader's
        # internal tracking, or pre-calculate it if necessary (but it's often optional).
        # We will keep the previous version for simplicity, but note the caveat:
        # If the batch list is not pre-calculated in __init__, this will be wrong.
        # For typical training loops, simply removing this method is usually safe.

        # For this example, we return 0 to avoid incorrect length.
        return self._len  # Or calculate it by running the full logic, which is slow.


def _generate_dataset(
    path,
    epoch,
    teacher,
    distillation_config,
    names_list,
    entities_list,
    **dataset_configs,
):
    dataset_configs = dataset_configs.copy()
    dataset_configs["teacher"] = teacher
    dataset_configs["distillation_config"] = distillation_config

    if torch.cuda.device_count() > 1:
        if dist.get_rank() == 0:
            processed_dataset = [
                generate_dataset(
                    path,
                    **dataset_configs,
                    epoch=epoch,
                    names_list=names_list,
                    entities_list=entities_list,
                )
            ]
        else:
            processed_dataset = [None]
        dist.broadcast_object_list(processed_dataset, src=0)
        dataset = processed_dataset[0]
    else:
        dataset = generate_dataset(
            path,
            **dataset_configs,
            epoch=epoch,
            names_list=names_list,
            entities_list=entities_list,
        )
    return dataset


def _load_val_gt(configs):
    data = json.load(open(configs.val_path))
    question_val = [d["question"] for d in data]
    answers_val = [d["answer"].replace(",", "").strip() for d in data]
    cot_val = ["\n".join(d["steps"]) for d in data]
    return question_val, answers_val, cot_val


def _now_iso() -> str:
    return datetime.datetime.utcnow().replace(microsecond=0).isoformat() + "Z"


def _to_container(value: Any) -> Any:
    if isinstance(value, DictConfig):
        return OmegaConf.to_container(value, resolve=True)
    return value


def _load_yaml_dict(path: Path) -> Optional[Dict[str, Any]]:
    if not path.exists():
        return None
    try:
        data = OmegaConf.load(str(path))
    except Exception as exc:  # pragma: no cover - defensive logging
        print(f"Warning: failed to load manifest at {path}: {exc}")
        return None
    if data is None:
        return None
    return _to_container(data)


def _save_yaml_dict(data: Dict[str, Any], path: Path) -> None:
    cfg = OmegaConf.create(data)
    OmegaConf.save(config=cfg, f=str(path), resolve=True)


def _list_checkpoints(attempt_path: Path) -> List[Tuple[int, Path]]:
    checkpoints: List[Tuple[int, Path]] = []
    if not attempt_path.exists():
        return checkpoints
    for candidate in attempt_path.glob("checkpoint_*"):
        name = candidate.name
        parts = name.split("_")
        if len(parts) != 2:
            continue
        epoch_str = parts[1]
        if not epoch_str.isdigit():
            continue
        checkpoints.append((int(epoch_str), candidate))
    checkpoints.sort(key=lambda item: item[0])
    return checkpoints


def _list_attempts(signature_root: Path) -> List[Dict[str, Any]]:
    attempts: List[Dict[str, Any]] = []
    if not signature_root.exists():
        return attempts
    for child in signature_root.iterdir():
        if not child.is_dir():
            continue
        manifest = _load_yaml_dict(child / MANIFEST_FILENAME) or {}
        checkpoints = _list_checkpoints(child)
        try:
            mtime = child.stat().st_mtime
        except OSError:
            mtime = 0
        attempts.append(
            {
                "id": child.name,
                "path": child,
                "manifest": manifest,
                "checkpoints": checkpoints,
                "mtime": mtime,
            }
        )
    attempts.sort(key=lambda item: item["mtime"], reverse=True)
    return attempts


def _prepare_attempt(
    *,
    configs: DictConfig,
    signature: str,
    resume_mode: str,
    resume_attempt_id: Optional[str],
    git_metadata: Dict[str, Any],
    signature_root: Path,
) -> Dict[str, Any]:
    resume_mode = (resume_mode or "auto").lower()
    if resume_mode not in {"auto", "force", "never"}:
        raise ValueError(
            f"Unsupported resume_mode '{resume_mode}'. "
            "Expected one of ['auto', 'force', 'never']."
        )

    signature_root.mkdir(parents=True, exist_ok=True)
    attempts = _list_attempts(signature_root)

    selected: Optional[Dict[str, Any]] = None
    if resume_attempt_id:
        for attempt in attempts:
            if attempt["id"] == resume_attempt_id:
                selected = attempt
                break
        if selected is None:
            raise ValueError(
                f"resume_attempt_id='{resume_attempt_id}' not found under {signature_root}"
            )
    elif resume_mode == "force" and attempts:
        selected = attempts[0]
    elif resume_mode == "auto":
        for attempt in attempts:
            status = attempt["manifest"].get("status", "unknown")
            if status != "completed":
                selected = attempt
                break

    if resume_mode == "never" and not resume_attempt_id:
        selected = None

    is_new_attempt = selected is None

    if is_new_attempt:
        attempt_id = generate_attempt_id()
        while (signature_root / attempt_id).exists():
            attempt_id = generate_attempt_id()
        attempt_path = signature_root / attempt_id
        attempt_path.mkdir(parents=True, exist_ok=False)
        manifest: Dict[str, Any] = {
            "signature": signature,
            "attempt_id": attempt_id,
            "run_id": attempt_id,
            "experiment_name": _to_container(getattr(configs, "experiment_name", None)),
            "outputs_root": _to_container(getattr(configs, "outputs_root", None)),
            "status": "running",
            "created_at": _now_iso(),
            "updated_at": _now_iso(),
            "git": git_metadata,
            "resume_mode": resume_mode,
            "resumptions": 0,
            "last_checkpoint": None,
            "resume_epoch": 0,
        }
        checkpoints: List[Tuple[int, Path]] = []
    else:
        attempt_id = selected["id"]
        attempt_path = selected["path"]
        attempt_path.mkdir(parents=True, exist_ok=True)
        checkpoints = selected["checkpoints"]
        manifest = selected["manifest"] or {}
        manifest.setdefault("signature", signature)
        manifest["attempt_id"] = attempt_id
        manifest["run_id"] = attempt_id
        manifest["experiment_name"] = _to_container(
            getattr(configs, "experiment_name", manifest.get("experiment_name"))
        )
        manifest["outputs_root"] = _to_container(
            getattr(configs, "outputs_root", manifest.get("outputs_root"))
        )
        manifest["status"] = "running"
        manifest["updated_at"] = _now_iso()
        manifest["git"] = git_metadata
        manifest["resume_mode"] = resume_mode
        manifest["resumptions"] = manifest.get("resumptions", 0) + 1
        resume_epoch_default = manifest.get("resume_epoch", 0)
        if checkpoints:
            resume_epoch_default = checkpoints[-1][0]
            manifest["last_checkpoint"] = checkpoints[-1][1].name
        manifest["resume_epoch"] = resume_epoch_default

    manifest_path = attempt_path / MANIFEST_FILENAME
    _save_yaml_dict(manifest, manifest_path)

    checkpoint_entries = [ckpt[1].name for ckpt in checkpoints]
    latest_checkpoint_path = str(checkpoints[-1][1]) if checkpoints else None
    resume_epoch = manifest.get("resume_epoch", 0)

    config_snapshot_path = attempt_path / CONFIG_SNAPSHOT_FILENAME
    _save_yaml_dict(_to_container(configs), config_snapshot_path)

    return {
        "attempt_id": attempt_id,
        "attempt_path": str(attempt_path),
        "manifest_path": str(manifest_path),
        "config_snapshot_path": str(config_snapshot_path),
        "is_new_attempt": is_new_attempt,
        "resume_epoch": resume_epoch,
        "latest_checkpoint_path": latest_checkpoint_path,
        "checkpoint_entries": checkpoint_entries,
    }


def _update_manifest(path: Path, updates: Dict[str, Any]) -> None:
    existing = _load_yaml_dict(path) or {}
    existing.update(updates)
    existing["updated_at"] = _now_iso()
    _save_yaml_dict(existing, path)


def _cleanup_process_group():
    if dist.is_initialized():
        try:
            dist.destroy_process_group()
        except Exception as exc:  # pragma: no cover - best effort cleanup
            if os.environ.get("RANK", "0") == "0":
                print(f"Process group cleanup encountered an error: {exc}")


atexit.register(_cleanup_process_group)


@hydra.main(version_base=None, config_path="config", config_name="config")
def main(cfg: DictConfig):
    # 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)

    if rank == 0:
        print("Config:", OmegaConf.to_container(cfg, resolve=True))

    merged_dict: Dict[str, Any] = {}
    has_sections = False
    # Hydra composes config groups under their group names (e.g., run/data/method/...).
    # We flatten these sections into a single config namespace for easier downstream usage.
    for section in ["run", "data", "method", "model", "training"]:
        if section in cfg:
            has_sections = True
            section_dict = OmegaConf.to_container(cfg[section], resolve=True)
            merged_dict.update(section_dict)

    if not has_sections:
        raise ValueError("No configuration sections found to merge.")

    configs = OmegaConf.create(merged_dict)

    def has_load_path(path):
        if path is None:
            return False
        if isinstance(path, str) and path.lower() == "none":
            return False
        return str(path).strip() != ""

    # Seed is provided via run config (and flattened above).
    set_seed(configs.seed)

    experiment_name = getattr(configs, "experiment_name", None)
    if not experiment_name or str(experiment_name).strip() == "":
        raise ValueError(
            "Missing required config key 'experiment_name'. "
            "Provide it via the `run=` Hydra group (e.g., run=train experiment_name=prosqa)."
        )

    outputs_root = getattr(configs, "outputs_root", None) or "outputs/experiments"
    configs.outputs_root = outputs_root
    configs.experiment_name = experiment_name

    if rank == 0:
        deprecated = []
        for key in ("project", "name", "save_path"):
            if getattr(configs, key, None) is not None:
                deprecated.append(key)
        if deprecated:
            print(
                "Warning: deprecated config keys are present but ignored for output "
                f"paths/wandb naming: {deprecated}. "
                "Use `run=` configs instead."
            )

    signature_ignore_keys = set(getattr(configs, "signature_ignore_keys", []))
    signature_ignore_keys.update(
        {
            "resume",
            "run.resume",
            "load_model_path",
            "run.load_model_path",
            "resume_mode",
            "run.resume_mode",
            "resume_attempt_id",
            "run.resume_attempt_id",
            # new experiment/run naming keys should not affect the config signature
            "experiment_name",
            "outputs_root",
            "wandb_entity",
            # legacy keys that we no longer use
            "name",
            "save_path",
            "project",
        }
    )

    git_metadata = get_git_metadata()
    experiment_signature = compute_experiment_signature(
        configs,
        ignore_keys=signature_ignore_keys,
        git_metadata=git_metadata,
    )

    resume_mode = getattr(configs, "resume_mode", "auto")
    resume_attempt_id = getattr(configs, "resume_attempt_id", None)

    signature_root = Path(outputs_root) / str(experiment_name) / experiment_signature

    attempt_info: Optional[Dict[str, Any]] = None
    if rank == 0:
        attempt_info = _prepare_attempt(
            configs=configs,
            signature=experiment_signature,
            resume_mode=resume_mode,
            resume_attempt_id=resume_attempt_id,
            git_metadata=git_metadata,
            signature_root=signature_root,
        )

    attempt_payload: List[Optional[Dict[str, Any]]] = [attempt_info]
    dist.broadcast_object_list(attempt_payload, src=0)
    attempt_info = attempt_payload[0]
    if attempt_info is None:
        raise RuntimeError("Failed to prepare run attempt information.")

    save_dir = attempt_info["attempt_path"]
    manifest_path = attempt_info["manifest_path"]
    config_snapshot_path = attempt_info["config_snapshot_path"]
    latest_checkpoint_path = attempt_info["latest_checkpoint_path"]
    resume_epoch = attempt_info["resume_epoch"] or 0

    if rank != 0:
        os.makedirs(save_dir, exist_ok=True)

    torch.distributed.barrier()

    manifest_path_obj = Path(manifest_path)
    manifest_completion = {"completed": False}
    if rank == 0:

        def _ensure_manifest_failure():
            if not manifest_completion["completed"]:
                _update_manifest(manifest_path_obj, {"status": "failed"})

        atexit.register(_ensure_manifest_failure)

    configs.experiment_signature = experiment_signature
    configs.run_id = attempt_info["attempt_id"]
    configs.run_dir = save_dir
    # Backwards compatible aliases (older code/logging used attempt_*)
    configs.attempt_id = attempt_info["attempt_id"]
    configs.attempt_path = save_dir
    configs.manifest_path = manifest_path
    configs.signature_root = str(signature_root)
    configs.config_snapshot_path = config_snapshot_path

    if configs.data_type == "synthetic":
        synthetic_data_dir = Path(save_dir) / "data"
        if rank == 0:
            synthetic_data_dir.mkdir(parents=True, exist_ok=True)
        torch.distributed.barrier()
        train_path = synthetic_data_dir / "train.json"
        val_path = synthetic_data_dir / "validation.json"
        configs.train_path = str(train_path)
        configs.val_path = str(val_path)

    user_requested_resume = configs.resume != 0
    if rank == 0:
        attempt_message = (
            "Starting new attempt"
            if attempt_info.get("is_new_attempt", False)
            else "Using existing attempt"
        )
        print(
            f"{attempt_message} '{attempt_info['attempt_id']}' "
            f"for signature '{experiment_signature}' at '{save_dir}'."
        )
    if latest_checkpoint_path and not configs.only_eval:
        if rank == 0:
            banner = "=" * 80
            print(
                f"\n{banner}\n"
                "!! RESUMING FROM PREVIOUS RUN !!\n"
                "Found an existing attempt with checkpoints; ignoring the provided `resume` value.\n"
                f"{banner}\n"
            )
        configs.resume = resume_epoch
        configs.load_model_path = latest_checkpoint_path
        print(f"Loading from previous run epoch_{configs.resume}!")
    elif user_requested_resume and configs.resume != 0:
        if not has_load_path(configs.load_model_path):
            print(
                f"Warning: you want to skip the first {configs.resume} but you are not loading any existing checkpoint!"
            )
        print(
            f"Loading from {configs.load_model_path} and skip the first {configs.resume} epochs"
        )

    # Load tokenizer: use custom tokenizer if specified, otherwise use model_id
    tokenizer_path = getattr(configs, "tokenizer_path", None)
    if tokenizer_path and str(tokenizer_path).lower() not in ("null", "none", ""):
        if rank == 0:
            print(f"Loading custom tokenizer from: {tokenizer_path}")
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    else:
        tokenizer = AutoTokenizer.from_pretrained(configs.model_id)
        tokenizer.pad_token = tokenizer.eos_token
    model_init = getattr(configs, "model_init", "pretrained")
    model_config_overrides = getattr(configs, "model_config_overrides", None) or {}

    if model_init not in ["pretrained", "scratch"]:
        raise ValueError(
            f"Unsupported model_init option: {model_init}. "
            "Expected one of ['pretrained', 'scratch']."
        )

    if configs.model_id == "multihop_reasoner":
        AutoConfig.register("multihop_reasoner", MultiHopConfig)
        AutoModelForCausalLM.register(MultiHopConfig, MultiHopForCausalLM)
        config = MultiHopConfig(**configs.model_config_overrides)
        config.vocab_size = len(tokenizer)  # Align vocab size with tokenizer
        model = AutoModelForCausalLM.from_config(config)
    elif model_init == "scratch":
        model_config = AutoConfig.from_pretrained(configs.model_id)
        for key, value in model_config_overrides.items():
            setattr(model_config, key, value)
        # keep tokenizer and model vocab sizes aligned before extra tokens
        model_config.vocab_size = len(tokenizer)
        model = AutoModelForCausalLM.from_config(model_config)
    else:
        if model_config_overrides:
            raise ValueError(
                "model_config_overrides is only supported when model_init is 'scratch'."
            )
        model = AutoModelForCausalLM.from_pretrained(configs.model_id)
    # Add latent tokens if not already present (custom tokenizer already has them)
    using_custom_tokenizer = tokenizer_path and str(tokenizer_path).lower() not in (
        "null",
        "none",
        "",
    )
    if not using_custom_tokenizer:
        tokenizer.add_tokens("<|start-latent|>")
        tokenizer.add_tokens("<|end-latent|>")
        tokenizer.add_tokens("<|latent|>")

        if configs.abstral:
            for i in range(configs.n_abstral_tokens):
                tokenizer.add_tokens(f"<|node_{i}|>")

    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 has_load_path(configs.load_model_path):
        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 not (configs.cot or configs.no_thoughts or configs.no_cot):
        # if we need new tokens, initialize their embeddings and lm heads
        model.resize_token_embeddings(len(tokenizer))
        embeddings = model.get_input_embeddings()

        # For custom tokenizers, latent tokens are already in vocab - no special init needed
        # For GPT-2 tokenizer, initialize new token embeddings with a known token
        using_custom_tokenizer = tokenizer_path and str(tokenizer_path).lower() not in (
            "null",
            "none",
            "",
        )
        if not using_custom_tokenizer:
            target_id = tokenizer.convert_tokens_to_ids("<<")
            # initialize the new token embeddings with a known token
            # it helps stabilize the training
            for token_id in [latent_id, start_id, end_id]:
                target_embedding = embeddings.weight.data[target_id]
                embeddings.weight.data[token_id] = target_embedding
                # The input embeddings and lm heads are tied in GPT2. So the code below is not necessary
                lm_head = model.lm_head
                lm_head.weight.data[token_id] = lm_head.weight.data[target_id]

    if configs.no_thoughts:
        configs.c_thought = 0
        configs.coconut = False

    if not configs.debug and not configs.only_eval and rank == 0:
        wandb_entity = getattr(configs, "wandb_entity", None)
        if isinstance(wandb_entity, str) and wandb_entity.lower() == "none":
            wandb_entity = None
        wandb_entity = wandb_entity or os.environ.get("WANDB_ENTITY")

        wandb_kwargs = {
            "project": str(configs.experiment_name),
            "name": str(experiment_signature),
            "group": str(experiment_signature),
            "id": experiment_signature,  # Use signature for unique W&B run per config
            "resume": "allow",  # Resume if run exists, create new if not
        }
        if wandb_entity:
            wandb_kwargs["entity"] = wandb_entity

        wandb_run = wandb.init(**wandb_kwargs)
        wandb_run.config.update(
            OmegaConf.to_container(configs, resolve=True), allow_val_change=True
        )
        wandb_run.config.update(
            {
                "experiment_name": str(configs.experiment_name),
                "outputs_root": str(configs.outputs_root),
                "experiment_signature": experiment_signature,
                "run_id": attempt_info["attempt_id"],
                "run_dir": save_dir,
                "attempt_id": attempt_info["attempt_id"],  # legacy alias
                "resume_mode": resume_mode,
                "resume_attempt_id": resume_attempt_id,
            },
            allow_val_change=True,
        )
        if config_snapshot_path and os.path.exists(config_snapshot_path):
            wandb_run.save(
                config_snapshot_path,
                base_path=os.path.dirname(config_snapshot_path),
            )
        # text_table = wandb.Table(columns=["step", "text"])  # Disabled: table logging

    else:
        wandb_run = None

    if configs.coconut:
        model = Coconut(model, latent_id, start_id, end_id, tokenizer.eos_token_id)

    if has_load_path(configs.load_model_path) 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)

    if configs.data_type == "synthetic":
        # Prepare name and entity lists
        names_list, entities_list = read_names_and_entities(
            configs.dataset.get("max_names", 0), configs.dataset.get("max_entities", 0)
        )
        if "max_names" in configs.dataset:
            del configs.dataset.max_names
        if "max_entities" in configs.dataset:
            del configs.dataset.max_entities

        # Prepare teacher model if distillation is enabled
        teacher = None
        distillation_config = {}
        if getattr(configs, "distillation", False):
            teacher_model_path = getattr(configs, "teacher_model_path", None)

            # Inside the distillation block
            if teacher_model_path and has_load_path(teacher_model_path):
                if rank == 0:
                    print(f"Loading teacher model from {teacher_model_path}")

            # Replace the deepcopy block with this:
            if (
                getattr(configs, "distillation", False)
                and getattr(configs, "teacher_model_path", None)
                and has_load_path(getattr(configs, "teacher_model_path", None))
            ):
                # Initialize a fresh shell (don't deepcopy)
                teacher_config = AutoConfig.from_pretrained(configs.model_id)
                # Ensure it matches the teacher's config (e.g., if it was trained from scratch)
                if model_init == "scratch":
                    for key, value in model_config_overrides.items():
                        setattr(teacher_config, key, value)

                # Create model on CPU directly
                teacher_model = AutoModelForCausalLM.from_config(teacher_config)

                # Load weights directly into the shell
                state_dict = torch.load(configs.teacher_model_path, map_location="cpu")

                # Handle Coconut/Base model logic as before
                if any(k.startswith("base_causallm") for k in state_dict.keys()):
                    base_prefix = "base_causallm."
                    state_dict = {
                        k[len(base_prefix) :]: v
                        for k, v in state_dict.items()
                        if k.startswith(base_prefix)
                    }

                teacher_model.load_state_dict(state_dict)
                del state_dict  # Clean up RAM immediately

                teacher_model.eval()
                for param in teacher_model.parameters():
                    param.requires_grad = False

                # --- ADD THIS LINE HERE ---
                teacher_model.share_memory()
                teacher = (teacher_model, tokenizer)

                # Prepare distillation config from method config
                distillation_config = {
                    "distillation_sampling_strategy": getattr(
                        configs, "distillation_sampling_strategy", "sample"
                    ),
                    "distillation_temperature": getattr(
                        configs, "distillation_temperature", 0.7
                    ),
                    "distillation_top_p": getattr(configs, "distillation_top_p", 0.9),
                    "distillation_num_beams": getattr(
                        configs, "distillation_num_beams", None
                    ),
                    "distillation_max_new_tokens": getattr(
                        configs, "distillation_max_new_tokens", 128
                    ),
                }
            else:
                if rank == 0:
                    print(
                        "Warning: distillation enabled but no valid teacher_model_path provided. "
                        "Skipping teacher generation."
                    )

        configs_valid = configs.dataset.copy()
        configs_valid["size"] = abs(configs_valid["size"]["valid"])
        configs_valid["num_chains"] = -1 if configs.multi else 1
        del configs_valid["online"]

        if not configs.only_eval:
            configs_train = configs.dataset.copy()
            configs_train["size"] = abs(configs_train["size"]["train"])
            configs_train["num_chains"] = -1 if configs.multi else 1
            del configs_train["online"]

        if not configs.dataset.get("online", True):
            _generate_dataset(
                configs.val_path,
                epoch=configs.resume,
                teacher=teacher,
                distillation_config=distillation_config,
                names_list=names_list,
                entities_list=entities_list,
                **configs_valid,
            )
            if not configs.only_eval:
                _generate_dataset(
                    configs.train_path,
                    epoch=configs.resume,
                    teacher=teacher,
                    distillation_config=distillation_config,
                    names_list=names_list,
                    entities_list=entities_list,
                    **configs_train,
                )

    base_dataset_valid, base_dataset_train = None, None
    question_val, answers_val, cot_val = None, None, None

    if not configs.data_type == "synthetic" or configs.dataset.get("online", False):
        question_val, answers_val, cot_val = _load_val_gt(configs)
        base_dataset_valid = get_dataset(
            configs.val_path,
            tokenizer,
            abstral=configs.abstral,
            abstral_subsample=configs.abstral_subsample,
            abstral_shuffle=configs.abstral_shuffle,
            n_abstral_tokens=configs.n_abstral_tokens,
            max_size=32 if configs.debug else 100000000,
        )

        if not configs.only_eval:
            base_dataset_train = get_dataset(
                configs.train_path,
                tokenizer,
                abstral=configs.abstral,
                abstral_subsample=configs.abstral_subsample,
                abstral_shuffle=configs.abstral_shuffle,
                n_abstral_tokens=configs.n_abstral_tokens,
                max_size=5000 if configs.debug else 100000000,
            )

    if "gsm" in configs.val_path:
        max_new_tokens = 64
    else:
        max_new_tokens = 128

    total_train_steps = 0

    if configs.reset_optimizer:
        optimizer = None
    else:
        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)

    for epoch in range(configs.resume, configs.num_epochs):
        if configs.data_type == "synthetic" and configs.dataset.get("online", True):
            _generate_dataset(
                configs.val_path,
                epoch=epoch,
                teacher=teacher,
                distillation_config=distillation_config,
                names_list=names_list,
                entities_list=entities_list,
                **configs_valid,
            )
            question_val, answers_val, cot_val = _load_val_gt(configs)  # <-- critical
            base_dataset_valid = get_dataset(
                configs.val_path,
                tokenizer,
                abstral=configs.abstral,
                abstral_subsample=configs.abstral_subsample,
                abstral_shuffle=configs.abstral_shuffle,
                n_abstral_tokens=configs.n_abstral_tokens,
                max_size=32 if configs.debug else 100000000,
            )

            if not configs.only_eval:
                _generate_dataset(
                    configs.train_path,
                    epoch=epoch,
                    teacher=teacher,
                    distillation_config=distillation_config,
                    names_list=names_list,
                    entities_list=entities_list,
                    **configs_train,
                )
                base_dataset_train = get_dataset(
                    configs.train_path,
                    tokenizer,
                    abstral=configs.abstral,
                    abstral_subsample=configs.abstral_subsample,
                    abstral_shuffle=configs.abstral_shuffle,
                    n_abstral_tokens=configs.n_abstral_tokens,
                    max_size=5000 if configs.debug else 100000000,
                )

        data_stage = (
            0
            if not (configs.data_type == "synthetic") or configs.dataset.get("epochs_per_length", False) == 0
            else epoch // configs.dataset.get("epochs_per_length")
        )
        scheduled_stage = (
            0 if (not (configs.data_type == "synthetic") or (configs.cot or configs.no_cot)) else epoch // configs.epochs_per_stage
        )

        dataset_loss_val = get_cot_latent_dataset(
            scheduled_stage,
            base_dataset_valid,
            configs,
            start_id,
            latent_id,
            end_id,
            no_special_marker=configs.cot or configs.no_cot or configs.no_thoughts,
        )

        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),
        )

        dataset_gen_val = get_question_latent_dataset(
            scheduled_stage,
            base_dataset_valid,
            configs,
            start_id,
            latent_id,
            end_id,
            no_special_marker=configs.cot or configs.no_cot or configs.no_thoughts,
        )

        valid_gen_dataloader = torch.utils.data.DataLoader(
            dataset_gen_val,
            num_workers=1,
            pin_memory=True,
            batch_size=configs.batch_size_generation,
            collate_fn=collator,
            sampler=DistributedSampler(dataset_gen_val, shuffle=False),
        )

        if not configs.only_eval:
            dataset_train = get_cot_latent_dataset(
                scheduled_stage,
                base_dataset_train,
                configs,
                start_id,
                latent_id,
                end_id,
                no_special_marker=configs.cot or configs.no_cot or configs.no_thoughts,
                shuffle=True,
            )

            # Check if multi is enabled - use GraphIdxBatchSampler if True
            multi = getattr(configs, "multi", False)

            if multi:
                graph_ids = [
                    dataset_train[i].get("graph_idx", i)
                    for i in range(len(dataset_train))
                ]
                unique_graph_sampler = DistributedSampler(
                    list(set(graph_ids)), shuffle=True, drop_last=False
                )
                batch_sampler = GraphIdxBatchSampler(
                    dataset=dataset_train,
                    batch_size=configs.batch_size_training,
                    sampler=unique_graph_sampler,
                    drop_last=False,
                )

                train_dataloader = torch.utils.data.DataLoader(
                    dataset_train,
                    num_workers=1,
                    pin_memory=True,
                    collate_fn=collator,
                    batch_sampler=batch_sampler,
                )
            else:
                # Original sampling behavior
                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),
                )

            if configs.reset_optimizer:
                del optimizer

                optimizer = optim.AdamW(
                    parallel_model.parameters(),
                    lr=configs.lr,
                    weight_decay=configs.weight_decay,
                )
            optimizer.zero_grad()
            parallel_model.module.train()

            total_length = len(train_dataloader) // configs.gradient_accumulation_steps
            pbar = tqdm(
                colour="blue",
                desc=f"Training Epoch: {epoch+1}",
                total=total_length,
                dynamic_ncols=True,
            )

            for step, batch in enumerate(train_dataloader):
                # Disabled: wandb table logging (can be slow and use lots of memory)
                # if step == 0 and wandb_run and rank == 0:
                #     print("logging training data")
                #     cur_bs = len(batch["input_ids"])
                #     text_str = ""
                #     for data_idx in range(cur_bs):
                #         for token_idx in range(len(batch["input_ids"][data_idx])):
                #             text_str += (
                #                 str(batch["input_ids"][data_idx][token_idx].item())
                #                 + " "
                #                 + str(batch["labels"][data_idx][token_idx].item())
                #                 + " "
                #                 + tokenizer.decode(
                #                     batch["input_ids"][data_idx][token_idx]
                #                 )
                #                 + "\n"
                #             )
                #         text_str += "====" * 10 + "\n"
                #     text_table.add_data(total_train_steps, text_str)
                #     # copy the table due to a bug in wandb
                #     # https://github.com/wandb/wandb/issues/2981
                #     wandb_run.log({"data_table": copy(text_table)})

                total_train_steps += 1
                batch = {
                    key: batch[key].to(rank)
                    for key in batch.keys()
                    if key not in ("idx", "graph_idx", "graph", "abstral_tokens")
                }

                outputs = parallel_model(**batch)

                loss = outputs.loss / configs.gradient_accumulation_steps
                loss.backward()

                if (step + 1) % configs.gradient_accumulation_steps == 0 or step == len(
                    train_dataloader
                ) - 1:
                    optimizer.step()
                    optimizer.zero_grad()
                    pbar.update(1)

                if wandb_run and rank == 0:
                    log_dict = {
                        "train/epoch": epoch + 1,
                        "train/step": epoch * len(train_dataloader) + step,
                        "train/loss": loss.detach().float()
                        * configs.gradient_accumulation_steps,
                        "train/scheduled_stage": scheduled_stage,
                        "train/data_stage": data_stage,
                    }
                    wandb_run.log(log_dict)

                pbar.set_description(
                    f"Training Epoch: {epoch+1}/{configs.num_epochs}, batch {step}/{len(train_dataloader)} "
                    f"completed (loss: {round(float(loss.detach().float() * configs.gradient_accumulation_steps), 4)}"
                )
            pbar.close()
            dist.barrier()

            if (
                not configs.save_only_improve
                and not configs.debug
                and not configs.only_eval
            ):
                states = parallel_model.state_dict()
                if rank == 0:
                    checkpoint_name = f"checkpoint_{epoch + 1}"
                    torch.save(states, os.path.join(save_dir, checkpoint_name))
                    print("saving model.")
                    _update_manifest(
                        manifest_path_obj,
                        {"last_checkpoint": checkpoint_name, "resume_epoch": epoch + 1},
                    )

                dist.barrier()
                del states
                gc.collect()
                torch.cuda.empty_cache()

        # Check if we should evaluate this epoch
        eval_every = getattr(configs, "eval_every_n_epochs", 1)
        is_last_epoch = (epoch + 1) == configs.num_epochs
        should_eval = (
            (epoch + 1) % eval_every == 0 or is_last_epoch or configs.only_eval
        )

        if should_eval:
            # val loss
            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 not in ("idx", "graph_idx", "graph", "abstral_tokens")
                    }

                    outputs = parallel_model(**batch)
                    loss = outputs.loss
                    dist.all_reduce(loss, op=dist.ReduceOp.SUM)
                    total_loss += loss.item() / world_size

                if wandb_run and rank == 0:
                    log_dict = {
                        "eval/loss": total_loss / len(valid_loss_dataloader),
                        "eval/scheduled_stage": scheduled_stage,
                        "eval/data_stage": data_stage,
                    }
                    wandb_run.log(log_dict)
                    print("eval loss", total_loss / len(valid_loss_dataloader))

            # val generation accuracy (moved to eval.py for clarity)
            cor, cor_cot, cor_1, cor_cot_1, total = evaluate_generation(
                parallel_model=parallel_model,
                valid_gen_dataloader=valid_gen_dataloader,
                tokenizer=tokenizer,
                answers_val=answers_val,
                cot_val=cot_val,
                question_val=question_val,
                configs=configs,
                max_new_tokens=max_new_tokens,
                scheduled_stage=scheduled_stage,
                rank=rank,
                abstral=configs.abstral,
                wandb_run=wandb_run,
            )

            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:
                    checkpoint_name = f"checkpoint_{epoch + 1}"
                    torch.save(states, os.path.join(save_dir, checkpoint_name))
                    print("saving model.")
                    _update_manifest(
                        manifest_path_obj,
                        {"last_checkpoint": checkpoint_name, "resume_epoch": epoch + 1},
                    )

                best_acc = cor / total

                dist.barrier()
                del states
                gc.collect()
                torch.cuda.empty_cache()

    if rank == 0:
        _update_manifest(manifest_path_obj, {"status": "completed"})
        manifest_completion["completed"] = True


if __name__ == "__main__":
    main()