experiment/training/continue_probe_gen.py

"""Warm-start continuation of probe training from an existing probes_gen checkpoint.

Sibling of ``train_probe_gen.py``: identical generation/forward/SAE/probe loop,
but instead of fresh-init weights it loads ``--probe_resume`` and trains
``--probe_epochs`` more epochs with a fresh optimizer and a pure cosine LR
(``probe_lr/100`` → ``probe_lr/10000``, no hold phase). Optimizer/scheduler
state is intentionally not persisted by the original trainer, so this is a
warm-start rather than a bit-exact resume.

Per-epoch checkpoints are written as
``probes_gen_{relation}_resume_epoch{NN}.pt`` into a new
``probe_run_resume_{run_id}/`` directory; the final un-suffixed file matches
the original naming so downstream tooling keeps working.

Usage:
    python -m experiment.training.continue_probe_gen --config experiment/adv_config.json \\
        --probe_resume /path/to/probes_gen_kitchen_oven.pt --probe_epochs 5
"""
from __future__ import annotations

import argparse
import os
import sys
from datetime import datetime

import torch
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 experiment.training.finetune_adv import (
    FrozenSAEEncoder,
    LayerProbes,
    HiddenStateCapture,
    count_lm_layers,
    probe_labels,
    probe_bce_loss_logits,
    parse_args,
)
from experiment.training.gen_features import (
    build_position_masks,
    left_pad_collate,
    masked_max_pool,
    masked_mean_pool,
)
from experiment.training.train_probe_gen import _generate_captions, _gen_features


def parse_extra_args_and_strip():
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--probe_epochs", type=int, default=5,
                        help="Additional epochs to train (default: 5).")
    parser.add_argument("--probe_resume", type=str, required=True,
                        help="Path to a probes_gen_{relation}.pt to warm-start from.")
    parser.add_argument("--probe_output", type=str, default=None,
                        help="Final probe path. Defaults to {output_dir}/probe_run_resume_{run_id}/probes_gen_{relation}.pt")
    parser.add_argument("--pool", choices=["max", "mean"], default="max")
    extra, remaining = parser.parse_known_args()
    sys.argv = [sys.argv[0]] + remaining
    return extra


def main():
    extra = parse_extra_args_and_strip()
    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)

    assert adv_cfg.sae_checkpoint, "adv.sae_checkpoint must be set in config"
    assert os.path.isfile(extra.probe_resume), f"--probe_resume not found: {extra.probe_resume}"

    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]
    run_dir = os.path.join(config.output_dir, f"probe_run_resume_{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

    if is_main:
        print(f"[continue_probe_gen] relation={config.relation} sae={adv_cfg.sae_checkpoint}")
        print(f"[continue_probe_gen] resume_from={extra.probe_resume}")
        print(f"[continue_probe_gen] +epochs={extra.probe_epochs} pool={extra.pool}")

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

    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)
    resume_sd = torch.load(extra.probe_resume, map_location=device, weights_only=True)
    missing, unexpected = probes.load_state_dict(resume_sd, strict=False)
    if is_main:
        print(f"[continue_probe_gen] loaded {len(resume_sd)} tensors "
              f"(missing={len(missing)}, unexpected={len(unexpected)})")
    if use_ddp:
        probes = DDP(probes, device_ids=[local_rank], find_unused_parameters=False)
    probes_module = probes.module if use_ddp else probes

    # Warm-start LR: pure cosine from probe_lr/10 → probe_lr/1000, no hold.
    # Sits roughly where epochs 6-7 of the original (full hold+cosine) schedule
    # were — enough to keep learning, not so high it undoes convergence.
    start_lr = adv_cfg.probe_lr / 10.0
    eta_min = adv_cfg.probe_lr / 1000.0
    probe_opt = torch.optim.AdamW(
        probes_module.parameters(),
        lr=start_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

    total_probe_steps = max(1, len(dataloader) * extra.probe_epochs)
    probe_sched = torch.optim.lr_scheduler.CosineAnnealingLR(
        probe_opt, T_max=total_probe_steps, eta_min=eta_min,
    )
    if is_main:
        print(f"[continue_probe_gen] LR schedule: cosine {start_lr:.2e} → {eta_min:.2e} "
              f"over {total_probe_steps} steps")

    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 = []
        pbar = tqdm(dataloader, desc=f"probe-resume {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)

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

            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)

            with torch.no_grad():
                with capture:
                    raw_model(
                        pixel_values=batch["pixel_values"],
                        input_ids=tf_ids,
                        attention_mask=tf_attn,
                        use_cache=False,
                    )

            S = next(iter(capture.hidden_states.values())).shape[1]
            if n_image_patches is None:
                L_max = tf_ids.shape[1]
                n_image_patches = S - L_max + 1
                if is_main:
                    print(f"[continue_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)
            _, gen_mask = build_position_masks(
                prompt_real_lens_t, gen_lens_t, n_image_patches, S
            )

            keep = gen_lens_t > 0
            if not keep.any():
                continue
            features = _gen_features(capture.hidden_states, sae, gen_mask, pool_fn)
            features = {l: f[keep] for l, f in features.items()}
            y = y_probe[keep]

            z = probes_module.forward_logits(features)
            loss = probe_bce_loss_logits(z, y, adv_cfg.probe_label_smoothing)

            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}",
                    "lr": f"{probe_sched.get_last_lr()[0]:.2e}",
                })

        if is_main and epoch_losses:
            print(f"[continue_probe_gen] epoch {epoch+1}: avg_bce="
                  f"{sum(epoch_losses)/len(epoch_losses):.4f}")

        if is_main:
            epoch_path = f"{out_stem}_resume_epoch{epoch+1:02d}{out_ext}"
            torch.save(probes_module.state_dict(), epoch_path)
            print(f"[continue_probe_gen] saved {epoch_path}")

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

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


if __name__ == "__main__":
    main()