From a595c493bea627acba4aed101f6fc2cc5e5cbb1e Mon Sep 17 00:00:00 2001 From: 0xEljh Date: Fri, 13 Nov 2026 08:10:48 +0000 Subject: [PATCH] Add explicit scalar calibration reuse --- AGENTS.md | 15 +++ README.md | 53 +++++++++ modulo/agent.py | 88 +++++++++++++++ tests/test_calibration.py | 226 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 382 insertions(+) create mode 100644 tests/test_calibration.py diff --git a/AGENTS.md b/AGENTS.md index 5ea5536..a0bce1a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -223,6 +223,9 @@ vec = extract_vector(model, tokenizer, ### SteeringAgent methods - `agent.step(messages, max_new_tokens=300, temperature=0.7, top_p=0.9, tools=None) -> StepResult` +- `agent.export_calibration(key=...) -> dict` — independent JSON-compatible snapshot of already measured scalar norms; no generation +- `agent.import_calibration(record, key=...)` — replace all norms from a matching record; `None` clears as a miss +- `agent.clear_calibration()` — forget imported and locally measured norms without disabling steering - `agent.build_system_prompt() -> str` — system prompt describing the steer tool - `agent.status() -> str` — active vector status - `agent.apply(name, strength) -> (bool, str)` — manually apply a vector @@ -232,6 +235,18 @@ vec = extract_vector(model, tokenizer, - `agent.debrief() -> str | None` — end the session's steering; send the returned status as a `tool` message for one last turn - `agent.menu_order() -> list[str]` — the order vectors are listed in this session (from `menu_seed`) +**Calibration reuse is opt-in startup reuse, not decode acceleration.** The caller's +nonempty key asserts complete checkpoint/tokenizer/weight/adapter/execution identity; +the library's schema, layer, dtype/device and runtime checks are sanity checks, not +automatic mutation detection. Before changing weights, adapters, tokenizer, dtype, +device or numerical settings, clear calibration and use a new key. Missing layers +calibrate normally. Invalid imports raise `ValueError` and leave all norms empty. +Export cannot relabel a bound key until clear; import itself clears/replaces. +Only canonical zero-based layers are reusable; negative-index memos cannot export. +Scalar records hold no live owners and are separate from `export_state()` steering +checkpoints. Normal tracebacks can retain owners until the caller releases errors. +See [Reuse calibration across identical model owners](README.md#reuse-calibration-across-identical-model-owners) for the complete record and operator contract. + ### StepResult - `prose` — model output with tool calls stripped diff --git a/README.md b/README.md index 5dc17ae..300abbd 100644 --- a/README.md +++ b/README.md @@ -50,6 +50,59 @@ messages += result.tool_messages() # tool results back to the model If your harness has tools of its own, pass them with `agent.step(messages, tools=[...])`. Calls to them come back unrun in `result.other_tool_calls`. See [AGENTS.md](AGENTS.md#1-as-a-library-for-existing-agent-harnesses). +#### Reuse calibration across identical model owners + +Calibration reuse is opt-in and saves startup work, not decoding work. By default, +each agent measures and memoizes its own residual norms. An operator can keep a +JSON-compatible scalar record after releasing an agent and its model: + +```python +key = "my-pinned-model-and-tokenizer; pristine; no-adapters; bf16-4090-sdpa-profile-v1" +agent.residual_norm(layer) # ordinary calibration, or let the first steered step do it +record = agent.export_calibration(key=key) # no extra calibration; json.dumps works +# Release the old agent/model, then load the exact same model and tokenizer. +new_agent = SteeringAgent(new_model, new_tokenizer) +new_agent.import_calibration(record, key=key) +``` + +**Your key asserts identity.** Include exact checkpoint and tokenizer revisions +(including the chat template), weight mutations, adapter revisions/order/scales +and enabled/merged state, generation configuration, custom code/hooks, and +numerical/backend/device settings. Distinguish physical GPUs, not just `cuda:0`. +Use a stable key for identical pristine reloads and a new unique key for mutations. +Matching keys are not proof of numerical equivalence: validate your execution +profile before reuse. The library does not hash weights or detect mutations. + +The library checks schema/recipe, key equality, model name/class, layer count, +hidden size, parameter/buffer dtype/device sets, input device, and Torch/Transformers +versions. Import/export require every model module in evaluation mode. These are +sanity checks, not complete identity checks. Records are trusted operator data, +not input from the model or community. A plausible but incorrect scalar cannot +be detected automatically. + +- `export_calibration(key=...)` returns an independent record of measured layers + only, possibly empty. Exporting under a different key requires a clear. +- `import_calibration(record, key=...)` **replaces** the memo. Unknown fields, + mismatches, invalid layer IDs, booleans, and nonpositive/nonfinite values raise + `ValueError`. Failure leaves the entire memo empty, with no partial import. + Norms accept finite positive Python integers/floats with canonical layer strings. + Import clears first, so a different matching record/key needs no separate clear. +- Reusable layers are zero-based (`0` through the last layer). An existing memo + measured via Python negative indexing cannot be exported. For reuse, configure + vectors with equivalent zero-based indices, clear, and recalibrate those layers. +- Missing layers still calibrate normally. `import_calibration(None, key=...)` + clears the memo as an explicit miss; its key is ignored. +- Before changing model, tokenizer, adapters, dtype, device, or execution settings, + call `clear_calibration()` and use a new key. It clears imported **and** locally + measured norms. Recreate the agent when replacing its model. There is no + automatic invalidation, concurrency support, or mid-step import/export. +- To disable reuse, omit imports; clear an existing memo if necessary. Steering + remains enabled. Calibration records are separate from steering checkpoints. + +Records contain no tensors, models, vectors, KV, or GPU state, and perform no file +I/O. Ordinary exception tracebacks remain intact and can retain owners/activations; +release caught errors and their tracebacks before expecting model/GPU collection. + ### As a standalone demo ```bash diff --git a/modulo/agent.py b/modulo/agent.py index 5a285a1..581e957 100644 --- a/modulo/agent.py +++ b/modulo/agent.py @@ -547,6 +547,7 @@ def __init__(self, model, tokenizer, registry: Optional[VectorRegistry] = None, self._turn_generated = False self._unit_vectors: dict[str, torch.Tensor] = {} self._residual_norms: dict[int, float] = {} + self._calibration_key: Optional[str] = None # name -> (strength, turn) of the most recent generation each vector steered self._last_steered: dict[str, tuple[float, int]] = {} # name -> turn of the latest report on it (a vector is reportable again once re-steered) @@ -925,6 +926,93 @@ def hook(module, inp, output): self._residual_norms[layer] = h.norm(dim=-1).median().item() return self._residual_norms[layer] + def clear_calibration(self) -> None: + """Forget all measured/imported norms. Subsequent use calibrates normally.""" + self._residual_norms.clear() + self._calibration_key = None + + def _calibration_profile(self) -> dict: + # Deliberately not a weight fingerprint. The operator key covers everything + # not checked here, including tokenizer, adapters and numerical settings. + import transformers + + if any(module.training for module in self.model.modules()): + raise ValueError("Calibration reuse requires all model modules in eval mode.") + return { + "recipe": [1, "median-generated-l2-full-replay", CALIBRATION_PROMPT, CALIBRATION_TOKENS], + "model": self.model_name, + "class": f"{type(self.model).__module__}.{type(self.model).__qualname__}", + "layers": len(decoder_layers(self.model)), + "hidden_size": getattr(self.model.config, "hidden_size", None), + "input_device": str(self.model.device), + "placements": [list(p) for p in sorted({(str(t.dtype), str(t.device)) + for t in (*self.model.parameters(), *self.model.buffers())})], + "torch": str(torch.__version__), + "transformers": transformers.__version__, + } + + def export_calibration(self, *, key: str) -> dict: + """Return a detached JSON-compatible record of already measured norms. + + key is the operator's identity for exact weights/tokenizer/adapters and + execution settings. It is an assertion, not automatic mutation detection. + Does not calibrate missing layers. See README.md for the reuse contract. + """ + if self._calibration_key is not None and key != self._calibration_key: + raise ValueError("Clear calibration before changing its operator key.") + record = {"schema": 1, "key": key, "profile": self._calibration_profile(), + "norms": {str(k): v for k, v in self._residual_norms.items()}} + self._validate_calibration(record, key=key, profile=record["profile"]) + self._calibration_key = key + return record + + @staticmethod + def _validate_calibration(record, *, key, profile) -> dict[int, float]: + if not isinstance(key, str) or not key.strip(): + raise ValueError("Calibration key must be a nonempty operator identity string.") + if (not isinstance(record, dict) or set(record) != {"schema", "key", "profile", "norms"} + or type(record["schema"]) is not int or record["schema"] != 1): + raise ValueError("Invalid calibration record schema.") + if type(record["key"]) is not str or record["key"] != key: + raise ValueError("Calibration operator key mismatch.") + try: + matches = json.dumps(record["profile"], sort_keys=True, allow_nan=False) == json.dumps(profile, sort_keys=True) + except (TypeError, ValueError, RecursionError) as error: + raise ValueError("Invalid calibration profile.") from error + if not matches: + raise ValueError("Calibration profile mismatch.") + norms = record["norms"] + if not isinstance(norms, dict): + raise ValueError("Calibration norms must be a layer-to-scalar dictionary.") + checked = {} + valid_layers = {str(i) for i in range(profile["layers"])} + for layer, value in norms.items(): + if (not isinstance(layer, str) or layer not in valid_layers + or type(value) not in (int, float)): + raise ValueError("Calibration norms require valid layer strings and positive finite numbers.") + try: + scalar = float(value) + except OverflowError as error: + raise ValueError("Calibration norm is outside the finite scalar range.") from error + if not math.isfinite(scalar) or scalar <= 0: + raise ValueError("Calibration norms must be positive and finite.") + checked[int(layer)] = scalar + return checked + + def import_calibration(self, record: Optional[dict], *, key: str) -> None: + """Replace all norms with a matching operator-owned record; None is a miss. + + Malformed/mismatched records raise ValueError and leave all norms empty. + Missing layers calibrate normally. Call only between steps on unchanged + models; clear before any mutation and use a new key. Tracebacks are not + redacted and can retain owners until the caller releases them. + """ + self.clear_calibration() + if record is not None: + self._residual_norms.update(self._validate_calibration( + record, key=key, profile=self._calibration_profile())) + self._calibration_key = key + def _build_steering_context(self, prompt_len: int): if not self.active or self.control == "sham": return nullcontext() diff --git a/tests/test_calibration.py b/tests/test_calibration.py new file mode 100644 index 0000000..a60c6d1 --- /dev/null +++ b/tests/test_calibration.py @@ -0,0 +1,226 @@ +"""Explicit scalar reuse across owners; no checkpoints or CUDA required.""" +import gc +import json +import traceback +import weakref + +import pytest +import torch +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from transformers import PreTrainedTokenizerFast, Qwen2Config, Qwen2ForCausalLM + +from modulo import SteeringAgent, VectorRegistry + + +KEY = "checkpoint@revision; tokenizer@revision; pristine; adapters=none; cpu-fp32-sdpa-v1" + + +def make_agent(tmp_path): + torch.manual_seed(17) + model = Qwen2ForCausalLM(Qwen2Config( + vocab_size=12, hidden_size=16, intermediate_size=32, num_hidden_layers=2, + num_attention_heads=2, num_key_value_heads=2, head_dim=8, + eos_token_id=None, pad_token_id=0, attn_implementation="sdpa", + )).eval() + raw = Tokenizer(WordLevel({"[PAD]": 0, "[UNK]": 1, "user": 2, "assistant": 3, + "Tell": 4, "me": 5, "a": 6, "little": 7, "about": 8, + "how": 9, "you": 10, "questions": 11}, unk_token="[UNK]")) + raw.pre_tokenizer = Whitespace() + tokenizer = PreTrainedTokenizerFast(tokenizer_object=raw, pad_token="[PAD]", unk_token="[UNK]") + tokenizer.model_input_names = ["input_ids", "attention_mask"] + tokenizer.chat_template = ("{% for m in messages %}{{ m['role'] + ' ' + m['content'] + ' ' }}" + "{% endfor %}{% if add_generation_prompt %}assistant {% endif %}") + registry = VectorRegistry(tmp_path) + return SteeringAgent(model, tokenizer, registry=registry, model_name="tiny", menu_seed=23) + + +@pytest.fixture +def agent(tmp_path): + return make_agent(tmp_path) + + +def test_cold_import_roundtrip_skips_calibration_and_copies_scalars(agent, tmp_path, monkeypatch): + expected = agent.residual_norm(0) + record = json.loads(json.dumps(agent.export_calibration(key=KEY))) + other = make_agent(tmp_path) + other.import_calibration(record, key=KEY) + monkeypatch.setattr(other, "_generate_ids", lambda *a, **k: pytest.fail("hit generated calibration")) + record["norms"]["0"] = expected * 2 + assert other.residual_norm(0) == expected + assert other.export_calibration(key=KEY)["norms"] == {"0": expected} + + +def test_default_missing_layer_none_and_reset_use_ordinary_calibration(agent, tmp_path, monkeypatch): + other = make_agent(tmp_path) + expected = agent.residual_norm(0) + record = agent.export_calibration(key=KEY) + calls = [] + generate = other._generate_ids + def counted(*args, **kwargs): + calls.append(1) + return generate(*args, **kwargs) + monkeypatch.setattr(other, "_generate_ids", counted) + assert other.export_calibration(key=KEY)["norms"] == {} # export is lazy + assert other.residual_norm(0) == expected + other.import_calibration(record, key=KEY) + assert other.residual_norm(0) == expected + assert other.residual_norm(1) == agent.residual_norm(1) + assert len(calls) == 2 # ordinary first call plus missing layer + other.clear_calibration() + assert other.residual_norm(0) == expected + other.import_calibration(None, key=KEY) + assert other.residual_norm(0) == expected + assert len(calls) == 4 + + +@pytest.mark.parametrize("change", [ + lambda r: r.update(schema=2), lambda r: r.update(schema=True), + lambda r: r.update(extra=1), lambda r: r.update(key="other-model"), + lambda r: r["profile"].update(placements=[["torch.float32", "cuda:0"]]), + lambda r: r.update(norms={"0": True}), lambda r: r.update(norms={"0": 0}), + lambda r: r.update(norms={"0": -1}), lambda r: r.update(norms={"0": float("nan")}), + lambda r: r.update(norms={"0": float("inf")}), lambda r: r.update(norms={"0": "3"}), + lambda r: r.update(norms={"0": 1.0, "1": False}), + lambda r: r.update(norms={"02": 1.0}), lambda r: r.update(norms={"2": 1.0}), + lambda r: r.update(norms={"-1": 1.0}), lambda r: r.update(norms={0: 1.0}), + lambda r: r.update(norms=[]), +]) +def test_invalid_import_clears_previous_norms_atomically(agent, change): + agent.residual_norm(0) + record = agent.export_calibration(key=KEY) + change(record) + with pytest.raises(ValueError): + agent.import_calibration(record, key=KEY) + assert agent.export_calibration(key=KEY)["norms"] == {} + + +def test_operator_identity_dtype_and_training_reject_stale_record(agent): + agent.residual_norm(0) + record = agent.export_calibration(key=KEY) + for key in ("different-checkpoint", "adapter-scale=2", "tokenizer-v2", "weights:mutation-uuid", ""): + with pytest.raises(ValueError): + agent.import_calibration(record, key=key) + agent.model.to(dtype=torch.bfloat16) + with pytest.raises(ValueError, match="profile"): + agent.import_calibration(record, key=KEY) + agent.model.float().train() + with pytest.raises(ValueError, match="eval"): + agent.import_calibration(record, key=KEY) + + +def test_export_cannot_relabel_an_imported_or_exported_memo(agent): + agent.residual_norm(0) + record = agent.export_calibration(key=KEY) + with pytest.raises(ValueError, match="Clear"): + agent.export_calibration(key="new-adapter") + agent.import_calibration(record, key=KEY) + with pytest.raises(ValueError, match="Clear"): + agent.export_calibration(key="new-adapter") + agent.clear_calibration() + assert agent.export_calibration(key="new-adapter")["norms"] == {} + + +def test_unserializable_recursive_profile_raises_validation_error_and_clears(agent): + agent.residual_norm(0) + record = agent.export_calibration(key=KEY) + profile = record["profile"] + for _ in range(1500): + profile = {"nested": profile} + record["profile"] = profile + with pytest.raises(ValueError, match="profile"): + agent.import_calibration(record, key=KEY) + assert agent.export_calibration(key=KEY)["norms"] == {} + + +def test_import_replaces_key_and_layers_without_a_separate_clear(agent, tmp_path): + agent.residual_norm(0) + agent.export_calibration(key=KEY) + other = make_agent(tmp_path) + expected = other.residual_norm(1) + record = other.export_calibration(key="new-profile") + agent.import_calibration(record, key="new-profile") + assert agent.export_calibration(key="new-profile")["norms"] == {"1": expected} + + +def test_negative_index_calibration_remains_default_only(agent): + expected = agent.residual_norm(-1) + assert expected == agent.residual_norm(1) + with pytest.raises(ValueError, match="valid layer strings"): + agent.export_calibration(key=KEY) + assert agent.residual_norm(-1) == expected + agent.clear_calibration() + assert agent.export_calibration(key=KEY)["norms"] == {} + + +def test_record_does_not_keep_model_or_agent_alive(tmp_path): + agent = make_agent(tmp_path) + agent.residual_norm(0) + record = agent.export_calibration(key=KEY) + refs = [weakref.ref(agent), weakref.ref(agent.model), weakref.ref(next(agent.model.parameters()))] + del agent + gc.collect() + assert all(ref() is None for ref in refs) + assert json.loads(json.dumps(record)) == record + + +def test_original_calibration_exception_keeps_traceback_and_removes_hook(agent, monkeypatch): + inputs = agent._encode([{"role": "user", "content": "hello"}]) + monkeypatch.setattr(agent, "_generate_ids", lambda *a, **k: inputs.input_ids) + def fail_forward(*args, **kwargs): + raise RuntimeError("original diagnostic") + monkeypatch.setattr(agent.model, "forward", fail_forward) + with pytest.raises(RuntimeError, match="original diagnostic") as error: + agent.residual_norm(0) + assert "fail_forward" in "".join(traceback.format_exception(error.value)) + assert agent.export_calibration(key=KEY)["norms"] == {} + assert not agent.model.model.layers[0]._forward_hooks + + +def test_retained_error_preserves_diagnostics_until_caller_releases_it(tmp_path): + def fail_import(): + agent = make_agent(tmp_path) + record = agent.export_calibration(key=KEY) + refs = [weakref.ref(agent), weakref.ref(agent.model)] + try: + agent.import_calibration(record, key="different-weight-state") + except ValueError as error: + return error, refs + error, refs = fail_import() + gc.collect() + assert all(ref() is not None for ref in refs) # ordinary traceback ownership + assert "fail_import" in "".join(traceback.format_exception(error)) + del error + gc.collect() + assert all(ref() is None for ref in refs) + + +def test_real_adapter_transitions_require_new_operator_keys(agent): + peft = pytest.importorskip("peft") + agent.model = peft.get_peft_model(agent.model, peft.LoraConfig( + task_type="CAUSAL_LM", r=2, lora_alpha=2, target_modules=["q_proj", "v_proj"], + )).eval() + with torch.no_grad(): + for name, parameter in agent.model.named_parameters(): + if "lora_B" in name: + parameter.fill_(0.2) + key = KEY + ";lora@revision;scale=1;enabled" + agent.residual_norm(0) + record = agent.export_calibration(key=key) + # The cheap profile is deliberately not an adapter-state detector. + for module in agent.model.modules(): + if isinstance(getattr(module, "scaling", None), dict): + module.scaling["default"] = 2.0 + agent.import_calibration(record, key=key) # caller's unchanged assertion is trusted + with pytest.raises(ValueError, match="key"): + agent.import_calibration(record, key=key + ";scale=2") + with agent.model.disable_adapter(): + with pytest.raises(ValueError, match="key"): + agent.import_calibration(record, key=key + ";disabled") + assert agent.export_calibration(key=key + ";disabled")["norms"] == {} + agent.clear_calibration() + agent.model = agent.model.merge_and_unload().eval() + with pytest.raises(ValueError): + agent.import_calibration(record, key=key + ";merged") + assert agent.export_calibration(key=key + ";merged")["norms"] == {}