diff --git a/cli/extract_text_embeddings.py b/cli/extract_text_embeddings.py index 37069dc..0047fce 100644 --- a/cli/extract_text_embeddings.py +++ b/cli/extract_text_embeddings.py @@ -147,14 +147,24 @@ def _encode( f"Embedding inputs {indices[:5]} exceed max_length={max_length}; " "shorten data or explicitly pass --allow_truncation" ) - encoded = tokenizer( - batch_texts, - add_special_tokens=True, - padding=True, - truncation=True, - max_length=max_length, - return_tensors="pt", - ) + if not any(batch_truncated): + # Nothing exceeds max_length, so padding the ids already produced + # yields exactly what a second `truncation=True` tokenization would, + # without re-tokenizing the batch. + encoded = tokenizer.pad( + {"input_ids": untruncated}, + padding=True, + return_tensors="pt", + ) + else: + encoded = tokenizer( + batch_texts, + add_special_tokens=True, + padding=True, + truncation=True, + max_length=max_length, + return_tensors="pt", + ) encoded = {key: value.to(model_device) for key, value in encoded.items()} with torch.inference_mode(): output = model(**encoded) diff --git a/cli/generate_task_rollouts.py b/cli/generate_task_rollouts.py index f899f6e..9d5247e 100644 --- a/cli/generate_task_rollouts.py +++ b/cli/generate_task_rollouts.py @@ -192,6 +192,26 @@ def main() -> None: generated_count = 0 with output.open("a", encoding="utf-8") as handle: for scenario_index, scenario in enumerate(scenarios): + # The prompt encoding depends only on the scenario, not the replicate + # or seed, so build it once per scenario. Sampling randomness is set + # by set_seed() immediately before each generate() call below, and + # generate() does not mutate its inputs, so replicates remain + # independent and their outputs are unchanged. + encoded = tokenizer.apply_chat_template( + scenario.messages, + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, + ) + if isinstance(encoded, torch.Tensor): + encoded = { + "input_ids": encoded, + "attention_mask": torch.ones_like(encoded), + } + encoded = { + key: value.to(model_device) for key, value in encoded.items() + } for replicate in range(args.num_rollouts): rollout_seed = ( args.seed + scenario_index * args.num_rollouts + replicate @@ -202,21 +222,6 @@ def main() -> None: if rollout_id in completed: continue set_seed(rollout_seed) - encoded = tokenizer.apply_chat_template( - scenario.messages, - tokenize=True, - add_generation_prompt=True, - return_tensors="pt", - return_dict=True, - ) - if isinstance(encoded, torch.Tensor): - encoded = { - "input_ids": encoded, - "attention_mask": torch.ones_like(encoded), - } - encoded = { - key: value.to(model_device) for key, value in encoded.items() - } kwargs = { "max_new_tokens": args.max_new_tokens, "do_sample": args.temperature > 0.0, diff --git a/cli/run_llm_judge_baselines.py b/cli/run_llm_judge_baselines.py index d10175d..267278e 100644 --- a/cli/run_llm_judge_baselines.py +++ b/cli/run_llm_judge_baselines.py @@ -238,13 +238,17 @@ def score_bundle( max_length = int(self.spec["max_length"]) for start in range(0, len(rendered), self.batch_size): batch_texts = rendered[start : start + self.batch_size] - untruncated = self.tokenizer( + encoded = self.tokenizer( batch_texts, add_special_tokens=False, - padding=False, + padding=True, truncation=False, - )["input_ids"] - lengths = [len(ids) for ids in untruncated] + return_tensors="pt", + ) + # Truncation is disabled, so every real token is retained and the + # attention mask sums to the exact untruncated length. This avoids a + # redundant second full tokenization pass over each batch. + lengths = [int(value) for value in encoded["attention_mask"].sum(dim=1).tolist()] too_long = [ start + i for i, length in enumerate(lengths) if length > max_length ] @@ -253,25 +257,33 @@ def score_bundle( f"LLM-judge prompts {too_long[:5]} exceed registered max_length={max_length}; " "truncation is prohibited" ) - encoded = self.tokenizer( - batch_texts, - add_special_tokens=False, - padding=True, - truncation=False, - return_tensors="pt", - ) encoded = { key: value.to(self.model_device) for key, value in encoded.items() } - with self.torch.inference_mode(): - output = self.model(**encoded) - mask = encoded["attention_mask"] - positions = self.torch.arange(mask.shape[1], device=mask.device).unsqueeze( - 0 + supports_keep = getattr(self.model, "_supports_logits_to_keep", None) + left_padded = self.tokenizer.padding_side == "left" + use_last_only = ( + left_padded and callable(supports_keep) and supports_keep() ) - last_positions = (positions * mask.to(dtype=self.torch.long)).argmax(dim=1) - batch_indices = self.torch.arange(mask.shape[0], device=mask.device) - next_logits = output.logits[batch_indices, last_positions] + forward_kwargs = {"logits_to_keep": 1} if use_last_only else {} + with self.torch.inference_mode(): + output = self.model(**encoded, **forward_kwargs) + if use_last_only: + # Left padding places every real final token in the last column, so + # the single retained logit row is exactly the forced-choice + # next-token distribution. Requesting only that row skips the LM-head + # projection over the whole (up to max_length) padded sequence. + next_logits = output.logits[:, -1, :] + else: + mask = encoded["attention_mask"] + positions = self.torch.arange( + mask.shape[1], device=mask.device + ).unsqueeze(0) + last_positions = ( + positions * mask.to(dtype=self.torch.long) + ).argmax(dim=1) + batch_indices = self.torch.arange(mask.shape[0], device=mask.device) + next_logits = output.logits[batch_indices, last_positions] probabilities = pairwise_positive_probability( next_logits, self.negative_token_id, self.positive_token_id ) diff --git a/cli/run_task_sweep.py b/cli/run_task_sweep.py index fb3e3cd..e781ca8 100644 --- a/cli/run_task_sweep.py +++ b/cli/run_task_sweep.py @@ -357,8 +357,15 @@ def _load_cached(features_dir: str, split: str, layer: int, suffix: str) -> Dict return bundle_cache[key] calibration_dir = args.calibration_dir or args.source_dir + summary_rows_written = 0 + summary_fsync_every = 128 with out_file.open("a", encoding="utf-8") as summary_handle: for layer in layers: + # Feature bundles are keyed by (dir, split, layer, suffix) and never + # reused across layers, so release the previous layer's arrays before + # loading this one. This bounds host/unified memory to a single + # layer's bundles instead of accumulating every layer's. + bundle_cache.clear() for probe_name in probe_names: probe_cls = TASK_PROBE_REGISTRY[probe_name] suffix = _bundle_suffix(probe_cls) @@ -468,9 +475,20 @@ def _load_cached(features_dir: str, split: str, layer: int, suffix: str) -> Dict failed += 1 summary_handle.write(json.dumps(row, sort_keys=True) + "\n") + # Flush every row so a resumed process sees every + # completed run, but only force the summary to + # stable storage periodically: prediction files + # are written atomically-and-fsynced before their + # summary row, and any summary tail lost to a hard + # crash is recomputed on resume. summary_handle.flush() - os.fsync(summary_handle.fileno()) + summary_rows_written += 1 + if summary_rows_written % summary_fsync_every == 0: + os.fsync(summary_handle.fileno()) existing_run_ids.add(run_id) + # Durably persist the full summary once every run has been recorded. + summary_handle.flush() + os.fsync(summary_handle.fileno()) print( f"completed {completed} valid runs; failed {failed}; " diff --git a/data/generation_confidence.py b/data/generation_confidence.py index e5dbb1c..1aa3b1e 100644 --- a/data/generation_confidence.py +++ b/data/generation_confidence.py @@ -68,26 +68,44 @@ def build_generation_confidence_trace( entropies: list[float] = [] margins: list[float] = [] selected_is_top1: list[bool] = [] - for index, (score, token_id) in enumerate(zip(scores, token_ids.tolist())): - values = torch.as_tensor(score).detach().float() - if values.ndim == 2 and values.shape[0] == 1: - values = values[0] - if values.ndim != 1 or not 0 <= int(token_id) < len(values): - raise ValueError( - f"Invalid generation score tensor at response token {index}" - ) - log_probs = torch.log_softmax(values, dim=-1) + token_id_list = token_ids.tolist() + # Summaries are computed in batches over the step-score matrix. Every stored + # value (selected log-prob, entropy, top-1/top-2 margin, is-top-1) is + # identical to a per-token computation up to floating-point reduction order; + # batching collapses the per-token Python/kernel dispatch that dominates + # wall-clock on MPS. The chunk bounds peak memory for large vocabularies. + chunk_size = 64 + for start in range(0, len(scores), chunk_size): + chunk_scores = scores[start : start + chunk_size] + chunk_token_ids = token_id_list[start : start + chunk_size] + rows: list[Any] = [] + for offset, score in enumerate(chunk_scores): + values = torch.as_tensor(score).detach().float() + if values.ndim == 2 and values.shape[0] == 1: + values = values[0] + if values.ndim != 1 or not 0 <= int(chunk_token_ids[offset]) < len(values): + raise ValueError( + f"Invalid generation score tensor at response token {start + offset}" + ) + rows.append(values) + stacked = torch.stack(rows, dim=0) + log_probs = torch.log_softmax(stacked, dim=-1) probabilities = torch.exp(log_probs) entropy_terms = torch.where( probabilities > 0, probabilities * torch.nan_to_num(log_probs, neginf=0.0), torch.zeros_like(probabilities), ) - top_probabilities, top_indices = torch.topk(probabilities, k=2) - selected_logprobs.append(float(log_probs[int(token_id)].item())) - entropies.append(float((-entropy_terms.sum()).item())) - margins.append(float((top_probabilities[0] - top_probabilities[1]).item())) - selected_is_top1.append(bool(int(top_indices[0].item()) == int(token_id))) + chunk_entropy = -entropy_terms.sum(dim=-1) + top_probabilities, top_indices = torch.topk(probabilities, k=2, dim=-1) + chunk_margin = top_probabilities[:, 0] - top_probabilities[:, 1] + selection = torch.as_tensor(chunk_token_ids, dtype=torch.long) + chunk_selected = log_probs[torch.arange(stacked.shape[0]), selection] + chunk_is_top1 = top_indices[:, 0] == selection + selected_logprobs.extend(float(value) for value in chunk_selected.tolist()) + entropies.extend(float(value) for value in chunk_entropy.tolist()) + margins.extend(float(value) for value in chunk_margin.tolist()) + selected_is_top1.extend(bool(value) for value in chunk_is_top1.tolist()) trace: dict[str, Any] = { "schema_version": GENERATION_CONFIDENCE_SCHEMA_VERSION, diff --git a/evaluation/hierarchical_statistics.py b/evaluation/hierarchical_statistics.py index 8a5dbc6..0f8938b 100644 --- a/evaluation/hierarchical_statistics.py +++ b/evaluation/hierarchical_statistics.py @@ -49,38 +49,43 @@ def hierarchical_paired_mean_difference( "Hierarchical inference requires at least two groups and two seeds" ) - cell_a: dict[tuple[str, str], float] = {} - cell_b: dict[tuple[str, str], float] = {} - for seed_id in unique_seeds: - for group_id in unique_groups: + # Collapse the rows to a dense seed-by-group cell-mean matrix once. The + # matrix stores exactly the per-cell means the previous dict held, in the + # same row-major (seed, group) order returned by ``np.unique``. + n_seeds = len(unique_seeds) + n_groups = len(unique_groups) + cell_mean_a = np.empty((n_seeds, n_groups), dtype=float) + cell_mean_b = np.empty((n_seeds, n_groups), dtype=float) + for seed_index, seed_id in enumerate(unique_seeds): + for group_index, group_id in enumerate(unique_groups): indices = np.flatnonzero((seeds == seed_id) & (groups == group_id)) if not len(indices): raise ValueError( f"Incomplete seed-by-group prediction grid at seed={seed_id}, group={group_id}" ) - cell_a[(seed_id, group_id)] = float(np.mean(a[indices])) - cell_b[(seed_id, group_id)] = float(np.mean(b[indices])) + cell_mean_a[seed_index, group_index] = float(np.mean(a[indices])) + cell_mean_b[seed_index, group_index] = float(np.mean(b[indices])) # Equal weight per seed-by-scenario cell prevents groups with more repeated # rollouts from masquerading as additional independent evidence. - observed_a = float(np.mean(list(cell_a.values()))) - observed_b = float(np.mean(list(cell_b.values()))) + observed_a = float(cell_mean_a.mean()) + observed_b = float(cell_mean_b.mean()) observed_diff = observed_a - observed_b rng = np.random.default_rng(seed) diffs = np.empty(n_boot, dtype=float) for iteration in range(n_boot): - sampled_seeds = rng.choice(unique_seeds, size=len(unique_seeds), replace=True) - sampled_groups = rng.choice( - unique_groups, size=len(unique_groups), replace=True + # Draw seed indices then group indices in this order: Generator.choice + # over an integer population consumes the random stream identically to + # choosing from the id arrays directly, so the resampled cells are the + # same as before. ``np.ix_`` reproduces the row-major seed-by-group + # ordering, keeping every reported statistic bit-identical while + # replacing ~n_boot * n_seeds * n_groups dict lookups with two means. + seed_choice = rng.choice(n_seeds, size=n_seeds, replace=True) + group_choice = rng.choice(n_groups, size=n_groups, replace=True) + selection = np.ix_(seed_choice, group_choice) + diffs[iteration] = float( + cell_mean_a[selection].mean() - cell_mean_b[selection].mean() ) - sample_a: list[float] = [] - sample_b: list[float] = [] - for sampled_seed in sampled_seeds: - for sampled_group in sampled_groups: - key = (str(sampled_seed), str(sampled_group)) - sample_a.append(cell_a[key]) - sample_b.append(cell_b[key]) - diffs[iteration] = float(np.mean(sample_a) - np.mean(sample_b)) sign_flip = min(float(np.mean(diffs <= 0.0)), float(np.mean(diffs >= 0.0))) return { diff --git a/evaluation/metrics.py b/evaluation/metrics.py index 1497658..d8d86a9 100644 --- a/evaluation/metrics.py +++ b/evaluation/metrics.py @@ -225,12 +225,19 @@ def paired_group_bootstrap_metric_diff( if len(unique_groups) < 2: raise ValueError("Grouped bootstrap requires at least two independent groups") + # Precompute each group's row indices once. The resampling draws and their + # order are unchanged, so results are identical to rescanning ``groups`` + # every iteration, but the per-iteration cost drops from O(n_boot * G * N) + # array scans to a dictionary lookup plus concatenation. + group_to_indices = { + group: np.flatnonzero(groups == group) for group in unique_groups + } observed_diff = float(metric_fn(y, a) - metric_fn(y, b)) rng = np.random.default_rng(seed) diffs: list[float] = [] for _ in range(n_boot): sampled_groups = rng.choice(unique_groups, size=len(unique_groups), replace=True) - sampled_indices = np.concatenate([np.flatnonzero(groups == group) for group in sampled_groups]) + sampled_indices = np.concatenate([group_to_indices[group] for group in sampled_groups]) sample_y = y[sampled_indices] if len(np.unique(sample_y)) < 2: continue diff --git a/extraction/task_extractor.py b/extraction/task_extractor.py index a5ead04..697102d 100644 --- a/extraction/task_extractor.py +++ b/extraction/task_extractor.py @@ -50,7 +50,7 @@ def __init__(self, cfg: TaskExtractionConfig): self.cfg = cfg if cfg.missing_view_policy not in {"error", "drop"}: raise ValueError("missing_view_policy must be 'error' or 'drop'") - from transformers import AutoModelForCausalLM, AutoTokenizer + from transformers import AutoModel, AutoTokenizer from cli.common import resolve_torch_device tokenizer_revision = cfg.tokenizer_revision or cfg.model_revision @@ -67,7 +67,10 @@ def __init__(self, cfg: TaskExtractionConfig): } if resolved_device == "auto": model_kwargs["device_map"] = "auto" - self.model = AutoModelForCausalLM.from_pretrained( + # Only the transformer-block hidden states are needed; loading the base + # model (no language-modeling head) yields identical ``hidden_states`` + # while skipping the full-vocabulary output projection on every forward. + self.model = AutoModel.from_pretrained( cfg.model_name, **model_kwargs, ) @@ -80,6 +83,7 @@ def __init__(self, cfg: TaskExtractionConfig): else None ) self._validate_layers() + self._truncate_unused_blocks() def _validate_layers(self) -> None: n_layers = int(getattr(self.model.config, "num_hidden_layers", -1)) @@ -89,6 +93,30 @@ def _validate_layers(self) -> None: f"Configured transformer-block layers {invalid} are invalid for a model with {n_layers} blocks" ) + def _truncate_unused_blocks(self) -> None: + """Drop transformer blocks above the highest requested layer. + + Extraction only reads ``hidden_states`` up to ``max(layers)``; any blocks + above it are computed and thrown away. For a standard decoder stack whose + single final norm produces ``last_hidden_state`` (e.g. ``Qwen3Model``), + removing the unused tail and neutralising that final norm leaves every + retained block's captured hidden state bit-identical — the top captured + entry is post-norm, so with an identity norm it equals the raw block + output, matching the untruncated model's intermediate state at that + index — while skipping the upper-network forward entirely. Models that do + not expose the expected ``.layers``/``.norm`` layout are left untouched. + """ + base = getattr(self.model, "model", self.model) + if not (hasattr(base, "layers") and hasattr(base, "norm")): + return + n_layers = int(getattr(base.config, "num_hidden_layers", len(base.layers))) + n_needed = max(self.cfg.layers) + 1 + if n_needed >= n_layers: + return + base.layers = base.layers[:n_needed] + base.norm = torch.nn.Identity() + base.config.num_hidden_layers = n_needed + def _prepared_segments(self, example: TaskExample) -> Dict[str, str]: if self.cfg.require_model_generated and example.metadata.get("data_origin") != "on_policy_generation": raise ValueError( @@ -277,25 +305,35 @@ def extract_example(self, example: TaskExample) -> Dict[int, Dict[str, np.ndarra input_tensor = torch.tensor([input_ids], device=model_device) attention_mask = torch.ones_like(input_tensor) - with torch.no_grad(): + with torch.inference_mode(): outputs = self.model( input_ids=input_tensor, attention_mask=attention_mask, output_hidden_states=True, ) - hidden_states = outputs.hidden_states + hidden_states = outputs.hidden_states + for block_index in self.cfg.layers: + hidden_state_index = block_index + 1 + if hidden_state_index >= len(hidden_states): + raise ValueError( + f"Block {block_index} maps to hidden state {hidden_state_index}, but model returned {len(hidden_states)} states" + ) + # Gather every requested block and copy to host in one device->host + # transfer instead of one per layer, minimising MPS synchronisations. + block_activations = ( + torch.stack( + [hidden_states[block_index + 1][0] for block_index in self.cfg.layers] + ) + .float() + .cpu() + .numpy() + ) wanted_views = self.cfg.views or ["full_text"] result: Dict[int, Dict[str, np.ndarray]] = {} - for block_index in self.cfg.layers: - hidden_state_index = block_index + 1 - if hidden_state_index >= len(hidden_states): - raise ValueError( - f"Block {block_index} maps to hidden state {hidden_state_index}, but model returned {len(hidden_states)} states" - ) - activations = hidden_states[hidden_state_index][0].detach().float().cpu().numpy() + for offset, block_index in enumerate(self.cfg.layers): spans = {name: token_spans[name] for name in wanted_views} result[block_index] = pool_named_spans( - activations, spans, mode=self.cfg.pooling_mode + block_activations[offset], spans, mode=self.cfg.pooling_mode ) return result diff --git a/probes/sae_probe.py b/probes/sae_probe.py index 33331ec..40c3aae 100644 --- a/probes/sae_probe.py +++ b/probes/sae_probe.py @@ -8,6 +8,12 @@ from .base import Probe +# Loaded SAEs are immutable, read-only feature encoders. The sweep constructs a +# fresh probe per run, so cache the loaded module by its pinned identity to avoid +# re-downloading/re-instantiating it on every run. +_SAE_CACHE: dict[tuple[str, str, str], object] = {} + + class SAEProbe(Probe): name = "P5_sae" scores_are_probabilities = True @@ -33,8 +39,13 @@ def load_sae(self): if not self.sae_release or not self.sae_id: raise RuntimeError("SAE release and id must be configured before using P5_sae") - self.sae = SAE.from_pretrained(self.sae_release, self.sae_id)[0].to(self.device) - self.sae.eval() + cache_key = (self.sae_release, self.sae_id, str(self.device)) + cached = _SAE_CACHE.get(cache_key) + if cached is None: + cached = SAE.from_pretrained(self.sae_release, self.sae_id)[0].to(self.device) + cached.eval() + _SAE_CACHE[cache_key] = cached + self.sae = cached def _encode(self, X: np.ndarray) -> np.ndarray: import torch