social-data-attribution-w2/src/unlearning/trainer/ngdiff.py

# pyright: reportPrivateImportUsage=false, reportAttributeAccessIssue=false, reportOptionalMemberAccess=false, reportCallIssue=false, reportArgumentType=false
"""
NGDiff trainer based on Algorithm 1 from Bu & Xu (NAACL 2025).

Normalized Gradient Difference:
    g = g_R / ||g_R|| - g_F / ||g_F||

Two extra forward passes on retain every AUTO_LR_INTERVAL optimizer steps
fit a quadratic to find the optimal learning rate automatically.

This implementation also supports length-normalized loss, periodic retain
resampling, optional retain-only cooldown phases, and periodic PPL checks.
"""

from __future__ import annotations

import json
import logging
import math
import os
import time
from typing import Any, cast

import torch

from unlearning.data.dolma_pool import ForgetRetainDataset

from .base import UnlearnTrainer
from .cooldown import CooldownConfig, CooldownManager
from .utils import compute_perplexity, per_document_loss

logger = logging.getLogger(__name__)

AUTO_LR_INTERVAL = 10
LOG_INTERVAL = 50
PPL_CHECK_INTERVAL = 200
PPL_LOG_INTERVAL = 200
MAX_STEPS = 5000
MMLU_SAFETY_THRESHOLD = 0.90
EPS = 1e-8


class NGDiff(UnlearnTrainer):
    """Normalized Gradient Difference unlearning trainer.

    Core constructor args:
        auto_lr, lr_delta, mmlu_baseline, mmlu_eval_samples,
        mmlu_choice_tokens, forget_eval_loader, max_walltime_minutes,
        ppl_stopping_threshold

    Optional method args:
        retain_pool             - RetainPool for periodic resampling
        length_normalized_loss  - per-doc mean then batch mean
        cooldown_config         - CooldownConfig for retain-only phases
    """

    def __init__(
        self,
        *args,
        auto_lr: bool = True,
        lr_delta: float = 1e-5,
        mmlu_baseline: float | None = None,
        mmlu_eval_samples=None,
        mmlu_choice_tokens=None,
        forget_eval_loader=None,
        retain_eval_loader=None,
        max_walltime_minutes: float | None = None,
        ppl_stopping_threshold: float | None = None,
        retain_pool=None,
        length_normalized_loss: bool = True,
        cooldown_config: CooldownConfig | None = None,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.auto_lr = auto_lr
        self.lr_delta = lr_delta
        self.mmlu_baseline = mmlu_baseline
        self.mmlu_eval_samples = mmlu_eval_samples or []
        self.mmlu_choice_tokens = mmlu_choice_tokens
        self.forget_eval_loader = forget_eval_loader
        self.retain_eval_loader = retain_eval_loader
        self.max_walltime_minutes = max_walltime_minutes
        self.ppl_stopping_threshold = ppl_stopping_threshold
        self.retain_pool = retain_pool
        self.length_normalized_loss = length_normalized_loss
        self.cooldown_mgr = CooldownManager(cooldown_config or CooldownConfig())
        self._step_count = 0
        self._optimizer_steps = 0
        self._job_start_time = time.time()
        self._ppl_log: dict[int, float] = {}
        self._resumed = False
        self._stopped_permanently = False

    def _signal_stop(self) -> None:
        if hasattr(self, "control") and self.control is not None:
            self.control.should_training_stop = True
            if self._stopped_permanently:
                self.control.should_save = True

    def _sync_step_count_on_resume(self) -> None:
        if self._resumed:
            return
        self._resumed = True
        if (
            hasattr(self, "state")
            and self.state is not None
            and self.state.global_step > 0
        ):
            grad_accum = self.args.gradient_accumulation_steps
            self._step_count = self.state.global_step * grad_accum
            logger.info(
                "Resumed from checkpoint: global_step=%d, _step_count restored to %d",
                self.state.global_step,
                self._step_count,
            )

    def _compute_loss(self, model_out, inputs) -> torch.Tensor:
        if self.length_normalized_loss:
            return per_document_loss(model_out.logits, inputs["labels"])
        return model_out.loss

    def training_step(
        self,
        model: torch.nn.Module,
        inputs: dict[str, Any],
        num_items_in_batch: torch.Tensor | None = None,
    ) -> torch.Tensor:
        self._sync_step_count_on_resume()
        model.train()

        forget_inputs = {
            k: v.to(model.device) if hasattr(v, "to") else v
            for k, v in inputs["forget"].items()
        }
        retain_inputs = {
            k: v.to(model.device) if hasattr(v, "to") else v
            for k, v in inputs["retain"].items()
        }

        in_cooldown = self.cooldown_mgr.in_cooldown()

        norm_F = 0.0
        norm_R = 0.0

        if in_cooldown:
            model.zero_grad()
            retain_out = model(**retain_inputs)
            retain_loss = self._compute_loss(retain_out, retain_inputs)
            self.accelerator.backward(retain_loss)
            g_R = {
                n: p.grad.clone()
                for n, p in model.named_parameters()
                if p.grad is not None
            }
            norm_R = (
                math.sqrt(
                    sum((g**2).sum().item() for g in g_R.values()),
                )
                + EPS
            )
            model.zero_grad()
            for n, p in model.named_parameters():
                if n in g_R:
                    p.grad = g_R[n] / norm_R
            forget_loss = retain_loss.new_zeros(())
            self.cooldown_mgr.step()
        else:
            model.zero_grad()
            forget_out = model(**forget_inputs)
            forget_loss = self._compute_loss(forget_out, forget_inputs)
            self.accelerator.backward(forget_loss)
            g_F = {
                n: p.grad.clone()
                for n, p in model.named_parameters()
                if p.grad is not None
            }

            model.zero_grad()
            retain_out = model(**retain_inputs)
            retain_loss = self._compute_loss(retain_out, retain_inputs)
            self.accelerator.backward(retain_loss)
            g_R = {
                n: p.grad.clone()
                for n, p in model.named_parameters()
                if p.grad is not None
            }

            norm_F = (
                math.sqrt(
                    sum((g**2).sum().item() for g in g_F.values()),
                )
                + EPS
            )
            norm_R = (
                math.sqrt(
                    sum((g**2).sum().item() for g in g_R.values()),
                )
                + EPS
            )

            model.zero_grad()
            for n, p in model.named_parameters():
                if n in g_F and n in g_R:
                    p.grad = g_R[n] / norm_R - g_F[n] / norm_F
                elif n in g_R:
                    p.grad = g_R[n] / norm_R

        self._step_count += 1
        grad_accum = self.args.gradient_accumulation_steps
        optimizer_steps = self._step_count // grad_accum
        is_optimizer_step = self._step_count % grad_accum == 0

        if optimizer_steps >= MAX_STEPS:
            logger.info("Hard ceiling %d optimizer steps reached.", MAX_STEPS)
            self._signal_stop()

        if self.max_walltime_minutes is not None:
            elapsed = (time.time() - self._job_start_time) / 60.0
            remaining = self.max_walltime_minutes - elapsed
            if remaining < 30.0:
                logger.info(
                    "Wall-time guard: %.1f min elapsed, %.1f min remaining.",
                    elapsed,
                    remaining,
                )
                self._signal_stop()

        if is_optimizer_step and not in_cooldown:
            if self.cooldown_mgr.should_cooldown(optimizer_steps):
                self.cooldown_mgr.enter_cooldown(optimizer_steps)

        if (
            is_optimizer_step
            and self.retain_pool is not None
            and self.retain_pool.should_resample(optimizer_steps)
        ):
            new_retain = self.retain_pool.resample(optimizer_steps)
            if isinstance(self.train_dataset, ForgetRetainDataset):
                self.train_dataset.retain = new_retain
            else:
                logger.warning("RetainPool configured for unsupported train dataset.")
            logger.info(
                "Retain resampled at optimizer_step=%d (%d docs)",
                optimizer_steps,
                len(new_retain),
            )

        if (
            self.auto_lr
            and is_optimizer_step
            and optimizer_steps > 0
            and optimizer_steps % AUTO_LR_INTERVAL == 0
        ):
            self._auto_lr_step(model, retain_inputs)

        if self._step_count % LOG_INTERVAL == 0:
            mmlu_acc = self._run_mmlu_eval(model)
            self._log_metrics(
                model,
                forget_loss,
                retain_loss,
                mmlu_acc=mmlu_acc,
                norm_F=norm_F,
                norm_R=norm_R,
                in_cooldown=in_cooldown,
            )
            if not math.isnan(mmlu_acc) and self.mmlu_baseline is not None:
                threshold = self.mmlu_baseline * MMLU_SAFETY_THRESHOLD
                if mmlu_acc < threshold:
                    logger.warning(
                        "MMLU safety gate at step %d: acc %.4f < %.4f",
                        self._step_count,
                        mmlu_acc,
                        threshold,
                    )
                    self._stopped_permanently = True
                    self._signal_stop()

        if (
            self._step_count % PPL_CHECK_INTERVAL == 0
            and self.forget_eval_loader is not None
        ):
            device = next(model.parameters()).device
            forget_ppl = compute_perplexity(
                model,
                self.forget_eval_loader,
                device,
            )
            retain_ppl = float("nan")
            if self.retain_eval_loader is not None:
                retain_ppl = compute_perplexity(
                    model,
                    self.retain_eval_loader,
                    device,
                )
            model.train()
            logger.info(
                "PPL at optimizer_step=%d: forget=%.4f  retain=%.4f",
                optimizer_steps,
                forget_ppl,
                retain_ppl if not math.isnan(retain_ppl) else -1,
            )
            self._ppl_log[optimizer_steps] = round(forget_ppl, 4)
            self._save_ppl_log()
            self._log_wandb_ppl(
                optimizer_steps,
                forget_ppl,
                retain_ppl,
            )
            if (
                self.ppl_stopping_threshold is not None
                and forget_ppl >= self.ppl_stopping_threshold
            ):
                logger.info(
                    "PPL stopping triggered at optimizer_step=%d: "
                    "PPL=%.4f >= threshold=%.4f",
                    optimizer_steps,
                    forget_ppl,
                    self.ppl_stopping_threshold,
                )
                self._stopped_permanently = True
                self._signal_stop()

        model.train()
        return (forget_loss + retain_loss).detach()

    def _save_ppl_log(self) -> None:
        try:
            path = os.path.join(str(self.args.output_dir or "."), "forget_ppl_log.json")
            with open(path, "w") as f:
                json.dump(self._ppl_log, f, indent=2)
        except Exception as e:
            logger.warning("Could not save forget_ppl_log.json: %s", e)

    def _auto_lr_step(
        self,
        model: torch.nn.Module,
        retain_inputs: dict,
    ) -> None:
        optimizer = cast(Any, self.optimizer)
        if optimizer is None:
            return

        current_lr = optimizer.param_groups[0]["lr"]
        delta = self.lr_delta

        probe_inputs = {
            k: v[:1] if isinstance(v, torch.Tensor) else v
            for k, v in retain_inputs.items()
        }
        param_backup = {n: p.data.clone() for n, p in model.named_parameters()}

        def probe_loss(lr_probe: float) -> float:
            torch.cuda.empty_cache()
            for pg in optimizer.param_groups:
                pg["lr"] = lr_probe
            optimizer.step()
            model.eval()
            with torch.no_grad():
                out = model(**probe_inputs)
                loss = out.loss.item()
            model.train()
            with torch.no_grad():
                for n, p in model.named_parameters():
                    p.copy_(param_backup[n])
            return loss

        try:
            l_minus = probe_loss(current_lr - delta)
            l_plus = probe_loss(current_lr + delta)
            l_curr = probe_loss(current_lr)
            a = (l_minus - 2 * l_curr + l_plus) / (2 * delta**2)
            b = (l_plus - l_minus) / (2 * delta)
            if abs(a) < 1e-30:
                return
            optimal_lr = max(1e-7, min(1e-3, -b / (2 * a)))
            for pg in optimizer.param_groups:
                pg["lr"] = optimal_lr
        except Exception as e:
            logger.warning("AutoLR probe failed at step %d: %s", self._step_count, e)
            for pg in optimizer.param_groups:
                pg["lr"] = current_lr
        finally:
            torch.cuda.empty_cache()

    def _run_mmlu_eval(self, model) -> float:
        if not self.mmlu_eval_samples or self.mmlu_choice_tokens is None:
            return float("nan")
        device = next(model.parameters()).device
        choice_tokens = self.mmlu_choice_tokens.to(device)
        model.eval()
        correct = 0
        with torch.no_grad():
            for input_ids, attn_mask, answer_idx in self.mmlu_eval_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()
        return correct / len(self.mmlu_eval_samples)

    def _log_wandb_ppl(
        self,
        optimizer_steps: int,
        forget_ppl: float,
        retain_ppl: float,
    ) -> None:
        if not (self.args.report_to and "wandb" in self.args.report_to):
            return
        try:
            import wandb

            payload: dict = {
                "eval/forget_ppl": forget_ppl,
                "train/optimizer_step": optimizer_steps,
            }
            if not math.isnan(retain_ppl):
                payload["eval/retain_ppl"] = retain_ppl
            if self.ppl_stopping_threshold is not None:
                payload["eval/ppl_target"] = self.ppl_stopping_threshold
            wandb.log(payload)
        except Exception:
            pass

    def _log_metrics(
        self,
        model,
        forget_loss: torch.Tensor,
        retain_loss: torch.Tensor,
        mmlu_acc: float = float("nan"),
        norm_F: float = 0.0,
        norm_R: float = 0.0,
        in_cooldown: bool = False,
    ):
        optimizer = cast(Any, self.optimizer)
        lr = optimizer.param_groups[0]["lr"] if optimizer else float("nan")
        grad_accum = self.args.gradient_accumulation_steps
        optimizer_steps = self._step_count // grad_accum
        mmlu_str = f"{mmlu_acc:.4f}" if not math.isnan(mmlu_acc) else "n/a"
        logger.info(
            "step=%d  opt_step=%d  f_loss=%.4f  r_loss=%.4f  mmlu=%s  lr=%.2e",
            self._step_count,
            optimizer_steps,
            forget_loss.item(),
            retain_loss.item(),
            mmlu_str,
            lr,
        )
        if self.args.report_to and "wandb" in self.args.report_to:
            try:
                import wandb

                payload: dict = {
                    "train/forget_loss": forget_loss.item(),
                    "train/retain_loss": retain_loss.item(),
                    "train/lr": lr,
                    "train/training_step": self._step_count,
                    "train/optimizer_step": optimizer_steps,
                    "train/grad_norm_forget": norm_F,
                    "train/grad_norm_retain": norm_R,
                    "train/in_cooldown": int(in_cooldown),
                }
                if not math.isnan(mmlu_acc):
                    payload["train/mmlu_acc"] = mmlu_acc
                wandb.log(payload)
            except Exception:
                pass