social-data-attribution-w2/src/unlearning/train.py

# pyright: reportPrivateImportUsage=false, reportCallIssue=false, reportArgumentType=false
"""
Entry point for NGDiff unlearning.

Usage:
    PYTHONPATH=src .venv/bin/python -m unlearning.train \
        topic_bin=entertainment \
        trainer.max_steps=10 \
        data.max_forget_docs=20
"""

from __future__ import annotations

import json
import logging
import math
import os
import random
import sys

import hydra
import torch
from omegaconf import DictConfig, OmegaConf
from peft import LoraConfig, get_peft_model
from safetensors import safe_open
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from transformers.trainer_utils import get_last_checkpoint

from unlearning.data.dolma_pool import build_forget_retain_loaders
from unlearning.data.retain_pool import RetainPool, RetainPoolConfig
from unlearning.trainer import TRAINER_REGISTRY
from unlearning.trainer.cooldown import CooldownConfig
from unlearning.trainer.utils import (
    cross_entropy_loss,
    shuffle_segments,
    shuffle_tokens,
)

logger = logging.getLogger(__name__)


def _checkpoint_step(path: str) -> int:
    name = os.path.basename(path.rstrip(os.sep))
    try:
        return int(name.removeprefix("checkpoint-"))
    except ValueError:
        return -1


def _valid_lora_checkpoint(path: str) -> bool:
    required = (
        "trainer_state.json",
        "adapter_config.json",
        "adapter_model.safetensors",
    )
    missing = [name for name in required if not os.path.exists(os.path.join(path, name))]
    if missing:
        logger.warning("Ignoring incomplete checkpoint %s; missing %s", path, missing)
        return False
    adapter_path = os.path.join(path, "adapter_model.safetensors")
    try:
        with safe_open(adapter_path, framework="pt", device="cpu") as handle:
            list(handle.keys())
    except Exception as exc:
        logger.warning("Ignoring unreadable checkpoint %s: %s", path, exc)
        return False
    return True


def _get_last_valid_checkpoint(output_dir: str) -> str | None:
    last_checkpoint = get_last_checkpoint(output_dir)
    if last_checkpoint is None:
        return None
    candidates = [
        os.path.join(output_dir, name)
        for name in os.listdir(output_dir)
        if name.startswith("checkpoint-")
        and os.path.isdir(os.path.join(output_dir, name))
    ]
    for checkpoint in sorted(candidates, key=_checkpoint_step, reverse=True):
        if _valid_lora_checkpoint(checkpoint):
            return checkpoint
    logger.warning("No valid checkpoint found in %s; starting fresh.", output_dir)
    return None


def _load_mmlu_samples(tokenizer, n_questions=100):
    try:
        from datasets import load_dataset

        logger.info(
            "Loading MMLU subset (%d questions) for safety guard ...", n_questions
        )
        ds = load_dataset(
            "cais/mmlu",
            "all",
            split=f"validation[:{n_questions}]",
            trust_remote_code=True,
        )
    except Exception as e:
        logger.warning("Could not load MMLU for safety guard: %s", e)
        return [], None

    choices = ["A", "B", "C", "D"]
    choice_token_ids = torch.tensor(
        [tokenizer.encode(f" {c}", add_special_tokens=False)[0] for c in choices],
        dtype=torch.long,
    )
    samples = []
    for row in ds:
        prompt = (
            f"Question: {row['question']}\n"
            f"A. {row['choices'][0]}\nB. {row['choices'][1]}\n"
            f"C. {row['choices'][2]}\nD. {row['choices'][3]}\nAnswer:"
        )
        enc = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
        samples.append((enc["input_ids"], enc["attention_mask"], row["answer"]))
    return samples, choice_token_ids


def _compute_shuffled_ppl(
    model,
    forget_eval_loader,
    device,
    max_batches: int = 20,
    shuffle_mode: str = "segments",
    n_factor: int = 10,
) -> float:
    """Compute base model PPL on shuffled forget-set tokens.

    shuffle_mode="segments" uses GRACE-style segment shuffle (larger segments).
    shuffle_mode="chunks" uses the original NGDiff chunk shuffle.
    """
    pad_id = getattr(model.config, "pad_token_id", 1) or 1
    model.eval()
    total_loss, count = 0.0, 0
    with torch.no_grad():
        for i, batch in enumerate(forget_eval_loader):
            if i >= max_batches:
                break
            input_ids = batch["input_ids"]
            if shuffle_mode == "segments":
                input_ids = shuffle_segments(input_ids, pad_id, n_factor)
            else:
                input_ids = shuffle_tokens(input_ids, pad_id)
            input_ids = input_ids.to(device)
            attention_mask = batch["attention_mask"].to(device)
            out = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
            )
            loss = cross_entropy_loss(out.logits, input_ids)
            total_loss += loss.item()
            del out, loss
            count += 1
    model.train()
    if count == 0:
        return float("inf")
    ppl = math.exp(total_loss / count)
    logger.info(
        "Base model shuffled-token PPL (forget set, mode=%s): %.4f",
        shuffle_mode,
        ppl,
    )
    return ppl


def _compute_mmlu_baseline(model, samples, choice_tokens, device) -> float:
    if not samples or choice_tokens is None:
        return float("nan")
    choice_tokens = choice_tokens.to(device)
    model.eval()
    correct = 0
    with torch.no_grad():
        for input_ids, attn_mask, answer_idx in samples:
            out = model(
                input_ids=input_ids.to(device),
                attention_mask=attn_mask.to(device),
            )
            last_logits = out.logits[0, -1, :]
            if last_logits[choice_tokens].argmax().item() == answer_idx:
                correct += 1
    model.train()
    acc = correct / len(samples)
    logger.info("Original model MMLU accuracy (n=%d): %.4f", len(samples), acc)
    return acc


@hydra.main(config_path="configs", config_name="train", version_base=None)
def main(cfg: DictConfig) -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    logger.info("Config:\n%s", OmegaConf.to_yaml(cfg))

    if cfg.get("wandb_project"):
        os.environ.setdefault("WANDB_PROJECT", cfg.wandb_project)
    if cfg.get("wandb_entity"):
        os.environ.setdefault("WANDB_ENTITY", cfg.wandb_entity)

    seed = cfg.get("seed", 42)
    random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    if cfg.get("topic_bins", None) is not None:
        target_topics = list(cfg.topic_bins)
        topic_label = "+".join(target_topics)
    else:
        target_topics = cfg.topic_bin
        topic_label = str(cfg.topic_bin)
    logger.info("Target topic(s): %s", topic_label)

    model_id = cfg.model.model_id
    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    output_dir = cfg.output_dir
    os.makedirs(output_dir, exist_ok=True)

    retain_topics = cfg.get("retain_topics", None)
    if retain_topics is not None:
        logger.info("Retain restricted to topic(s): %s", retain_topics)

    train_loader, forget_eval_loader, retain_eval_loader, filtered_ds = (
        build_forget_retain_loaders(
            target_topic=target_topics,
            tokenizer=tokenizer,
            batch_size=cfg.trainer.per_device_train_batch_size,
            max_forget_docs=cfg.data.get("max_forget_docs", None),
            max_retain_docs=cfg.data.get("max_retain_docs", 9_000),
            docs_per_retain_bin=cfg.data.get("docs_per_retain_bin", None),
            min_tokens=cfg.data.get("min_tokens", 0),
            max_length=cfg.data.get("max_length", 2048),
            num_workers=cfg.data.get("num_workers", 0),
            seed=seed,
            retain_topics=retain_topics,
            output_dir=output_dir,
            forget_manifest_path=cfg.data.get("forget_manifest_path", None),
            forget_texts_path=(
                cfg.data.get("forget_texts_path", None)
                or cfg.get("forget_texts_file", None)
            ),
        )
    )

    resample_interval = cfg.data.get("resample_interval", 0)
    if resample_interval > 0 and filtered_ds is not None:
        topics_set = set(
            target_topics
            if isinstance(target_topics, list)
            else [t.strip() for t in str(target_topics).split(",")]
        )
        retain_pool = RetainPool(
            ds=filtered_ds,
            tokenizer=tokenizer,
            target_topics=topics_set,
            config=RetainPoolConfig(
                docs_per_bin=cfg.data.get("docs_per_retain_bin", 1_000),
                resample_interval=resample_interval,
                max_length=cfg.data.get("max_length", 2048),
            ),
            base_seed=seed,
        )
        logger.info(
            "RetainPool: resample every %d optimizer steps",
            resample_interval,
        )
    else:
        retain_pool = None

    last_checkpoint = _get_last_valid_checkpoint(output_dir)
    thresh_path = os.path.join(output_dir, "ppl_stopping_threshold.json")
    if last_checkpoint:
        state_path = os.path.join(last_checkpoint, "trainer_state.json")
        if os.path.exists(state_path):
            with open(state_path) as f:
                saved_state = json.load(f)
            if saved_state.get("global_step", 0) >= cfg.trainer.max_steps:
                logger.info(
                    "Training already complete at step %d. Exiting.",
                    saved_state["global_step"],
                )
                return
        if os.path.exists(thresh_path) and os.path.exists(
            os.path.join(output_dir, "forget_ppl_log.json"),
        ):
            with open(thresh_path) as f:
                saved_thresh = json.load(f).get("threshold")
            with open(os.path.join(output_dir, "forget_ppl_log.json")) as f:
                ppl_hist = json.load(f)
            if saved_thresh is not None and ppl_hist:
                max_ppl = max(float(v) for v in ppl_hist.values())
                if math.isinf(max_ppl) or max_ppl >= saved_thresh:
                    logger.info(
                        "PPL stopping was triggered in a previous job. Exiting.",
                    )
                    return
        logger.info("Will resume training from: %s", last_checkpoint)
    else:
        logger.info("Starting fresh training in: %s", output_dir)

    logger.info("Loading model %s ...", model_id)
    dtype_map = {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }
    torch_dtype = dtype_map.get(
        cfg.model.get("torch_dtype", "bfloat16"),
        torch.bfloat16,
    )
    model_kwargs: dict = {
        "torch_dtype": torch_dtype,
        "device_map": "auto",
        "trust_remote_code": True,
        "low_cpu_mem_usage": False,
    }
    if attn_impl := cfg.model.get("attn_implementation", None):
        model_kwargs["attn_implementation"] = attn_impl
    model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)

    device = next(model.parameters()).device
    shuffle_mode = cfg.trainer.get("ppl_shuffle_mode", "segments")
    n_factor = cfg.trainer.get("ppl_n_factor", 10)

    if last_checkpoint:
        logger.info("Resuming. Skipping pre-training baseline eval.")
        mmlu_samples, mmlu_choice_tokens, mmlu_baseline = [], None, float("nan")
        if os.path.exists(thresh_path):
            with open(thresh_path) as f:
                ppl_stopping_threshold = json.load(f).get("threshold")
            logger.info("Restored PPL threshold: %.4f", ppl_stopping_threshold)
        else:
            ppl_stopping_threshold = None
    else:
        mmlu_samples, mmlu_choice_tokens = _load_mmlu_samples(tokenizer)
        mmlu_baseline = _compute_mmlu_baseline(
            model,
            mmlu_samples,
            mmlu_choice_tokens,
            device,
        )
        ppl_stopping_threshold = _compute_shuffled_ppl(
            model,
            forget_eval_loader,
            device,
            shuffle_mode=shuffle_mode,
            n_factor=n_factor,
        )
        with open(thresh_path, "w") as f:
            json.dump({"threshold": ppl_stopping_threshold}, f)

    model.enable_input_require_grads()
    model.gradient_checkpointing_enable()

    lora_cfg = cfg.model.lora
    peft_config = LoraConfig(
        r=lora_cfg.r,
        target_modules=list(lora_cfg.target_modules),
        lora_alpha=lora_cfg.lora_alpha,
        lora_dropout=lora_cfg.lora_dropout,
        bias=lora_cfg.bias,
        task_type=lora_cfg.get("task_type", "CAUSAL_LM"),
    )
    model = get_peft_model(model, peft_config)
    model.print_trainable_parameters()

    t = cfg.trainer
    training_args = TrainingArguments(
        output_dir=output_dir,
        max_steps=t.max_steps,
        per_device_train_batch_size=t.per_device_train_batch_size,
        gradient_accumulation_steps=t.gradient_accumulation_steps,
        learning_rate=t.learning_rate,
        lr_scheduler_type=t.lr_scheduler_type,
        warmup_steps=t.warmup_steps,
        weight_decay=t.weight_decay,
        adam_beta1=t.adam_beta1,
        adam_beta2=t.adam_beta2,
        adam_epsilon=t.adam_epsilon,
        max_grad_norm=t.max_grad_norm,
        logging_steps=t.logging_steps,
        save_strategy=t.get("save_strategy", "steps"),
        save_steps=t.save_steps,
        save_total_limit=t.get("save_total_limit", 1),
        bf16=t.get("bf16", True),
        fp16=t.get("fp16", False),
        report_to=list(t.get("report_to", [])),
        run_name=f"ngdiff-{topic_label}",
        remove_unused_columns=False,
        dataloader_pin_memory=True,
        seed=seed,
        gradient_checkpointing=t.get("gradient_checkpointing", True),
    )

    cooldown_config = CooldownConfig(
        enabled=t.get("cooldown_enabled", False),
        every_k=t.get("cooldown_every_k", 200),
        steps=t.get("cooldown_steps", 50),
    )

    trainer_cls = TRAINER_REGISTRY["ngdiff"]
    trainer = trainer_cls(
        model=model,
        args=training_args,
        train_dataset=train_loader.dataset,
        auto_lr=t.get("auto_lr", True),
        lr_delta=t.get("lr_delta", 1e-5),
        mmlu_baseline=mmlu_baseline if not math.isnan(mmlu_baseline) else None,
        mmlu_eval_samples=mmlu_samples,
        mmlu_choice_tokens=mmlu_choice_tokens,
        forget_eval_loader=forget_eval_loader,
        retain_eval_loader=retain_eval_loader,
        max_walltime_minutes=t.get("max_walltime_minutes", None),
        ppl_stopping_threshold=ppl_stopping_threshold,
        retain_pool=retain_pool,
        length_normalized_loss=t.get("length_normalized_loss", True),
        cooldown_config=cooldown_config,
        data_collator=lambda x: x,
    )
    trainer.get_train_dataloader = lambda: train_loader

    ppl_log_path = os.path.join(output_dir, "forget_ppl_log.json")
    if last_checkpoint and os.path.exists(ppl_log_path):
        with open(ppl_log_path) as f:
            trainer._ppl_log = {int(k): v for k, v in json.load(f).items()}
        logger.info("Loaded %d existing PPL log entries.", len(trainer._ppl_log))

    logger.info("Starting NGDiff unlearning for topic(s) '%s'...", topic_label)
    if not math.isnan(mmlu_baseline):
        logger.info(
            "MMLU safety baseline: %.4f  (threshold: %.4f)",
            mmlu_baseline,
            mmlu_baseline * 0.90,
        )

    if t.get("report_to") and "wandb" in list(t.get("report_to", [])):
        try:
            import wandb

            baselines: dict = {
                "config/seed": seed,
                "config/topic": topic_label,
                "config/max_forget_docs": cfg.data.get("max_forget_docs", 2000),
                "config/docs_per_retain_bin": cfg.data.get("docs_per_retain_bin", None),
                "config/min_tokens": cfg.data.get("min_tokens", 0),
                "config/resample_interval": resample_interval,
                "config/length_normalized_loss": t.get("length_normalized_loss", True),
                "config/ppl_shuffle_mode": shuffle_mode,
                "config/cooldown_enabled": cooldown_config.enabled,
            }
            if ppl_stopping_threshold is not None:
                baselines["eval/ppl_target"] = ppl_stopping_threshold
            if not math.isnan(mmlu_baseline):
                baselines["config/mmlu_baseline"] = mmlu_baseline
                baselines["config/mmlu_threshold"] = mmlu_baseline * 0.90
            wandb.config.update(baselines, allow_val_change=True)
        except Exception:
            pass

    trainer.train(resume_from_checkpoint=last_checkpoint)

    stopped_for_good = trainer.state.global_step >= t.max_steps or getattr(
        trainer, "_stopped_permanently", False
    )
    if stopped_for_good:
        adapter_dir = os.path.join(output_dir, "adapter")
        model.save_pretrained(adapter_dir)
        tokenizer.save_pretrained(adapter_dir)
        logger.info(
            "LoRA adapter saved to %s (step %d)",
            adapter_dir,
            trainer.state.global_step,
        )
    else:
        logger.info(
            "Training stopped at step %d/%d (wall-time guard). "
            "Checkpoint saved; next job will resume.",
            trainer.state.global_step,
            t.max_steps,
        )
    logger.info("Done.")


if __name__ == "__main__":
    main()