From 26c77ce83bad89552b9ec1bb3b7058a416970ad0 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 1 Nov 2026 18:06:49 +0800 Subject: [PATCH 1/3] Pilot agent-authored normalization checks on native API traces --- tests/test_agent_authored_rule.py | 94 +++++ tests/test_agent_invariant_observation.py | 203 ++++++++++ tests/test_agent_invariant_trace_check.py | 334 +++++++++++++++++ tools/agent_invariants/AUTHORING.md | 121 ++++++ tools/agent_invariants/__init__.py | 1 + tools/agent_invariants/authored_rule.py | 133 +++++++ tools/agent_invariants/check.py | 431 ++++++++++++++++++++++ tools/agent_invariants/observation.py | 65 ++++ tools/distributed_probe/workload.py | 31 +- 9 files changed, 1412 insertions(+), 1 deletion(-) create mode 100644 tests/test_agent_authored_rule.py create mode 100644 tests/test_agent_invariant_observation.py create mode 100644 tests/test_agent_invariant_trace_check.py create mode 100644 tools/agent_invariants/AUTHORING.md create mode 100644 tools/agent_invariants/__init__.py create mode 100644 tools/agent_invariants/authored_rule.py create mode 100644 tools/agent_invariants/check.py create mode 100644 tools/agent_invariants/observation.py diff --git a/tests/test_agent_authored_rule.py b/tests/test_agent_authored_rule.py new file mode 100644 index 00000000..ec598c46 --- /dev/null +++ b/tests/test_agent_authored_rule.py @@ -0,0 +1,94 @@ +"""Validate the frozen author artifact without editing its predicate.""" + +from dataclasses import replace + +import pytest + +from tools.agent_invariants.authored_rule import Observation, evaluate + + +def observation() -> Observation: + return Observation(0, 2, 1, 0, 2, 10, 32, 0.0625) + + +@pytest.mark.parametrize( + "world_size,local_tokens,global_tokens,accumulation", + [(1, 5, 5, 1), (2, 5, 16, 1), (2, 10, 32, 2), (2, 15, 48, 3), (4, 7, 37, 4)], +) +def test_rule_uses_window_counts_instead_of_a_fixed_scale( + world_size: int, local_tokens: int, global_tokens: int, accumulation: int +) -> None: + item = Observation( + 0, + world_size, + 1, + 0, + accumulation, + local_tokens, + global_tokens, + world_size / global_tokens, + ) + assert evaluate(item).passed + assert not evaluate(replace(item, applied_scale=item.applied_scale / 2)).passed + + +@pytest.mark.parametrize("scale", [0.1, 1.0 / 22, 0.0, -0.0625, 0.0625 / 2]) +def test_finite_wrong_multiplier_is_a_violation(scale: float) -> None: + result = evaluate(replace(observation(), applied_scale=scale)) + assert not result.passed + assert result.observed == scale + assert result.expected == 0.0625 + + +@pytest.mark.parametrize( + "field,value", + [ + ("rank", True), + ("world_size", "2"), + ("step", 1.0), + ("microbatch", False), + ("accumulation", None), + ("local_tokens", 10.0), + ("global_tokens", True), + ("applied_scale", 1), + ("applied_scale", True), + ], +) +def test_wrong_types_are_errors_instead_of_passes(field: str, value: object) -> None: + # Simulate a malformed runtime payload without weakening the public signature. + item = observation() + object.__setattr__(item, field, value) + with pytest.raises(TypeError): + evaluate(item) + + +@pytest.mark.parametrize( + "field,value", + [ + ("rank", -1), + ("rank", 2), + ("world_size", 0), + ("step", 0), + ("microbatch", 2), + ("accumulation", 0), + ("local_tokens", 0), + ("global_tokens", 10), + ("global_tokens", 10**400), + ("applied_scale", float("nan")), + ("applied_scale", float("inf")), + ], +) +def test_unsupported_domains_are_errors(field: str, value: object) -> None: + item = observation() + object.__setattr__(item, field, value) + with pytest.raises(ValueError): + evaluate(item) + + +def test_single_rank_and_relative_tolerance_boundary() -> None: + with pytest.raises(ValueError, match="single-rank"): + evaluate(replace(observation(), world_size=1)) + assert evaluate(replace(observation(), applied_scale=0.0625 * (1 + 1e-13))).passed + assert not evaluate( + replace(observation(), applied_scale=0.0625 * (1 + 1e-10)) + ).passed diff --git a/tests/test_agent_invariant_observation.py b/tests/test_agent_invariant_observation.py new file mode 100644 index 00000000..1197e9f9 --- /dev/null +++ b/tests/test_agent_invariant_observation.py @@ -0,0 +1,203 @@ +"""Native scalar observations preserve the operation they describe.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +from tools.distributed_probe.workload import Options, parse_options, run + +ROOT = Path(__file__).resolve().parents[1] +FUNCTION = "tools.agent_invariants.observation.apply_loss_scale" +SCALES = (0.0625, 0.1, 0.0, -0.25) + +WORKER = """ +import json +import sys +from pathlib import Path + +import torch + +from tools.agent_invariants.observation import wrapped_loss_scale +from tools.distributed_probe.tracing import start_tracing + +torch.set_num_threads(1) +torch.manual_seed(733) +model = torch.nn.Linear(2, 1) +session = start_tracing('api', Path(sys.argv[1]), model) +operation = wrapped_loss_scale() +for step, scale in enumerate((0.0625, 0.1, 0.0, -0.25), start=1): + with session.suppress(): + source = torch.tensor(3.75, requires_grad=True) + reference = source.detach().clone().requires_grad_(True) + expected = reference * scale + expected.backward() + rng_before = torch.get_rng_state().clone() + with session.step(step): + actual = operation( + source, rank=0, world_size=2, step=step, microbatch=0, + accumulation=2, local_tokens=10, global_tokens=32, + applied_scale=scale, + ) + assert isinstance(actual, torch.Tensor) + actual.backward() + with session.suppress(): + assert torch.equal(actual, expected) + assert source.grad is not None and reference.grad is not None + assert torch.equal(source.grad, reference.grad) + assert torch.equal(torch.get_rng_state(), rng_before) + +with session.step(5): + with session.suppress(): + skipped = operation( + torch.tensor(2.0), rank=0, world_size=2, step=5, microbatch=0, + accumulation=2, local_tokens=10, global_tokens=32, + applied_scale=0.0625, + ) + assert isinstance(skipped, torch.Tensor) and skipped.item() == 0.125 + +try: + with session.step(6): + operation( + 3.75, rank=0, world_size=2, step=6, microbatch=0, + accumulation=2, local_tokens=10, global_tokens=32, + applied_scale=0.0625, + ) +except TypeError as error: + assert str(error) == 'Loss multiplication must return a Tensor' +else: + raise AssertionError('Operation exception was swallowed') + +report = session.close() +assert report['completed_trace_steps'] == [1, 2, 3, 4, 5] +Path(sys.argv[2]).write_text(json.dumps(report)) +""" + + +def mapping(value: object) -> dict[str, object]: + if not isinstance(value, dict): + raise TypeError("Expected a JSON object") + result: dict[str, object] = {} + for key, child in value.items(): + if not isinstance(key, str): + raise TypeError("Expected string keys") + result[key] = child + return result + + +def integer(value: object) -> int: + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError("Expected an integer, excluding bool") + return value + + +def trace_records(root: Path) -> list[dict[str, object]]: + records: list[dict[str, object]] = [] + for path in sorted(root.glob("trace_API_*.log")): + for line in path.read_text().splitlines(): + record = mapping(json.loads(line)) + if record["function"] == FUNCTION: + records.append(record) + return records + + +def test_actual_native_records_preserve_scalars_gradients_rng_and_exceptions( + tmp_path: Path, +) -> None: + trace_root = tmp_path / "traces" + report_path = tmp_path / "report.json" + environment = dict(os.environ) + environment["PYTHONPATH"] = str(ROOT) + environment["PYTHONDONTWRITEBYTECODE"] = "1" + result = subprocess.run( + [sys.executable, "-c", WORKER, str(trace_root), str(report_path)], + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=90, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + report = mapping(json.loads(report_path.read_text())) + assert report["completed_trace_steps"] == [1, 2, 3, 4, 5] + records = trace_records(trace_root) + assert len(records) == 10 + observed_steps = [ + integer(mapping(record["meta_vars"])["step"]) for record in records + ] + assert observed_steps == [1, 1, 2, 2, 3, 3, 4, 4, 6, 6] + identifiers: set[tuple[int, int, str]] = set() + for offset, step in enumerate((1, 2, 3, 4, 6)): + pre, post = records[2 * offset : 2 * offset + 2] + assert pre["type"] == "function_call (pre)" + assert post["type"] == ( + "function_call (post) (exception)" if step == 6 else "function_call (post)" + ) + assert integer(pre["time"]) <= integer(post["time"]) + for field in ("func_call_id", "thread_id", "process_id", "meta_vars"): + assert pre[field] == post[field] + call_id = pre["func_call_id"] + assert isinstance(call_id, str) + identity = (integer(pre["process_id"]), integer(pre["thread_id"]), call_id) + assert identity not in identifiers + identifiers.add(identity) + assert pre["is_bound_method"] is False + assert pre["obj_id"] is None + assert pre["args"] == {} + scalar = SCALES[step - 1] if step < 5 else SCALES[0] + expected_kwargs: dict[str, object] = { + "rank": {"int": 0}, + "world_size": {"int": 2}, + "step": {"int": step}, + "microbatch": {"int": 0}, + "accumulation": {"int": 2}, + "local_tokens": {"int": 10}, + "global_tokens": {"int": 32}, + "applied_scale": {"float": scalar}, + } + assert pre["kwargs"] == expected_kwargs + expected_post_fields = { + "func_call_id", + "thread_id", + "process_id", + "meta_vars", + "type", + "function", + "is_bound_method", + "obj_id", + "time", + } + if step == 6: + expected_post_fields.update(("exception", "exception_msg")) + assert post["exception_msg"] == "Loss multiplication must return a Tensor" + assert post.keys() == expected_post_fields + + +@pytest.mark.parametrize("mode", ["api", "subclass"]) +def test_observation_option_requires_traced_mode(mode: str) -> None: + options = parse_options( + ["--output", "/unused", "--instrumentation", mode, "--observe-normalization"] + ) + assert options.observe_normalization is True + assert options.instrumentation == mode + + +def test_plain_observation_cli_is_rejected() -> None: + with pytest.raises(SystemExit): + parse_options(["--output", "/unused", "--observe-normalization"]) + + +def test_plain_observation_programmatic_run_is_rejected_before_output( + tmp_path: Path, +) -> None: + output = tmp_path / "must-not-exist" + options = Options(output, None, 1, 1, "fp32", "none", "plain", "cpu", 2, True) + with pytest.raises(ValueError, match="require API or subclass tracing"): + run(options) + assert not output.exists() diff --git a/tests/test_agent_invariant_trace_check.py b/tests/test_agent_invariant_trace_check.py new file mode 100644 index 00000000..46dabdd4 --- /dev/null +++ b/tests/test_agent_invariant_trace_check.py @@ -0,0 +1,334 @@ +"""Coverage errors must remain distinct from completed rule evaluations.""" + +from __future__ import annotations + +import copy +import hashlib +import json +from pathlib import Path + +import pytest + +from tools.agent_invariants.check import ( + EXCEPTION, + FUNCTION, + POST, + PRE, + ExpectedRun, + check_run, + main, +) + + +def records(rank: int, scale: float) -> list[dict[str, object]]: + events: list[dict[str, object]] = [] + for step in (3, 4): + for microbatch in (0, 1): + index = 2 * (step - 3) + microbatch + keywords: dict[str, object] = { + "rank": {"int": rank}, + "world_size": {"int": 2}, + "step": {"int": step}, + "microbatch": {"int": microbatch}, + "accumulation": {"int": 2}, + "local_tokens": {"int": 10 if rank == 0 else 22}, + "global_tokens": {"int": 32}, + "applied_scale": {"float": scale}, + } + base: dict[str, object] = { + "function": FUNCTION, + "func_call_id": f"call-{rank}-{index}", + "process_id": 100 + rank, + "thread_id": 42, + "meta_vars": {"step": step, "epoch": 0}, + "is_bound_method": False, + "obj_id": None, + } + events.append( + { + **base, + "time": 1000 + 2 * index, + "type": PRE, + "args": {}, + "kwargs": keywords, + } + ) + events.append({**base, "time": 1001 + 2 * index, "type": POST}) + return events + + +def write_run( + root: Path, rank0: list[dict[str, object]], rank1: list[dict[str, object]] +) -> None: + for rank, events in enumerate((rank0, rank1)): + directory = root / f"rank{rank}" / "trace" + directory.mkdir(parents=True) + (directory / "trace_API_100_42.log").write_text( + "".join(json.dumps(event, allow_nan=False) + "\n" for event in events) + ) + + +def expected() -> ExpectedRun: + return ExpectedRun(world_size=2, start_step=3, steps=2, accumulation=2) + + +def keywords(event: dict[str, object]) -> dict[str, object]: + raw = event["kwargs"] + assert isinstance(raw, dict) + return raw + + +def test_complete_check_preserves_identity_and_provenance(tmp_path: Path) -> None: + rank0, rank1 = records(0, 0.0625), records(1, 0.0625) + unrelated: dict[str, object] = { + "function": "torch.example", + "arbitrary_metadata": [2, 2, 1], + } + rank0.insert(0, unrelated) + write_run(tmp_path, rank0, rank1) + report = check_run(tmp_path, expected()) + assert report["status"] == "passed" + assert report["expected_observations"] == report["completed_observations"] == 8 + assert report["violations"] == 0 + assert report["decisions"][0]["pre"]["record"] == rank0[1] + assert report["decisions"][0]["post"]["record"] == rank0[2] + assert report["decisions"][0]["pre"]["line"] == 2 + assert report["decisions"][0]["observation"]["step"] == 3 + for source in (*report["trace_files"], report["rule_source"]): + data = Path(source["path"]).read_bytes() + assert source["bytes"] == len(data) + assert source["sha256"] == hashlib.sha256(data).hexdigest() + + +def test_completed_bad_scale_is_a_violation_not_a_coverage_error( + tmp_path: Path, +) -> None: + write_run(tmp_path, records(0, 0.1), records(1, 1.0 / 22)) + report = check_run(tmp_path, expected()) + assert report["status"] == "violated" + assert report["completed_observations"] == report["violations"] == 8 + assert all(not decision["passed"] for decision in report["decisions"]) + + +def test_different_valid_window_size_does_not_memorize_scale(tmp_path: Path) -> None: + rows = (records(0, 0.05), records(1, 0.05)) + for rank, events in enumerate(rows): + for event in events: + if event["type"] == PRE: + keywords(event)["local_tokens"] = {"int": 13 if rank == 0 else 27} + keywords(event)["global_tokens"] = {"int": 40} + write_run(tmp_path, *rows) + assert check_run(tmp_path, expected())["status"] == "passed" + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_missing_rank_trace_is_not_success(tmp_path: Path, rank: int) -> None: + write_run(tmp_path, records(0, 0.0625), records(1, 0.0625)) + (tmp_path / f"rank{rank}" / "trace" / "trace_API_100_42.log").unlink() + with pytest.raises(ValueError, match="No API trace"): + check_run(tmp_path, expected()) + + +def test_missing_entire_run_is_not_success(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="rank directories"): + check_run(tmp_path, expected()) + + +def test_empty_native_logs_are_not_success(tmp_path: Path) -> None: + write_run(tmp_path, [], []) + with pytest.raises(ValueError, match="coverage"): + check_run(tmp_path, expected()) + + +def test_missing_complete_call_is_not_success(tmp_path: Path) -> None: + write_run(tmp_path, records(0, 0.0625)[2:], records(1, 0.0625)) + with pytest.raises(ValueError, match="coverage"): + check_run(tmp_path, expected()) + + +@pytest.mark.parametrize("drop", [0, 1]) +def test_unpaired_events_are_not_success(tmp_path: Path, drop: int) -> None: + rank0 = records(0, 0.0625) + del rank0[drop] + write_run(tmp_path, rank0, records(1, 0.0625)) + with pytest.raises(ValueError, match="Incomplete observation"): + check_run(tmp_path, expected()) + + +@pytest.mark.parametrize("index", [0, 1]) +def test_duplicate_native_events_are_not_success(tmp_path: Path, index: int) -> None: + rank0 = records(0, 0.0625) + rank0.append(copy.deepcopy(rank0[index])) + write_run(tmp_path, rank0, records(1, 0.0625)) + with pytest.raises(ValueError, match="Duplicate observation event"): + check_run(tmp_path, expected()) + + +def test_different_call_id_does_not_hide_duplicate_observation(tmp_path: Path) -> None: + rank0 = records(0, 0.0625) + duplicated = copy.deepcopy(rank0[:2]) + for event in duplicated: + event["func_call_id"] = "extra-call" + rank0.extend(duplicated) + write_run(tmp_path, rank0, records(1, 0.0625)) + with pytest.raises(ValueError, match="Duplicate observation for"): + check_run(tmp_path, expected()) + + +def test_exception_post_is_not_a_completed_evaluation(tmp_path: Path) -> None: + rank0 = records(0, 0.0625) + rank0[1]["type"] = EXCEPTION + rank0[1]["exception"] = "RuntimeError" + rank0[1]["exception_msg"] = "failure" + write_run(tmp_path, rank0, records(1, 0.0625)) + with pytest.raises(ValueError, match="raised an exception"): + check_run(tmp_path, expected()) + + +@pytest.mark.parametrize( + "field,value,match", + [ + ("process_id", 400, "Incomplete observation"), + ("thread_id", 9, "Incomplete observation"), + ("func_call_id", "changed", "Incomplete observation"), + ("meta_vars", {"step": 8}, "metadata differs"), + ("time", 1, "precedes"), + ("kwargs", {"applied_scale": {"float": 1.0}}, "must not repeat"), + ], +) +def test_changed_post_contract_is_rejected( + tmp_path: Path, field: str, value: object, match: str +) -> None: + rank0 = records(0, 0.0625) + rank0[1][field] = value + write_run(tmp_path, rank0, records(1, 0.0625)) + with pytest.raises(ValueError, match=match): + check_run(tmp_path, expected()) + + +@pytest.mark.parametrize( + "field,value,error", + [ + ("world_size", {"int": True}, TypeError), + ("local_tokens", {"int": "10"}, TypeError), + ("step", {"float": 3.0}, TypeError), + ("rank", {"int": 1}, ValueError), + ("accumulation", {"int": 3}, ValueError), + ("global_tokens", {"int": 0}, ValueError), + ("applied_scale", {"float": 1}, TypeError), + ("applied_scale", {"float": True}, TypeError), + ("applied_scale", {"float": "0.0625"}, TypeError), + ("applied_scale", {"float": 0.0625, "extra": 1}, TypeError), + ], +) +def test_malformed_required_scalar_is_rejected( + tmp_path: Path, field: str, value: object, error: type[Exception] +) -> None: + rank0 = records(0, 0.0625) + keywords(rank0[0])[field] = value + write_run(tmp_path, rank0, records(1, 0.0625)) + with pytest.raises(error): + check_run(tmp_path, expected()) + + +@pytest.mark.parametrize( + "field", ["applied_scale", "global_tokens", "local_tokens", "world_size"] +) +def test_missing_required_scalar_is_rejected(tmp_path: Path, field: str) -> None: + rank0 = records(0, 0.0625) + del keywords(rank0[0])[field] + write_run(tmp_path, rank0, records(1, 0.0625)) + with pytest.raises(ValueError, match="required fields"): + check_run(tmp_path, expected()) + + +def test_inconsistent_microbatch_window_counts_are_rejected(tmp_path: Path) -> None: + rank0 = records(0, 0.0625) + keywords(rank0[0])["local_tokens"] = {"int": 11} + write_run(tmp_path, rank0, records(1, 0.0625)) + with pytest.raises(ValueError, match="within an accumulation window"): + check_run(tmp_path, expected()) + + +@pytest.mark.parametrize("both_ranks", [False, True]) +def test_wrong_cross_rank_totals_are_rejected(tmp_path: Path, both_ranks: bool) -> None: + rows = records(0, 0.0625), records(1, 0.0625) + for events in rows if both_ranks else rows[:1]: + for event in events: + if event["type"] == PRE: + keywords(event)["global_tokens"] = {"int": 33} + write_run(tmp_path, *rows) + with pytest.raises(ValueError, match="Cross-rank token counts"): + check_run(tmp_path, expected()) + + +@pytest.mark.parametrize("value", ["NaN", "Infinity", "-Infinity"]) +def test_nonfinite_json_is_rejected(tmp_path: Path, value: str) -> None: + write_run(tmp_path, records(0, 0.0625), records(1, 0.0625)) + path = tmp_path / "rank0" / "trace" / "trace_API_100_42.log" + path.write_text(path.read_text().replace('"float": 0.0625', f'"float": {value}', 1)) + with pytest.raises(ValueError, match="Nonfinite JSON"): + check_run(tmp_path, expected()) + + +def test_duplicate_json_fields_are_rejected(tmp_path: Path) -> None: + write_run(tmp_path, records(0, 0.0625), records(1, 0.0625)) + path = tmp_path / "rank0" / "trace" / "trace_API_100_42.log" + path.write_text(path.read_text().replace('"int": 10', '"int": 10, "int": 11', 1)) + with pytest.raises(ValueError, match="Duplicate JSON field"): + check_run(tmp_path, expected()) + + +@pytest.mark.parametrize( + "config", + [ExpectedRun(0, 1, 1, 1), ExpectedRun(1, 0, 1, 1), ExpectedRun(1, 1, 0, 1)], +) +def test_nonpositive_expected_dimensions_are_rejected( + tmp_path: Path, config: ExpectedRun +) -> None: + with pytest.raises(ValueError, match="positive"): + check_run(tmp_path, config) + + +@pytest.mark.parametrize( + "scenario,code,status", + [("healthy", 0, "passed"), ("fault", 1, "violated"), ("missing", 2, "error")], +) +def test_cli_exit_status_keeps_error_separate( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + scenario: str, + code: int, + status: str, +) -> None: + root = tmp_path / "run" + scale = 0.0625 if scenario == "healthy" else 0.1 + write_run(root, records(0, scale), records(1, scale)) + if scenario == "missing": + (root / "rank0" / "trace" / "trace_API_100_42.log").unlink() + output = tmp_path / "report.json" + monkeypatch.setattr( + "sys.argv", + [ + "check", + "--root", + str(root), + "--world-size", + "2", + "--start-step", + "3", + "--steps", + "2", + "--accumulation", + "2", + "--output", + str(output), + ], + ) + assert main() == code + result: object = json.loads(output.read_text()) + assert isinstance(result, dict) + assert result["status"] == status + if scenario == "missing": + assert result["violations"] is None + assert result["completed_observations"] is None diff --git a/tools/agent_invariants/AUTHORING.md b/tools/agent_invariants/AUTHORING.md new file mode 100644 index 00000000..7506fda5 --- /dev/null +++ b/tools/agent_invariants/AUTHORING.md @@ -0,0 +1,121 @@ +# Authored objective-normalization invariant + +## Contract used + +Training uses PyTorch DistributedDataParallel with its default averaging of +gradients over `world_size` ranks. An optimizer-update window contains exactly +`accumulation` microbatches on every rank. Each microbatch's unscaled loss is the +SUM of per-token cross entropy for valid, unmasked tokens. The application +multiplies that scalar by the recorded built-in Python `float` `applied_scale` +and invokes backward. `DDP.no_sync` applies to every microbatch except the final +one. No GradScaler is present; any gradient clipping happens after backward. +The intended objective is mean loss over all valid tokens on every rank and in +every microbatch of that optimizer-update window. All windows are complete, +global token counts are positive, and this probe supports positive local +window totals on every rank. + +`step` uses positive, one-based optimizer-update numbers; `microbatch` uses +zero-based indexing. This convention follows the capture-adapter protocol +clarification supplied before freezing; it contains no faulty-run information. + +## Derivation and focused invariant + +Let W denote the number of ranks, N the valid-token count for the entire update +across all ranks, and L[r,m] the token-SUM loss on rank r, microbatch m. The +intended loss is: + + J = (1 / N) * sum_r sum_m L[r,m]. + +With scalar multipliers s[r,m], accumulation followed by default DDP averaging +produces, before clipping: + + G = (1 / W) * sum_r sum_m s[r,m] * gradient(L[r,m]). + +Matching each microbatch's coefficient to the intended objective gives the +single invariant checked at every observed loss application: + + applied_scale = world_size / global_tokens. + +Consequently G equals gradient(J), subject to ordinary floating-point gradient +arithmetic and the stated DDP/accumulation assumptions. The full-window token +total already accounts for accumulation, so an extra accumulation divisor is +not part of this invariant. Unequal local token counts do not change the rule +because the supplied unscaled losses are token sums, not local means. + +This checks objective coefficients for arbitrary gradients. Special cases such +as zero gradients or accidental cancellation can make a wrong multiplier yield +the same numerical gradient; such cases still fail the normalization contract. + +## Requested fields and validation + +`Observation` is a frozen dataclass with these required fields: + +| Field | Contract and use | +| --- | --- | +| rank: int | Zero-based rank; 0 <= rank < world_size. | +| world_size: int | Positive number of DDP ranks; gradient-average denominator. | +| step: int | Positive, one-based optimizer-update number; external grouping. | +| microbatch: int | Zero-based position; 0 <= microbatch < accumulation. | +| accumulation: int | Positive number of microbatches in each complete window. | +| local_tokens: int | Positive valid-token total on this rank over the whole window. | +| global_tokens: int | Positive valid-token total across every rank over the whole window. | +| applied_scale: float | Actual finite built-in Python float used at this loss application. | + +Integers explicitly exclude booleans. No values are coerced or defaulted. +Single-rank local/global totals must agree; for multiple ranks the global total +must leave at least one valid token for every other rank. Required factors that +underflow to zero or otherwise cannot be represented as finite positive Python +floats are unsupported and raise `ValueError`. Wrong types raise `TypeError`; +invalid domain values raise `ValueError`. The caller must report these as +invalid/unsupported observations, not passing decisions or successful coverage. +A finite zero or negative applied multiplier is an observed normalization error +and returns a failed `Decision`. + +`Decision` is frozen and returns a stable `rule_id`, `passed`, the actual +`observed` scalar, computed `expected` scalar, and an explanation. The rule is +stateless; it does not infer missing observations or maintain a cross-rank cache. + +## Tolerance rationale + +The comparison is `math.isclose(observed, expected, rel_tol=1e-12, +abs_tol=0.0)`. These values concern Python scalar multipliers, not mixed-precision +tensor reductions. A 1e-12 relative tolerance permits small binary64 rounding +differences from algebraically equivalent scalar arithmetic, while requiring +approximately twelve decimal digits of agreement. Zero absolute tolerance +prevents small expected factors from passing solely because their absolute +magnitude is small; a zero multiplier cannot pass. Deviations within the stated +tolerance are intentionally accepted, so this is not an exact symbolic proof. +At the extreme subnormal boundary the floating-point comparison may effectively +require exact equality. No tolerance is tied to a batch shape or rank count. + +## Assumptions and limits + +A pass establishes only that this recorded scalar agrees with the coefficient +required by the supplied contract and supplied metadata. It does not establish +that the record reflects actual execution. In particular, this rule cannot +independently detect incorrect token masks or counts, a loss that is actually a +mean or weighted sum, a detached multiplier, extra backward scaling, DDP +communication hooks that change reduction, incorrect no_sync use, absent or +repeated backward calls, gradients carried over from an earlier update, +premature optimizer steps, optimizer/clipping defects, or incorrect model +behavior. Later clipping can change the final update even when pre-clipping +objective normalization is correct. + +The evaluator must independently establish truthful loss-site capture, exactly +one relevant observation for every rank/microbatch, coherent step identities, +consistent world_size/accumulation/global_tokens within a window, stable +local_tokens within each rank's window, and sum(local window totals) equal to +global_tokens. Repeated per-microbatch window totals must not be added repeatedly. +A single passing observation is not evidence of complete execution or coverage. +Zero-local-token windows and partial accumulation windows are outside this +probe's contract. Mixed precision reduction error is not assessed by this +scalar-only invariant. + +## Authoring provenance + +This artifact was derived only from the supplied task contract. No TrainCheck +repository, history, failure implementation, trace, diagnosis, prior output +directory, or faulty-run evaluation was inspected during authoring. The access +restriction is instructional, not a filesystem sandbox; this is an authoring +feasibility probe and is not claimed to be formal blinding. Only this file and +its adjacent authored_rule.py were written/read. No distributed training was run. diff --git a/tools/agent_invariants/__init__.py b/tools/agent_invariants/__init__.py new file mode 100644 index 00000000..9f4a8deb --- /dev/null +++ b/tools/agent_invariants/__init__.py @@ -0,0 +1 @@ +"""Experimental agent-authored predicates over explicit TrainCheck observations.""" diff --git a/tools/agent_invariants/authored_rule.py b/tools/agent_invariants/authored_rule.py new file mode 100644 index 00000000..ae462931 --- /dev/null +++ b/tools/agent_invariants/authored_rule.py @@ -0,0 +1,133 @@ +"""Check normalization of token-summed losses accumulated under default DDP.""" + +from dataclasses import dataclass +from math import isclose, isfinite + +RULE_ID: str = "ddp_accumulated_global_token_mean" +RELATIVE_TOLERANCE: float = 1e-12 +ABSOLUTE_TOLERANCE: float = 0.0 + + +@dataclass(frozen=True) +class Observation: + """One loss application; step is one-based and microbatch is zero-based. + + local_tokens and global_tokens cover the complete optimizer-update window, + including every accumulation microbatch. applied_scale is the actual Python + float multiplying this microbatch's token-summed loss before backward. + """ + + rank: int + world_size: int + step: int + microbatch: int + accumulation: int + local_tokens: int + global_tokens: int + applied_scale: float + + +@dataclass(frozen=True) +class Decision: + rule_id: str + passed: bool + observed: float + expected: float + explanation: str + + +def _require_integer(name: str, value: object) -> int: + """Accept integer values without treating booleans as integer metadata.""" + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an int, excluding bool") + return value + + +def _require_finite_float(name: str, value: object) -> float: + """Require an actual built-in Python float, not a coerced numeric value.""" + value_type: type[object] = type(value) + if not isinstance(value, float) or value_type is not float: + raise TypeError(f"{name} must be a Python float") + if not isfinite(value): + raise ValueError(f"{name} must be finite") + return value + + +def evaluate(observation: Observation) -> Decision: + """Check the multiplier needed for a complete-window global token mean. + + Malformed or unsupported observations raise TypeError or ValueError. A + finite but mathematically incorrect multiplier produces a failed Decision. + The caller must treat validation exceptions as invalid observations, never + as passes or as evidence that this rule executed successfully. + """ + if not isinstance(observation, Observation): + raise TypeError("observation must be an Observation") + + rank: int = _require_integer("rank", observation.rank) + world_size: int = _require_integer("world_size", observation.world_size) + step: int = _require_integer("step", observation.step) + microbatch: int = _require_integer("microbatch", observation.microbatch) + accumulation: int = _require_integer("accumulation", observation.accumulation) + local_tokens: int = _require_integer("local_tokens", observation.local_tokens) + global_tokens: int = _require_integer("global_tokens", observation.global_tokens) + applied_scale: float = _require_finite_float( + "applied_scale", observation.applied_scale + ) + + if world_size <= 0: + raise ValueError("world_size must be positive") + if rank < 0 or rank >= world_size: + raise ValueError("rank must satisfy 0 <= rank < world_size") + if step < 1: + raise ValueError("step must be a positive, one-based update index") + if accumulation <= 0: + raise ValueError("accumulation must be positive") + if microbatch < 0 or microbatch >= accumulation: + raise ValueError("microbatch must satisfy 0 <= microbatch < accumulation") + if local_tokens <= 0: + raise ValueError("local_tokens must be positive for this probe") + if global_tokens <= 0: + raise ValueError("global_tokens must be positive") + if world_size == 1 and global_tokens != local_tokens: + raise ValueError("single-rank local_tokens and global_tokens must agree") + if global_tokens < local_tokens + world_size - 1: + raise ValueError( + "global_tokens must include local_tokens and positive totals for " + "every other rank" + ) + + try: + expected: float = world_size / global_tokens + except OverflowError as error: + raise ValueError( + "required multiplier is not representable as a float" + ) from error + if not isfinite(expected) or expected <= 0.0: + raise ValueError( + "required multiplier must be representable as a finite positive float" + ) + + passed: bool = isclose( + applied_scale, + expected, + rel_tol=RELATIVE_TOLERANCE, + abs_tol=ABSOLUTE_TOLERANCE, + ) + status: str = "matches" if passed else "does not match" + explanation: str = ( + f"Applied multiplier {applied_scale!r} {status} world_size/global_tokens " + f"= {expected!r}. Default DDP averages gradients across ranks, so this " + "factor makes each token-summed microbatch contribute to the mean over " + "all valid tokens in the complete optimizer-update window. No additional " + "accumulation divisor is needed because global_tokens already covers " + f"the full window. Comparison uses rel_tol={RELATIVE_TOLERANCE!r}, " + f"abs_tol={ABSOLUTE_TOLERANCE!r}." + ) + return Decision( + rule_id=RULE_ID, + passed=passed, + observed=applied_scale, + expected=expected, + explanation=explanation, + ) diff --git a/tools/agent_invariants/check.py b/tools/agent_invariants/check.py new file mode 100644 index 00000000..225e3a38 --- /dev/null +++ b/tools/agent_invariants/check.py @@ -0,0 +1,431 @@ +"""Evaluate a frozen agent-authored predicate on native TrainCheck API records. + +This experiment checker is not TrainCheck's built-in inferred-invariant engine. +It requires complete, successful observation calls for an explicitly declared run. +Missing coverage and malformed data raise errors; neither becomes a passing rule. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, TypedDict + +from tools.agent_invariants import authored_rule +from tools.agent_invariants.authored_rule import Observation + +FUNCTION = "tools.agent_invariants.observation.apply_loss_scale" +PRE = "function_call (pre)" +POST = "function_call (post)" +EXCEPTION = "function_call (post) (exception)" +INTEGER_FIELDS = ( + "rank", + "world_size", + "step", + "microbatch", + "accumulation", + "local_tokens", + "global_tokens", +) +FIELDS = frozenset((*INTEGER_FIELDS, "applied_scale")) + + +@dataclass(frozen=True) +class ExpectedRun: + world_size: int + start_step: int + steps: int + accumulation: int + + def validate(self) -> None: + for name, value in self.to_dict().items(): + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError(f"expected {name} must be an integer") + if value < 1: + raise ValueError(f"expected {name} must be positive") + + def to_dict(self) -> dict[str, int]: + return { + "world_size": self.world_size, + "start_step": self.start_step, + "steps": self.steps, + "accumulation": self.accumulation, + } + + +@dataclass(frozen=True) +class Event: + source: str + line: int + rank_directory: int + process_id: int + thread_id: int + func_call_id: str + timestamp: int + step: int + kind: str + raw: dict[str, object] + observation: Observation | None + + def identity(self) -> tuple[int, int, int, str]: + return ( + self.rank_directory, + self.process_id, + self.thread_id, + self.func_call_id, + ) + + +class TraceFile(TypedDict): + path: str + bytes: int + sha256: str + + +class EventEvidence(TypedDict): + path: str + line: int + record: dict[str, object] + + +class DecisionRecord(TypedDict): + rule_id: str + passed: bool + observed: float + expected: float + explanation: str + observation: dict[str, int | float] + pre: EventEvidence + post: EventEvidence + + +class CheckReport(TypedDict): + status: Literal["passed", "violated"] + scope: str + expected_run: dict[str, int] + expected_observations: int + completed_observations: int + violations: int + rule_source: TraceFile + trace_files: list[TraceFile] + decisions: list[DecisionRecord] + + +def mapping(value: object, context: str) -> dict[str, object]: + if not isinstance(value, dict): + raise TypeError(f"{context} must be an object") + result: dict[str, object] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise TypeError(f"{context} must have string keys") + result[key] = item + return result + + +def integer(value: object, context: str) -> int: + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError(f"{context} must be an integer") + return value + + +def text(value: object, context: str) -> str: + if not isinstance(value, str) or not value: + raise TypeError(f"{context} must be a nonempty string") + return value + + +def unique_json_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"Duplicate JSON field: {key}") + result[key] = value + return result + + +def reject_json_constant(value: str) -> object: + raise ValueError(f"Nonfinite JSON constant: {value}") + + +def observation_from_record(record: dict[str, object]) -> Observation: + positional = mapping(record["args"], "args") + if positional: + raise ValueError("Observation calls must record only named scalar arguments") + keywords = mapping(record["kwargs"], "kwargs") + if set(keywords) != FIELDS: + raise ValueError("Observation kwargs do not match the required fields") + values: dict[str, int] = {} + for field in INTEGER_FIELDS: + payload = mapping(keywords[field], f"kwargs.{field}") + if set(payload) != {"int"}: + raise TypeError(f"kwargs.{field} must have the native int payload") + values[field] = integer(payload["int"], f"kwargs.{field}.int") + scale_payload = mapping(keywords["applied_scale"], "kwargs.applied_scale") + if set(scale_payload) != {"float"}: + raise TypeError("kwargs.applied_scale must have the native float payload") + scale = scale_payload["float"] + if not isinstance(scale, float): + raise TypeError("kwargs.applied_scale.float must be a float") + if not math.isfinite(scale): + raise ValueError("kwargs.applied_scale.float must be finite") + if values["local_tokens"] < 1 or values["global_tokens"] < values["local_tokens"]: + raise ValueError("Observation requires positive, ordered token counts") + return Observation( + rank=values["rank"], + world_size=values["world_size"], + step=values["step"], + microbatch=values["microbatch"], + accumulation=values["accumulation"], + local_tokens=values["local_tokens"], + global_tokens=values["global_tokens"], + applied_scale=scale, + ) + + +def parse_event(record: dict[str, object], path: Path, line: int, rank: int) -> Event: + kind = text(record["type"], "type") + if kind not in (PRE, POST, EXCEPTION): + raise ValueError(f"Unexpected observation event type: {kind}") + process_id = integer(record["process_id"], "process_id") + thread_id = integer(record["thread_id"], "thread_id") + timestamp = integer(record["time"], "time") + if process_id < 1 or thread_id < 1 or timestamp < 0: + raise ValueError("Invalid observation event identity or timestamp") + meta = mapping(record["meta_vars"], "meta_vars") + step = integer(meta["step"], "meta_vars.step") + if record["is_bound_method"] is not False or record["obj_id"] is not None: + raise ValueError("Observation helper must be an unbound function") + observation = observation_from_record(record) if kind == PRE else None + if kind != PRE and ("args" in record or "kwargs" in record): + raise ValueError("Observation post record must not repeat argument fields") + if observation is not None and ( + observation.rank != rank or observation.step != step + ): + raise ValueError("Observation rank/step disagrees with its trace context") + return Event( + source=str(path.resolve()), + line=line, + rank_directory=rank, + process_id=process_id, + thread_id=thread_id, + func_call_id=text(record["func_call_id"], "func_call_id"), + timestamp=timestamp, + step=step, + kind=kind, + raw=record, + observation=observation, + ) + + +def file_identity(path: Path) -> TraceFile: + data = path.read_bytes() + return { + "path": str(path.resolve()), + "bytes": len(data), + "sha256": hashlib.sha256(data).hexdigest(), + } + + +def read_events( + root: Path, expected: ExpectedRun +) -> tuple[list[Event], list[TraceFile]]: + if not root.is_dir(): + raise FileNotFoundError(root) + rank_paths = sorted(root.glob("rank*")) + if {path.name for path in rank_paths} != { + f"rank{rank}" for rank in range(expected.world_size) + } or any(not path.is_dir() for path in rank_paths): + raise ValueError("Run rank directories do not match expected world_size") + events: list[Event] = [] + files: list[TraceFile] = [] + for rank in range(expected.world_size): + paths = sorted((root / f"rank{rank}" / "trace").glob("trace_API_*.log")) + if not paths: + raise ValueError(f"No API trace files for rank {rank}") + for path in paths: + data = path.read_bytes() + files.append( + { + "path": str(path.resolve()), + "bytes": len(data), + "sha256": hashlib.sha256(data).hexdigest(), + } + ) + for line, content in enumerate(data.decode("utf-8").splitlines(), 1): + decoded: object = json.loads( + content, + object_pairs_hook=unique_json_object, + parse_constant=reject_json_constant, + ) + record = mapping(decoded, f"{path}:{line}") + function = text(record["function"], "function") + if function == FUNCTION: + events.append(parse_event(record, path, line, rank)) + return events, files + + +def complete_calls(events: list[Event]) -> list[tuple[Event, Event]]: + calls: dict[tuple[int, int, int, str], dict[str, Event]] = {} + for event in events: + call = calls.setdefault(event.identity(), {}) + if event.kind in call: + raise ValueError(f"Duplicate observation event: {event.identity()}") + call[event.kind] = event + paired: list[tuple[Event, Event]] = [] + for identity, call in calls.items(): + if EXCEPTION in call: + raise ValueError(f"Observation call raised an exception: {identity}") + if set(call) != {PRE, POST}: + raise ValueError(f"Incomplete observation call: {identity}") + before, after = call[PRE], call[POST] + if before.timestamp > after.timestamp: + raise ValueError("Observation post timestamp precedes its pre timestamp") + if before.raw["meta_vars"] != after.raw["meta_vars"]: + raise ValueError("Observation pre/post metadata differs") + paired.append((before, after)) + return paired + + +def validate_coverage(calls: list[tuple[Event, Event]], expected: ExpectedRun) -> None: + required = { + (rank, step, microbatch) + for rank in range(expected.world_size) + for step in range(expected.start_step, expected.start_step + expected.steps) + for microbatch in range(expected.accumulation) + } + observed: dict[tuple[int, int, int], Observation] = {} + windows: dict[tuple[int, int], tuple[int, int]] = {} + for before, _ in calls: + sample = before.observation + assert sample is not None + if ( + sample.world_size != expected.world_size + or sample.accumulation != expected.accumulation + ): + raise ValueError("Observation configuration differs from expected run") + key = sample.rank, sample.step, sample.microbatch + if key in observed: + raise ValueError(f"Duplicate observation for rank/step/microbatch: {key}") + observed[key] = sample + window = sample.rank, sample.step + counts = sample.local_tokens, sample.global_tokens + if window in windows and windows[window] != counts: + raise ValueError("Token counts differ within an accumulation window") + windows[window] = counts + if set(observed) != required: + raise ValueError( + f"Incomplete or unexpected observation coverage: " + f"missing={sorted(required - set(observed))}, extra={sorted(set(observed) - required)}" + ) + for step in range(expected.start_step, expected.start_step + expected.steps): + rank_counts = [windows[(rank, step)] for rank in range(expected.world_size)] + globals_seen = {global_count for _, global_count in rank_counts} + if ( + len(globals_seen) != 1 + or sum(local for local, _ in rank_counts) != rank_counts[0][1] + ): + raise ValueError( + "Cross-rank token counts do not sum to the shared global count" + ) + + +def observation_dict(sample: Observation) -> dict[str, int | float]: + return { + "rank": sample.rank, + "world_size": sample.world_size, + "step": sample.step, + "microbatch": sample.microbatch, + "accumulation": sample.accumulation, + "local_tokens": sample.local_tokens, + "global_tokens": sample.global_tokens, + "applied_scale": sample.applied_scale, + } + + +def check_run(root: Path, expected: ExpectedRun) -> CheckReport: + expected.validate() + events, trace_files = read_events(root, expected) + calls = complete_calls(events) + validate_coverage(calls, expected) + decisions: list[DecisionRecord] = [] + for before, after in calls: + sample = before.observation + assert sample is not None + decision = authored_rule.evaluate(sample) + decisions.append( + { + "rule_id": decision.rule_id, + "passed": decision.passed, + "observed": decision.observed, + "expected": decision.expected, + "explanation": decision.explanation, + "observation": observation_dict(sample), + "pre": { + "path": before.source, + "line": before.line, + "record": before.raw, + }, + "post": {"path": after.source, "line": after.line, "record": after.raw}, + } + ) + violations = sum(not decision["passed"] for decision in decisions) + return { + "status": "violated" if violations else "passed", + "scope": "custom agent-authored predicate on TrainCheck API traces; not built-in inferred invariants", + "expected_run": expected.to_dict(), + "expected_observations": expected.world_size + * expected.steps + * expected.accumulation, + "completed_observations": len(decisions), + "violations": violations, + "rule_source": file_identity(Path(authored_rule.__file__)), + "trace_files": trace_files, + "decisions": decisions, + } + + +class Options(argparse.Namespace): + root: Path + world_size: int + start_step: int + steps: int + accumulation: int + output: Path + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", type=Path, required=True) + parser.add_argument("--world-size", type=int, required=True) + parser.add_argument("--start-step", type=int, required=True) + parser.add_argument("--steps", type=int, required=True) + parser.add_argument("--accumulation", type=int, required=True) + parser.add_argument("--output", type=Path, required=True) + options = parser.parse_args(namespace=Options()) + expected = ExpectedRun( + options.world_size, options.start_step, options.steps, options.accumulation + ) + try: + report = check_run(options.root, expected) + except (OSError, UnicodeError, ValueError, TypeError, KeyError) as error: + failure: dict[str, object] = { + "status": "error", + "scope": "custom agent-authored predicate on TrainCheck API traces", + "expected_run": expected.to_dict(), + "completed_observations": None, + "violations": 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 1 if report["violations"] else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/agent_invariants/observation.py b/tools/agent_invariants/observation.py new file mode 100644 index 00000000..59ada409 --- /dev/null +++ b/tools/agent_invariants/observation.py @@ -0,0 +1,65 @@ +"""Capture the actual loss multiplier with native TrainCheck API records. + +This helper adds scalar observations for one experimental predicate. It does not +infer invariants or register a new built-in TrainCheck relation. Window counts +come from the workload; they are not independently recounted here. +""" + +from __future__ import annotations + +from collections.abc import Callable + +import torch + +OBSERVATION_FIELDS = ( + "rank", + "world_size", + "step", + "microbatch", + "accumulation", + "local_tokens", + "global_tokens", + "applied_scale", +) + + +def apply_loss_scale( + loss: torch.Tensor, + *, + rank: int, + world_size: int, + step: int, + microbatch: int, + accumulation: int, + local_tokens: int, + global_tokens: int, + applied_scale: float, +) -> torch.Tensor: + """Apply the supplied multiplier, retaining the original autograd operation.""" + result: object = loss * applied_scale + if not isinstance(result, torch.Tensor): + raise TypeError("Loss multiplication must return a Tensor") + return result + + +def wrapped_loss_scale() -> Callable[..., object]: + """Use the existing wrapper and capture only the primitive keyword arguments. + + The positional loss tensor and return tensor are intentionally unobserved; + this predicate uses the scalar multiplier, not tensor values or metadata. + No predicate runs in the training process. TraceSession owns tracing gates + and buffer drainage, as it does for ordinary PyTorch API wrappers. + """ + from traincheck.instrumentor.tracer import wrapper + + argument_config: dict[str, object] = {key: {} for key in OBSERVATION_FIELDS} + return wrapper( + apply_loss_scale, + is_bound_method=False, + scan_proxy_in_args=False, + dump_stack_trace=False, + dump_args_config=argument_config, + dump_ret=False, + handle_proxy=False, + trigger_var_dump=False, + ) diff --git a/tools/distributed_probe/workload.py b/tools/distributed_probe/workload.py index 5baff6f5..200b92b9 100644 --- a/tools/distributed_probe/workload.py +++ b/tools/distributed_probe/workload.py @@ -62,6 +62,7 @@ class Options: instrumentation: Instrumentation device: Literal["cpu", "cuda"] accumulation: int + observe_normalization: bool = False @dataclass(frozen=True) @@ -370,10 +371,13 @@ def parse_options(argv: list[str] | None = None) -> Options: "--instrumentation", choices=("plain", "api", "subclass"), default="plain" ) parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda") + parser.add_argument("--observe-normalization", action="store_true") args = parser.parse_args(argv) stop_after = args.steps if args.stop_after is None else args.stop_after if not 1 <= stop_after <= args.steps or not 1 <= args.accumulation <= 4: parser.error("Require 1 <= stop-after <= steps and 1 <= accumulation <= 4") + if args.observe_normalization and args.instrumentation == "plain": + parser.error("--observe-normalization requires API or subclass tracing") return Options( args.output, args.resume, @@ -384,12 +388,16 @@ def parse_options(argv: list[str] | None = None) -> Options: args.instrumentation, args.device, args.accumulation, + args.observe_normalization, ) def run(options: Options) -> None: + from tools.agent_invariants.observation import wrapped_loss_scale from tools.distributed_probe.tracing import start_tracing + if options.observe_normalization and options.instrumentation == "plain": + raise ValueError("Normalization observations require API or subclass tracing") 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")) @@ -420,6 +428,9 @@ def run(options: Options) -> None: ) session = start_tracing(options.instrumentation, rank_output / "trace", model) try: + scale_operation = ( + wrapped_loss_scale() if options.observe_normalization else None + ) with session.suppress(): wrapped: nn.Module = model if world_size > 1: @@ -488,7 +499,24 @@ def run(options: Options) -> None: logits: object = wrapped(batch.tokens) if not isinstance(logits, torch.Tensor): raise TypeError("Training output must be a Tensor") - loss = token_loss(logits, batch) * scale + summed_loss = token_loss(logits, batch) + if scale_operation is None: + loss = summed_loss * scale + else: + observed_loss = scale_operation( + summed_loss, + rank=rank, + world_size=world_size, + step=step, + microbatch=microbatch, + accumulation=options.accumulation, + local_tokens=local_tokens, + global_tokens=global_tokens, + applied_scale=scale, + ) + if not isinstance(observed_loss, torch.Tensor): + raise TypeError("Observed loss must be a Tensor") + loss = observed_loss loss.backward() losses.append(loss.detach()) if step == 1: @@ -549,6 +577,7 @@ def run(options: Options) -> None: "precision": options.precision, "fault": options.fault, "instrumentation": options.instrumentation, + "normalization_observations": options.observe_normalization, "start_step": start_step, "completed_steps": options.stop_after, "total_steps": options.steps, From fe482640d2979481419040fe3395b894cb8600fd Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 1 Nov 2026 18:10:57 +0800 Subject: [PATCH 2/3] Preserve temporary tensor lifetimes in observation adapter --- tests/test_agent_invariant_observation.py | 99 ++++++++++++++++++++++- tools/agent_invariants/observation.py | 30 ++++++- tools/distributed_probe/workload.py | 10 +-- 3 files changed, 128 insertions(+), 11 deletions(-) diff --git a/tests/test_agent_invariant_observation.py b/tests/test_agent_invariant_observation.py index 1197e9f9..71551e66 100644 --- a/tests/test_agent_invariant_observation.py +++ b/tests/test_agent_invariant_observation.py @@ -63,7 +63,10 @@ try: with session.step(6): - operation( + # Exercise malformed runtime input across an untyped external boundary. + malformed_operation: object = operation + assert callable(malformed_operation) + malformed_operation( 3.75, rank=0, world_size=2, step=6, microbatch=0, accumulation=2, local_tokens=10, global_tokens=32, applied_scale=0.0625, @@ -78,6 +81,72 @@ Path(sys.argv[2]).write_text(json.dumps(report)) """ +LIFETIME_WORKER = """ +import importlib.util +import sys +import weakref +from pathlib import Path + +import torch + +from tools.distributed_probe import tracing + +spec = importlib.util.spec_from_file_location( + 'tools.distributed_probe.workload', sys.argv[1] +) +assert spec is not None and spec.loader is not None +workload = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = workload +spec.loader.exec_module(workload) + +pending: list[weakref.ReferenceType[torch.Tensor]] = [] +active = False +checked = 0 +original_token_loss = workload.token_loss +original_start_tracing = tracing.start_tracing + + +def tracked_token_loss(logits: torch.Tensor, batch: workload.Batch) -> torch.Tensor: + value = original_token_loss(logits, batch) + if active: + pending.append(weakref.ref(value)) + return value + + +def install_lifetime_check( + mode: tracing.Mode, output: Path, model: torch.nn.Module +) -> tracing.TraceSession: + global active + session = original_start_tracing(mode, output, model) + original_backward = torch.Tensor.backward + + def checked_backward( + tensor: torch.Tensor, *args: object, **kwargs: object + ) -> None: + global checked + assert len(pending) == 1, 'Expected one unscaled training loss per backward' + reference = pending.pop() + assert reference() is None, 'Unscaled loss remains live at backward entry' + checked += 1 + original_backward(tensor, *args, **kwargs) + + torch.Tensor.backward = checked_backward + active = True + return session + + +workload.token_loss = tracked_token_loss +tracing.start_tracing = install_lifetime_check +mode = sys.argv[2] +options = workload.Options( + Path(sys.argv[3]), None, 2, 2, 'fp32', 'none', mode, 'cpu', 2, + observe_normalization=mode != 'plain', +) +workload.run(options) +assert checked == 4 and not pending +print('Four temporary loss lifetimes verified') +""" + def mapping(value: object) -> dict[str, object]: if not isinstance(value, dict): @@ -201,3 +270,31 @@ def test_plain_observation_programmatic_run_is_rejected_before_output( with pytest.raises(ValueError, match="require API or subclass tracing"): run(options) assert not output.exists() + + +@pytest.mark.parametrize("mode", ["plain", "api"]) +def test_workload_releases_unscaled_loss_before_backward( + tmp_path: Path, mode: str +) -> None: + """Observe Python lifetime during actual training without keeping tensors alive.""" + environment = dict(os.environ) + environment["PYTHONPATH"] = str(ROOT) + environment["PYTHONDONTWRITEBYTECODE"] = "1" + result = subprocess.run( + [ + sys.executable, + "-c", + LIFETIME_WORKER, + str(ROOT / "tools/distributed_probe/workload.py"), + mode, + str(tmp_path / "run"), + ], + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=90, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "Four temporary loss lifetimes verified" in result.stdout diff --git a/tools/agent_invariants/observation.py b/tools/agent_invariants/observation.py index 59ada409..68cc8bf6 100644 --- a/tools/agent_invariants/observation.py +++ b/tools/agent_invariants/observation.py @@ -7,7 +7,7 @@ from __future__ import annotations -from collections.abc import Callable +from typing import Protocol import torch @@ -23,6 +23,22 @@ ) +class LossScaleOperation(Protocol): + def __call__( + self, + loss: torch.Tensor, + *, + rank: int, + world_size: int, + step: int, + microbatch: int, + accumulation: int, + local_tokens: int, + global_tokens: int, + applied_scale: float, + ) -> torch.Tensor: ... + + def apply_loss_scale( loss: torch.Tensor, *, @@ -42,7 +58,7 @@ def apply_loss_scale( return result -def wrapped_loss_scale() -> Callable[..., object]: +def wrapped_loss_scale() -> LossScaleOperation: """Use the existing wrapper and capture only the primitive keyword arguments. The positional loss tensor and return tensor are intentionally unobserved; @@ -53,7 +69,7 @@ def wrapped_loss_scale() -> Callable[..., object]: from traincheck.instrumentor.tracer import wrapper argument_config: dict[str, object] = {key: {} for key in OBSERVATION_FIELDS} - return wrapper( + wrapped = wrapper( apply_loss_scale, is_bound_method=False, scan_proxy_in_args=False, @@ -63,3 +79,11 @@ def wrapped_loss_scale() -> Callable[..., object]: handle_proxy=False, trigger_var_dump=False, ) + + def checked(*args: object, **kwargs: object) -> torch.Tensor: + result = wrapped(*args, **kwargs) + if not isinstance(result, torch.Tensor): + raise TypeError("Observed loss must be a Tensor") + return result + + return checked diff --git a/tools/distributed_probe/workload.py b/tools/distributed_probe/workload.py index 200b92b9..c48046b1 100644 --- a/tools/distributed_probe/workload.py +++ b/tools/distributed_probe/workload.py @@ -499,12 +499,11 @@ def run(options: Options) -> None: logits: object = wrapped(batch.tokens) if not isinstance(logits, torch.Tensor): raise TypeError("Training output must be a Tensor") - summed_loss = token_loss(logits, batch) if scale_operation is None: - loss = summed_loss * scale + loss = token_loss(logits, batch) * scale else: - observed_loss = scale_operation( - summed_loss, + loss = scale_operation( + token_loss(logits, batch), rank=rank, world_size=world_size, step=step, @@ -514,9 +513,6 @@ def run(options: Options) -> None: global_tokens=global_tokens, applied_scale=scale, ) - if not isinstance(observed_loss, torch.Tensor): - raise TypeError("Observed loss must be a Tensor") - loss = observed_loss loss.backward() losses.append(loss.detach()) if step == 1: From 37303114ade059e436c663f3904df4c8607d1dc9 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 1 Nov 2026 18:18:11 +0800 Subject: [PATCH 3/3] Report verified agent-authored invariant GPU pilot results --- .../agent-authored-normalization-20261101.md | 183 ++++++++++++++++++ 1 file changed, 183 insertions(+) create mode 100644 docs/measurements/agent-authored-normalization-20261101.md diff --git a/docs/measurements/agent-authored-normalization-20261101.md b/docs/measurements/agent-authored-normalization-20261101.md new file mode 100644 index 00000000..50ea4819 --- /dev/null +++ b/docs/measurements/agent-authored-normalization-20261101.md @@ -0,0 +1,183 @@ +# Agent-authored normalization invariant: targeted feasibility pilot + +This experiment asks whether an agent can turn a stated training objective into +an executable invariant that detects the distributed normalization defect from +the [Ring18 probe](ring18-normalization-coverage-20261031.md). It evaluates one +targeted bug family. It is not an RSI score, an autonomous bug-discovery result, +or an evaluation of TrainCheck's built-in invariant inference. + +## Authoring and execution boundary + +A separate agent, with no inherited conversation history, received the intended +global valid-token mean objective, token-summed microbatch losses, default DDP +gradient averaging, complete accumulation windows, and a scalar observation +interface. The parent selected normalization as the topic and supplied the +observation fields. The author received no faulty implementation, trace, +diagnostic result, numeric fixture, or expected equation. Access isolation was +instructional, not a filesystem sandbox or a formal blinded evaluation. + +The agent derived and implemented: + +```text +applied_scale approximately equals world_size / global_tokens +``` + +`global_tokens` covers every rank and every microbatch in the optimizer update. +The rule compares Python scalar multipliers with relative tolerance `1e-12` and +zero absolute tolerance. Its complete derivation, validation, and unsupported +cases are in [AUTHORING.md](../../tools/agent_invariants/AUTHORING.md). + +The predicate was frozen before training/native-trace evaluation. A subsequent +lint-only rewrite preserved exact-float validation; original and final freeze +receipts are retained. Unit and synthetic-record tests had run before that +rewrite; the equation and tolerance did not change. The final rule SHA-256 is +`f8af836ff635cf5a8de94c0bc7e2457a3397687066c8799af89a5d414aa8df4b`. + +The probe's opt-in `--observe-normalization` flag records the scalar arguments of +the actual loss multiplication through TrainCheck's existing native API wrapper. +It captures rank, update, microbatch, accumulation, local/global window counts, +and the multiplier actually applied. It adds no tensor-content serialization and +does not enable the model tracker. Count computation remains outside tracing. + +[check.py](../../tools/agent_invariants/check.py) runs the authored Python +predicate offline. It requires exactly one successful pre/post pair per expected +rank/update/microbatch, validates native scalar payloads and cross-rank window +counts, and retains original event identities, timestamps, source locations, +records, and file hashes. Malformed/incomplete evidence reports an error with +unknown violation count, never a passing check. CLI exit codes are 0 for complete +pass, 1 for complete violations, and 2 for an error. Training completion is +verified separately through successful worker/supervisor exits and checkpoints. + +## Matched GPU experiment + +Final source revision: `fe482640d2979481419040fe3395b894cb8600fd`, based on PR #167 +revision `1df0ee5d9f3c8d31c58111215dd69f11b35dfba8`. Ring18 NVIDIA A16 GPUs 0/1; +Python 3.10, PyTorch 2.2.2+cu121. The 5,504-parameter Transformer uses FP32, +dropout=0, AdamW, deterministic data, one Torch CPU thread, and four complete +optimizer updates per run. Observed runs select `--instrumentation api`; the +adapter also supports subclass mode, which this matrix does not exercise. +This is a correctness pilot, not a performance run. + +The matrix contains six observed two-GPU runs: healthy/global-token and faulty +local-mean normalization at accumulation lengths 1, 2, and 3. Global window token +counts are respectively 16, 32, and 48. A seventh observed run uses one GPU and +local-mean normalization, which is mathematically correct in that configuration. +Two additional plain runs at accumulation 2 provide matched healthy/faulty +checkpoint controls. Each configuration is run once. + +The first nine-run matrix used `26c77ce`. Review found extra persistent loss +locals could change temporary tensor lifetimes, including with observation off. +The final source restores the original disabled expression and avoids additional +persistent tensor aliases in the observed path. Both real-training lifetime +regressions fail with `workload.py` from `26c77ce` in the same test environment and +pass after correction. Preliminary runs are +preserved in their own directory; final claims use the corrected rerun only. + +## Corrected-run results + +All nine training phases exited successfully with unchanged source. The seven +observed runs completed all **104 expected rule applications**, with no missing, +duplicate, or exceptional observation calls. Counts below include both ranks +where applicable; one application corresponds to one loss multiplication. + +| GPUs | Accumulation | Normalization | Completed checks | Violations | First violating update | +| --- | --- | --- | ---: | ---: | --- | +| 2 | 1 | global-token mean | 8 | 0 | — | +| 2 | 1 | local mean | 8 | 8 | 1 | +| 2 | 2 | global-token mean | 16 | 0 | — | +| 2 | 2 | local mean | 16 | 16 | 1 | +| 2 | 3 | global-token mean | 24 | 0 | — | +| 2 | 3 | local mean | 24 | 24 | 1 | +| 1 | 2 | local mean, correct on one rank | 8 | 0 | — | + +Thus all three tested faulty two-GPU configurations triggered the same authored +rule, while all 56 checks in the four mathematically healthy runs passed. These +are repeated applications of one invariant to one bug family, not 104 independent +bug examples or an estimate of general detection accuracy. The separate +first-update global-batch gradient oracle agreed with the healthy/faulty +classification in every run. Both healthy and faulty runs still had identical +model replicas across ranks. + +For healthy and faulty accumulation-2 runs, the complete plain/observed checkpoint +comparisons passed: **4 rank-file pairs, 548 tensor pairs**, including model, +gradients, optimizer, scheduler, CPU/CUDA RNG and provenance. Corresponding whole +checkpoint file SHA-256 values also matched. No fields were dropped or normalized +to establish equality. Plain controls were unobserved by this rule. + +## Validation + +The composed focused suite passed **113 tests** on macOS/Python 3.10 in 19.77 s +and **113 tests** on Ring18/Python 3.10 in 45.43 s. These are test-run durations, +not tracing/checking overhead measurements. Coverage includes actual native +records, gradients and RNG preservation, exception propagation, suppression, +temporary tensor lifetimes, malformed input, and incomplete/duplicate evidence. +The local PyTorch runs required elevated execution for OpenMP shared-memory +access; the first sandbox-only integration attempt failed on that environment +restriction, not a rule assertion. + +Configured pre-commit hooks passed. Explicit strict mypy checked the eight +changed/new Python modules because the hook does not select tools/tests. Pinned +Ruff/Black/isort also checked tests, which the hooks exclude. Configured core +mypy passed 79 unchanged source modules; its legacy untyped-body coverage limit +remains. No remote CI success is claimed. + +```bash +PYTHONDONTWRITEBYTECODE=1 python -m pytest -q -p no:cacheprovider \ + tests/test_agent_authored_rule.py tests/test_agent_invariant_observation.py \ + tests/test_agent_invariant_trace_check.py tests/test_distributed_probe_workload.py \ + tests/test_distributed_probe_tracing.py tests/test_distributed_probe_compare.py +``` + +## Interpretation limits + +- This tests one human-selected contract and observation interface. It does not + establish that an agent can discover which invariants to write on an unknown + task, choose trustworthy observation sites, or improve a model autonomously. +- Counts originate in the workload. Cross-rank agreement does not independently + establish that masks, count computation, loss reduction, or DDP reduction are + correct. A consistently wrong count can escape this rule. +- It does not check missing/repeated backward calls, altered communication hooks, + zero-local-token or partial windows, stochastic training, or mixed precision. +- First violating update refers to an offline trace location, not online + detection latency. No checking-overhead or general detection-rate claim is made. +- This custom predicate bypasses the built-in relation-template limitation. It + neither adds a relation family nor tests what existing inference could learn + from these newly exposed scalars. Existing identity-related false alarms are + outside this pilot. +- Complete checkpoint comparison preserves every stored value and field; it maps + storage to CPU and does not reconstruct device placement or storage aliasing. + +## Reproduction and artifacts + +From a checkout at the final source revision, in the existing compatible +environment, run each case in its own fresh output directory. Example: + +```bash +CUDA_VISIBLE_DEVICES=0,1 CUBLAS_WORKSPACE_CONFIG=:4096:8 \ + python -m torch.distributed.run --standalone --nproc-per-node=2 \ + tools/distributed_probe/workload.py --output /scratch/run \ + --steps 4 --device cuda --accumulation 2 --fault none \ + --instrumentation api --observe-normalization +python -m tools.agent_invariants.check --root /scratch/run \ + --world-size 2 --start-step 1 --steps 4 --accumulation 2 \ + --output /scratch/check.json +``` + +For the faulty case select `--fault local-mean`; for plain controls omit +`--observe-normalization` and select `--instrumentation plain`. A plain run has no +rule observations and must not be interpreted as passing the predicate. + +Final remote artifacts: +`/home/yuxuan/experiments/traincheck-agent-invariants-20261101-v2`. +Preliminary artifacts remain at the same path without the `-v2` suffix. +Local evidence root: +`/Users/matt/gitrepos/traincheck/outputs/ring18-agent-invariants-20261101`. +The operator's exact nine commands, exit codes, environment, source hashes, +frozen rule receipts, raw API traces, checkpoints, metrics, and reports are kept +with the experiment. No credentials or training artifacts are committed. + +The final evidence archive contains 266 manifest-verified files, including raw +API traces, and is 10,950,069 bytes. Its SHA-256 is +`c138f5e7e050939182ade7f60b79baad9a2c5833ca42ad4c36a237d6a549f09b`. +The source manifest covers 124 files at the stated revision. The archive receipt, +manifest, and source identities are also checked during local retrieval.