diff --git a/docs/measurements/continuation-observation-pilot-20261101.md b/docs/measurements/continuation-observation-pilot-20261101.md new file mode 100644 index 00000000..f55f8f86 --- /dev/null +++ b/docs/measurements/continuation-observation-pilot-20261101.md @@ -0,0 +1,266 @@ +# Training continuation with optional TrainCheck observations + +## Question and result + +Can an agent use TrainCheck as an optional observation tool while continuing a +training task, without being told to find a bug or being supplied a target +invariant? This development pilot exercises that workflow with two fresh agents: +ordinary development tools versus the same tools plus TrainCheck instrumentation. +The inference engine is outside the experiment. + +Both agents independently diagnosed the same extra gradient-normalization +factor from source inspection and selected the same one-line revision at the +original learning rate. The treatment agent subsequently used a native +TrainCheck probe and its own Python checker to corroborate that diagnosis. This +establishes a working observation workflow; it does **not** demonstrate an agent +quality, discovery, speed, or cost advantage from TrainCheck. + +This is one calibrated synthetic development case, not an official RSI benchmark +result, a blinded experiment, or an estimate of success rates across tasks. + +## Neutral task and experimental boundary + +Both agents received byte-identical copies of 11 common files: the previous and +current pipeline, shared runner, actual training/validation curves, run profile, +engineering guidance, and [TASK.md](../../tools/continuation_pilot/TASK.md). The +task says that a change to microbatch execution reduced training quality and +asks the agent to decide what to investigate or try next. It permits learning +rate changes, diagnostics, and rerunning either version while preserving the +effective batch, intended mean-token objective, and public interfaces. It does +not say that the implementation is faulty or supply a diagnosis or invariant. + +The treatment additionally received [TRAINCHECK.md](../../tools/continuation_pilot/TRAINCHECK.md) +and 32 instrumenter/support Python files. The guide explains selected callables, +argument fields, explicit probes for local state, native event identities, and +arbitrary offline checkers. It makes no unmeasured low-overhead promise. There is +no inferred rule, targeted normalization checker, or inference engine in either +packet. Prior conversations, operator tests/protocol, other arm files, and private +evaluation were excluded by instructions. This was not an OS access boundary. + +Both agents used the same inherited model/reasoning configuration with no parent +history. Exact operator dispatch messages are saved in each packet. Each arm +could request four two-GPU runs of at most 24 optimizer updates, including failed +runs, and perform tiny CPU operation checks. The 15-minute reasoning target was +advisory and excluded operator queue/setup time; tokens and actual reasoning +time were not independently measured or enforced. + +## Workload and calibration + +Ring18, two NVIDIA A16 GPUs, Python 3.10.14, PyTorch 2.2.2+cu121, FP32. The model +is a dropout-free, 5,504-parameter Transformer: width 16, two layers, two attention +heads, feed-forward width 32, vocabulary 32, sequence length 8. Synthetic targets +are `(token + 1) % 32`. SGD uses learning rate 0.15 and momentum 0.9. There is no +gradient clipping. This is a new small development workload, not the earlier +6.2M-parameter TorchTitan Llama inference trace. + +Each rank generates its complete eight-example batch before splitting it into +microbatches. Accumulation changes therefore retain the same data and 128 global +tokens per optimizer update. All tokens are valid. The reference concatenates +microbatches and performs one full-batch backward pass; the current pipeline +uses DDP `no_sync()` except on the last microbatch. Training and independent +validation curves report unscaled mean-token cross entropy, not the scalar +passed to backward. Evaluation preserves RNG, model state, and module modes. + +The operator-known initial fixture divides by accumulation in both `current.py` +and the shared `backward` helper. SGD was chosen because adaptive optimizers can +cancel uniform gradient scaling. This information was not in the agent prompt. +The committed `training_task/current.py` deliberately retains the initial task; +the agents' selected revisions are frozen separately in the evidence archive. + +Calibration is part of the evidence: + +- A first qualification launch failed before training because the supervisor + hides CUDA by default. The corrected command explicitly supplies + `CUDA_VISIBLE_DEVICES=0,1` to its child. The failure log is retained. +- The first completed qualification used 64 updates. Both versions nearly + converged: the final-three-validation-point means differed by 0.0452539 nats, + and the current/reference improvement ratio was 0.987573. This failed the + preselected gap threshold of 0.05 and improvement-ratio threshold of 0.8. +- Before launching either agent, the development budget was changed to 24 + updates using these calibration curves. A fresh qualification passed the + same criteria: final-three means 0.663464 versus 2.142884 nats, ratio 0.503334. + Final validation loss was 0.00896090380 for reference and 1.18585254252 for + current. All attempts are retained; this is not a first-attempt preregistered + success. No held-out seed or solver outcome was used to select the task. + +Protocol revision 2 SHA-256: +`0a71f6e452fee57862bb9a8cfa2b1d685efe5414a7dd22b75b1e36f12e073861`. +Implementation/evaluator revision: `8d511938eb5940cbc02023b9096df317768ae1bc`. +The workload was first committed as `13550b75cdc5ee14d1556a27e516ddb191ebc257`; +the runtime workload files are identical across those revisions. + +## Agent decisions and development evidence + +Both agents removed the same extra division, retaining LR 0.15: + +```diff +- backward(loss / accumulation, accumulation) ++ backward(loss, accumulation) +``` + +The control agent wrote ordinary autograd regressions on a tiny embedding +operation. The treatment wrote a tiny CPU `FixedLogits` diagnostic, inserted a +`loss_input_probe` immediately before the shared backward helper, wrapped that +probe with TrainCheck, and parsed the native trace with its own checker. It did +not wrap native backward execution or instrument the GPU training runs. + +Both CPU diagnostics measured gradient ratios 1, 1/2, 1/4, and 1/8 before the +revision for accumulation 1, 2, 4, and 8, then matching gradients after it. The +treatment's two native traces each contain 15 unique, complete pre/post probe +pairs. Its observed and unobserved post-revision diagnostic JSON outputs are +byte-identical. That covers the selected probe, not every backward event or +arbitrary training state. + +Development seed 20261101, evaluation seed 777000, 24 updates per run: + +| Arm | Run | Accumulation | Final validation loss | +| --- | --- | ---: | ---: | +| Control | selected current | 4 | 0.00896090368042 | +| Control | reference rerun | 4 | 0.00896090379683 | +| Control | selected current | 8 | 0.00896090053720 | +| Control | selected current | 1 | 0.00896090379683 | +| TrainCheck | selected current | 4 | 0.00896090368042 | +| TrainCheck | reference rerun | 4 | 0.00896090379683 | +| TrainCheck | selected current | 8 | 0.00896090053720 | + +All seven requests completed, with unchanged snapshotted sources. Each agent's +complete checkpoint comparison against its reference rerun found maximum model +and momentum differences of 2.38419e-7 and 5.96046e-8, respectively, for +accumulation 4/8. CPU/CUDA RNG states and optimizer parameter groups matched +exactly. Control's accumulation-1 tensors/curves match reference exactly; +accumulation metadata differences are recorded explicitly. These are numerical +comparisons of every saved tensor, not a claim that all checkpoint files match +byte-for-byte. + +## Frozen independent evaluation + +Both final selections and sources were frozen at **2026-11-01 10:57:36 UTC**, before +scoring. All immutable packet files were rechecked. The 117-file freeze manifest +has SHA-256 `0fdf1f84cb14e15e0b97692b3a6877a2bfdda253ba212f0908eaaec7837662b4`. +Selected `current.py` hashes: + +- Control: `c5d452fbfbd94dbd85a601e6e6d453a1dd59e34bb2cf2f59b6fc78f5db4db62e`. +- TrainCheck: `f3a29fdcae9fd25ecc3de21478b61d92d8ec0e5d810092ec585a811665d25bb9`. + +The files differ in their explanatory comment, not the executed operations. +The private evaluator first computes frozen reference outcomes, then imports the +selected candidate in a fresh process. Quality and contract outcomes are +reported separately. Held-out quality uses training seeds 20261117/20261129, +evaluation seed 777001, 24 updates, accumulation 4, and two GPUs. Candidate order +is reversed for the second seed. Independent CPU contract scenarios cover both +seeds and accumulation 1/2/4, comparing all parameters, gradients at the actual +optimizer step, momentum, RNG, scalar/container structure, and the microbatch +bound. FP32 tolerance is `atol=2e-6, rtol=2e-4`; exact content hashes are retained. + +All eight GPU runs completed with frozen sources unchanged and matching rank +curves. Final independent validation losses: + +| Pipeline | Seed 20261117 | Seed 20261129 | Mean | +| --- | ---: | ---: | ---: | +| Reference | 0.011329145054 | 0.008691648662 | 0.010010396858 | +| Original current | 1.258648961782 | 1.070365890861 | 1.164507426322 | +| Control selection | 0.011329147383 | 0.008691646799 | 0.010010397091 | +| TrainCheck selection | 0.011329147383 | 0.008691646799 | 0.010010397091 | + +Both selected candidates pass all six CPU scenarios, each comparing 82 tensor +pairs and scalar/container structure, with one optimizer step and the required +microbatch bound. The original current passes accumulation 1 but fails +accumulation 2/4 on both seeds, with 81 tensor mismatches per failing scenario. +The checker thus executed and rejected a known negative, rather than merely +reporting no violations. + +An additional independent comparison of complete held-out GPU checkpoints +passes for reference/control, reference/treatment, and control/treatment across +both seeds and ranks: **12 checkpoint pairs, 672 tensor pairs**, without removing +any saved field. All control/treatment tensor hashes match exactly, along with +all scalar/structure checks. Reference/candidate comparisons use the stated FP32 +tolerance. The two selected solutions have identical held-out curves; this is a +tie for this paired task, not a general equivalence result. + +## Costs and limitations + +Control used 4/4 training requests and treatment used 3/4. Their summed client +request times were 260.321 s and 205.142 s, including 79.021 s and 70.987 s of +queueing. These are operational receipts, not causal efficiency estimates: +control requested an extra boundary configuration, early uploads included type +checker caches, and setup/queue conditions differed. No GPU training request +retained TrainCheck instrumentation, and CPU probe overhead was not measured. +No tracing slowdown or low-overhead result follows from this experiment. + +This simple case exposes the duplicate divisor directly in readable source; +both agents reached the diagnosis before the runtime confirmation. Two held-out +seeds validate the selected implementations on this task, but do not create two +independent agent trials. The CPU arithmetic check does not certify distributed +collectives, long-run stability, or checkpoint resumption. The GPU runs exercise +standard two-rank DDP only. More realistic continuation tasks should require +cross-component state observations before investing in a larger benchmark. + +## Validation and reproduction + +The 40 focused tests pass locally (1.29 s) and on Ring18 (3.58 s). They cover +deterministic batches, state-preserving unscaled evaluation, malformed run +options, the known negative, a correct revision, learning-rate compensation, +full-batch escape, multiple optimizer steps, exceptions, missing/invalid +interfaces, frozen-reference isolation, import-time tracing setup, and source +mutation. Pinned pre-commit checks pass. Explicit strict mypy covers eight +Python modules, including the tools/tests that the configured core-only mypy +hook does not select; tests also passed explicit pinned lint/format checks. +These checks concern the committed harness/evaluator. The archived one-off +operator transport scripts were not included in that static-check claim. No +remote CI result is claimed. + +```bash +python -m pytest -q -p no:cacheprovider \ + tests/test_continuation_workload.py tests/test_continuation_evaluator.py +mypy --explicit-package-bases --check-untyped-defs --disallow-untyped-defs \ + --disallow-any-generics --ignore-missing-imports --follow-imports=silent \ + tools/continuation_pilot tests/test_continuation_workload.py \ + tests/test_continuation_evaluator.py +``` + +Private evidence roots: + +- Local: `/Users/matt/gitrepos/traincheck/outputs/ring18-continuation-20261101`. +- Ring18: `/home/yuxuan/experiments/traincheck-continuation-20261101`. + +Archive in both roots: `continuation-pilot-evidence.tar.gz`, 19,997,742 bytes, +442 individually verified evidence files plus the manifest. SHA-256: +`9f72d6e42c78e5ceb9f12439dadfaac0b5c0a21e3185a5f8aad2c7c5d7cc5d79`. +The archive was reopened and every member's size and content hash verified. + +The evidence includes exact prompts, packet hashes, all qualification attempts, +requests, original input archives, raw native probe traces, metrics, complete +checkpoints, frozen selections, and independent evaluation reports. Large +artifacts and machine-specific transport scripts remain outside Git. Inspect +`independent-results.json`, `frozen/freeze.json`, `heldout/contracts/`, each +`phase-result.json`, and the per-rank metrics/checkpoints. Individual commands +and environment are preserved in `invocation.json`. The initial SSH launch +wrapper timed out while launching held-out scoring; the detached process +continued and completed. It was not resubmitted. + +To repeat the fixed candidate quality check after restoring the archive, use +the same Python/Torch environment, reserve two GPUs, set a new output directory, +and run from a frozen candidate directory: + +```bash +CUDA_VISIBLE_DEVICES=0,1 CUBLAS_WORKSPACE_CONFIG=:4096:8 \ + python -m torch.distributed.run --standalone --nproc-per-node=2 \ + -m training_task.engine --pipeline current --output /new/output \ + --steps 24 --accumulation 4 --seed 20261117 --eval-seed 777001 \ + --lr 0.15 --device cuda +``` + +For the separate CPU arithmetic check, run from the committed repository (or +the archive's `frozen/operator` directory), in a fresh process per candidate: + +```bash +python -m tools.continuation_pilot.evaluate_candidate \ + --candidate /restored/frozen/control/training_task/current.py \ + --lr 0.15 --output /existing/report-directory/control.json +``` + +Exit status 0 means pass, 1 means a contract violation, and 2 means an evaluation +error with an unknown outcome. This reproduces the implementation check, not a +new independent agent trial. Repeating the agent experiment requires fresh +agents and freshly prepared neutral packets; never expose the operator protocol, +this report, or selected solutions in those packets. diff --git a/tests/test_continuation_evaluator.py b/tests/test_continuation_evaluator.py new file mode 100644 index 00000000..79e44801 --- /dev/null +++ b/tests/test_continuation_evaluator.py @@ -0,0 +1,300 @@ +"""Operator-only checks must separate numerical violations from invalid evidence.""" + +from __future__ import annotations + +import importlib +import json +import os +import shutil +import sys +from pathlib import Path + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from tools.continuation_pilot import evaluate_candidate as evaluator +from tools.continuation_pilot.training_task import current, engine, reference + + +def candidate_packet(root: Path, *, corrected: bool) -> Path: + package = root / "training_task" + package.mkdir() + source = Path(engine.__file__).parent + for name in ("__init__.py", "engine.py", "reference.py", "current.py"): + shutil.copyfile(source / name, package / name) + path = package / "current.py" + if corrected: + text = path.read_text() + original = "backward(loss / accumulation, accumulation)" + assert text.count(original) == 1 + path.write_text(text.replace(original, "backward(loss, accumulation)")) + return path + + +def test_corrected_candidate_matches_full_batch_update_and_preserves_rng( + tmp_path: Path, +) -> None: + path = candidate_packet(tmp_path, corrected=True) + rng = torch.get_rng_state().clone() + threads = torch.get_num_threads() + report = evaluator.evaluate_candidate(path, 0.15) + assert report["status"] == "passed" + assert report["contract_passed"] is True + assert len(report["scenarios"]) == 6 + assert torch.equal(torch.get_rng_state(), rng) + assert torch.get_num_threads() == threads + for case in report["scenarios"]: + assert case["mean_loss_matches"] is True + assert case["state_matches"] is True + assert case["optimizer_steps"] == 1 + assert ( + case["forward_batch_sizes"] + == [8 // case["accumulation"]] * case["accumulation"] + ) + assert case["microbatch_bound_satisfied"] is True + tensors = [item for item in case["comparisons"] if item["kind"] == "tensor"] + assert ( + len(tensors) == 82 + ) # 27 parameters, gradients, momentum buffers + CPU RNG. + assert any("applied_gradients" in item["path"] for item in tensors) + assert any("optimizer" in item["path"] for item in tensors) + assert any("model" in item["path"] for item in tensors) + assert any("cpu_rng" in item["path"] for item in tensors) + assert all(len(item["candidate_sha256"]) == 64 for item in tensors) + for seed in evaluator.HELD_OUT_SEEDS: + cases = [item for item in report["scenarios"] if item["seed"] == seed] + assert all( + item["full_batch_sha256"] == cases[0]["full_batch_sha256"] for item in cases + ) + + +def test_original_candidate_has_real_gradient_and_update_discrepancies( + tmp_path: Path, +) -> None: + path = candidate_packet(tmp_path, corrected=False) + report = evaluator.evaluate_candidate(path, 0.15) + assert report["status"] == "violated" + for case in report["scenarios"]: + assert case["mean_loss_matches"] is True + assert case["microbatch_bound_satisfied"] is True + assert case["contract_passed"] is (case["accumulation"] == 1) + if case["accumulation"] > 1: + failures = [item for item in case["comparisons"] if not item["passed"]] + for component in ("model", "optimizer", "applied_gradients"): + assert any(component in item["path"] for item in failures) + + +def test_reverting_to_full_batch_does_not_pass_microbatch_contract() -> None: + cases = evaluator.evaluate_update(reference.run_update, 0.15) + for case in cases: + assert case["state_matches"] is True + assert case["forward_batch_sizes"] == [8] + assert case["microbatch_bound_satisfied"] is (case["accumulation"] == 1) + + +def test_post_step_gradient_cleanup_does_not_hide_applied_gradient_evidence() -> None: + def update( + model: nn.Module, + optimizer: torch.optim.Optimizer, + batches: list[engine.Batch], + ) -> float: + value = reference.run_update(model, optimizer, batches) + optimizer.zero_grad(set_to_none=True) + return value + + model = engine.TinyTransformer() + batch = engine.make_batch(3, 1, 0, torch.device("cpu")) + result = evaluator.run_observed_update(update, model, batch, 1, 0.15) + assert all(parameter.grad is None for parameter in model.parameters()) + assert isinstance(result.state, dict) + gradients = result.state["applied_gradients"] + assert isinstance(gradients, list) and len(gradients) == 1 + assert isinstance(gradients[0], dict) and len(gradients[0]) == 27 + + +def test_missing_gradient_is_an_error_not_zero_discrepancy() -> None: + def update( + model: nn.Module, + optimizer: torch.optim.Optimizer, + batches: list[engine.Batch], + ) -> float: + optimizer.step() + return 0.0 + + model = engine.TinyTransformer() + batch = engine.make_batch(3, 1, 0, torch.device("cpu")) + with pytest.raises(ValueError, match="Missing applied gradient"): + evaluator.run_observed_update(update, model, batch, 1, 0.15) + + +def test_nonfinite_gradient_is_an_error() -> None: + def update( + model: nn.Module, + optimizer: torch.optim.Optimizer, + batches: list[engine.Batch], + ) -> float: + for parameter in model.parameters(): + parameter.grad = torch.full_like(parameter, float("nan")) + optimizer.step() + return 0.0 + + model = engine.TinyTransformer() + batch = engine.make_batch(3, 1, 0, torch.device("cpu")) + with pytest.raises(ValueError, match="nonfinite tensor"): + evaluator.run_observed_update(update, model, batch, 1, 0.15) + + +def test_no_optimizer_step_is_an_error() -> None: + def update( + model: nn.Module, + optimizer: torch.optim.Optimizer, + batches: list[engine.Batch], + ) -> float: + return 0.0 + + with pytest.raises(ValueError, match="Expected one optimizer step"): + evaluator.evaluate_update(update, 0.15) + + +@pytest.mark.parametrize("value", [None, True, 1, torch.tensor(1.0), float("nan")]) +def test_invalid_return_is_rejected(value: object) -> None: + def invalid( + model: nn.Module, + optimizer: torch.optim.Optimizer, + batches: list[engine.Batch], + ) -> object: + return value + + with pytest.raises((TypeError, ValueError)): + evaluator.validated_update(invalid)( + engine.TinyTransformer(), + torch.optim.SGD([nn.Parameter(torch.ones(1))], lr=0.1), + [], + ) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), 0.0, -1.0]) +def test_invalid_learning_rate_is_rejected(value: float) -> None: + with pytest.raises(ValueError, match="Learning rate"): + evaluator.evaluate_update(current.run_update, value) + + +def test_state_comparison_retains_dtype_shape_order_and_scalar_bits() -> None: + left = evaluator.clone_state({"x": torch.tensor([1.0]), "y": [0.0]}, "state") + right = evaluator.clone_state( + {"y": [-0.0], "x": torch.tensor([1.0], dtype=torch.float64)}, "state" + ) + records = evaluator.compare_states(left, right, "state") + failures = [item for item in records if not item["passed"]] + assert {item["path"] for item in failures} == { + "state.keys", + "state['x']", + "state['y'][0]", + } + tensor = next(item for item in records if item["kind"] == "tensor") + assert tensor["max_abs_error"] is None + assert tensor["reference_dtype"] == "torch.float32" + assert tensor["candidate_dtype"] == "torch.float64" + + +def test_loader_restores_existing_package_binding_even_after_error( + tmp_path: Path, +) -> None: + package = importlib.import_module("tools.continuation_pilot.training_task") + previous = sys.modules.get("training_task") + original_path = list(sys.path) + sys.modules["training_task"] = package + path = tmp_path / "current.py" + path.write_text("run_update = 3\n") + try: + with pytest.raises(TypeError, match="callable"): + with evaluator.load_candidate(path): + raise AssertionError("Malformed callable unexpectedly loaded") + assert sys.modules["training_task"] is package + assert sys.path == original_path + finally: + if previous is None: + del sys.modules["training_task"] + else: + sys.modules["training_task"] = previous + + +def test_cli_records_unknown_outcome_for_candidate_exception( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + path = tmp_path / "current.py" + path.write_text("raise ZeroDivisionError('candidate import failed')\n") + output = tmp_path / "report.json" + monkeypatch.setattr( + sys, + "argv", + ["check", "--candidate", str(path), "--lr", "0.15", "--output", str(output)], + ) + assert evaluator.main() == 2 + report: object = json.loads(output.read_text()) + assert isinstance(report, dict) + assert report["status"] == "error" + assert report["contract_passed"] is None + assert report["error_type"] == "ZeroDivisionError" + + +def test_reference_is_frozen_before_candidate_installs_global_wrapper( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + path = candidate_packet(tmp_path, corrected=True) + # Keep pytest's restoration registered before packet import changes F. + monkeypatch.setattr(F, "cross_entropy", F.cross_entropy) + path.write_text( + path.read_text() + + """ +import torch.nn.functional as _functional +_original_cross_entropy = _functional.cross_entropy +def _scaled_cross_entropy(*args: object, **kwargs: object) -> torch.Tensor: + result: object = _original_cross_entropy(*args, **kwargs) + if not isinstance(result, torch.Tensor): + raise TypeError("Cross entropy must return a tensor") + return result * 0.5 +_functional.cross_entropy = _scaled_cross_entropy +""" + ) + report = evaluator.evaluate_candidate(path, 0.15, tmp_path) + assert report["contract_passed"] is False + assert all(not case["mean_loss_matches"] for case in report["scenarios"]) + + +def test_changed_source_is_an_error_instead_of_wrong_provenance(tmp_path: Path) -> None: + path = candidate_packet(tmp_path, corrected=True) + path.write_text( + path.read_text() + + """ +from pathlib import Path as _Path +_source_path = _Path(__file__) +_source_path.write_text(_source_path.read_text() + "\\n") +""" + ) + with pytest.raises(ValueError, match="source changed"): + evaluator.evaluate_candidate(path, 0.15, tmp_path) + + +def test_optional_diagnostics_get_unique_output_before_import( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + path = candidate_packet(tmp_path, corrected=True) + path.write_text( + path.read_text() + + """ +import os as _os +from pathlib import Path as _Path +_trace_directory = _Path(_os.environ["TRAINCHECK_OUTPUT_DIR"]) +(_trace_directory / "imported.txt").write_text("native diagnostics may initialize here") +""" + ) + monkeypatch.setenv("TRAINCHECK_OUTPUT_DIR", "/previous-value") + first = evaluator.evaluate_candidate(path, 0.15, tmp_path) + second = evaluator.evaluate_candidate(path, 0.15, tmp_path) + assert first["contract_passed"] and second["contract_passed"] + assert first["diagnostic_trace_directory"] != second["diagnostic_trace_directory"] + assert Path(first["diagnostic_trace_directory"], "imported.txt").is_file() + assert os.environ["TRAINCHECK_OUTPUT_DIR"] == "/previous-value" diff --git a/tests/test_continuation_workload.py b/tests/test_continuation_workload.py new file mode 100644 index 00000000..74d7ce0e --- /dev/null +++ b/tests/test_continuation_workload.py @@ -0,0 +1,192 @@ +"""Operator-only qualification tests; excluded from the agent's task packet.""" + +from __future__ import annotations + +import copy +import json +import math +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +from tools.continuation_pilot.training_task.current import run_update as current_update +from tools.continuation_pilot.training_task.engine import ( + Batch, + Options, + TinyTransformer, + evaluate, + make_batch, + parse_options, + run, + split_batch, +) +from tools.continuation_pilot.training_task.reference import ( + run_update as reference_update, +) + + +def test_batch_is_reproducible_independent_of_accumulation_and_global_rng() -> None: + torch.manual_seed(17) + before = torch.get_rng_state().clone() + device = torch.device("cpu") + original = make_batch(123, 7, 0, device) + repeated = make_batch(123, 7, 0, device) + assert torch.equal(original.tokens, repeated.tokens) + assert torch.equal(original.labels, repeated.labels) + assert torch.equal(torch.get_rng_state(), before) + assert not torch.equal(original.tokens, make_batch(123, 7, 1, device).tokens) + for accumulation in (1, 2, 4, 8): + parts = split_batch(original, accumulation) + assert torch.equal(torch.cat([part.tokens for part in parts]), original.tokens) + assert torch.equal(torch.cat([part.labels for part in parts]), original.labels) + + +@pytest.mark.parametrize("accumulation", [1, 2, 4]) +def test_current_gradient_differs_by_one_accumulation_factor( + accumulation: int, +) -> None: + torch.set_num_threads(1) + torch.manual_seed(31) + reference = TinyTransformer().train() + current = copy.deepcopy(reference) + batch = make_batch(444, 1, 0, torch.device("cpu")) + parts = split_batch(batch, accumulation) + reference_optimizer = torch.optim.SGD(reference.parameters(), lr=0.01, momentum=0.9) + current_optimizer = torch.optim.SGD(current.parameters(), lr=0.01, momentum=0.9) + reference_loss = reference_update(reference, reference_optimizer, parts) + current_loss = current_update(current, current_optimizer, parts) + assert current_loss == pytest.approx(reference_loss, rel=1e-6, abs=1e-6) + nonzero = False + for reference_parameter, current_parameter in zip( + reference.parameters(), current.parameters(), strict=True + ): + assert ( + reference_parameter.grad is not None and current_parameter.grad is not None + ) + assert torch.allclose( + current_parameter.grad * accumulation, + reference_parameter.grad, + rtol=1e-5, + atol=2e-6, + ) + nonzero = nonzero or bool(reference_parameter.grad.abs().max() > 1e-4) + assert nonzero + if accumulation > 1: + assert any( + not torch.equal(left, right) + for left, right in zip( + reference.parameters(), current.parameters(), strict=True + ) + ) + + +def test_evaluation_preserves_rng_parameters_gradients_and_module_modes() -> None: + torch.set_num_threads(1) + torch.manual_seed(43) + model = TinyTransformer().train() + model.embedding.eval() + model(torch.tensor([[0, 1]])).sum().backward() + state = {key: value.clone() for key, value in model.state_dict().items()} + gradients = [ + parameter.grad.clone() + for parameter in model.parameters() + if parameter.grad is not None + ] + modes = [module.training for module in model.modules()] + before = torch.get_rng_state().clone() + score = evaluate(model, torch.device("cpu"), 127) + assert math.isfinite(score) and score > 0 + assert evaluate(model, torch.device("cpu"), 127) == score + assert torch.equal(torch.get_rng_state(), before) + assert [module.training for module in model.modules()] == modes + for name, parameter in model.state_dict().items(): + assert torch.equal(parameter, state[name]) + actual_gradients = [ + parameter.grad for parameter in model.parameters() if parameter.grad is not None + ] + assert len(actual_gradients) == len(gradients) + for actual, expected in zip(actual_gradients, gradients, strict=True): + assert torch.equal(actual, expected) + + +class UniformModel(nn.Module): + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + return torch.zeros((*tokens.shape, 32), device=tokens.device) + + +def test_evaluation_reports_unscaled_token_cross_entropy() -> None: + value = evaluate(UniformModel(), torch.device("cpu")) + assert value == pytest.approx(math.log(32), abs=1e-6) + + +@pytest.mark.parametrize( + "arguments", + [ + ["--steps", "129"], + ["--steps", "0"], + ["--accumulation", "3"], + ["--accumulation", "0"], + ["--seed", "-1"], + ["--eval-seed", "-1"], + ["--lr", "nan"], + ["--lr", "0"], + ], +) +def test_cli_rejects_invalid_experiment_configuration(arguments: list[str]) -> None: + with pytest.raises(SystemExit): + parse_options(["--pipeline", "current", "--output", "/unused", *arguments]) + + +def test_invalid_complete_batch_shape_is_rejected() -> None: + malformed = Batch(torch.ones((7, 8)), torch.ones((7, 8))) + with pytest.raises(ValueError, match="shape"): + split_batch(malformed, 2) + + +@pytest.mark.parametrize("pipeline", ["reference", "current"]) +def test_short_cpu_run_writes_real_curve_and_checkpoint( + tmp_path: Path, pipeline: str +) -> None: + output = tmp_path / pipeline + options = parse_options( + [ + "--pipeline", + pipeline, + "--output", + str(output), + "--steps", + "2", + "--device", + "cpu", + "--accumulation", + "4", + "--seed", + "117", + "--eval-seed", + "719", + ] + ) + run(options) + raw: object = json.loads((output / "rank0/metrics.json").read_text()) + assert isinstance(raw, dict) + assert raw["optimizer_updates"] == 2 + assert raw["tokens_per_global_update"] == 64 + assert raw["eval_seed"] == 719 + curve = raw["validation_curve"] + assert isinstance(curve, list) + assert [item["step"] for item in curve] == [0, 2] + checkpoint: object = torch.load(output / "rank0/checkpoint.pt", weights_only=True) + assert isinstance(checkpoint, dict) + assert checkpoint["completed_steps"] == 2 + assert isinstance(checkpoint["model"], dict) + assert isinstance(checkpoint["optimizer"], dict) + + +def test_programmatic_boolean_count_is_rejected_before_output(tmp_path: Path) -> None: + options = Options("current", tmp_path / "run", device="cpu") + object.__setattr__(options, "steps", True) + with pytest.raises(TypeError, match="steps"): + run(options) + assert not options.output.exists() diff --git a/tools/continuation_pilot/TASK.md b/tools/continuation_pilot/TASK.md new file mode 100644 index 00000000..0fefc762 --- /dev/null +++ b/tools/continuation_pilot/TASK.md @@ -0,0 +1,42 @@ +We changed this training pipeline to process each per-rank batch as configurable +microbatches with gradient accumulation, reducing activation memory while +retaining two-GPU training. The attached curves show the previous and current +pipelines under the same training budget. The current result is below expectations. + +Continue this task: decide what to investigate or try next, implement the most +justified revision, and report the evidence for your decision. You may rerun +either pipeline, adjust the learning rate, and add diagnostics. Preserve the +microbatch capability, the effective batch, the intended mean-token objective, +and the public interfaces. Keep the model, dataset, SGD momentum, and evaluator +unchanged. You may also conclude that further investigation is preferable to a +code change if the evidence supports that decision. + +Your editable implementation is `training_task/current.py`. You may add local +Python diagnostic helpers. `training_task/reference.py` and `engine.py` are the +fixed previous implementation and shared runner. The learning rate is a run +option. All new functions should have explicit types; do not suppress errors or +replace missing required data with a successful result. + +You have four training requests, each of at most 24 optimizer updates, on two +GPUs. Requests for either pipeline and diagnostic training runs use the same +allowance; failed requests also count. Tiny CPU checks on individual operations +are allowed without training a model. The supplied curves cost neither arm a +request. Aim to spend about 15 minutes deciding and experimenting, excluding +time waiting for the shared training service. + +To request a run, write `request-N.json` containing `pipeline` (`reference` or +`current`), `steps`, `accumulation`, `learning_rate`, and a short `purpose`. Send +the operator the file path. It will snapshot the current task files, run the +request, and return the complete log, metrics, curves, and checkpoint location. +Use unique request names. Do not access the training host directly. + +Finish by writing `selection.json` with required fields `learning_rate`, +`rationale`, and `evidence` (a list of relative evidence file paths). The final +implementation is the current contents of `training_task/current.py` and any +helpers it imports. Explain what you changed or retained, why, what actually ran, +and what remains uncertain. Your work will be evaluated on fresh training seeds. + +Only inspect and modify your assigned packet. Do not inspect the parent +repository, operator workspace, other agents, previous study outputs, or private +evaluation materials. This scope rule is an experimental protocol, not a claim +of operating-system isolation. diff --git a/tools/continuation_pilot/TRAINCHECK.md b/tools/continuation_pilot/TRAINCHECK.md new file mode 100644 index 00000000..26c46225 --- /dev/null +++ b/tools/continuation_pilot/TRAINCHECK.md @@ -0,0 +1,57 @@ +You additionally have TrainCheck's runtime observation facility. Its source is +available in the supplied `traincheck/` package. Use it if it would help your +next decision; it is optional. + +TrainCheck can wrap selected Python/PyTorch callables and record API entry, +successful return, or exception events, with process/thread/call identities and +timestamps. You can choose argument fields, enable or suppress recording for +selected scopes, and write your own analysis or checker in any useful format. +There is no requirement to use an invariant inference engine or a predefined +relation language. Selective observation lets you limit the amount of recorded +data; actual cost depends on the observations and should be measured when relevant. + +The basic callable interface is: + +```python +from traincheck.config import config +from traincheck.instrumentor.tracer import wrapper + +# callable_to_observe is a function you selected from your program. +observed = wrapper( + callable_to_observe, + is_bound_method=False, + scan_proxy_in_args=False, + dump_stack_trace=False, + dump_ret=False, + handle_proxy=False, +) +config.DISABLE_WRAPPER = False +# Use observed(...) at the same call site with the original arguments. +``` + +The training service sets `TRAINCHECK_OUTPUT_DIR` to a unique per-rank directory +before importing your pipeline. Do not share an output file between ranks. +Native JSON lines are written as `trace_API_*.log` and become available after the +worker exits. Inspect actual records before writing a parser; primitive values +have tagged payloads such as `{"int": 3}`. Post records do not repeat arguments. +Preserve identities when joining entry/return records. + +For selective argument recording, `dump_args_config` maps a positional argument +index as a string, or a keyword name, to the desired metadata selection; an empty +selection value records that argument normally. Omitted arguments are not +recorded. For example, `{"0": {}}` records the first positional argument. +`dump_ret=False` omits the return value. The current implementation does not +apply the advertised `dump_ret_config` field filter, so do not rely on it. + +`traincheck.instrumentor.tracer.suppress_tracing()` excludes recorder-owned work +within a scope. `traincheck.instrumentor.control.start_step()` can maintain step +metadata and apply the configured sampling interval. Arbitrary function locals +are not automatically visible: a small explicit probe or helper at the actual +execution site can expose a value you decide is useful. Live configuration-file +changes do not automatically reconfigure already-installed wrappers. + +Recordings should capture actual execution state without changing computation, +RNG, or object lifetimes. A passing predicate needs evidence that the intended +observations actually occurred; missing records, failed calls, or checker errors +are distinct from successful checks. You may analyze raw traces offline or +manage your own checker; no problem-specific rule is supplied. diff --git a/tools/continuation_pilot/evaluate_candidate.py b/tools/continuation_pilot/evaluate_candidate.py new file mode 100644 index 00000000..574bfb22 --- /dev/null +++ b/tools/continuation_pilot/evaluate_candidate.py @@ -0,0 +1,537 @@ +"""Operator-only accumulation contract evaluation for a trusted candidate. + +This CPU evaluator is separate from held-out two-GPU training quality. It checks +one SGD update against the frozen full-batch implementation, with matched initial +state, data and learning rate. It does not certify distributed communication or +isolate arbitrary candidate code. Run each selected candidate in a fresh process. +""" + +from __future__ import annotations + +import argparse +import copy +import hashlib +import importlib.util +import json +import math +import os +import struct +import sys +import tempfile +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType +from typing import Literal, Protocol, TypeAlias, TypedDict + +import torch +import torch.nn as nn + +from tools.continuation_pilot.training_task import engine, reference + +HELD_OUT_SEEDS = (20261117, 20261129) +ACCUMULATIONS = (1, 2, 4) +ATOL = 2e-6 +RTOL = 2e-4 + +State: TypeAlias = ( + torch.Tensor + | None + | bool + | int + | float + | str + | list["State"] + | tuple["State", ...] + | dict[str | int, "State"] +) + + +class Update(Protocol): + def __call__( + self, + model: nn.Module, + optimizer: torch.optim.Optimizer, + batches: list[engine.Batch], + ) -> float: ... + + +class TensorComparison(TypedDict): + kind: Literal["tensor"] + path: str + passed: bool + reference_shape: list[int] + candidate_shape: list[int] + reference_dtype: str + candidate_dtype: str + reference_sha256: str + candidate_sha256: str + max_abs_error: float | None + + +class ScalarComparison(TypedDict): + kind: Literal["scalar"] + path: str + passed: bool + reference: str + candidate: str + + +Comparison: TypeAlias = TensorComparison | ScalarComparison + + +class Scenario(TypedDict): + seed: int + accumulation: int + full_batch_sha256: dict[str, str] + reference_loss: float + candidate_loss: float + mean_loss_matches: bool + optimizer_steps: int + forward_batch_sizes: list[int] + permitted_microbatch_size: int + microbatch_bound_satisfied: bool | None + state_matches: bool + contract_passed: bool + comparisons: list[Comparison] + + +class Evaluation(TypedDict): + status: Literal["passed", "violated"] + contract_passed: bool + scope: str + learning_rate: float + atol: float + rtol: float + candidate_path: str + candidate_sha256: str + reference_sha256: str + diagnostic_trace_directory: str + scenarios: list[Scenario] + limitations: list[str] + + +@dataclass +class Observation: + loss: float + state: State + forward_batch_sizes: list[int] + optimizer_steps: int + + +def clone_state(value: object, path: str) -> State: + """Snapshot every accepted leaf; reject unsupported or nonfinite values.""" + if isinstance(value, torch.Tensor): + if value.layout != torch.strided or value.is_quantized: + raise TypeError(f"{path}: only dense nonquantized tensors are supported") + if value.is_floating_point() or value.is_complex(): + if not bool(torch.isfinite(value).all()): + raise ValueError(f"{path}: nonfinite tensor") + return value.detach().cpu().clone() + if value is None: + return None + if type(value) in (bool, int, float, str): + if isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"{path}: nonfinite scalar") + if isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, list): + return [ + clone_state(item, f"{path}[{index}]") for index, item in enumerate(value) + ] + if isinstance(value, tuple): + return tuple( + clone_state(item, f"{path}[{index}]") for index, item in enumerate(value) + ) + if isinstance(value, dict): + result: dict[str | int, State] = {} + for key, item in value.items(): + if type(key) not in (str, int) or not isinstance(key, (str, int)): + raise TypeError(f"{path}: invalid state mapping key") + result[key] = clone_state(item, f"{path}[{key!r}]") + return result + raise TypeError(f"{path}: unsupported state value {type(value).__name__}") + + +def tensor_sha256(value: torch.Tensor) -> str: + data = value.detach().cpu().contiguous().reshape(-1).view(torch.uint8) + return hashlib.sha256(data.numpy().tobytes()).hexdigest() + + +def scalar_comparison(path: str, left: object, right: object) -> ScalarComparison: + same_type = type(left) is type(right) + equal = left == right + if isinstance(left, float) and isinstance(right, float): + equal = struct.pack("!d", left) == struct.pack("!d", right) + return { + "kind": "scalar", + "path": path, + "passed": same_type and equal is True, + "reference": repr(left), + "candidate": repr(right), + } + + +def compare_states(left: State, right: State, path: str) -> list[Comparison]: + """Compare all leaves, retaining mapping/sequence order and tensor metadata.""" + reports: list[Comparison] = [] + if isinstance(left, torch.Tensor) and isinstance(right, torch.Tensor): + comparable = left.shape == right.shape and left.dtype == right.dtype + maximum: float | None = None + if comparable and (left.is_floating_point() or left.is_complex()): + passed = bool(torch.allclose(left, right, atol=ATOL, rtol=RTOL)) + maximum = float((left - right).abs().max().item()) if left.numel() else 0.0 + else: + passed = comparable and torch.equal(left, right) + reports.append( + { + "kind": "tensor", + "path": path, + "passed": passed, + "reference_shape": list(left.shape), + "candidate_shape": list(right.shape), + "reference_dtype": str(left.dtype), + "candidate_dtype": str(right.dtype), + "reference_sha256": tensor_sha256(left), + "candidate_sha256": tensor_sha256(right), + "max_abs_error": maximum, + } + ) + elif isinstance(left, dict) and isinstance(right, dict): + left_keys = [(type(key).__name__, key) for key in left] + right_keys = [(type(key).__name__, key) for key in right] + reports.append(scalar_comparison(path + ".keys", left_keys, right_keys)) + for key in left: + if key in right and (type(key).__name__, key) in right_keys: + reports.extend( + compare_states(left[key], right[key], f"{path}[{key!r}]") + ) + elif isinstance(left, (list, tuple)) and isinstance(right, (list, tuple)): + reports.append( + scalar_comparison(path + ".type", type(left).__name__, type(right).__name__) + ) + reports.append(scalar_comparison(path + ".length", len(left), len(right))) + for index, (a, b) in enumerate(zip(left, right)): + reports.extend(compare_states(a, b, f"{path}[{index}]")) + elif isinstance(left, torch.Tensor) or isinstance(right, torch.Tensor): + reports.append( + scalar_comparison(path + ".type", type(left).__name__, type(right).__name__) + ) + else: + reports.append(scalar_comparison(path, left, right)) + return reports + + +def validated_update(value: object) -> Update: + if not callable(value): + raise TypeError("Candidate run_update must be callable") + + def invoke( + model: nn.Module, + optimizer: torch.optim.Optimizer, + batches: list[engine.Batch], + ) -> float: + result: object = value(model, optimizer, batches) + if not isinstance(result, float): + raise TypeError("Candidate run_update must return a float") + if not math.isfinite(result): + raise ValueError("Candidate returned nonfinite mean loss") + return result + + return invoke + + +@contextmanager +def load_candidate(path: Path) -> Iterator[Update]: + """Load trusted packet code with relative imports, restoring import bindings. + + The packet's engine is verified immutable by the operator. This loader is an + import adapter, not a security sandbox or rollback of arbitrary side effects. + """ + candidate = path.resolve(strict=True) + if not candidate.is_file() or candidate.suffix != ".py": + raise ValueError("Candidate must be a Python source file") + package_name = "training_task" + module_name = package_name + ".current" + previous_path = list(sys.path) + previous_modules = { + name: module + for name, module in sys.modules.items() + if name == package_name or name.startswith(package_name + ".") + } + package = ModuleType(package_name) + package.__path__ = [str(candidate.parent)] + spec = importlib.util.spec_from_file_location(module_name, candidate) + if spec is None or spec.loader is None: + raise ImportError("Cannot construct candidate module loader") + module = importlib.util.module_from_spec(spec) + try: + for name in previous_modules: + del sys.modules[name] + sys.modules[package_name] = package + sys.modules[module_name] = module + sys.path[:0] = [str(candidate.parent.parent), str(candidate.parent)] + spec.loader.exec_module(module) + if "run_update" not in vars(module): + raise ValueError("Candidate does not define run_update") + yield validated_update(vars(module)["run_update"]) + finally: + sys.path[:] = previous_path + for name in list(sys.modules): + if name == package_name or name.startswith(package_name + "."): + del sys.modules[name] + sys.modules.update(previous_modules) + + +def run_observed_update( + update: Update, + model: nn.Module, + batch: engine.Batch, + accumulation: int, + learning_rate: float, +) -> Observation: + optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9) + applied_gradients: list[dict[str, torch.Tensor]] = [] + forward_sizes: list[int] = [] + + def observe_forward(module: nn.Module, args: tuple[object, ...]) -> None: + if len(args) != 1 or not isinstance(args[0], torch.Tensor) or args[0].ndim != 2: + raise ValueError("Expected a single two-dimensional token batch") + forward_sizes.append(args[0].shape[0]) + + def observe_step( + observed_optimizer: torch.optim.Optimizer, + args: tuple[object, ...], + kwargs: dict[str, object], + ) -> None: + if observed_optimizer is not optimizer: + raise ValueError("Unexpected optimizer observed") + gradients: dict[str, torch.Tensor] = {} + for name, parameter in model.named_parameters(): + if parameter.grad is None: + raise ValueError(f"Missing applied gradient for {name}") + snapshot = clone_state(parameter.grad, f"gradients.{name}") + assert isinstance(snapshot, torch.Tensor) + gradients[name] = snapshot + if not gradients: + raise ValueError("No model parameter gradients observed") + applied_gradients.append(gradients) + + forward_hook = model.register_forward_pre_hook(observe_forward) + step_hook = optimizer.register_step_pre_hook(observe_step) + try: + loss = validated_update(update)( + model, optimizer, engine.split_batch(batch, accumulation) + ) + finally: + forward_hook.remove() + step_hook.remove() + if len(applied_gradients) != 1: + raise ValueError( + f"Expected one optimizer step, observed {len(applied_gradients)}" + ) + state = clone_state( + { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "applied_gradients": applied_gradients, + "cpu_rng": torch.get_rng_state(), + }, + "update", + ) + return Observation(loss, state, forward_sizes, len(applied_gradients)) + + +@dataclass(frozen=True) +class ReferenceCase: + seed: int + accumulation: int + initial_model: nn.Module + rng: torch.Tensor + batch: engine.Batch + observation: Observation + + +def prepare_reference_cases(learning_rate: float) -> list[ReferenceCase]: + """Freeze reference outcomes before candidate import or instrumentation.""" + if isinstance(learning_rate, bool) or not isinstance(learning_rate, (float, int)): + raise TypeError("Learning rate must be numeric") + if not math.isfinite(learning_rate) or learning_rate <= 0: + raise ValueError("Learning rate must be finite and positive") + cases: list[ReferenceCase] = [] + previous_threads = torch.get_num_threads() + torch.set_num_threads(1) + try: + with torch.random.fork_rng(devices=[]): + for seed in HELD_OUT_SEEDS: + torch.manual_seed(seed) + initial = engine.TinyTransformer().train() + rng = torch.get_rng_state().clone() + batch = engine.make_batch(seed, 1, 0, torch.device("cpu")) + for accumulation in ACCUMULATIONS: + torch.set_rng_state(rng) + expected = run_observed_update( + reference.run_update, + copy.deepcopy(initial), + engine.Batch(batch.tokens.clone(), batch.labels.clone()), + accumulation, + learning_rate, + ) + cases.append( + ReferenceCase(seed, accumulation, initial, rng, batch, expected) + ) + finally: + torch.set_num_threads(previous_threads) + return cases + + +def compare_candidate_cases( + update: Update, learning_rate: float, cases: list[ReferenceCase] +) -> list[Scenario]: + if not cases: + raise ValueError("No reference cases supplied") + scenarios: list[Scenario] = [] + previous_threads = torch.get_num_threads() + torch.set_num_threads(1) + try: + with torch.random.fork_rng(devices=[]): + for case in cases: + expected = case.observation + torch.set_rng_state(case.rng) + actual = run_observed_update( + update, + copy.deepcopy(case.initial_model), + engine.Batch(case.batch.tokens.clone(), case.batch.labels.clone()), + case.accumulation, + learning_rate, + ) + comparisons = compare_states(expected.state, actual.state, "update") + state_matches = all(item["passed"] for item in comparisons) + loss_matches = math.isclose( + expected.loss, actual.loss, abs_tol=ATOL, rel_tol=RTOL + ) + maximum_batch = engine.EXAMPLES_PER_RANK_UPDATE // case.accumulation + batch_bound = ( + all( + 0 < size <= maximum_batch for size in actual.forward_batch_sizes + ) + if actual.forward_batch_sizes + else None + ) + scenarios.append( + { + "seed": case.seed, + "accumulation": case.accumulation, + "full_batch_sha256": { + "tokens": tensor_sha256(case.batch.tokens), + "labels": tensor_sha256(case.batch.labels), + }, + "reference_loss": expected.loss, + "candidate_loss": actual.loss, + "mean_loss_matches": loss_matches, + "optimizer_steps": actual.optimizer_steps, + "forward_batch_sizes": actual.forward_batch_sizes, + "permitted_microbatch_size": maximum_batch, + "microbatch_bound_satisfied": batch_bound, + "state_matches": state_matches, + "contract_passed": state_matches + and loss_matches + and batch_bound is True, + "comparisons": comparisons, + } + ) + finally: + torch.set_num_threads(previous_threads) + return scenarios + + +def evaluate_update(update: Update, learning_rate: float) -> list[Scenario]: + """Evaluate an already-loaded trusted callable; imports are the caller's scope.""" + return compare_candidate_cases( + update, learning_rate, prepare_reference_cases(learning_rate) + ) + + +def evaluate_candidate( + path: Path, learning_rate: float, trace_parent: Path | None = None +) -> Evaluation: + candidate_sha = hashlib.sha256(path.read_bytes()).hexdigest() + reference_path = Path(reference.__file__) + reference_sha = hashlib.sha256(reference_path.read_bytes()).hexdigest() + cases = prepare_reference_cases(learning_rate) + trace_directory = Path( + tempfile.mkdtemp(prefix="continuation-check-", dir=trace_parent) + ) + previous_trace_directory = os.environ.get("TRAINCHECK_OUTPUT_DIR") + os.environ["TRAINCHECK_OUTPUT_DIR"] = str(trace_directory) + try: + with load_candidate(path) as update: + scenarios = compare_candidate_cases(update, learning_rate, cases) + finally: + if previous_trace_directory is None: + os.environ.pop("TRAINCHECK_OUTPUT_DIR", None) + else: + os.environ["TRAINCHECK_OUTPUT_DIR"] = previous_trace_directory + if hashlib.sha256(path.read_bytes()).hexdigest() != candidate_sha: + raise ValueError("Candidate source changed during evaluation") + if hashlib.sha256(reference_path.read_bytes()).hexdigest() != reference_sha: + raise ValueError("Reference source changed during evaluation") + passed = all(item["contract_passed"] for item in scenarios) + return { + "status": "passed" if passed else "violated", + "contract_passed": passed, + "scope": "CPU one-update accumulation arithmetic and microbatch bound; training quality evaluated separately", + "learning_rate": learning_rate, + "atol": ATOL, + "rtol": RTOL, + "candidate_path": str(path.resolve()), + "candidate_sha256": candidate_sha, + "reference_sha256": reference_sha, + "diagnostic_trace_directory": str(trace_directory), + "scenarios": scenarios, + "limitations": [ + "single-process CPU check does not certify DDP collective behavior", + "one update per seed/configuration does not certify long-run optimizer or checkpoint behavior", + "floating tensor comparisons allow stated FP32 tolerance; every tensor retains exact content hashes", + "stored state is compared without reconstructing storage aliasing or device placement", + "trusted candidate code runs in-process; use a fresh process per candidate", + "optional trace writer artifacts may finish at process exit; this report does not certify their coverage", + "CPU numerical evaluation does not establish CPU-only compatibility of candidate instrumentation dependencies", + ], + } + + +class Options(argparse.Namespace): + candidate: Path + lr: float + output: Path + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--candidate", type=Path, required=True) + parser.add_argument("--lr", type=float, required=True) + parser.add_argument("--output", type=Path, required=True) + options = parser.parse_args(namespace=Options()) + try: + report = evaluate_candidate( + options.candidate, options.lr, options.output.parent + ) + except Exception as error: + # Trusted candidate code can raise arbitrary exception classes. This is + # an explicit failed evaluation, with unknown outcome and nonzero exit. + failure: dict[str, object] = { + "status": "error", + "contract_passed": None, + "error_type": type(error).__name__, + "error": str(error), + } + options.output.write_text(json.dumps(failure, indent=2, allow_nan=False) + "\n") + return 2 + options.output.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n") + return 0 if report["contract_passed"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/continuation_pilot/protocol.json b/tools/continuation_pilot/protocol.json new file mode 100644 index 00000000..df4d9343 --- /dev/null +++ b/tools/continuation_pilot/protocol.json @@ -0,0 +1,79 @@ +{ + "schema_version": 1, + "study": "continuation_observation_pilot_20261101", + "scope": "One paired synthetic training-continuation task; not an official RSI benchmark or a general effect estimate", + "agent_context": "Fresh agents without parent history; same inherited model and reasoning configuration", + "public_setup": { + "change": "Process the effective per-rank batch as configurable microbatches while retaining two-GPU training", + "world_size": 2, + "per_rank_batch": 8, + "sequence_length": 8, + "optimizer": "SGD", + "momentum": 0.9, + "learning_rate": 0.15, + "steps": 24, + "accumulation": 4, + "training_seed": 20261101, + "evaluation_seed": 777000, + "evaluation_interval": 8 + }, + "qualification": { + "parameters": "Analyst-chosen screening thresholds frozen before qualification", + "reference_validation_loss_drop_at_least_nats": 0.2, + "current_minus_reference_last_three_validation_points_mean_at_least_nats": 0.05, + "current_validation_improvement_less_than_fraction_of_reference": 0.8, + "if_not_met": "Do not launch solvers with an asserted underperformance story; preserve and report all qualification attempts", + "metric": "Independent unscaled mean token cross entropy on the same fixed development data" + }, + "operator_known_change": "The microbatch pipeline applies an additional accumulation divisor. SGD is used because adaptive optimizers may cancel uniform gradient scaling. This description is excluded from solver packets.", + "arms": { + "control": "Ordinary code inspection, logging, assertions, custom analyses, and reruns", + "traincheck": "Same facilities plus generic TrainCheck native observation capability; no inferred rules or targeted checker supplied" + }, + "solver_budget": { + "training_requests_per_arm": 4, + "maximum_updates_per_request": 24, + "maximum_seconds_per_request": 180, + "requested_reasoning_minutes": 15, + "reasoning_limit_enforcement": "Advisory; actual wall span and overruns reported; native subagent token budget not independently enforced", + "failed_training_requests_count": true, + "cpu_microchecks": "Allowed on tiny tensors; no full training outside the request budget", + "initial_curves_free_to_both_arms": true + }, + "packet_exclusions": [ + "operator protocol and known change", + "prior conversation and diagnosis", + "existing authored normalization invariant", + "private evaluator and operator tests", + "other arm files and final scores" + ], + "held_out_evaluation": { + "training_seeds": [ + 20261117, + 20261129 + ], + "evaluation_seed": 777001, + "quality_updates": 24, + "quality_world_size": 2, + "quality_accumulation": 4, + "accumulation_contract_values": [ + 1, + 2, + 4 + ], + "contract_scope": "Separate CPU accumulation-equivalence and microbatch-size checks; not a distributed communication proof", + "freeze_selection_before_scoring": true, + "report_quality_and_contract_separately": true + }, + "access_isolation": "Explicit per-agent packet scope and no inherited context; not an OS security boundary or formal blinding guarantee", + "provenance": "Retain prompts, packet hashes, all requests, source snapshots, full curves/checkpoints, actual tool use, and measured costs", + "protocol_revision": 2, + "calibration_history": { + "previous_training_budget": 64, + "previous_protocol_sha256": "606085adfd0b266984bfdc88c67a63dd35efb519ac48f094ea4f4bffa4d72907", + "screen_passed": false, + "last_three_gap_nats": 0.04525390006650317, + "current_to_reference_improvement_ratio": 0.9875729563949539, + "revision_reason": "Both versions nearly converge by64; use24updates for a development case with a meaningful early-budget gap. Selected using qualification curves before solver access, never held-out seeds or solver outcomes. Preserve all calibration artifacts." + } +} diff --git a/tools/continuation_pilot/run_worker.py b/tools/continuation_pilot/run_worker.py new file mode 100644 index 00000000..ec0efe59 --- /dev/null +++ b/tools/continuation_pilot/run_worker.py @@ -0,0 +1,28 @@ +"""Operator bootstrap; the solver packet contains only the training_task package.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + + +def main() -> None: + workspace = Path(os.environ["CONTINUATION_WORKSPACE"]).resolve(strict=True) + trace_root = Path(os.environ["CONTINUATION_TRACE_ROOT"]).resolve() + rank = int(os.environ["RANK"]) + if rank < 0: + raise ValueError("RANK must be nonnegative") + trace = trace_root / f"rank{rank}" + trace.mkdir(parents=True, exist_ok=False) + os.environ["TRAINCHECK_OUTPUT_DIR"] = str(trace) + sys.path.insert(0, str(workspace)) + # Import after setting per-rank output so optional diagnostics see the correct + # sink even when they install wrappers during module initialization. + from training_task.engine import main as train + + train() + + +if __name__ == "__main__": + main() diff --git a/tools/continuation_pilot/training_task/__init__.py b/tools/continuation_pilot/training_task/__init__.py new file mode 100644 index 00000000..b75f4584 --- /dev/null +++ b/tools/continuation_pilot/training_task/__init__.py @@ -0,0 +1 @@ +"""Small distributed token-prediction training task.""" diff --git a/tools/continuation_pilot/training_task/current.py b/tools/continuation_pilot/training_task/current.py new file mode 100644 index 00000000..18be07d5 --- /dev/null +++ b/tools/continuation_pilot/training_task/current.py @@ -0,0 +1,33 @@ +"""Microbatch training update.""" + +from __future__ import annotations + +from contextlib import nullcontext + +import torch +import torch.nn as nn +from torch.nn.parallel import DistributedDataParallel + +from .engine import Batch, backward, loss_for_batch + + +def run_update( + model: nn.Module, optimizer: torch.optim.Optimizer, batches: list[Batch] +) -> float: + if not batches: + raise ValueError("An update requires at least one batch") + optimizer.zero_grad(set_to_none=True) + accumulation = len(batches) + total = torch.zeros((), device=batches[0].tokens.device) + for index, batch in enumerate(batches): + context = ( + model.no_sync() + if isinstance(model, DistributedDataParallel) and index + 1 < accumulation + else nullcontext() + ) + with context: + loss = loss_for_batch(model, batch) + total += loss.detach() + backward(loss / accumulation, accumulation) + optimizer.step() + return float((total / accumulation).item()) diff --git a/tools/continuation_pilot/training_task/engine.py b/tools/continuation_pilot/training_task/engine.py new file mode 100644 index 00000000..8870a6aa --- /dev/null +++ b/tools/continuation_pilot/training_task/engine.py @@ -0,0 +1,306 @@ +"""Run a reproducible token-prediction training experiment.""" + +from __future__ import annotations + +import argparse +import datetime +import json +import math +import os +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel + +VOCAB_SIZE = 32 +SEQUENCE_LENGTH = 8 +EXAMPLES_PER_RANK_UPDATE = 8 +EVALUATION_BATCHES = 8 + + +@dataclass(frozen=True) +class Batch: + tokens: torch.Tensor + labels: torch.Tensor + + +@dataclass(frozen=True) +class Options: + pipeline: Literal["reference", "current"] + output: Path + steps: int = 64 + accumulation: int = 4 + seed: int = 20261101 + lr: float = 0.15 + device: Literal["cpu", "cuda"] = "cuda" + eval_seed: int = 777000 + + def validate(self) -> None: + if self.pipeline not in ("reference", "current"): + raise ValueError("pipeline must be reference or current") + if self.device not in ("cpu", "cuda"): + raise ValueError("device must be cpu or cuda") + for name, value in ( + ("steps", self.steps), + ("accumulation", self.accumulation), + ("seed", self.seed), + ("eval_seed", self.eval_seed), + ): + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError(f"{name} must be an integer") + if not 1 <= self.steps <= 128: + raise ValueError("steps must be between 1 and 128") + if self.accumulation < 1 or EXAMPLES_PER_RANK_UPDATE % self.accumulation: + raise ValueError("accumulation must be a positive divisor of 8") + if not 0 <= self.seed < 2**31 or not 0 <= self.eval_seed < 2**31: + raise ValueError("seed and eval_seed must be in [0, 2**31)") + if isinstance(self.lr, bool) or not isinstance(self.lr, (int, float)): + raise TypeError("lr must be numeric") + if not math.isfinite(self.lr) or self.lr <= 0: + raise ValueError("lr must be finite and positive") + + +class TinyTransformer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.embedding = nn.Embedding(VOCAB_SIZE, 16) + layer = nn.TransformerEncoderLayer( + d_model=16, + nhead=2, + dim_feedforward=32, + dropout=0.0, + batch_first=True, + ) + self.encoder = nn.TransformerEncoder(layer, num_layers=2) + self.readout = nn.Linear(16, VOCAB_SIZE) + + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + result: object = self.readout(self.encoder(self.embedding(tokens))) + if not isinstance(result, torch.Tensor): + raise TypeError("Model output must be a tensor") + return result + + +def make_batch(seed: int, step: int, rank: int, device: torch.device) -> Batch: + """Regenerate a complete rank batch without consuming the training RNG.""" + if seed < 0 or step < 0 or rank < 0: + raise ValueError("seed, step, and rank must be nonnegative") + generator = torch.Generator(device="cpu") + generator.manual_seed(seed + 10000 * step + 100 * rank) + tokens = torch.randint( + VOCAB_SIZE, + (EXAMPLES_PER_RANK_UPDATE, SEQUENCE_LENGTH), + generator=generator, + ) + labels = (tokens + 1) % VOCAB_SIZE + return Batch(tokens.to(device), labels.to(device)) + + +def split_batch(batch: Batch, accumulation: int) -> list[Batch]: + if accumulation < 1 or EXAMPLES_PER_RANK_UPDATE % accumulation: + raise ValueError("accumulation must be a positive divisor of 8") + if batch.tokens.shape != batch.labels.shape or tuple(batch.tokens.shape) != ( + EXAMPLES_PER_RANK_UPDATE, + SEQUENCE_LENGTH, + ): + raise ValueError("Unexpected complete batch shape") + size = EXAMPLES_PER_RANK_UPDATE // accumulation + return [ + Batch(batch.tokens[start : start + size], batch.labels[start : start + size]) + for start in range(0, EXAMPLES_PER_RANK_UPDATE, size) + ] + + +def loss_for_batch(model: nn.Module, batch: Batch) -> torch.Tensor: + result: object = model(batch.tokens) + if not isinstance(result, torch.Tensor): + raise TypeError("Model output must be a tensor") + loss: object = F.cross_entropy( + result.reshape(-1, VOCAB_SIZE), batch.labels.reshape(-1), reduction="mean" + ) + if not isinstance(loss, torch.Tensor): + raise TypeError("Cross entropy must return a tensor") + return loss + + +def backward(loss: torch.Tensor, accumulation: int) -> None: + """Accumulate one microbatch's contribution to the current update.""" + if accumulation < 1: + raise ValueError("accumulation must be positive") + (loss / accumulation).backward() + + +def evaluate(model: nn.Module, device: torch.device, seed: int = 777000) -> float: + """Compute mean held-out token loss and restore every module's mode. + + The fixed evaluation data uses private generators. Evaluation does not + consume CPU/CUDA RNG streams or update the model/optimizer. + """ + modes = [(module, module.training) for module in model.modules()] + model.eval() + try: + with torch.no_grad(): + total = torch.zeros((), dtype=torch.float64, device=device) + for index in range(EVALUATION_BATCHES): + batch = make_batch(seed, index, 0, device) + total += loss_for_batch(model, batch).to(torch.float64) + value = float((total / EVALUATION_BATCHES).item()) + finally: + for module, training in modes: + module.training = training + if not math.isfinite(value): + raise ValueError("Evaluation loss must be finite") + return value + + +def parse_options(argv: list[str] | None = None) -> Options: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--pipeline", choices=("reference", "current"), required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--steps", type=int, default=64) + parser.add_argument("--accumulation", type=int, default=4) + parser.add_argument("--seed", type=int, default=20261101) + parser.add_argument("--eval-seed", type=int, default=777000) + parser.add_argument("--lr", type=float, default=0.15) + parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda") + args = parser.parse_args(argv) + options = Options( + args.pipeline, + args.output, + args.steps, + args.accumulation, + args.seed, + args.lr, + args.device, + args.eval_seed, + ) + try: + options.validate() + except (TypeError, ValueError) as error: + parser.error(str(error)) + return options + + +def run(options: Options) -> None: + from .current import run_update as current_update + from .reference import run_update as reference_update + + options.validate() + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + if rank < 0 or world_size < 1 or rank >= world_size or local_rank < 0: + raise ValueError("Invalid distributed rank metadata") + torch.set_num_threads(1) + device = ( + torch.device("cuda", local_rank) + if options.device == "cuda" + else torch.device("cpu") + ) + if device.type == "cuda": + torch.cuda.set_device(device) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.use_deterministic_algorithms(True) + torch.manual_seed(options.seed) + if world_size > 1: + dist.init_process_group( + "nccl" if device.type == "cuda" else "gloo", + timeout=datetime.timedelta(seconds=60), + ) + try: + output = options.output / f"rank{rank}" + output.mkdir(parents=True, exist_ok=False) + model = TinyTransformer().to(device).train() + wrapped: nn.Module = model + if world_size > 1: + wrapped = DistributedDataParallel( + model, device_ids=[local_rank] if device.type == "cuda" else None + ) + optimizer = torch.optim.SGD(model.parameters(), lr=options.lr, momentum=0.9) + update = reference_update if options.pipeline == "reference" else current_update + curve: list[dict[str, int | float]] = [ + {"step": 0, "validation_loss": evaluate(model, device, options.eval_seed)} + ] + steps: list[dict[str, int | float]] = [] + started = time.perf_counter() + for step in range(1, options.steps + 1): + batch = make_batch(options.seed, step, rank, device) + microbatches = split_batch(batch, options.accumulation) + local_loss = update(wrapped, optimizer, microbatches) + global_loss = torch.tensor(local_loss, dtype=torch.float64, device=device) + if world_size > 1: + dist.all_reduce(global_loss) + global_loss /= world_size + loss = float(global_loss.item()) + if not math.isfinite(loss): + raise ValueError("Training loss must be finite") + steps.append({"step": step, "train_token_loss": loss}) + if step % 8 == 0 or step == options.steps: + curve.append( + { + "step": step, + "validation_loss": evaluate(model, device, options.eval_seed), + } + ) + if device.type == "cuda": + torch.cuda.synchronize(device) + seconds = time.perf_counter() - started + torch.save( + { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "cpu_rng": torch.get_rng_state(), + "cuda_rng": ( + torch.cuda.get_rng_state(device) if device.type == "cuda" else None + ), + "completed_steps": options.steps, + "rank": rank, + "world_size": world_size, + "seed": options.seed, + "accumulation": options.accumulation, + }, + output / "checkpoint.pt", + ) + metrics: dict[str, object] = { + "pipeline": options.pipeline, + "rank": rank, + "world_size": world_size, + "device": str(device), + "torch_version": str(torch.__version__), + "seed": options.seed, + "eval_seed": options.eval_seed, + "lr": options.lr, + "accumulation": options.accumulation, + "examples_per_rank_update": EXAMPLES_PER_RANK_UPDATE, + "tokens_per_global_update": world_size + * EXAMPLES_PER_RANK_UPDATE + * SEQUENCE_LENGTH, + "optimizer_updates": options.steps, + "parameter_count": sum( + parameter.numel() for parameter in model.parameters() + ), + "training_seconds_including_validation": seconds, + "steps": steps, + "validation_curve": curve, + } + (output / "metrics.json").write_text( + json.dumps(metrics, indent=2, allow_nan=False) + "\n" + ) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def main() -> None: + run(parse_options()) + + +if __name__ == "__main__": + main() diff --git a/tools/continuation_pilot/training_task/reference.py b/tools/continuation_pilot/training_task/reference.py new file mode 100644 index 00000000..5acbadf9 --- /dev/null +++ b/tools/continuation_pilot/training_task/reference.py @@ -0,0 +1,24 @@ +"""Full-batch training update.""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from .engine import Batch, backward, loss_for_batch + + +def run_update( + model: nn.Module, optimizer: torch.optim.Optimizer, batches: list[Batch] +) -> float: + if not batches: + raise ValueError("An update requires at least one batch") + optimizer.zero_grad(set_to_none=True) + combined = Batch( + torch.cat([batch.tokens for batch in batches]), + torch.cat([batch.labels for batch in batches]), + ) + loss = loss_for_batch(model, combined) + backward(loss, 1) + optimizer.step() + return float(loss.detach().item())