training_method/train_probe_gen.py

"""Offline probe pre-trainer on generation-time SAE features (frozen base model).

Loads the base LLaVA model (no LoRA), iterates the FinetuneDataset, and for
each batch:
  1. Generates a caption per row (no_grad, greedy by default).
  2. Builds a left-padded teacher-forced batch of [prompt + caption].
  3. Forwards through the frozen base with hooks; captures decoder hidden states.
  4. SAE-encodes every position, then pools latents over generated positions
     (gen_mask) with --pool ∈ {max, mean}. Encode-then-pool matches the per-token
     distribution the SAE was trained on, so the JumpReLU threshold is
     in-distribution.
  5. Trains LayerProbes via BCEWithLogits on those features.

Output: ``probes_gen.pt`` (LayerProbes state_dict) reusable by
``finetune_adv_gen.py`` and downstream experiments.

Usage:
    python -m experiment.training.train_probe_gen --config experiment/adv_config.json
    torchrun --nproc_per_node=4 -m experiment.training.train_probe_gen --config ...
"""
from __future__ import annotations

import argparse
import os
import sys
from datetime import datetime

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from transformers import AutoProcessor, AutoModelForPreTraining
from tqdm import tqdm

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../.."))

from sae.Training_Utils import str_to_torch_dtype
from experiment.config.train_config import TrainConfig
from experiment.config.relation_config import get_relation_config
from experiment.data.datasets import FinetuneDataset, finetune_dataset_extra_kwargs
from training_method.finetune_adv import (
    FrozenSAEEncoder,
    LayerProbes,
    HiddenStateCapture,
    count_lm_layers,
    probe_labels,
    probe_bce_loss_logits,
    parse_args,
)
from experiment.evaluation.metrics import KeywordMentionDetector
from experiment.training.gen_features import (
    build_position_masks,
    build_scope_mask,
    gather_gen_positions,
    gather_masked_positions,
    left_pad_collate,
    masked_max_pool,
    masked_mean_pool,
)
from training_method.sequence_probe import SequenceLayerProbes

import wandb

wandb.init(project="multilayer-sae", name=f"train_probe_gen_{datetime.now().strftime('%Y%m%d_%H%M%S')}")

def _strip_pad(input_ids_row: torch.Tensor, attn_row: torch.Tensor) -> torch.Tensor:
    """Return the unpadded prompt tokens for a single row (1-D LongTensor)."""
    return input_ids_row[attn_row.bool()]


def _generate_captions(
    raw_model,
    pixel_values: torch.Tensor,         # (B, ...)
    input_ids: torch.Tensor,            # (B, L)
    attention_mask: torch.Tensor,       # (B, L)
    max_new_tokens: int,
    do_sample: bool,
    temperature: float,
    pad_token_id: int,
) -> tuple[list[torch.Tensor], list[torch.Tensor], list[int], list[int]]:
    """Per-row generation. Returns (prompt_unpadded, full_seq, prompt_lens, gen_lens)."""
    B = input_ids.shape[0]
    prompt_unpadded: list[torch.Tensor] = []
    full_seq: list[torch.Tensor] = []
    prompt_lens: list[int] = []
    gen_lens: list[int] = []
    # Disable GC while generating to avoid PEFT+GC interaction warnings.
    was_gc = getattr(raw_model, "is_gradient_checkpointing", False)
    if was_gc:
        raw_model.gradient_checkpointing_disable()
    try:
        for i in range(B):
            prompt_real = _strip_pad(input_ids[i], attention_mask[i])
            with torch.no_grad():
                out = raw_model.generate(
                    pixel_values=pixel_values[i : i + 1],
                    input_ids=prompt_real.unsqueeze(0),
                    attention_mask=torch.ones_like(prompt_real).unsqueeze(0),
                    max_new_tokens=max_new_tokens,
                    do_sample=do_sample,
                    temperature=temperature if do_sample else 1.0,
                    use_cache=True,
                    pad_token_id=pad_token_id,
                )
            full = out[0]
            gen = full[prompt_real.shape[0] :]
            prompt_unpadded.append(prompt_real)
            full_seq.append(full)
            prompt_lens.append(int(prompt_real.shape[0]))
            gen_lens.append(int(gen.shape[0]))
    finally:
        if was_gc:
            raw_model.gradient_checkpointing_enable(
                gradient_checkpointing_kwargs={"use_reentrant": False}
            )
    return prompt_unpadded, full_seq, prompt_lens, gen_lens


def _gen_features(
    capture_hidden: dict[int, torch.Tensor],
    sae: FrozenSAEEncoder,
    pool_mask: torch.Tensor,
    pool_fn,
) -> dict[int, torch.Tensor]:
    """SAE-encode every position then pool latents over pool_mask positions.

    The SAE was trained on per-token hidden states; encode-then-pool keeps inputs
    in-distribution so the JumpReLU threshold fires sensibly. ``pool_fn`` selects
    the reduction (``masked_max_pool`` or ``masked_mean_pool``). Per-layer (B, d_sae).
    ``pool_mask`` can cover generated-only positions or all real positions depending
    on the --pool_tokens setting.
    """
    out: dict[int, torch.Tensor] = {}
    for l, h in capture_hidden.items():
        latents = sae(h)                               # (B, S, d_sae)
        out[l] = pool_fn(latents, pool_mask)           # (B, d_sae)
    return out


def reserve_gpu_memory(device: torch.device, gib: float = 35.0, verbose: bool = True) -> None:
    """Pre-claim a fixed amount of VRAM so competing processes can't steal it."""
    if device.type != "cuda":
        return
    try:
        buf = torch.empty(int(gib * 2**30) // 2, dtype=torch.int16, device=device)
        del buf
        if verbose:
            reserved = torch.cuda.memory_reserved(device)
            total = torch.cuda.get_device_properties(device).total_memory
            print(f"  [mem-reserve] {reserved/2**30:.1f}/{total/2**30:.1f} GB reserved")
    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        if verbose:
            print(f"  [mem-reserve] WARNING: could not reserve {gib:.0f} GiB — GPU already loaded")


def parse_extra_args_and_strip():
    """Parse train_probe_gen-specific args, then REMOVE them from sys.argv so the
    shared parse_args() in finetune_adv doesn't try to push them into apply_overrides
    (which would AttributeError on unknown TrainConfig fields).
    """
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--probe_epochs", type=int, default=10)
    parser.add_argument("--probe_output", type=str, default=None,
                        help="Path to save probes_gen.pt; defaults to {output_dir}/probe_run_{run_id}/probes_gen_{relation}.pt")
    parser.add_argument("--pool", choices=["max", "mean"], default="max",
                        help="Pooling over generated positions for SAE latents. "
                             "'max' = peak-firing (gameable via redistribution); "
                             "'mean' = non-redistributable, harder to evade under adversarial finetuning.")
    parser.add_argument("--pool_tokens",
                        choices=["gen", "vision", "prompt", "prompt_gen", "all"], default="gen",
                        help="Which token positions the probe reads. "
                             "'gen' = generated tokens only (default); "
                             "'vision' = image-patch tokens only; "
                             "'prompt' = prompt text only; "
                             "'prompt_gen' = prompt text + generated; "
                             "'all' = image patches + prompt text + generated tokens.")
    parser.add_argument("--label_from_mention", action="store_true",
                        help="Label each row by whether the MODEL'S GENERATED caption "
                             "mentions the object (within the pooled token window), instead "
                             "of ground-truth presence. Makes labels consistent with the "
                             "generated-token activations (fixes the ~24%% contradictory-label "
                             "ceiling on relations the model hallucinates/misses heavily).")
    parser.add_argument("--mention_keywords", type=str, default=None,
                        help="Comma-separated keywords for --label_from_mention "
                             "(default: relation_config.mention_keywords).")
    extra, remaining = parser.parse_known_args()
    sys.argv = [sys.argv[0]] + remaining
    return extra


def main():
    extra = parse_extra_args_and_strip()       # must run BEFORE parse_args
    args, overrides = parse_args()

    config = TrainConfig.load(args.config)
    if args.relation:
        config.relation = args.relation
    if overrides:
        config.apply_overrides(overrides)
    config.resolve_from_relation()

    relation_config = get_relation_config(config.relation)
    adv_cfg = config.adv
    model_dtype = str_to_torch_dtype(config.dtype)

    mention_detector = None
    if extra.label_from_mention:
        if extra.mention_keywords:
            kw = [k.strip() for k in extra.mention_keywords.split(",") if k.strip()]
        else:
            kw = list(relation_config.mention_keywords)
        mention_detector = KeywordMentionDetector(keywords=kw)

    if not adv_cfg.raw_activation_probe:
        assert adv_cfg.sae_checkpoint, "adv.sae_checkpoint must be set in config"

    use_ddp = "LOCAL_RANK" in os.environ
    if use_ddp:
        local_rank = int(os.environ["LOCAL_RANK"])
        dist.init_process_group(backend="nccl")
        torch.cuda.set_device(local_rank)
        device = torch.device(f"cuda:{local_rank}")
        is_main = local_rank == 0
    else:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        local_rank = 0
        is_main = True

    run_id = os.environ.get("RUN_ID") or datetime.now().strftime("%Y%m%d_%H%M%S")
    if use_ddp:
        run_id_list = [run_id]
        dist.broadcast_object_list(run_id_list, src=0)
        run_id = run_id_list[0]
    # Only materialise run_dir if we actually save into it (i.e., no explicit --probe_output).
    run_dir = os.path.join(config.output_dir, f"probe_run_{run_id}")
    if is_main and not extra.probe_output:
        os.makedirs(run_dir, exist_ok=True)
    if use_ddp:
        dist.barrier()

    pool_fn = masked_mean_pool if extra.pool == "mean" else masked_max_pool

    # Reserve GPU memory before loading model so competing processes can't steal VRAM.
    _mem_gib = float(os.environ.get("GPU_MEM_RESERVE_GIB", "35.0"))
    if _mem_gib > 0:
        reserve_gpu_memory(device, gib=_mem_gib, verbose=is_main)

    if is_main:
        print(f"[train_probe_gen] relation={config.relation} sae={adv_cfg.sae_checkpoint}")
        print(f"[train_probe_gen] epochs={extra.probe_epochs} pool={extra.pool} "
              f"pool_tokens={extra.pool_tokens} "
              f"out={extra.probe_output or os.path.join(run_dir, f'probes_gen_{config.relation}.pt')}")
        if mention_detector is not None:
            print(f"[train_probe_gen] LABEL=mention (generated-caption), "
                  f"keywords={mention_detector.keywords}")
        else:
            print(f"[train_probe_gen] LABEL=ground-truth ({adv_cfg.probe_label_mode})")

    # Frozen base model (no LoRA).
    model = AutoModelForPreTraining.from_pretrained(
        config.model_name,
        torch_dtype=model_dtype,
        device_map={"": device},
    )
    processor = AutoProcessor.from_pretrained(config.model_name)
    pad_token_id = processor.tokenizer.pad_token_id
    if pad_token_id is None:
        pad_token_id = processor.tokenizer.eos_token_id
    for p in model.parameters():
        p.requires_grad_(False)
    model.eval()
    raw_model = model

    n_layers = count_lm_layers(model)
    probe_layers = adv_cfg.probe_layers if adv_cfg.probe_layers else list(range(n_layers))

    image_token_id = int(getattr(model.config, "image_token_index", 32000))

    sae = None
    if adv_cfg.raw_activation_probe:
        try:
            d_model = int(model.config.text_config.hidden_size)
        except AttributeError:
            d_model = int(model.config.hidden_size)
        probes = SequenceLayerProbes(
            probe_layers, d_model,
            d_probe=adv_cfg.probe_dim, n_heads=adv_cfg.probe_heads,
            n_ctx_blocks=adv_cfg.probe_ctx_blocks,
            spectral_norm=adv_cfg.probe_spectral_norm,
        ).to(device)
        if is_main:
            print(f"[train_probe_gen] SequenceLayerProbes (raw activations) "
                  f"d_model={d_model} d_probe={adv_cfg.probe_dim} layers={len(probe_layers)}")
    else:
        sae = FrozenSAEEncoder.from_checkpoint(adv_cfg.sae_checkpoint, device)
        d_sae = sae.encoder.weight.shape[0]
        probes = LayerProbes(
            probe_layers, d_sae, spectral_norm=adv_cfg.probe_spectral_norm
        ).to(device)
    if use_ddp:
        probes = DDP(probes, device_ids=[local_rank], find_unused_parameters=False)
    probes_module = probes.module if use_ddp else probes

    probe_opt = torch.optim.AdamW(
        probes_module.parameters(),
        lr=adv_cfg.probe_lr,
        weight_decay=adv_cfg.probe_weight_decay,
    )

    dataset = FinetuneDataset(

        processor=processor,
        prompt_config=config.prompts,
        dataset_id=config.dataset_id,
        scene_col=relation_config.scene_key,
        object_col=relation_config.object_key,
        csv_path=config.csv_path,
        image_dir=config.image_dir,
        max_samples=config.max_train_samples,
        split="train",
        upsample_categories=None,
        **finetune_dataset_extra_kwargs(config),
    )
    if use_ddp:
        sampler = DistributedSampler(dataset, shuffle=True)
        dataloader = DataLoader(
            dataset, batch_size=config.batch_size, sampler=sampler,
            num_workers=config.num_workers, pin_memory=True, drop_last=True,
        )
    else:
        dataloader = DataLoader(
            dataset, batch_size=config.batch_size, shuffle=True,
            num_workers=config.num_workers, pin_memory=True, drop_last=True,
        )

    capture = HiddenStateCapture(raw_model, probe_layers)
    n_image_patches = None  # filled on first batch

    # Delayed cosine LR: hold probe_lr for the first half, then cosine-decay to probe_lr/100.
    total_probe_steps = max(1, len(dataloader) * extra.probe_epochs)
    hold_steps = max(1, total_probe_steps // 2)
    decay_steps = max(1, total_probe_steps - hold_steps)
    eta_min = adv_cfg.probe_lr / 100.0
    sched_hold = torch.optim.lr_scheduler.ConstantLR(probe_opt, factor=1.0, total_iters=hold_steps)
    sched_cos = torch.optim.lr_scheduler.CosineAnnealingLR(probe_opt, T_max=decay_steps, eta_min=eta_min)
    probe_sched = torch.optim.lr_scheduler.SequentialLR(
        probe_opt, schedulers=[sched_hold, sched_cos], milestones=[hold_steps],
    )
    if is_main:
        print(f"[train_probe_gen] LR schedule: hold {adv_cfg.probe_lr:.2e} for {hold_steps} steps, "
              f"then cosine → {eta_min:.2e} over {decay_steps} steps")

    # Resolve checkpoint paths once. Per-epoch saves get an _epoch{N:02d} suffix;
    # the canonical (un-suffixed) path is also rewritten at the end so downstream
    # tools that look for probes_gen_{relation}.pt keep working.
    final_out_path = extra.probe_output or os.path.join(
        run_dir, f"probes_gen_{config.relation}.pt"
    )
    out_dir = os.path.dirname(final_out_path)
    if is_main and out_dir:
        os.makedirs(out_dir, exist_ok=True)
    out_stem, out_ext = os.path.splitext(final_out_path)

    for epoch in range(extra.probe_epochs):
        if use_ddp:
            sampler.set_epoch(epoch)
        epoch_losses = []
        epoch_accs = []
        epoch_pos = []
        pbar = tqdm(dataloader, desc=f"probe-epoch {epoch+1}/{extra.probe_epochs}", disable=not is_main)
        for batch in pbar:
            has_object = batch.pop("has_object").to(device)
            is_scene = batch.pop("is_scene").to(device)
            batch = {k: v.to(device) for k, v in batch.items()}
            y_probe = probe_labels(is_scene, has_object, adv_cfg.probe_label_mode)

            # 1. Generate per-row.
            _, full_seqs, prompt_lens, gen_lens = _generate_captions(
                raw_model,
                pixel_values=batch["pixel_values"],
                input_ids=batch["input_ids"],
                attention_mask=batch["attention_mask"],
                max_new_tokens=adv_cfg.max_new_tokens_train,
                do_sample=adv_cfg.gen_do_sample,
                temperature=adv_cfg.gen_temperature,
                pad_token_id=pad_token_id,
            )

            # 1b. Optional: relabel by whether the GENERATED caption mentions the
            # object (within this same gen window). Keeps labels consistent with the
            # generated-token activations the probe reads (vs ground-truth presence,
            # which contradicts activations on hallucinated/missed rows).
            if mention_detector is not None:
                mlabels = []
                for i in range(len(full_seqs)):
                    gen_ids = full_seqs[i][prompt_lens[i]:]
                    txt = processor.tokenizer.decode(gen_ids, skip_special_tokens=True)
                    mlabels.append(1.0 if mention_detector.mentions_object(txt) else 0.0)
                y_probe = torch.tensor(mlabels, device=device, dtype=y_probe.dtype)

            # 2. Left-pad collate teacher-forced batch.
            tf_ids, tf_attn = left_pad_collate(full_seqs, pad_id=pad_token_id)
            tf_ids = tf_ids.to(device)
            tf_attn = tf_attn.to(device)

            # 3. Forward (no grad through model; probe is the only thing with grad).
            with torch.no_grad():
                with capture:
                    raw_model(
                        pixel_values=batch["pixel_values"],
                        input_ids=tf_ids,
                        attention_mask=tf_attn,
                        use_cache=False,
                    )

            # 4. Build masks. Infer n_image_patches on first batch.
            S = next(iter(capture.hidden_states.values())).shape[1]
            if n_image_patches is None:
                # S = N_patches + L_max - 1 ⇒ N_patches = S - L_max + 1.
                L_max = tf_ids.shape[1]
                n_image_patches = S - L_max + 1
                if is_main:
                    print(f"[train_probe_gen] inferred n_image_patches={n_image_patches} "
                          f"(S={S}, L_max={L_max})")

            prompt_real_lens_t = torch.tensor(prompt_lens, device=device)
            gen_lens_t = torch.tensor(gen_lens, device=device)
            prompt_mask, gen_mask = build_position_masks(
                prompt_real_lens_t, gen_lens_t, n_image_patches, S
            )

            # Select probe token scope (gen | vision | prompt | prompt_gen | all).
            vision_mask = (tf_ids == image_token_id)
            pool_mask = build_scope_mask(extra.pool_tokens, prompt_mask, gen_mask, vision_mask)

            # Drop rows with no selectable positions.
            keep = pool_mask.any(dim=1)
            # DDP-safe skip: all ranks must agree, else a rank that skips its
            # backward() desyncs the gradient all-reduce (NCCL collective timeout).
            skip = not bool(keep.any())
            if use_ddp:
                flag = torch.tensor([1.0 if skip else 0.0], device=device)
                dist.all_reduce(flag, op=dist.ReduceOp.MAX)
                skip = flag.item() > 0.5
            if skip:
                continue
            if adv_cfg.raw_activation_probe:
                # General gather (handles non-tail scopes like vision patches).
                feats_seq, valid = gather_masked_positions(capture.hidden_states, pool_mask)
                features = {l: f[keep] for l, f in feats_seq.items()}
                z = probes_module.forward_logits(features, ~valid[keep])
            else:
                features = _gen_features(capture.hidden_states, sae, pool_mask, pool_fn)
                features = {l: f[keep] for l, f in features.items()}
                z = probes_module.forward_logits(features)
            y = y_probe[keep]

            # 5. Probe BCE.
            loss = probe_bce_loss_logits(z, y, adv_cfg.probe_label_smoothing)
            with torch.no_grad():
                p_avg = sum(torch.sigmoid(zz) for zz in z) / len(z)
                acc = (((p_avg > 0.5).float() == y).float().mean().item())
                epoch_accs.append(acc)
                epoch_pos.append(y.float().mean().item())

            probe_opt.zero_grad()
            loss.backward()
            probe_opt.step()
            probe_sched.step()

            epoch_losses.append(loss.item())
            if is_main:
                pbar.set_postfix({
                    "bce": f"{loss.item():.4f}",
                    "acc": f"{acc:.3f}",
                    "lr": f"{probe_sched.get_last_lr()[0]:.2e}",
                })

        if is_main and epoch_losses:
            print(f"[train_probe_gen] epoch {epoch+1}: avg_bce="
                  f"{sum(epoch_losses)/len(epoch_losses):.4f}  "
                  f"avg_acc={sum(epoch_accs)/max(len(epoch_accs),1):.4f}  "
                  f"pos_rate={sum(epoch_pos)/max(len(epoch_pos),1):.3f}")

        # Per-epoch checkpoint (rank 0 only). Other ranks proceed straight into
        # the next epoch and naturally re-sync at the next NCCL allreduce.
        if is_main:
            epoch_path = f"{out_stem}_epoch{epoch+1:02d}{out_ext}"
            torch.save(probes_module.state_dict(), epoch_path)
            print(f"[train_probe_gen] saved {epoch_path}")

    if is_main:
        torch.save(probes_module.state_dict(), final_out_path)
        print(f"[train_probe_gen] saved {final_out_path} (final)")

    if use_ddp:
        dist.barrier()
        dist.destroy_process_group()


if __name__ == "__main__":
    main()