diff --git a/docs/WWPGD_PROJECTION_DOSE.md b/docs/WWPGD_PROJECTION_DOSE.md new file mode 100644 index 0000000..3232669 --- /dev/null +++ b/docs/WWPGD_PROJECTION_DOSE.md @@ -0,0 +1,51 @@ +# WW-PGD projection dose telemetry + +The main scientific runner’s `wwpgd_projection.csv` rows report the realized intervention size and distinguish scheduled events from actual per-layer applications. + +## Scope + +This schema is applied by the shared durable artifact writer used by `run_scientific_single`. It covers the scientific Level 0–2 WW-PGD, delayed-onset, adaptive, and norm-matched-sham arms that write the root experiment artifact `wwpgd_projection.csv`. + +It does not silently redefine the separate invalid-for-science `run_single` smoke artifact or the isolated `level_0_wwpgd` package’s independent projection schema. Those paths have different lifecycle contracts and require separate schema changes if they are ever promoted to scientific evidence. + +## Realized dose + +For each projected matrix, let `W_before` be the live matrix immediately before the WW-PGD projection and `W_after` the matrix after the applied, trust-region-limited projection. The reported dose is + +```text +dose_value = ||W_after - W_before||_F / max(||W_before||_F, 1e-12) +``` + +The row fields are: + +- `dose_definition = applied_projection_delta_frobenius_over_preprojection_weight_frobenius` +- `dose_value`: the realized relative Frobenius change +- `dose_relative_frobenius`: compatibility alias of `dose_value` +- `dose_applied`: true only when the live matrix actually changed and the realized dose is positive + +This is an observed dose. It is not `blend_eta`, requested hardness, or a configured trust-region cap. + +## Scheduled event indexing + +Within the main scientific runner, `projection_event` is zero-based. Rows also include: + +- `projection_event_index_base = 0` +- `projection_event_number = projection_event + 1` +- `is_first_scheduled_projection_event = (projection_event == 0)` + +The scheduled-event flag is intentionally not called `is_first_projection_event`. Warm-up, delayed onset, adaptive gates, zero hardness, unchanged candidates, and trust-region behavior can make the first scheduled event different from the first actual application. + +## First actual application + +Actual applications are counted independently for each `layer_name`: + +- `layer_application_index`: one-based application count for an applied row; empty when no dose was applied +- `is_first_applied_projection`: true exactly when `layer_application_index == 1` + +The projection CSV remains the source of truth. A small hidden sidecar caches only the current per-layer counts together with the CSV byte size and nanosecond modification time. Normal interval-one runs therefore do not rescan the complete projection history after every optimizer step. If the sidecar is absent, malformed, stale, or disagrees with a reconciled/truncated CSV, it is discarded and rebuilt from the durable CSV before the next append. + +Checkpoint recovery first reconciles the CSV to its committed prefix. The first-application annotation therefore remains correct after interruption or process restart without adding counters to the optimizer, controller, model, or checkpoint state. + +## Compatibility + +These fields extend the main scientific projection CSV schema but do not change model weights, optimizer state, projection mathematics, spectral targets, or defaults. An incomplete scientific run created with an older projection CSV header is rejected before append; start a fresh run under the new schema instead. diff --git a/src/wwgpt/checkpointing.py b/src/wwgpt/checkpointing.py index 9e85789..7386f97 100644 --- a/src/wwgpt/checkpointing.py +++ b/src/wwgpt/checkpointing.py @@ -1,5 +1,5 @@ from __future__ import annotations -import json, os, time, random, hashlib +import json, os, time, random, hashlib, math from pathlib import Path from typing import Any import numpy as np @@ -9,6 +9,12 @@ CODE_VERSION_COMPAT = ("git_commit", "git_dirty", "weightwatcher_version", "wwpgd_commit", "torch_version", "optimizer_implementation_version") REQUIRED_COMPAT=("configuration_hash","data_hash","tokenizer_hash","initialization_hash","model_configuration_hash","training_configuration_hash","wwpgd_configuration_hash","validation_probe_hash","training_probe_hash","scientific_schema_version","optimizer_fingerprint", *CODE_VERSION_COMPAT) +WWPGD_PROJECTION_DOSE_DEFINITION = ( + "applied_projection_delta_frobenius_over_preprojection_weight_frobenius" +) +WWPGD_PROJECTION_EVENT_INDEX_BASE = 0 +WWPGD_PROJECTION_APPLICATION_INDEX_SCHEMA_VERSION = 1 + def stable_hash(obj: Any) -> str: return hashlib.sha256(json.dumps(obj, sort_keys=True, default=str, separators=(",", ":")).encode()).hexdigest() @@ -65,12 +71,195 @@ def _sha256_file(path: Path) -> str: h.update(chunk) return h.hexdigest() + +def _projection_bool(value: Any, *, field: str) -> bool: + if isinstance(value, bool): + return value + token = str(value).strip().lower() + if token in {"true", "1"}: + return True + if token in {"false", "0", ""}: + return False + raise ValueError(f"{field} must be boolean-compatible, got {value!r}") + + +def _projection_event_index(value: Any) -> int: + if isinstance(value, bool): + raise ValueError("projection_event must be an integer") + try: + numeric = float(value) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("projection_event must be an integer") from exc + if not math.isfinite(numeric) or not numeric.is_integer(): + raise ValueError("projection_event must be an integer") + event = int(numeric) + if event < 0: + raise ValueError("projection_event must be non-negative") + return event + + +def _projection_application_index_path(path: Path) -> Path: + return path.with_name(f".{path.name}.application_counts.json") + + +def _scan_projection_application_counts(path: Path) -> dict[str, int]: + import csv + if not path.exists() or not path.stat().st_size: + return {} + with path.open(newline="") as handle: + reader = csv.DictReader(handle) + fields = set(reader.fieldnames or ()) + required = {"layer_name", "dose_applied", "layer_application_index"} + missing = sorted(required - fields) + if missing: + raise ValueError( + "existing WWPGD projection CSV uses an older telemetry schema; " + f"missing fields {missing}. Start a fresh run instead of appending." + ) + counts: dict[str, int] = {} + for row in reader: + layer_name = str(row.get("layer_name") or "").strip() + if not layer_name: + raise ValueError("existing WWPGD projection row is missing layer_name") + if _projection_bool(row.get("dose_applied", False), field="dose_applied"): + counts[layer_name] = counts.get(layer_name, 0) + 1 + return counts + + +def _valid_projection_application_counts(value: Any) -> dict[str, int] | None: + if not isinstance(value, dict): + return None + counts: dict[str, int] = {} + for raw_name, raw_count in value.items(): + name = str(raw_name).strip() + if not name or isinstance(raw_count, bool): + return None + try: + count = int(raw_count) + except (TypeError, ValueError, OverflowError): + return None + if count < 0 or count != raw_count: + return None + counts[name] = count + return counts + + +def _load_projection_application_counts(path: Path) -> dict[str, int]: + if not path.exists() or not path.stat().st_size: + return {} + stat = path.stat() + index_path = _projection_application_index_path(path) + try: + index = json.loads(index_path.read_text()) + except (OSError, UnicodeError, TypeError, json.JSONDecodeError): + index = None + if isinstance(index, dict): + counts = _valid_projection_application_counts(index.get("counts")) + try: + schema_version = int(index.get("schema_version", -1)) + source_size_bytes = int(index.get("source_size_bytes", -1)) + source_mtime_ns = int(index.get("source_mtime_ns", -1)) + except (TypeError, ValueError, OverflowError): + pass + else: + if ( + schema_version == WWPGD_PROJECTION_APPLICATION_INDEX_SCHEMA_VERSION + and source_size_bytes == stat.st_size + and source_mtime_ns == stat.st_mtime_ns + and counts is not None + ): + return counts + return _scan_projection_application_counts(path) + + +def _store_projection_application_counts(path: Path, counts: dict[str, int]) -> None: + stat = path.stat() + _atomic_write_json( + _projection_application_index_path(path), + { + "schema_version": WWPGD_PROJECTION_APPLICATION_INDEX_SCHEMA_VERSION, + "source_size_bytes": stat.st_size, + "source_mtime_ns": stat.st_mtime_ns, + "counts": dict(sorted(counts.items())), + }, + ) + + +def _enrich_wwpgd_projection_rows( + path: Path, + rows: list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], dict[str, int]]: + """Add exact dose and resume-safe first-application identity to projection rows. + + ``projection_event`` is zero-based. First actual application is tracked per + layer from a self-validating durable-history index. The index is rebuilt + from the reconciled CSV only after a crash, truncation, or metadata mismatch, + avoiding an O(n^2) history scan during interval-one experiments. + """ + application_counts = _load_projection_application_counts(path) + enriched: list[dict[str, Any]] = [] + for source in rows: + row = dict(source) + if "is_first_projection_event" in row: + raise ValueError( + "ambiguous is_first_projection_event is unsupported; use " + "is_first_scheduled_projection_event and is_first_applied_projection" + ) + if "projection_event" not in row: + raise ValueError("WWPGD projection row is missing projection_event") + event = _projection_event_index(row["projection_event"]) + + raw_dose = row.get("relative_frobenius_change_applied") + if raw_dose is None or raw_dose == "" or isinstance(raw_dose, bool): + raise ValueError( + "WWPGD projection row is missing relative_frobenius_change_applied" + ) + try: + dose = float(raw_dose) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError( + "relative_frobenius_change_applied must be numeric" + ) from exc + if not math.isfinite(dose) or dose < 0.0: + raise ValueError( + "relative_frobenius_change_applied must be finite and non-negative" + ) + + layer_name = str(row.get("layer_name") or "").strip() + if not layer_name: + raise ValueError("WWPGD projection row is missing layer_name") + changed = _projection_bool(row.get("changed", False), field="changed") + dose_applied = changed and dose > 0.0 + application_index = None + if dose_applied: + application_index = application_counts.get(layer_name, 0) + 1 + application_counts[layer_name] = application_index + row.update( + { + "dose_definition": WWPGD_PROJECTION_DOSE_DEFINITION, + "dose_value": dose, + "dose_relative_frobenius": dose, + "dose_applied": dose_applied, + "layer_application_index": application_index, + "is_first_applied_projection": application_index == 1, + "projection_event_index_base": WWPGD_PROJECTION_EVENT_INDEX_BASE, + "projection_event_number": event + 1, + "is_first_scheduled_projection_event": event == 0, + } + ) + enriched.append(row) + return enriched, application_counts + + def append_csv_records(path: Path, rows: list[dict[str, Any]]) -> None: """Durably append records. A flush boundary is a transaction boundary.""" import csv if not rows: return path = Path(path); path.parent.mkdir(parents=True, exist_ok=True) + projection_application_counts = None + if path.name == "wwpgd_projection.csv": + rows, projection_application_counts = _enrich_wwpgd_projection_rows(path, rows) # Cached-endpoint fast-relaxation rows already have the dedicated # wwpgd_endpoint_relaxation.csv stream. Do not duplicate them into # wwpgd_controller.csv, whose stable schema is the slow measurement record. @@ -117,6 +306,8 @@ def append_csv_records(path: Path, rows: list[dict[str, Any]]) -> None: writer.writeheader() writer.writerows(rows) f.flush(); os.fsync(f.fileno()) + if projection_application_counts is not None: + _store_projection_application_counts(path, projection_application_counts) def csv_commit(path: Path) -> dict[str, Any]: """Describe the exact durable prefix referenced by a checkpoint.""" diff --git a/tests/test_wwpgd_interval_config.py b/tests/test_wwpgd_interval_config.py index 8d6747c..b75f07a 100644 --- a/tests/test_wwpgd_interval_config.py +++ b/tests/test_wwpgd_interval_config.py @@ -51,12 +51,34 @@ def init_state(cfg): return {k: v.detach().clone() for k, v in m.state_dict().items()}, "init" +def _fake_projection_row(*, event_index, actual_step, actual_tokens_seen=None): + row = { + "projection_event": event_index, + "actual_step": actual_step, + "layer_name": "fake", + "relative_frobenius_change_applied": 0.01, + "changed": True, + } + if actual_tokens_seen is not None: + row["actual_tokens_seen"] = actual_tokens_seen + return row + + def test_extension_interval_cadences_and_skips(monkeypatch): calls = [] pre_calls = [] monkeypatch.setattr("wwgpt.train.weightwatcher_details", lambda model: pre_calls.append("pre") or object()) monkeypatch.setattr("wwgpt.train.build_stock_wwpgd_candidate", lambda *args, **kwargs: SimpleNamespace(internal_diagnostics=[])) - monkeypatch.setattr("wwgpt.train.apply_external_wwpgd", lambda *args, **kw: calls.append(kw) or [{"projection_event": kw["event_index"], "actual_step": kw["actual_step"], "actual_tokens_seen": kw["actual_tokens_seen"], "layer_name": "fake"}]) + monkeypatch.setattr( + "wwgpt.train.apply_external_wwpgd", + lambda *args, **kw: calls.append(kw) or [ + _fake_projection_row( + event_index=kw["event_index"], + actual_step=kw["actual_step"], + actual_tokens_seen=kw["actual_tokens_seen"], + ) + ], + ) ext = WWPGDExtension(WWPGDConfig(), interval=2) rows_by_step = [ext.after_optimizer_step(model=object(), optimizer_step=s, total_optimizer_steps=6, tokens_seen=s*10, collect_pre_details=True)[1] for s in range(1, 7)] @@ -82,7 +104,16 @@ def test_extension_rejects_non_positive_interval(): def test_training_interval_manifest_and_counts(monkeypatch, tmp_path): calls = [] monkeypatch.setattr("wwgpt.train.spectral_summary", lambda *a, **k: []) - monkeypatch.setattr("wwgpt.train.apply_external_wwpgd", lambda *a, **kw: calls.append(kw["actual_step"]) or [{"projection_event": kw["event_index"], "actual_step": kw["actual_step"], "actual_tokens_seen": kw["actual_tokens_seen"], "layer_name": "fake"}]) + monkeypatch.setattr( + "wwgpt.train.apply_external_wwpgd", + lambda *a, **kw: calls.append(kw["actual_step"]) or [ + _fake_projection_row( + event_index=kw["event_index"], + actual_step=kw["actual_step"], + actual_tokens_seen=kw["actual_tokens_seen"], + ) + ], + ) cfg = tiny_cfg(steps=6, interval=2) state, h = init_state(cfg) run = run_scientific_single(tmp_path, "adamw_wwpgd", 3, cfg, tiny_data(), "pair", state, h, 0, 1, device="cpu") @@ -100,7 +131,15 @@ def test_training_interval_manifest_and_counts(monkeypatch, tmp_path): def test_run_function_cli_override_precedes_config(monkeypatch, tmp_path): calls = [] monkeypatch.setattr("wwgpt.train.spectral_summary", lambda *a, **k: []) - monkeypatch.setattr("wwgpt.train.apply_external_wwpgd", lambda *a, **kw: calls.append(kw["actual_step"]) or [{"projection_event": kw["event_index"], "actual_step": kw["actual_step"], "layer_name": "fake"}]) + monkeypatch.setattr( + "wwgpt.train.apply_external_wwpgd", + lambda *a, **kw: calls.append(kw["actual_step"]) or [ + _fake_projection_row( + event_index=kw["event_index"], + actual_step=kw["actual_step"], + ) + ], + ) cfg = tiny_cfg(steps=8, interval=2) state, h = init_state(cfg) run_scientific_single(tmp_path, "adamw_wwpgd", 3, cfg, tiny_data(), "pair", state, h, 0, 1, device="cpu", ww_interval=4) diff --git a/tests/test_wwpgd_projection_dose.py b/tests/test_wwpgd_projection_dose.py new file mode 100644 index 0000000..bc641eb --- /dev/null +++ b/tests/test_wwpgd_projection_dose.py @@ -0,0 +1,232 @@ +"""Regression coverage for exact, resume-safe WW-PGD projection-dose telemetry.""" +from __future__ import annotations + +import csv +import math + +import pytest + +import wwgpt.checkpointing as checkpointing +from wwgpt.checkpointing import append_csv_records, csv_commit, reconcile_csv_artifacts + + +DOSE_DEFINITION = ( + "applied_projection_delta_frobenius_over_preprojection_weight_frobenius" +) + + +def _read_rows(path): + with path.open(newline="", encoding="utf-8") as handle: + return list(csv.DictReader(handle)) + + +def _projection_row( + *, + event: int | float, + step: int, + dose: float, + changed: bool, + layer: str = "blocks.0.attn.query", +): + return { + "layer_name": layer, + "projection_event": event, + "actual_step": step, + "relative_frobenius_change_applied": dose, + "changed": changed, + } + + +def test_projection_csv_reports_exact_realized_dose_and_first_application(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + source = _projection_row(event=0, step=10, dose=0.125, changed=True) + + append_csv_records(path, [source]) + + row = _read_rows(path)[0] + assert row["dose_definition"] == DOSE_DEFINITION + assert float(row["dose_value"]) == pytest.approx(0.125) + assert float(row["dose_relative_frobenius"]) == pytest.approx(0.125) + assert row["dose_value"] == row["relative_frobenius_change_applied"] + assert row["dose_applied"] == "True" + assert int(row["layer_application_index"]) == 1 + assert row["is_first_applied_projection"] == "True" + assert int(row["projection_event_index_base"]) == 0 + assert int(row["projection_event_number"]) == 1 + assert row["is_first_scheduled_projection_event"] == "True" + assert "is_first_projection_event" not in row + assert "dose_value" not in source # durable enrichment must not mutate callers + + +def test_first_scheduled_event_can_precede_first_actual_application(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + append_csv_records( + path, + [ + _projection_row(event=0, step=10, dose=0.0, changed=False), + _projection_row(event=1, step=20, dose=0.05, changed=True), + ], + ) + + rows = _read_rows(path) + assert rows[0]["is_first_scheduled_projection_event"] == "True" + assert rows[0]["dose_applied"] == "False" + assert rows[0]["layer_application_index"] == "" + assert rows[0]["is_first_applied_projection"] == "False" + assert float(rows[0]["dose_value"]) == 0.0 + + assert rows[1]["is_first_scheduled_projection_event"] == "False" + assert rows[1]["dose_applied"] == "True" + assert int(rows[1]["layer_application_index"]) == 1 + assert rows[1]["is_first_applied_projection"] == "True" + assert int(rows[1]["projection_event_number"]) == 2 + + first_applied = next(row for row in rows if row["is_first_applied_projection"] == "True") + assert int(first_applied["projection_event"]) == 1 + assert int(first_applied["actual_step"]) == 20 + + +def test_layer_application_index_survives_sequential_appends(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + append_csv_records(path, [_projection_row(event=0, step=10, dose=0.0, changed=False)]) + append_csv_records(path, [_projection_row(event=1, step=20, dose=0.05, changed=True)]) + append_csv_records(path, [_projection_row(event=2, step=30, dose=0.04, changed=True)]) + + rows = _read_rows(path) + assert len(rows) == 3 + assert rows[0]["layer_application_index"] == "" + assert int(rows[1]["layer_application_index"]) == 1 + assert rows[1]["is_first_applied_projection"] == "True" + assert int(rows[2]["layer_application_index"]) == 2 + assert rows[2]["is_first_applied_projection"] == "False" + assert int(rows[2]["projection_event_number"]) == 3 + + +def test_application_count_index_avoids_full_history_rescan(monkeypatch, tmp_path): + path = tmp_path / "wwpgd_projection.csv" + append_csv_records(path, [_projection_row(event=0, step=10, dose=0.05, changed=True)]) + index_path = path.with_name(f".{path.name}.application_counts.json") + assert index_path.is_file() + + def fail_scan(_path): + raise AssertionError("valid sidecar index should avoid a full CSV rescan") + + monkeypatch.setattr(checkpointing, "_scan_projection_application_counts", fail_scan) + append_csv_records(path, [_projection_row(event=1, step=20, dose=0.04, changed=True)]) + assert int(_read_rows(path)[-1]["layer_application_index"]) == 2 + + +def test_invalid_sidecar_metadata_rebuilds_from_csv(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + append_csv_records(path, [_projection_row(event=0, step=10, dose=0.05, changed=True)]) + index_path = path.with_name(f".{path.name}.application_counts.json") + index_path.write_text('{"schema_version": "broken", "counts": {}}') + + append_csv_records(path, [_projection_row(event=1, step=20, dose=0.04, changed=True)]) + assert int(_read_rows(path)[-1]["layer_application_index"]) == 2 + + +def test_reconciled_uncommitted_suffix_does_not_advance_application_count(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + append_csv_records(path, [_projection_row(event=0, step=10, dose=0.05, changed=True)]) + committed = csv_commit(path) + + # Simulate rows flushed after the last durable checkpoint, then a restart. + append_csv_records(path, [_projection_row(event=1, step=20, dose=0.04, changed=True)]) + assert int(_read_rows(path)[-1]["layer_application_index"]) == 2 + reconcile_csv_artifacts(tmp_path, {path.name: committed}) + + append_csv_records(path, [_projection_row(event=2, step=30, dose=0.03, changed=True)]) + rows = _read_rows(path) + assert len(rows) == 2 + assert int(rows[-1]["layer_application_index"]) == 2 + assert rows[-1]["is_first_applied_projection"] == "False" + + +def test_first_application_is_tracked_independently_per_layer(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + append_csv_records( + path, + [ + _projection_row( + event=0, + step=10, + dose=0.05, + changed=True, + layer="blocks.0.attn.query", + ), + _projection_row( + event=0, + step=10, + dose=0.03, + changed=True, + layer="blocks.0.attn.key", + ), + ], + ) + append_csv_records( + path, + [ + _projection_row( + event=1, + step=20, + dose=0.02, + changed=True, + layer="blocks.0.attn.query", + ) + ], + ) + + rows = _read_rows(path) + assert [int(rows[index]["layer_application_index"]) for index in (0, 1, 2)] == [1, 1, 2] + assert [rows[index]["is_first_applied_projection"] for index in (0, 1, 2)] == [ + "True", + "True", + "False", + ] + + +def test_non_projection_csv_rows_are_unchanged(tmp_path): + path = tmp_path / "metrics.csv" + append_csv_records(path, [{"step": 1, "loss": 2.5}]) + assert _read_rows(path) == [{"step": "1", "loss": "2.5"}] + + +def test_old_projection_schema_is_rejected_before_append(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter( + handle, + fieldnames=[ + "layer_name", + "projection_event", + "actual_step", + "relative_frobenius_change_applied", + "changed", + ], + ) + writer.writeheader() + writer.writerow(_projection_row(event=0, step=10, dose=0.05, changed=True)) + + before = path.read_bytes() + with pytest.raises(ValueError, match="older telemetry schema"): + append_csv_records(path, [_projection_row(event=1, step=20, dose=0.04, changed=True)]) + assert path.read_bytes() == before + + +@pytest.mark.parametrize( + "row, message", + [ + ({"relative_frobenius_change_applied": 0.1, "changed": True}, "projection_event"), + (_projection_row(event=-1, step=1, dose=0.1, changed=True), "non-negative"), + (_projection_row(event=0.5, step=1, dose=0.1, changed=True), "must be an integer"), + (_projection_row(event=0, step=1, dose=-0.1, changed=True), "finite and non-negative"), + (_projection_row(event=0, step=1, dose=math.nan, changed=True), "finite and non-negative"), + (_projection_row(event=0, step=1, dose=0.1, changed=True, layer=""), "layer_name"), + ], +) +def test_invalid_projection_telemetry_fails_before_writing(tmp_path, row, message): + path = tmp_path / "wwpgd_projection.csv" + with pytest.raises(ValueError, match=message): + append_csv_records(path, [row]) + assert not path.exists() diff --git a/tests/test_wwpgd_training_cadence.py b/tests/test_wwpgd_training_cadence.py index 30a5fab..2078da9 100644 --- a/tests/test_wwpgd_training_cadence.py +++ b/tests/test_wwpgd_training_cadence.py @@ -59,6 +59,19 @@ def _init_state(cfg: ExperimentConfig): return state, "tiny-init" +def _fake_projection_row(*, event_index, actual_step, actual_tokens_seen=None, layer_name="fake"): + row = { + "projection_event": event_index, + "actual_step": actual_step, + "layer_name": layer_name, + "relative_frobenius_change_applied": 0.01, + "changed": True, + } + if actual_tokens_seen is not None: + row["actual_tokens_seen"] = actual_tokens_seen + return row + + def test_wwpgd_runs_once_per_successful_optimizer_step(monkeypatch, tmp_path: Path): steps = 3 calls = [] @@ -67,8 +80,18 @@ def fake_apply(model, *, event_index, scheduled_token_fraction, actual_step, act assert stock_candidate is not None calls.append(actual_step) return [ - {"projection_event": event_index, "layer_name": "blocks.0.attn.key"}, - {"projection_event": event_index, "layer_name": "blocks.0.attn.value"}, + _fake_projection_row( + event_index=event_index, + actual_step=actual_step, + actual_tokens_seen=actual_tokens_seen, + layer_name="blocks.0.attn.key", + ), + _fake_projection_row( + event_index=event_index, + actual_step=actual_step, + actual_tokens_seen=actual_tokens_seen, + layer_name="blocks.0.attn.value", + ), ] monkeypatch.setattr("wwgpt.train.apply_external_wwpgd", fake_apply) @@ -136,7 +159,16 @@ def test_checkpoint_resume_is_deterministic_and_complete(monkeypatch, tmp_path: ) init_state, init_hash = _init_state(cfg) monkeypatch.setattr("wwgpt.train.spectral_summary", lambda *args, **kwargs: []) - monkeypatch.setattr("wwgpt.train.apply_external_wwpgd", lambda *args, **kwargs: [{"projection_event": kwargs["event_index"], "actual_step": kwargs["actual_step"], "actual_tokens_seen": kwargs["actual_tokens_seen"], "layer_name": "fake"}]) + monkeypatch.setattr( + "wwgpt.train.apply_external_wwpgd", + lambda *args, **kwargs: [ + _fake_projection_row( + event_index=kwargs["event_index"], + actual_step=kwargs["actual_step"], + actual_tokens_seen=kwargs["actual_tokens_seen"], + ) + ], + ) full = run_scientific_single(tmp_path / "full", "adamw_wwpgd", 7, cfg, _tiny_data(), "pair_tiny", init_state, init_hash, 0, 1, device="cpu")