"""Build a gen-only SAE feature cache (drop-in for Train_Probe_AllLayers output, but pooled over the assistant's generated caption tokens only). The existing ``training/multilayer_sae_ckpt/feature_cache/`` was written by ``training/Train_Probe_AllLayers.py``'s ``topk_sae_max_pool``, which max-pools over the *entire* sequence (image patches + prompt + caption + padding). That mismatches the probes trained by ``train_probe_gen.py``, which pool only over the assistant's generated tokens via ``gen_mask``. This script reuses the exact pipeline of ``train_probe_gen.py``: generate caption → teacher-force [prompt+caption] → capture hidden states → ``FrozenSAEEncoder`` → ``masked_max_pool`` (or ``masked_mean_pool``) over ``gen_mask``. The result is saved in the same ``feature_cache/`` directory layout (``labels.pt`` + ``layer_LL.pt`` + ``meta.json``) so ``Select_Features_F1_AllLayers.py`` and ``mechanistic_interp/select_features_activation.py`` consume it unchanged. Single-GPU or torchrun multi-GPU. Each rank writes its features into a shared tensor list, and rank 0 gathers + saves. Usage: python -m experiment.training.build_gen_feature_cache \\ --config experiment/adv_config.json \\ --cache_out training/multilayer_sae_ckpt/feature_cache_gen torchrun --nproc_per_node=8 -m experiment.training.build_gen_feature_cache \\ --config experiment/adv_config.json \\ --cache_out training/multilayer_sae_ckpt/feature_cache_gen """ from __future__ import annotations import argparse import json import os import sys import torch import torch.distributed as dist from torch.utils.data import DataLoader, DistributedSampler from transformers import AutoModelForPreTraining, AutoProcessor 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, HiddenStateCapture, count_lm_layers, ) from experiment.training.train_probe_gen import _generate_captions from experiment.training.gen_features import ( build_position_masks, left_pad_collate, masked_max_pool, masked_mean_pool, ) def parse_local_args(): p = argparse.ArgumentParser(add_help=False) p.add_argument("--config", required=True) p.add_argument("--cache_out", required=True, help="Output directory (will be created).") p.add_argument("--relation", default=None) p.add_argument("--pool", choices=("max", "mean"), default="max") p.add_argument("--split", default="train", help="HF split fed into FinetuneDataset.") p.add_argument("--max_samples", type=int, default=0, help="0 = all (overrides config.max_train_samples).") extra, remaining = p.parse_known_args() sys.argv = [sys.argv[0]] + remaining return extra def main(): extra = parse_local_args() config = TrainConfig.load(extra.config) if extra.relation: config.relation = extra.relation 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" 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 world = dist.get_world_size() else: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") local_rank, is_main, world = 0, True, 1 pool_fn = masked_max_pool if extra.pool == "max" else masked_mean_pool if is_main: print(f"[build_gen_cache] relation={config.relation} " f"sae={adv_cfg.sae_checkpoint} pool={extra.pool} " f"split={extra.split} out={extra.cache_out}") os.makedirs(extra.cache_out, exist_ok=True) 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 or processor.tokenizer.eos_token_id for p in model.parameters(): p.requires_grad_(False) model.eval() n_layers = count_lm_layers(model) layer_indices = adv_cfg.probe_layers if adv_cfg.probe_layers else list(range(n_layers)) if is_main: print(f"[build_gen_cache] layers: {layer_indices}") sae = FrozenSAEEncoder.from_checkpoint(adv_cfg.sae_checkpoint, device) d_sae = sae.encoder.weight.shape[0] if extra.max_samples > 0: max_samples = extra.max_samples else: max_samples = config.max_train_samples 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=max_samples, split=extra.split, upsample_categories=None, **finetune_dataset_extra_kwargs(config), ) if use_ddp: sampler = DistributedSampler(dataset, shuffle=False, drop_last=False) loader = DataLoader( dataset, batch_size=config.batch_size, sampler=sampler, num_workers=config.num_workers, pin_memory=True, drop_last=False, ) else: loader = DataLoader( dataset, batch_size=config.batch_size, shuffle=False, num_workers=config.num_workers, pin_memory=True, drop_last=False, ) capture = HiddenStateCapture(model, layer_indices) n_image_patches = None layer_feats: dict[int, list[torch.Tensor]] = {l: [] for l in layer_indices} has_object_list: list[torch.Tensor] = [] is_scene_list: list[torch.Tensor] = [] for batch in tqdm(loader, desc=f"rank{local_rank}", disable=not is_main): has_object = batch.pop("has_object").to(device) is_scene = batch.pop("is_scene").to(device) batch.pop("labels", None) batch = {k: v.to(device) for k, v in batch.items()} # 1) Free-generate captions per row. _, full_seqs, prompt_lens, gen_lens = _generate_captions( 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, ) # 2) Teacher-force [prompt + caption], capture hidden states. 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: model( pixel_values=batch["pixel_values"], input_ids=tf_ids, attention_mask=tf_attn, use_cache=False, ) # 3) Build gen_mask. Infer n_image_patches on first batch. S = next(iter(capture.hidden_states.values())).shape[1] if n_image_patches is None: n_image_patches = S - tf_ids.shape[1] + 1 if is_main: print(f"[build_gen_cache] inferred n_image_patches={n_image_patches} " f"(S={S}, L_max={tf_ids.shape[1]})") 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 # 4) SAE-encode every position, gen-pool, store. for l, h in capture.hidden_states.items(): latents = sae(h) # (B, S, d_sae) float32 (gated) feat = pool_fn(latents, gen_mask) # (B, d_sae) # masked_max_pool returns -inf where mask is empty; we already # dropped those rows via `keep`. Cast to float16 to mirror the # original cache (compact + fast to load). layer_feats[l].append(feat[keep].detach().to(torch.float16).cpu()) has_object_list.append(has_object[keep].detach().cpu().float()) is_scene_list.append(is_scene[keep].detach().cpu().float()) local_layer_feats = {l: torch.cat(layer_feats[l]) for l in layer_indices} local_has_object = torch.cat(has_object_list) if has_object_list else torch.zeros(0) local_is_scene = torch.cat(is_scene_list) if is_scene_list else torch.zeros(0) if use_ddp: # Gather across ranks (concatenate per-rank tensors). def _gather_concat(t: torch.Tensor) -> torch.Tensor: parts = [None] * world dist.all_gather_object(parts, t) return torch.cat(parts) if is_main else None all_has_object = _gather_concat(local_has_object) all_is_scene = _gather_concat(local_is_scene) all_layer_feats = {} for l in layer_indices: all_layer_feats[l] = _gather_concat(local_layer_feats[l]) else: all_has_object = local_has_object all_is_scene = local_is_scene all_layer_feats = local_layer_feats if not is_main: if use_ddp: dist.barrier() dist.destroy_process_group() return # ── Save in feature_cache format (labels=has_object, mirrors original) ── out = extra.cache_out os.makedirs(out, exist_ok=True) torch.save(all_has_object, os.path.join(out, "labels.pt")) torch.save(all_is_scene, os.path.join(out, "is_scene.pt")) for l in layer_indices: torch.save(all_layer_feats[l], os.path.join(out, f"layer_{l:02d}.pt")) meta = { "d_sae": int(d_sae), "n_layers": int(n_layers), "probe_label_mode": "object_only", "dataset_id": config.dataset_id, "pool_over": "gen", "pool": extra.pool, "sae_ckpt": adv_cfg.sae_checkpoint, "split": extra.split, "relation": config.relation, "n_samples": int(all_has_object.numel()), "n_pos": int(all_has_object.sum()), "n_neg": int((1 - all_has_object).sum()), "n_image_patches": int(n_image_patches), "source": "build_gen_feature_cache.py (gen_mask, FrozenSAEEncoder)", } with open(os.path.join(out, "meta.json"), "w") as f: json.dump(meta, f, indent=2) print(f"[build_gen_cache] wrote cache → {out}") print(f" labels: N={meta['n_samples']} pos={meta['n_pos']} neg={meta['n_neg']}") print(f" per-layer shape: ({meta['n_samples']}, {d_sae}) float16") if use_ddp: dist.barrier() dist.destroy_process_group() if __name__ == "__main__": main()