From 8ae8526c9d18369120d0741c361dae707e14f10b Mon Sep 17 00:00:00 2001 From: benthecarman <15256660+benthecarman@users.noreply.github.com> Date: Tue, 6 Oct 2026 03:10:51 -0400 Subject: [PATCH] fix: prevent empty layer range and too many neighbors in ARA (#478) * Cap ARA's neighbor count at the number of prompts The nearest neighbors are selected from the outputs for the good and bad prompts, which have one row per prompt. With fewer than 15 prompts in either dataset, the optimizer could pick a neighbor count larger than that, and torch.topk would crash with "selected index k out of range". Bounding the search range instead of clamping the value keeps the optimizer from wasting trials on counts that would all behave the same. * Render chat templates with a fixed date Some chat templates, including gpt-oss's, put the current date in the system prompt. That makes the residuals, and with them the modified model, depend on the day Heretic runs, so a model could not be reproduced on a different day. It also made the gpt-oss test's hash change daily once ARA started modifying the model. Prompts used for computing residuals and for evaluation are now rendered with a fixed date. Interactive chat keeps the real one. * Prevent ARA from sampling an empty layer range The start index was sampled from [0, L//2] and the exclusive end index from [L//2, L], so both could land on L//2 and leave no layers to modify, wasting the trial. Starting the end range at L//2 + 1 rules this out while keeping the bounds fixed, as multivariate TPE requires. On the 2-layer test models, every ARA trial had picked this empty range, so the seed-oss and gpt-oss tests never ran the optimization. They now modify layer 1, which changes their expected hashes. The seed-oss output also differs between the CPUs of GitHub's runners, because L-BFGS amplifies small floating-point differences, so it gets CI hashes as well. --- src/heretic/model.py | 10 ++++++++++ src/heretic/modifiers/ara.py | 11 +++++++++-- tests/gpt-oss/SHA256SUMS.linux | 2 +- tests/gpt-oss/SHA256SUMS.windows | 2 +- tests/seed-oss/SHA256SUMS.ci | 6 ++++++ tests/seed-oss/SHA256SUMS.ci2 | 6 ++++++ tests/seed-oss/SHA256SUMS.linux | 2 +- tests/seed-oss/SHA256SUMS.windows | 2 +- 8 files changed, 35 insertions(+), 6 deletions(-) create mode 100644 tests/seed-oss/SHA256SUMS.ci create mode 100644 tests/seed-oss/SHA256SUMS.ci2 diff --git a/src/heretic/model.py b/src/heretic/model.py index aaac268..d50b245 100644 --- a/src/heretic/model.py +++ b/src/heretic/model.py @@ -3,6 +3,7 @@ from collections.abc import Callable from contextlib import suppress +from datetime import datetime, timezone from typing import Any, TypeAlias, cast import torch @@ -43,6 +44,13 @@ def get_model_class( return AutoModelForCausalLM +# Some chat templates (e.g. gpt-oss) insert the current date into the prompt, +# which makes the residuals, and therefore the modified model, depend on the day +# Heretic is run. Rendering templates with a fixed date keeps results reproducible. +def strftime_fixed_date(format: str) -> str: + return datetime(2026, 1, 1, tzinfo=timezone.utc).strftime(format) + + # The list contains one element per layer. # Each element maps from the component name to a (possibly sparse) mapping # from the module index to an (input, output) tuple containing the I/O @@ -473,6 +481,8 @@ class Model: chats, add_generation_prompt=True, tokenize=False, + # Overrides the function that templates call to get the current date. + strftime_now=strftime_fixed_date, ), ) diff --git a/src/heretic/modifiers/ara.py b/src/heretic/modifiers/ara.py index a3b7aaf..123895b 100644 --- a/src/heretic/modifiers/ara.py +++ b/src/heretic/modifiers/ara.py @@ -189,6 +189,9 @@ class ARA(Modifier[Parameters]): bad_prompts = ctx.load_prompts(self.settings.bad_prompts) print(f"* [bold]{len(bad_prompts)}[/] prompts loaded") + self.good_prompt_count = len(good_prompts) + self.bad_prompt_count = len(bad_prompts) + print() print("Obtaining module I/O for good prompts...") self.good_module_io = model.get_module_io_batched(good_prompts) @@ -209,9 +212,11 @@ class ARA(Modifier[Parameters]): 0, layer_count // 2, ) + # The end index is exclusive, so starting its range one above the start range's + # upper bound ensures that the layer range is never empty. end_layer_index = trial.suggest_int( "end_layer_index", - layer_count // 2, + layer_count // 2 + 1, layer_count, ) preserve_good_behavior_weight = trial.suggest_float( @@ -230,10 +235,12 @@ class ARA(Modifier[Parameters]): 0.0, 1.3, ) + # The nearest neighbors are selected from the outputs for the good and bad prompts, + # which contain one vector per prompt, so there can't be more neighbors than prompts. neighbor_count = trial.suggest_int( "neighbor_count", 1, - 15, + min(15, self.good_prompt_count, self.bad_prompt_count), ) return Parameters( diff --git a/tests/gpt-oss/SHA256SUMS.linux b/tests/gpt-oss/SHA256SUMS.linux index 8e8ced3..aa0a0a0 100644 --- a/tests/gpt-oss/SHA256SUMS.linux +++ b/tests/gpt-oss/SHA256SUMS.linux @@ -1,6 +1,6 @@ f8d9255777615591a7cc1a7c932f5a69e181128902295e1b81221d20d983cac7 *chat_template.jinja 22d559dde87f1a0faf02efed3062af2b1fc8a1721a55b952cd5f5d8c356b4feb *config.json 523e422e425e7a7ba23d85cc0faeef143dc5da37b72bfc649df5a8d83ee95beb *generation_config.json -aefe8b9c4b4969f6d13c5d778760f3dce4e25134324b33677934550d9df02a7c *model.safetensors +8473b458481540948b46b4001f282073f9337bb3ca9c40b589d3a77add3ecb23 *model.safetensors fce342a4642cb8afc42d8d89cfa21198b64a43458ded7f6ff28d1151a08c9cda *tokenizer.json 9ba5fa877168e24823cb583c55b4c2e4df0331f30084953c7cf07de294640384 *tokenizer_config.json diff --git a/tests/gpt-oss/SHA256SUMS.windows b/tests/gpt-oss/SHA256SUMS.windows index 5422da6..3119497 100644 --- a/tests/gpt-oss/SHA256SUMS.windows +++ b/tests/gpt-oss/SHA256SUMS.windows @@ -1,6 +1,6 @@ e381f0b98be3c3714676ec4102bc75a885d75b9002feb02fb6f25ad0e641bce3 *chat_template.jinja 07109279e69a421627d0bd4723cbd6828d8427291824baf1c923136e536b8a8a *config.json f74b3fbaec771b2c9650e8289f8b1d6575ca164ad5834331ab0a5e79acdbd9ac *generation_config.json -aefe8b9c4b4969f6d13c5d778760f3dce4e25134324b33677934550d9df02a7c *model.safetensors +8473b458481540948b46b4001f282073f9337bb3ca9c40b589d3a77add3ecb23 *model.safetensors fce342a4642cb8afc42d8d89cfa21198b64a43458ded7f6ff28d1151a08c9cda *tokenizer.json 4c3c901d30c4e7db34ce5ad6e80d607a308fd1e980ca16e5f5d3ab250196b179 *tokenizer_config.json diff --git a/tests/seed-oss/SHA256SUMS.ci b/tests/seed-oss/SHA256SUMS.ci new file mode 100644 index 0000000..75e7be1 --- /dev/null +++ b/tests/seed-oss/SHA256SUMS.ci @@ -0,0 +1,6 @@ +39a23a1a28f68acd747373fa688bea630021ec4db2bc8dec2fd3962240fb0418 *chat_template.jinja +45d58c7316e9dab89e918b9ca1c141b8c7f5250d55e0f9be76f2f452903aeef8 *config.json +02b35dff3057b998f224a3a5599dc7737ba028e3b4181ba89ba74ba1e707a5dc *generation_config.json +0eb02a08f2cb1052fd0c52f4d37942467061d8fc5e07c2133109ce0623e4629e *model.safetensors +3cf3a6d9520f195638a36f0194239d817de7288710bca55f1f5753de226748f7 *tokenizer.json +388b47e61cb40f2fd51a89999053686ab4c45b40b43c0329d15645f5910d069e *tokenizer_config.json diff --git a/tests/seed-oss/SHA256SUMS.ci2 b/tests/seed-oss/SHA256SUMS.ci2 new file mode 100644 index 0000000..dc6f6de --- /dev/null +++ b/tests/seed-oss/SHA256SUMS.ci2 @@ -0,0 +1,6 @@ +39a23a1a28f68acd747373fa688bea630021ec4db2bc8dec2fd3962240fb0418 *chat_template.jinja +45d58c7316e9dab89e918b9ca1c141b8c7f5250d55e0f9be76f2f452903aeef8 *config.json +02b35dff3057b998f224a3a5599dc7737ba028e3b4181ba89ba74ba1e707a5dc *generation_config.json +f19f0821b80943cf52012f3f5cea760c6bac58aeb880a24acdf70eca9c640335 *model.safetensors +3cf3a6d9520f195638a36f0194239d817de7288710bca55f1f5753de226748f7 *tokenizer.json +388b47e61cb40f2fd51a89999053686ab4c45b40b43c0329d15645f5910d069e *tokenizer_config.json diff --git a/tests/seed-oss/SHA256SUMS.linux b/tests/seed-oss/SHA256SUMS.linux index bebf3f6..bfa2618 100644 --- a/tests/seed-oss/SHA256SUMS.linux +++ b/tests/seed-oss/SHA256SUMS.linux @@ -1,6 +1,6 @@ 39a23a1a28f68acd747373fa688bea630021ec4db2bc8dec2fd3962240fb0418 *chat_template.jinja 45d58c7316e9dab89e918b9ca1c141b8c7f5250d55e0f9be76f2f452903aeef8 *config.json 02b35dff3057b998f224a3a5599dc7737ba028e3b4181ba89ba74ba1e707a5dc *generation_config.json -0cdea9064dcbe6db666f9d42a283d664c133d4c67a2f6ea0fd863a54f522160c *model.safetensors +28997c7fe4a4dfaae94a7d587cd2dd0a0affed6567e95902dc52c4cd2195313e *model.safetensors 3cf3a6d9520f195638a36f0194239d817de7288710bca55f1f5753de226748f7 *tokenizer.json 388b47e61cb40f2fd51a89999053686ab4c45b40b43c0329d15645f5910d069e *tokenizer_config.json diff --git a/tests/seed-oss/SHA256SUMS.windows b/tests/seed-oss/SHA256SUMS.windows index 670efd6..cd7d282 100644 --- a/tests/seed-oss/SHA256SUMS.windows +++ b/tests/seed-oss/SHA256SUMS.windows @@ -1,6 +1,6 @@ 711867f04895d4947fb13ad8df6a675854b36cc43f124226d9fee92a306ebb24 *chat_template.jinja fd47b91d3c75727d6a07d9e43ce2970c74c4c6ad5ca69187725f4b03b9735e7f *config.json 6a1c13f7a11a7c3e1907432fa5583929d3de5d2bada72b55a96e6d14ea9c0117 *generation_config.json -0cdea9064dcbe6db666f9d42a283d664c133d4c67a2f6ea0fd863a54f522160c *model.safetensors +0eb02a08f2cb1052fd0c52f4d37942467061d8fc5e07c2133109ce0623e4629e *model.safetensors 3cf3a6d9520f195638a36f0194239d817de7288710bca55f1f5753de226748f7 *tokenizer.json ca510154d75bdc38492cb68114fa05a9bc5ce4777f8198f6ae63761926981cff *tokenizer_config.json