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