Speed up experiment on Apple Silicon/MPS (verified result-preserving)

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)

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,

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
             )

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}; "

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,

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 {

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