diff --git a/docs/measurements/model-observation-workflow-20261102.md b/docs/measurements/model-observation-workflow-20261102.md new file mode 100644 index 00000000..672afb23 --- /dev/null +++ b/docs/measurements/model-observation-workflow-20261102.md @@ -0,0 +1,367 @@ +# Model snapshots and a fresh optional-tool comparison + +Status: implementation, native snapshot compatibility and the fresh comparison +are complete. The optional-tool agent voluntarily used the new snapshot and +comparison APIs. Both agents first submitted a valid repair in turn two, and +both final sources passed all six frozen scenarios. This single interrupted +pair establishes working adoption and measured costs, not a causal speedup. + +## Problem and implementation + +In the [preceding FSDP pilot](fsdp-source-consistency-pilot-20261102.md), the +optional-tool agent read TrainCheck's callable-wrapper manual and then wrote a +79-line parameter/buffer fingerprint diagnostic. Its final diagnostic was 128 +lines. No attempt to activate TrainCheck was observed in the reviewed operation +record. The operation record supports +an interface/task mismatch, not an observed installation failure or a measured +advantage of handwritten code over an equivalent TrainCheck implementation. + +The signed candidate `498bf5f9f21c44c9d9487fb9b11b14585fd28304` adds +`traincheck.observation`: explicit model snapshots, validated file reading and +writing, selected-pair comparison and complete-rank comparison. It shares model +enumeration with the existing sampler and reuses the tensor tracker's exact +SHA-256 implementation. It preserves all materialized registered parameter and +buffer names, nonpersistent buffers and repeated object aliases. It does not +replace parameters, implicitly gather shards, infer observation points or run +the invariant engine. Missing ranks, duplicate observations and malformed input +are errors. A difference is an observation, not automatically a training bug. + +The [generic guide](../model-observation.md) and runnable CPU/DDP examples expose +the workflow without this case's diagnosis, repair or source provenance. The +agent selects its own boundaries and predicates. Complete fingerprints require +synchronous device-to-CPU transfer; this change does not promise negligible +overhead or implement asynchronous tracing. + +Importing the lightweight observation API now leaves PyTorch hooks unchanged. +The four existing top-level package exports load lazily and retain their object +identities; explicit instrumentor imports retain their prior activation +behavior. The legacy parallel tensor hash remains unchanged. Its quantization +and first-row scope are not used to establish complete-content fingerprint +agreement. + +A focused regression also exposed a legacy proxy observer ignoring the existing +tracing-suppression scope. Snapshot collection emitted four extra variable +records and changed parameter metadata/timestamps with the parent observer. +The narrow suppression guard produces zero extra records and unchanged +metadata in the same test. Ordinary calls still produce their two pre/post +records, and explicit nested application tracing restores those observations. +This is a correctness result, not a tracing-overhead speedup measurement. + +## Local validation + +The candidate's full local test suite reports 1,074 passed, nine skipped and two +failures in 112.28 seconds. Both failures reproduce with the same assertions in +an isolated archive of parent `55305542d195086f04526bf7ace8febe250bc0e1`: +`test_traincheck_internal_attr_hidden` and `test_inject_eval_loop`. No new +failure was observed. The focused overlapping selection reports 102 passed +and five CUDA-only skips. These counts must not be added together. + +Configured Ruff 0.3.3, mypy 1.9.0, Black 24.3.0 and isort 5.10.1 hooks pass. +New reusable modules and examples also pass scoped strict mypy. Tests cover +whole tensor contents, scalar/empty/noncontiguous/bfloat16/Boolean/complex +representations, signed zero and NaN payloads, cached-hash avoidance, alias/name +coverage, malformed inputs, exceptions, RNG/model/optimizer identity and object +lifetime, pure imports and coexistence with active tracking. + +The actual two-process CPU example completes two updates per rank and writes +six snapshots. The final complete-rank comparison covers five tensors per +rank and reports equal fingerprints. On this Mac it required an explicit +loopback rendezvous rather than the standalone hostname route. This does not +establish CUDA/FSDP compatibility; that is checked separately below. + +## Native Ring18 compatibility + +Two-update matched plain/snapshot runs passed on both the two-GPU current +configuration and the one-GPU reference. All 60 complete native state-file +comparisons passed, including model, all buffers, gradients, optimizer, inputs +and Torch/CUDA/Python/NumPy RNG. All seven explicit snapshots cover 76 parameters +and 19 buffers each. Their 665 tensor fingerprints, shapes, dtypes and complete +ordered name sets match independently saved native tensor values. + +The rank comparator reports differing loaded and pre-wrap states and matching +post-wrap states, consistent with the already qualified case. This also +illustrates why matching reconstructed replicas alone does not establish +coherent initialization. These observations are operator qualification evidence, +not information supplied to either solver. + +| New operator preflight | Full state-file comparisons | Snapshots | Allocated GPU-seconds | +| --- | ---: | ---: | ---: | +| Two-GPU current, plain/snapshot | 42 | 6 | 70.677 | +| One-GPU reference, plain/snapshot | 18 | 1 | 25.176 | +| Total | 60 | 7 | 95.853 | + +All four phase starts have matched completion and cleanup receipts, with no +unknown cost. Full tensor files remain on Ring18; 52 compact artifacts were +retrieved and hash-verified, and snapshot reading/rank comparisons were repeated +locally. The paired jobs include private observer/state I/O and have no timing +repeats, so their durations are not an overhead benchmark. They do not establish +live storage-alias preservation or compatibility with arbitrary sharded models. + +## Fresh comparison protocol + +The new pair reuses the preceding qualified 14,067,712-parameter native FSDP +case. All 3,784 public source files, the corpus and ordinary observations are +unchanged. Both agents get the same neutral readiness task, editable source and +request service. Treatment alone adds 38 explicitly allowlisted package, +documentation and example files. Use remains optional. Neither packet says +there is a bug or includes a case-specific observation recipe. + +Both native CLI conversations start fresh with the same pinned model/settings. +Budgets remain six requests, 1–16 updates per request, 240 seconds of worker +lifetime, 900 aggregate CLI seconds and eight decision turns. GPU requests are +serialized, with queue waiting recorded separately. New tool compatibility +must pass before packet review and launch freeze. + +The requested agent configuration is native CLI 0.160.0 with `gpt-6-astra`, +high reasoning effort and priority service. The exact binary hash and settings +are bound in `launch-plan.json`; the provider's underlying model weights are +not independently attested. The conversations begin without the parent chat's +task history. Packet equality and instructions are verified, but this is not a +claim of enforced filesystem isolation. + +Ring18 runs Python 3.10.14, PyTorch 2.2.2+cu121, CUDA 12.1 and the supplied +Transformers 4.47.0.dev0 source. The current profile uses two NVIDIA A16 GPUs +(physical devices 4/5, driver 535.309.01); the reference uses one. The model is +14,067,712-parameter Pythia with six layers, hidden size 128, four attention +heads and sequence length 64. Two cache domains are represented on one physical +host; this is not a multi-host hardware trial. The fixed corpus, seed 1234, +source manifests and original oracle states are retained with hashes. + +The six-scenario private criterion and original native oracle states are reused +read-only with hashes. The earlier failed strict cross-world numerical criterion +remains failed; acceptance continues to require coherent initial source and +complete same-topology trajectories. Every submitted source occurrence is +retained, and only exact complete source maps may share replay results. + +Report mechanism hypothesis, runtime confirmation, first independently valid +submitted repair, visible verification and terminal completion separately. +Full cache-version diagnosis remains distinct from finding inconsistent rank +state. Tokens include cached input once; within-turn milestones retain token +intervals rather than interpolation. GPU cost is allocated device count times +supervised worker lifetime after verified release. New qualification, solver +and private grading costs are separate; historical qualification is not charged +again. + +Tool-use accounting includes the new snapshot interface. Absence of API trace +files is not evidence of non-use. Documentation reads, preparation, actual +capture, complete rank coverage and successful comparison are separate facts. +Capture attribution must bind the submitted source, successful request and +validated runtime artifacts; a JSON filename alone is insufficient. + +This is a fresh workflow pilot on a previously examined case. One matched pair +cannot establish a general or causal efficiency improvement, especially when +agents select different validation work. + +### Transport interruption + +The first source upload in each arm exceeded the controller's 60-second SSH +timeout. Both request receipts explicitly record that remote execution had not +started and allocated GPU cost was zero. The completed first CLI turns remain +part of the experiment: control used 78.330 CLI seconds and treatment used +69.314 seconds. Their cumulative input-plus-output token counts were 371,462 +and 324,841, respectively, including cached input once. These are partial +costs, not time-to-diagnosis or repair results. + +The first recovery transferred the delta successfully but rejected the reusable +preflight directory because that earlier run had generated Python bytecode +there. The strict full-source check failed before either request launched GPU +workers or resumed an agent. The directory and failed receipts were preserved. +A separate clean base now contains only the 3,824 manifest-selected source +files; it is never executed. Both actual pending submissions passed a remote +dry reconstruction with matching complete byte and permission maps. + +Recovery retains the same native conversation IDs, submissions, logical request +numbers and remaining budgets. Its separate amendment replaces the full-source +upload with hash-verified reuse of this clean base and transfer of changed +contents, equally for both arms. The first delta packages contain seven blobs +each and approximately 440 KiB instead of approximately 20 MB. The frozen task, +optional package and acceptance criteria remain unchanged. Permissions are +preserved relative to the actual frozen submissions; the original protocol's +file-copy snapshots had already made those files mode 0644. This recovery does +not claim to preserve executable modes from the original agent workspace. + +All three physical attempts at each logical first request retain separate +receipts and queue intervals. The failed assembly duration is not classified as +operator suspension or charged as a second logical request. Report actual +elapsed time, queue-adjusted time and time additionally adjusted for measured +operator interruptions separately. The failed upload itself remains visible; +operator suspension is not a measurement of exact network downtime. This +interruption further limits any end-to-end efficiency comparison. + +The clean-base recovery subsequently completed both arms' first and second GPU +requests. Treatment's third upload then exceeded the same 60-second deadline +despite its 452,435-byte delta package. Its receipt explicitly records no remote +execution and zero GPU cost. Control continued independently and finished. An +append-only, same-thread recovery completed treatment's pending third request +with longer transport deadlines and unchanged worker, turn and CLI budgets. +Its accounting preserves the complete earlier composition, failed attempts +and original submissions. One prelaunch check also stopped before any service +started because the earlier timed-out upload subsequently appeared as a partial +manifest. That manifest was verified as the exact 603,648-byte prefix of the +2,129,045-byte original, with twelve other complete matching files. The partial +upload remains unused. A separate bound prelaunch amendment verified the +preserved partial map and absence of owned processes, then used a fresh complete +upload and the unchanged strict assembly checks. This abort consumed no agent +turn or GPU execution. These network interruptions preclude treating this as a +clean wall-time benchmark. + +## Observed adoption + +The treatment voluntarily imported `snapshot_model`, `write_snapshot` and +`compare_model_snapshots` in its first diagnostic. These calls executed on the +GPU workers. The two distributed requests returned four valid snapshots each; +the reference returned one. Every snapshot covers 95 registered tensors, for +855 tensor observations across nine snapshots. Turn three also successfully +executed `read_snapshot`, pair comparison and +`compare_ranks(expected_ranks=[0, 1])` locally with assertions. Turn four compared +both repaired ranks with the reference initialization. + +The observations changed the agent's explanation: it rejected its initial +RAM-efficient-loading hypothesis after observing that the flag was disabled. +It identified 76 differing pre-wrap tensor fingerprints and 38 changes to rank +zero's loaded state during wrapping, then selected synchronization. It also +executed the preservation predicate against the original and revised snapshots, +observing failure and success respectively. The actual cache-version origin +remained unresolved in both agents' final responses. + +No TrainCheck import, configuration or runtime integration error was observed. +The agent still selected observation boundaries, wrote an FSDP subclass adapter, +entered the full-parameter context and added a rank-zero snapshot broadcast. +Its first cross-rank inspection used manual JSON pairing before it adopted the +strict comparator. A local missing-`safetensors` error occurred before snapshot +reading; a subsequent checkpoint-reader retry and TrainCheck comparisons +succeeded. This is distinct from a TrainCheck integration failure. + +## Independent correctness and first valid repair + +All nine ordered submission occurrences were retained. Only exact +arm/content/permission identities shared replay results, yielding six source +maps and 36 actual scenario replays. Both initial sources passed the four +healthy/uniform controls and failed mixed and swapped generation. Every repaired +source passed all six scenarios. Thus 32 replays passed and four expected +baseline mismatches failed, with zero unknown outcomes. Control's final turn +five and treatment's final turn four both passed 6/6. + +The first independently valid submitted source is turn two in both arms. +This is a retrospective correctness judgment at a frozen submission boundary, +not the earlier file-edit event or the later moment when the agent saw its +own validation results. All preceding sources were checked. Costs below include +the completed first diagnostic request and the completed first two CLI turns; +they exclude the second GPU request, which followed that submission. + +| At the first valid submitted repair | Control | Optional TrainCheck | +| --- | ---: | ---: | +| Global decision turn | 2 | 2 | +| CLI execution time, seconds | 188.260 | 141.211 | +| Input plus output tokens | 981,070 | 715,584 | +| Cached input tokens, already included above | 901,376 | 641,024 | +| Uncached input plus output tokens | 79,694 | 74,560 | +| Allocated solver GPU-seconds | 35.031 | 25.289 | +| Actual elapsed time including interruptions, seconds | 2,113.871 | 2,075.498 | +| Elapsed time excluding queue and measured operator suspension, seconds | 434.263 | 365.628 | + +With control as denominator, this pair used 24.99% less CLI execution time, +27.06% fewer total tokens and 27.81% fewer allocated GPU-seconds to its first +valid submission in treatment. CLI time excludes remote service time and +queue/operator pauses. The adjusted elapsed clock still includes failed +transfers; neither clock establishes a clean, general wall-time speedup. +There are no timing repeats, and the agents wrote different diagnostics and +regression checks. + +Visible runtime confirmation of inconsistent initial state occurred in turn +two: 376.418 versus 338.619 seconds after removing queue and recorded operator +suspension. Visible repair verification followed in turn three: 568.975 versus +515.522 seconds on that same adjusted clock. These are received-event upper +bounds, distinct from submitted-source boundaries. Tokens for such within-turn +events are retained as completed-turn brackets, without interpolation. Neither +agent established the actual cache-version origin of the inconsistency. + +## Completed-trial costs + +| Completed-trial metric | Control | Optional TrainCheck | +| --- | ---: | ---: | +| Accepted logical execution requests | 4 | 3 | +| CLI execution time, seconds | 362.937 | 294.554 | +| Input plus output tokens | 1,852,226 | 1,269,560 | +| Cached input tokens, already included above | 1,750,784 | 1,181,184 | +| Uncached input plus output tokens | 101,442 | 88,376 | +| Allocated solver GPU-seconds | 118.342 | 72.822 | +| Actual wall time including interruptions, seconds | 2,729.242 | 4,171.665 | +| Wall time excluding recorded queue waits, seconds | 2,534.728 | 4,008.938 | +| Wall time additionally excluding measured operator suspension, seconds | 949.162 | 847.126 | + +These are original-start totals across the preserved native conversations and +recoveries. Tokens include cached input once; reasoning tokens are already part +of output. The last clock still includes failed transfers and service overhead; +it is not a clean counterfactual wall time. Core qualification and independent +grading GPU costs are separate from the solver totals. +Token volumes are not dollar prices: most of the total token difference is +cached input, and input, cached input and output can have different prices. +The uncached-input-plus-output volume is 12.88% lower; this also is not a +billing estimate. + +The terminal totals are 31.46% lower in tokens, 18.84% lower in CLI time and +38.46% lower in solver GPU-seconds for treatment, with control as denominator. +However, the agents chose different verification work. Both ran one current +update, sixteen current updates and sixteen reference updates. Control alone +then ran sixteen updates through the ordinary entry point to check whether its +diagnostic changed outputs. That extra request cost 35.627 GPU-seconds, +accounting for 78.27% of the 45.520 GPU-second terminal difference. Treatment did +not perform that extra noninterference check. These totals therefore do not +establish an intrinsic tool speedup. Restricting GPU accounting to the first +three matching request profiles gives 82.716 versus 72.822 GPU-seconds; this +still compares different agent-authored diagnostics without timing repeats. + +| New experiment cost category | Allocated GPU-seconds | +| --- | ---: | +| Operator snapshot compatibility preflight | 95.853 | +| Both solver conversations | 191.164 | +| Independent 36-scenario grading | 995.453 | +| Total new GPU use | 1,282.470 | + +All grading starts have known completion/cleanup costs, with no missing or +duplicate phase charges. Both GPUs were released after grading. Historical +qualification/oracle cost is not charged again. The grading jobs include full +private state observations and are not solver timing measurements. + +## Evidence + +Local evidence: `outputs/ring18-source-observation-20261102/` in the canonical +TrainCheck checkout. Remote execution: +`/home/yuxuan/experiments/traincheck-runtime-study-d-20261102` on Ring18. +Earlier oracles remain under `traincheck-runtime-study-c-20261102`. + +`CHECKPOINT.json`, `core-validation.json`, `optional-manifest.json`, +`qualification-reuse.json`, `protocol-reuse.json` and `README-operator.md` +record the implementation, exact inputs, validation and reproduction steps. +Large model states and raw transcripts are not committed. + +The final solver index is `recovery-3/solver-completion.json`; it references +control's preserved recovery-two composition and treatment's flattened +three-segment history. `independent-adoption-audit.json` binds the nine returned +snapshots, executed comparison operations and diagnostic sources. +`milestone-annotations-final.json`, `milestone-summary-final.json` and +`private-grading-join.json` bind the independent source grades to the original +turn timestamps, native usage counters and physical service attempts. Earlier +visible-only annotations remain unchanged. + +The complete grading metadata is under +`grading-results-v1/independent-grade-v1/`, including `summary.json` with SHA-256 +`6cc322509b79525b57f0078b8a5238bcbfe666021456491aed8274039b6f1989`. +All 203 exported member hashes were checked locally. The identical metadata +files were losslessly recompressed from a 27.6 MB gzip export to a 682 KB xz +export for transfer; no fields, sources or results were dropped. Native tensor +states stay on Ring18. The final milestone summary has SHA-256 +`d3beb3ae0b7575e85019a43992bd0c43644c15b46a29bedac82acfbb11550c8b`. +The separately frozen grader and its reproduction/review instructions are under +`evaluation/grading/`; post-freeze recovery and evidence-path adapters retain +their prior versions, hashes, focused regressions and explicit amendments. + +An independent export review checked all 203 files, complete source/permission +bindings, 36 phase receipts, outcomes and costs. Its record is +`evaluation/grading/export-review-v1.json`, SHA-256 +`8b59ecdc66ccd545a3f67922143687204986d3bdf5a299a4bb3690e87726547b`. +It validates the exported evidence without rerunning training or recomputing +native tensor oracles. The compact final accounting and adoption report is +`final-adjudication-report.json`, SHA-256 +`5239d04780f0d9475ab51d970d68f5503b827c74278da81827e83eac928d43f5`. diff --git a/docs/model-observation.md b/docs/model-observation.md new file mode 100644 index 00000000..9fde19c0 --- /dev/null +++ b/docs/model-observation.md @@ -0,0 +1,161 @@ +# Observe model state at a boundary you choose + +TrainCheck can capture a model's named parameters and buffers and compare the +result across ranks or execution phases. Use this when it helps answer your +next question; observation is optional. You choose the observation points and +the predicate. No inferred invariant or predefined diagnosis is required. + +The snapshot interface reuses TrainCheck's model-state enumeration and existing +tensor fingerprint implementation. It does not replace parameters, install +tensor subclasses, or enable full API tracing. Each snapshot includes every +registered parameter and buffer name, including nonpersistent buffers and +multiple names for an aliased tensor. It records shape, dtype and a SHA-256 +fingerprint of every tensor's complete logical contents. Snapshot comparison +checks names and metadata as well as fingerprints. + +Importing `traincheck` or `traincheck.observation` does not activate the legacy +instrumentor. Existing public exports such as `Instrumentor` remain available +and load their implementation when requested; explicit instrumentor imports +retain their existing behavior. + +## Add one observation + +```python +from pathlib import Path +from traincheck.observation import snapshot_model, write_snapshot + +# model is your existing torch.nn.Module; rank is its worker's actual rank. +snapshot = snapshot_model( + model, run_id="my-investigation", rank=rank, + phase="chosen-boundary", step=step, model_name="policy", +) +output = Path("observations") +output.mkdir(parents=True, exist_ok=True) +write_snapshot(snapshot, output / f"rank{rank}-chosen-boundary-{step}.json") +``` + +Call this at the actual execution site where the state is useful. Keep the model +quiescent while collecting it: this is a synchronous observation, not an atomic +snapshot against concurrent training or mutation. Use a unique output path for +every observation. `write_snapshot` refuses to overwrite an existing file; +missing parent directories must be created by the caller. In a managed training +service, place files under its supplied output directory and retain the actual +rank in the snapshot. No shared append-only file or logging configuration is +needed. + +`run_id` identifies your logical investigation. `model_name` identifies the model +within it. `phase` and nonnegative `step` identify a boundary in your program; +they do not infer training progress. Choose them consistently so independently +written files can be compared. Values are detached observations, so subsequent +training does not change a saved snapshot. + +## Compare ranks offline + +```python +from pathlib import Path +from traincheck.observation import compare_ranks, read_snapshot + +snapshots = [read_snapshot(Path(f"observations/rank{r}-chosen-boundary-0.json")) + for r in (0, 1)] +result = compare_ranks(snapshots, expected_ranks=[0, 1]) +print(result.equal) +for comparison in result.comparisons: + for difference in comparison.differences: + print(difference.name, difference.fields) +``` + +This requires at least two distinct expected ranks and exactly one observation +for every explicitly expected rank, with +the same run, model, phase and step. Missing or duplicate ranks, malformed +records and incompatible labels raise errors. They never become a successful +empty comparison. Every non-reference rank is compared with `reference_rank` +(zero by default). Physical device placement may differ across ranks and does +not itself make the model values different. + +Equality is evidence of matching complete fingerprints and metadata; it is not +a mathematical proof that hash collisions are impossible. A difference is an +observation, not automatically a bug. Whether two states should agree depends +on the training contract. Rank-specific buffers or legitimately different model +components need an appropriate predicate instead of unconditional equality. + +## Compare two explicitly selected phases + +```python +from traincheck.observation import compare_model_snapshots + +result = compare_model_snapshots(before, after) +print(result.equal, result.tensors_compared) +for difference in result.differences: + print(difference.name, difference.fields) +``` + +`before` and `after` must be snapshots of the same named model in the same +logical investigation. This explicit pair comparison permits different ranks, +phases and steps. For example, an optimizer update may legitimately change model +parameters. The result identifies changed fields; it does not prescribe which +changes are expected. Independently collected observations can use the same +logical run identifier when you deliberately intend to compare them. + +## Runnable examples + +The repository includes `tools/model_observation/capture_example.py` and +`tools/model_observation/compare.py`. In an exported agent packet, use the +provided copies of these examples from that packet's example directory. + +From a checkout with TrainCheck and its dependencies installed: + +```sh +python -m tools.model_observation.capture_example \ + --output /tmp/model-observation-demo --run-id demo --steps 2 +python -m tools.model_observation.compare pair \ + /tmp/model-observation-demo/rank0-ready-0.json \ + /tmp/model-observation-demo/rank0-update-2.json +``` + +The small model trains on fixed illustrative data. The pair command normally +reports changed parameter fingerprints and exits with status 1 because the +observed states differ; that is expected after training. An equal comparison +returns 0, while invalid or incomplete inputs return 2 with an error. + +On an allocated two-GPU machine, the same example supports ordinary DDP: + +```sh +torchrun --standalone --nproc-per-node=2 -m tools.model_observation.capture_example \ + --device cuda --output /tmp/model-observation-ddp --run-id demo-ddp --steps 2 +python -m tools.model_observation.compare ranks --expected-ranks 0 1 \ + --snapshots /tmp/model-observation-ddp/rank0-update-2.json \ + /tmp/model-observation-ddp/rank1-update-2.json +``` + +For a two-process CPU example, omit `--device cuda`. Use fresh output directories +for each invocation. These examples generate their own tiny workload; insert +the observation calls into your actual training code for your investigation. + +## Scope and cost + +- The API observes registered model parameters and buffers. It does not + automatically capture optimizer state, gradients, arbitrary locals, source + provenance, or application-specific objects. Add separate observations when + those are needed. +- Materialized dense tensors with the supported ordinary dtypes are accepted. + Sparse, quantized, meta and unsupported tensor subclasses are explicit errors. + Custom module state-enumeration overrides are also rejected because they can + hide registered state. These cases do not produce partial passing snapshots. +- It performs no distributed communication or implicit all-gather. If a model + is sharded, the caller must select and coordinate the appropriate full-model + observation context supplied by its framework. Comparing local shards as if + they were replicas is not a valid check. Read the API's supported-tensor errors + instead of treating an unsupported or absent observation as success. +- Capturing complete fingerprints reads every selected tensor and copies GPU + data to the CPU. It can synchronize device work. This makes occasional + observations convenient but does not promise negligible cost or suitability + for every training step. Choose useful boundaries and measure overhead for + your workload. +- Fingerprints cannot reconstruct tensor values or measure numerical distance. + They do distinguish bit patterns such as signed zero and NaN representations; + they do not apply floating-point tolerances. Matching tensor contents does not + establish identical storage aliasing or a complete application checkpoint. +- Completeness is relative to the registered model state at each chosen + boundary and the expected ranks you supplied. Equality alone does not prove + the intended program location executed; keep the call in the real execution + path and retain failures as failures. diff --git a/tests/test_model_observation_examples.py b/tests/test_model_observation_examples.py new file mode 100644 index 00000000..7f3ef7a8 --- /dev/null +++ b/tests/test_model_observation_examples.py @@ -0,0 +1,144 @@ +"""Exercise the agent-facing capture and comparison commands end to end.""" + +from __future__ import annotations + +import contextlib +import io +import os +import tempfile +import unittest +from dataclasses import replace +from pathlib import Path +from unittest.mock import patch + +import torch + +from tools.model_observation.capture_example import Settings, run +from tools.model_observation.compare import main as compare_main +from traincheck.observation import ( + compare_model_snapshots, + read_snapshot, + write_snapshot, +) + + +class ModelObservationExampleTests(unittest.TestCase): + def capture(self, directory: Path) -> None: + environment = { + name: value + for name, value in os.environ.items() + if name not in {"RANK", "LOCAL_RANK", "WORLD_SIZE"} + } + with ( + patch.dict(os.environ, environment, clear=True), + torch.random.fork_rng(), + contextlib.redirect_stdout(io.StringIO()), + ): + run(Settings(directory, "example-test", "cpu", 2)) + + def test_capture_observes_optimizer_updates_and_nonpersistent_buffer(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + directory = Path(temporary) + self.capture(directory) + before = read_snapshot(directory / "rank0-ready-0.json") + after = read_snapshot(directory / "rank0-update-2.json") + result = compare_model_snapshots(before, after) + self.assertFalse(result.equal) + self.assertEqual(result.tensors_compared, 5) + self.assertNotIn( + "example_counter", {item.name for item in result.differences} + ) + + def test_comparison_commands_report_equal_different_and_missing(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + directory = Path(temporary) + self.capture(directory) + before = str(directory / "rank0-ready-0.json") + after = str(directory / "rank0-update-2.json") + other = directory / "rank1-ready-0.json" + write_snapshot(replace(read_snapshot(Path(before)), rank=1), other) + with contextlib.redirect_stdout(io.StringIO()): + self.assertEqual(compare_main(["pair", before, before]), 0) + self.assertEqual(compare_main(["pair", before, after]), 1) + self.assertEqual( + compare_main( + [ + "ranks", + "--expected-ranks", + "0", + "1", + "--snapshots", + before, + str(other), + ] + ), + 0, + ) + with ( + contextlib.redirect_stderr(io.StringIO()), + self.assertRaises(SystemExit) as raised, + ): + compare_main( + ["ranks", "--expected-ranks", "0", "1", "--snapshots", before] + ) + self.assertEqual(raised.exception.code, 2) + + def test_duplicate_rank_files_are_not_accepted(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + directory = Path(temporary) + self.capture(directory) + before = str(directory / "rank0-ready-0.json") + with ( + contextlib.redirect_stderr(io.StringIO()), + self.assertRaises(SystemExit) as raised, + ): + compare_main( + [ + "ranks", + "--expected-ranks", + "0", + "1", + "--snapshots", + before, + before, + ] + ) + self.assertEqual(raised.exception.code, 2) + + def test_missing_file_is_an_error(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + missing = str(Path(temporary) / "absent.json") + with ( + contextlib.redirect_stderr(io.StringIO()), + self.assertRaises(SystemExit) as raised, + ): + compare_main(["pair", missing, missing]) + self.assertEqual(raised.exception.code, 2) + + def test_partial_distributed_environment_is_rejected(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + environment = { + name: value + for name, value in os.environ.items() + if name not in {"RANK", "LOCAL_RANK", "WORLD_SIZE"} + } + environment["RANK"] = "0" + with ( + patch.dict(os.environ, environment, clear=True), + self.assertRaisesRegex(ValueError, "requires RANK"), + ): + run(Settings(Path(temporary), "example-test", "cpu", 2)) + + def test_existing_observations_are_not_overwritten(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + directory = Path(temporary) + self.capture(directory) + path = directory / "rank0-ready-0.json" + original = path.read_bytes() + with self.assertRaises(FileExistsError): + self.capture(directory) + self.assertEqual(path.read_bytes(), original) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_model_snapshot.py b/tests/test_model_snapshot.py new file mode 100644 index 00000000..23683eea --- /dev/null +++ b/tests/test_model_snapshot.py @@ -0,0 +1,528 @@ +"""Complete, non-invasive model observations and strict offline comparisons.""" + +from __future__ import annotations + +import copy +import gc +import json +import subprocess +import sys +import weakref +from dataclasses import asdict, replace +from pathlib import Path + +import pytest +import torch +from torch import nn + +from traincheck._model_state import iter_model_tensors +from traincheck.config import config +from traincheck.instrumentor import tracer +from traincheck.instrumentor.proxy_wrapper.subclass import proxy_parameter +from traincheck.observation import ( + ModelSnapshot, + compare_model_snapshots, + compare_ranks, + read_snapshot, + snapshot_model, + write_snapshot, +) + + +def capture(model: nn.Module, rank: int = 0, phase: str = "selected") -> ModelSnapshot: + return snapshot_model(model, run_id="test", rank=rank, phase=phase) + + +def model_with_state() -> nn.Module: + model = nn.Linear(3, 2) + model.register_parameter("tied_weight", model.weight) + model.register_buffer("running", torch.tensor([0.0, -0.0, float("nan")])) + model.register_buffer("temporary", torch.tensor(True), persistent=False) + model.register_buffer("empty", torch.empty(2, 0)) + model.register_buffer("absent", None) + return model + + +def test_snapshot_keeps_all_names_buffers_aliases_and_registration_order() -> None: + model = model_with_state() + snapshot = capture(model) + assert [tensor.name for tensor in snapshot.tensors] == [ + "weight", + "bias", + "tied_weight", + "running", + "temporary", + "empty", + ] + assert snapshot.tensors[2].object_alias_of == "weight" + assert snapshot.tensors[2].sha256 == snapshot.tensors[0].sha256 + assert [tensor.name for tensor in snapshot.tensors if tensor.kind == "buffer"] == [ + "running", + "temporary", + "empty", + ] + assert [name for _, name, _ in iter_model_tensors(model)] == ["weight", "bias"] + + +@pytest.mark.parametrize( + "value", + [ + torch.tensor(4.0), + torch.tensor([True, False]), + torch.arange(6, dtype=torch.bfloat16).reshape(2, 3), + torch.tensor([1 + 2j, 3 - 4j]).conj(), + torch.arange(6.0).reshape(2, 3).T, + torch.empty(0, 2), + ], +) +def test_supported_representations_hash_exact_logical_bytes( + value: torch.Tensor, +) -> None: + model = nn.Module() + model.register_buffer("value", value) + first = capture(model) + model.value = value.resolve_conj().resolve_neg().contiguous().clone() + assert compare_model_snapshots(first, capture(model)).equal + + +def test_every_row_sub_quantization_changes_and_cached_hash_are_observed() -> None: + model = nn.Module() + value = torch.tensor([[0.0, 0.0], [0.0, 0.0]], dtype=torch.float64) + value.__dict__["_traincheck_tensor_hash"] = 123 + model.register_buffer("value", value) + first = capture(model) + value[1, 1] = 1e-12 + second = capture(model) + comparison = compare_model_snapshots(first, second) + assert not comparison.equal + assert comparison.differences[0].fields == ("sha256",) + assert value.__dict__["_traincheck_tensor_hash"] == 123 + + +def test_signed_zero_and_distinct_nan_payloads_are_not_normalized() -> None: + model = nn.Module() + model.register_buffer("value", torch.tensor([0.0])) + before = capture(model) + model.value[0] = -0.0 + assert not compare_model_snapshots(before, capture(model)).equal + model.value = torch.tensor([0x7FC00001], dtype=torch.int32).view(torch.float32) + before_nan = capture(model) + model.value = torch.tensor([0x7FC00002], dtype=torch.int32).view(torch.float32) + assert not compare_model_snapshots(before_nan, capture(model)).equal + + +def test_snapshot_preserves_rng_gradients_optimizer_identity_and_model_state() -> None: + model = nn.Linear(2, 2) + optimizer = torch.optim.AdamW(model.parameters()) + model(torch.ones(1, 2)).sum().backward() + optimizer.step() + tensor_ids = [id(parameter) for parameter in model.parameters()] + gradients = [parameter.grad.clone() for parameter in model.parameters()] + state = copy.deepcopy(model.state_dict()) + parameter_metadata = [dict(parameter.__dict__) for parameter in model.parameters()] + rng = torch.get_rng_state().clone() + captured = capture(model) + assert torch.equal(rng, torch.get_rng_state()) + assert tensor_ids == [id(parameter) for parameter in model.parameters()] + assert tensor_ids == [ + id(parameter) for parameter in optimizer.param_groups[0]["params"] + ] + assert parameter_metadata == [ + dict(parameter.__dict__) for parameter in model.parameters() + ] + for parameter, gradient in zip(model.parameters(), gradients): + assert torch.equal(parameter.grad, gradient) + for name, value in state.items(): + assert torch.equal(model.state_dict()[name], value) + with torch.no_grad(): + model.weight.add_(1) + assert not compare_model_snapshots(captured, capture(model)).equal + + +def test_snapshot_does_not_retain_model_or_tensors() -> None: + model = nn.Linear(2, 2) + model_reference = weakref.ref(model) + tensor_reference = weakref.ref(model.weight) + snapshot = capture(model) + del model + gc.collect() + assert model_reference() is None + assert tensor_reference() is None + assert len(snapshot.tensors) == 2 + + +def test_existing_subclass_tracker_works_without_promoting_new_parameters() -> None: + model = nn.Linear(2, 2) + before = capture(model) + proxy_parameter(model, parent_name="observed", should_dump_trace=False) + tracked = capture(model) + assert compare_model_snapshots(before, tracked).equal + ordinary = nn.Linear(2, 2) + capture(ordinary) + assert type(ordinary.weight) is nn.Parameter + + +def test_recorder_operations_do_not_emit_api_records( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = nn.Linear(2, 2) + records: list[object] = [] + monkeypatch.setattr(config, "DISABLE_WRAPPER", False) + monkeypatch.setattr(tracer, "dump_trace_API", records.append) + monkeypatch.setattr( + torch.Tensor, + "detach", + tracer.wrapper( + torch.Tensor.detach, + is_bound_method=True, + scan_proxy_in_args=False, + dump_stack_trace=False, + dump_ret=False, + ), + ) + capture(model) + assert records == [] + assert not tracer.tracing_suppressed() + + +def test_active_parameter_tracker_and_api_wrapper_do_not_observe_capture( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from traincheck.instrumentor.proxy_wrapper import subclass + from traincheck.instrumentor.proxy_wrapper.proxy_registry import get_global_registry + + model = nn.Linear(2, 2) + proxy_parameter(model, parent_name="combined", should_dump_trace=False) + metadata = [dict(parameter.__dict__) for parameter in model.parameters()] + registry = get_global_registry() + entries = { + name: (id(entry), entry.stale) for name, entry in registry.registry.items() + } + records: list[object] = [] + monkeypatch.setattr(config, "DISABLE_WRAPPER", False) + monkeypatch.setattr(tracer, "dump_trace_API", records.append) + monkeypatch.setattr(subclass, "dump_trace_VAR", records.append) + monkeypatch.setattr( + torch.Tensor, + "detach", + tracer.wrapper( + torch.Tensor.detach, + is_bound_method=True, + scan_proxy_in_args=True, + dump_stack_trace=False, + dump_ret=False, + ), + ) + capture(model) + assert records == [] + assert metadata == [dict(parameter.__dict__) for parameter in model.parameters()] + assert entries == { + name: (id(entry), entry.stale) for name, entry in registry.registry.items() + } + + +def test_legacy_observer_does_not_change_parameters_during_capture( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from traincheck.instrumentor.proxy_wrapper import subclass + from traincheck.instrumentor.proxy_wrapper.proxy_observer import ( + add_observer_to_func, + ) + + model = nn.Linear(2, 2) + proxy_parameter(model, parent_name="legacy", should_dump_trace=False) + metadata = [dict(parameter.__dict__) for parameter in model.parameters()] + records: list[object] = [] + monkeypatch.setattr(config, "DISABLE_WRAPPER", False) + monkeypatch.setattr(subclass, "dump_trace_VAR", records.append) + monkeypatch.setattr( + torch.Tensor, "detach", add_observer_to_func(torch.Tensor.detach) + ) + capture(model) + assert records == [] + assert metadata == [dict(parameter.__dict__) for parameter in model.parameters()] + + +def test_observation_import_and_capture_have_no_framework_patch_side_effects() -> None: + script = """ +import sys +import tokenize +import torch +import torch.optim.optimizer as optimizer +from torch._C._distributed_c10d import ProcessGroup +before = (torch.nn.Module.__setattr__, ProcessGroup.broadcast, ProcessGroup.allreduce, + ProcessGroup.allgather, tokenize._tokenize, optimizer._default_to_fused_or_foreach) +from traincheck.observation import snapshot_model +model = torch.nn.Linear(2, 2) +snapshot_model(model, run_id='isolated', rank=0, phase='chosen') +after = (torch.nn.Module.__setattr__, ProcessGroup.broadcast, ProcessGroup.allreduce, + ProcessGroup.allgather, tokenize._tokenize, optimizer._default_to_fused_or_foreach) +assert all(left is right for left,right in zip(before,after)) +assert not any(name.startswith('traincheck.instrumentor') for name in sys.modules) +assert all(not key.startswith('_tc_') for key in vars(model)) +assert type(model.weight) is torch.nn.Parameter +""" + subprocess.run( + [sys.executable, "-c", script], check=True, capture_output=True, text=True + ) + + +def test_lazy_public_exports_keep_existing_objects_and_errors() -> None: + script = """ +import traincheck +expected = {'Instrumentor','get_all_subclasses','annotate_stage','annotate_answer_start_token_ids'} +assert set(traincheck.__all__) == expected +assert expected <= set(dir(traincheck)) +try: + traincheck.nonexistent_export +except AttributeError: + pass +else: + raise AssertionError('unknown export must fail') +from traincheck import Instrumentor, get_all_subclasses +from traincheck.instrumentor.tracer import Instrumentor as implementation +assert Instrumentor is implementation +namespace = {} +exec('from traincheck import *', namespace) +assert all(namespace[name] is getattr(traincheck,name) for name in expected) +""" + subprocess.run( + [sys.executable, "-c", script], check=True, capture_output=True, text=True + ) + + +@pytest.mark.parametrize("kind", ["sparse", "meta", "quantized", "subclass", "flat"]) +def test_unsupported_tensors_fail_explicitly_without_partial_capture(kind: str) -> None: + model = nn.Module() + model.register_buffer("good", torch.ones(2)) + if kind == "sparse": + bad = torch.eye(2).to_sparse() + elif kind == "meta": + bad = torch.empty(2, device="meta") + elif kind == "quantized": + bad = torch.quantize_per_tensor(torch.ones(2), 0.1, 0, torch.qint8) + elif kind == "subclass": + + class CustomTensor(torch.Tensor): + pass + + bad = torch.ones(2).as_subclass(CustomTensor) + else: + bad = torch.ones(2) + bad.__dict__["_is_flat_param"] = True + model.register_buffer("bad", bad) + with pytest.raises((TypeError, ValueError)): + capture(model) + assert not tracer.tracing_suppressed() + + +def test_empty_model_is_not_reported_as_successful_observation() -> None: + with pytest.raises(ValueError, match="nonempty"): + capture(nn.Module()) + + +def test_custom_enumeration_and_cycles_fail_before_silently_omitting_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = nn.Linear(2, 2) + monkeypatch.setattr(model, "named_parameters", lambda **kwargs: iter(())) + with pytest.raises(TypeError, match="custom named_parameters"): + capture(model) + parent = nn.Sequential(model) + with pytest.raises(TypeError, match="custom named_parameters"): + capture(parent) + cyclic = nn.Module() + cyclic.child = cyclic + with pytest.raises(ValueError, match="cyclic"): + capture(cyclic) + + +def test_capture_failure_propagates_and_restores_tracing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from traincheck import observation + + failure = RuntimeError("snapshot hash failure") + + def fail_hash(tensor: torch.Tensor) -> int: + raise failure + + monkeypatch.setattr(observation, "exact_tensor_hash", fail_hash) + with pytest.raises(RuntimeError) as caught: + capture(nn.Linear(2, 2)) + assert caught.value is failure + assert not tracer.tracing_suppressed() + + +@pytest.mark.parametrize( + "module_name", + [ + "traincheck.instrumentor.tracer", + "traincheck.instrumentor.proxy_wrapper.subclass", + ], +) +def test_present_blocked_tracker_module_is_not_treated_as_absent( + monkeypatch: pytest.MonkeyPatch, + module_name: str, +) -> None: + model = nn.Linear(2, 2) + proxy_parameter(model, parent_name="blocked", should_dump_trace=False) + monkeypatch.setitem(sys.modules, module_name, None) + with pytest.raises(TypeError, match="Python module"): + capture(model) + + +def test_existing_exact_tensor_hash_keeps_digest_and_cached_hash_behavior() -> None: + from traincheck.instrumentor.proxy_wrapper.hash import tensor_hash + + value = torch.tensor([0, 1, 255], dtype=torch.uint8) + expected = int( + "26a66b061e8f48f39927c312f25293959729eee95978e2892d49d3512a5cc092", 16 + ) + assert tensor_hash(value, with_parallel=False) == expected + value.__dict__["_traincheck_tensor_hash"] = 99 + assert tensor_hash(value, with_parallel=False) == 99 + with pytest.raises(TypeError): + tensor_hash(torch.ones(2, dtype=torch.bfloat16), with_parallel=False) + + +def test_sampler_default_retains_custom_no_argument_parameter_iterator() -> None: + from collections.abc import Iterator + + class CustomModule(nn.Linear): + def named_parameters(self) -> Iterator[tuple[str, nn.Parameter]]: + yield "custom.weight", self.weight + + model = CustomModule(2, 2) + entries = tuple(iter_model_tensors(model)) + assert len(entries) == 1 + assert entries[0][1] == "custom.weight" + assert entries[0][2] is model.weight + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") +def test_cuda_snapshots_preserve_rng_and_match_cpu() -> None: + model = nn.Linear(2, 2, device="cuda") + cpu_rng = torch.get_rng_state().clone() + cuda_rng = torch.cuda.get_rng_state().clone() + gpu_snapshot = capture(model) + assert torch.equal(torch.get_rng_state(), cpu_rng) + assert torch.equal(torch.cuda.get_rng_state(), cuda_rng) + assert compare_model_snapshots( + gpu_snapshot, capture(copy.deepcopy(model).cpu()) + ).equal + + +def test_round_trip_overwrite_guard_and_required_fields(tmp_path: Path) -> None: + snapshot = capture(model_with_state()) + path = tmp_path / "snapshot.json" + write_snapshot(snapshot, path) + assert read_snapshot(path) == snapshot + with pytest.raises(FileExistsError): + write_snapshot(snapshot, path) + value = asdict(snapshot) + del value["phase"] + path.write_text(json.dumps(value)) + with pytest.raises(ValueError, match="fields"): + read_snapshot(path) + + +@pytest.mark.parametrize( + "field,value", + [ + ("rank", True), + ("step", -1), + ("phase", ""), + ("schema_version", 2), + ("tensors", []), + ], +) +def test_malformed_top_level_values_fail( + tmp_path: Path, field: str, value: object +) -> None: + data = asdict(capture(nn.Linear(2, 2))) + data[field] = value + path = tmp_path / "bad.json" + path.write_text(json.dumps(data)) + with pytest.raises(ValueError): + read_snapshot(path) + + +@pytest.mark.parametrize( + "field,value", + [ + ("shape", [True]), + ("dtype", "float"), + ("sha256", ""), + ("requires_grad", 1), + ("object_alias_of", "absent"), + ], +) +def test_malformed_tensor_values_fail( + tmp_path: Path, field: str, value: object +) -> None: + data = asdict(capture(nn.Linear(2, 2))) + data["tensors"][0][field] = value + path = tmp_path / "bad.json" + path.write_text(json.dumps(data)) + with pytest.raises(ValueError): + read_snapshot(path) + + +def test_duplicate_json_keys_and_tensor_names_fail(tmp_path: Path) -> None: + path = tmp_path / "bad.json" + path.write_text('{"rank":0,"rank":1}') + with pytest.raises(ValueError, match="duplicate JSON"): + read_snapshot(path) + snapshot = capture(nn.Linear(2, 2)) + with pytest.raises(ValueError, match="duplicate tensor"): + replace(snapshot, tensors=(snapshot.tensors[0], snapshot.tensors[0])) + + +def test_comparison_preserves_order_names_dtype_aliases_and_provenance() -> None: + original = capture(model_with_state()) + peer = replace( + original, + rank=1, + phase="after", + step=1, + time_ns=0, + tensors=tuple(replace(tensor, device="cuda:1") for tensor in original.tensors), + ) + assert compare_model_snapshots(original, peer).equal + assert not compare_model_snapshots( + original, replace(original, tensors=original.tensors[:-1]) + ).equal + # Untying an alias matters even when every element remains identical. + untied = replace( + original, + tensors=tuple( + replace(tensor, object_alias_of=None) for tensor in original.tensors + ), + ) + assert not compare_model_snapshots(original, untied).equal + renamed = replace( + original, + tensors=(*original.tensors[:-1], replace(original.tensors[-1], name="other")), + ) + assert not compare_model_snapshots(original, renamed).equal + with pytest.raises(ValueError, match="same run"): + compare_model_snapshots(original, replace(original, run_id="different")) + + +def test_rank_comparison_requires_exact_coverage_and_matching_observation() -> None: + first = capture(nn.Linear(2, 2)) + second = replace(first, rank=1) + assert compare_ranks([second, first], expected_ranks=[0, 1]).equal + with pytest.raises(ValueError, match="coverage"): + compare_ranks([first], expected_ranks=[0, 1]) + with pytest.raises(ValueError, match="duplicate"): + compare_ranks([first, first], expected_ranks=[0, 1]) + with pytest.raises(ValueError, match="phase and step"): + compare_ranks([first, replace(second, phase="later")], expected_ranks=[0, 1]) + with pytest.raises(ValueError, match="distinct"): + compare_ranks([first], expected_ranks=[0]) + with pytest.raises(ValueError, match="same run"): + compare_ranks( + [first, replace(second, model_name="other")], expected_ranks=[0, 1] + ) diff --git a/tools/model_observation/__init__.py b/tools/model_observation/__init__.py new file mode 100644 index 00000000..d6d2fa2d --- /dev/null +++ b/tools/model_observation/__init__.py @@ -0,0 +1 @@ +"""Runnable examples for explicit model observations and offline comparison.""" diff --git a/tools/model_observation/capture_example.py b/tools/model_observation/capture_example.py new file mode 100644 index 00000000..90d2ba00 --- /dev/null +++ b/tools/model_observation/capture_example.py @@ -0,0 +1,124 @@ +"""Capture selected model states during a small single-process or DDP run.""" + +from __future__ import annotations + +import argparse +import json +import os +from dataclasses import dataclass +from pathlib import Path + +import torch +from torch import distributed as dist +from torch import nn +from torch.nn.parallel import DistributedDataParallel + +from traincheck.observation import snapshot_model, write_snapshot + + +@dataclass(frozen=True) +class Settings: + output: Path + run_id: str + device: str + steps: int + + +def _environment_integer(name: str) -> int: + value = os.environ[name] + if not value.isdecimal(): + raise ValueError(f"{name} must be a nonnegative decimal integer") + return int(value) + + +def run(settings: Settings) -> None: + """Train an illustrative model; every explicit snapshot observes all its state.""" + if settings.steps < 1: + raise ValueError("steps must be positive") + if settings.device not in ("cpu", "cuda"): + raise ValueError("device must be cpu or cuda") + required = {"RANK", "LOCAL_RANK", "WORLD_SIZE"} + present = required.intersection(os.environ) + if present and present != required: + raise ValueError( + "distributed execution requires RANK, LOCAL_RANK and WORLD_SIZE" + ) + rank = _environment_integer("RANK") if present else 0 + local_rank = _environment_integer("LOCAL_RANK") if present else 0 + world_size = _environment_integer("WORLD_SIZE") if present else 1 + if world_size < 1 or rank >= world_size: + raise ValueError("rank must be in [0, WORLD_SIZE)") + device = ( + torch.device(settings.device, local_rank) + if settings.device == "cuda" + else torch.device("cpu") + ) + if settings.device == "cuda": + torch.cuda.set_device(device) + distributed = world_size > 1 + if distributed: + dist.init_process_group(backend="nccl" if settings.device == "cuda" else "gloo") + try: + torch.manual_seed(1234) + model = nn.Sequential(nn.Linear(4, 8), nn.Tanh(), nn.Linear(8, 2)).to(device) + model.register_buffer( + "example_counter", + torch.zeros((), dtype=torch.int64, device=device), + persistent=False, + ) + training_model: nn.Module = model + if distributed: + training_model = DistributedDataParallel( + model, device_ids=[local_rank] if settings.device == "cuda" else None + ) + optimizer = torch.optim.SGD(training_model.parameters(), lr=0.01) + settings.output.mkdir(parents=True, exist_ok=True) + + def observe(phase: str, step: int) -> None: + snapshot = snapshot_model( + model, run_id=settings.run_id, rank=rank, phase=phase, step=step + ) + write_snapshot( + snapshot, settings.output / f"rank{rank}-{phase}-{step}.json" + ) + + observe("ready", 0) + inputs = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4) / 12 + targets = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]], device=device) + for step in range(1, settings.steps + 1): + optimizer.zero_grad(set_to_none=True) + loss = nn.functional.mse_loss(training_model(inputs), targets) + loss.backward() + optimizer.step() + observe("update", step) + print( + json.dumps( + {"rank": rank, "step": step, "loss": float(loss.detach().cpu())} + ), + flush=True, + ) + finally: + if distributed: + dist.destroy_process_group() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--run-id", required=True) + parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu") + parser.add_argument("--steps", type=int, default=2) + args = parser.parse_args() + output: object = args.output + run_id: object = args.run_id + device: object = args.device + steps: object = args.steps + assert isinstance(output, Path) + assert isinstance(run_id, str) + assert isinstance(device, str) + assert isinstance(steps, int) + run(Settings(output, run_id, device, steps)) + + +if __name__ == "__main__": + main() diff --git a/tools/model_observation/compare.py b/tools/model_observation/compare.py new file mode 100644 index 00000000..fdff495b --- /dev/null +++ b/tools/model_observation/compare.py @@ -0,0 +1,61 @@ +"""Compare explicitly selected model snapshots, rejecting incomplete coverage.""" + +from __future__ import annotations + +import argparse +import json +from dataclasses import asdict +from pathlib import Path +from typing import Sequence + +from traincheck.observation import compare_model_snapshots, compare_ranks, read_snapshot + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="mode", required=True) + ranks = commands.add_parser( + "ranks", help="compare one phase across a complete rank set" + ) + ranks.add_argument("--expected-ranks", type=int, nargs="+", required=True) + ranks.add_argument("--reference-rank", type=int, default=0) + ranks.add_argument("--snapshots", type=Path, nargs="+", required=True) + pair = commands.add_parser( + "pair", help="compare two explicitly chosen states, including different phases" + ) + pair.add_argument("reference", type=Path) + pair.add_argument("candidate", type=Path) + args = parser.parse_args(argv) + try: + if args.mode == "ranks": + paths: object = args.snapshots + expected: object = args.expected_ranks + reference_rank: object = args.reference_rank + assert isinstance(paths, list) and all( + isinstance(path, Path) for path in paths + ) + assert isinstance(expected, list) and all( + isinstance(rank, int) for rank in expected + ) + assert isinstance(reference_rank, int) + result = compare_ranks( + [read_snapshot(path) for path in paths], + expected_ranks=expected, + reference_rank=reference_rank, + ) + print(json.dumps({"equal": result.equal, **asdict(result)}, indent=2)) + return 0 if result.equal else 1 + reference: object = args.reference + candidate: object = args.candidate + assert isinstance(reference, Path) and isinstance(candidate, Path) + comparison = compare_model_snapshots( + read_snapshot(reference), read_snapshot(candidate) + ) + print(json.dumps({"equal": comparison.equal, **asdict(comparison)}, indent=2)) + return 0 if comparison.equal else 1 + except (OSError, TypeError, ValueError) as error: + parser.error(str(error)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/traincheck/__init__.py b/traincheck/__init__.py index e8556f53..2c439565 100644 --- a/traincheck/__init__.py +++ b/traincheck/__init__.py @@ -1,12 +1,18 @@ -from traincheck.developer.annotations import ( - annotate_answer_start_token_ids, - annotate_stage, -) +"""TrainCheck public exports, loaded when requested rather than on package import. -from .instrumentor.tracer import ( # Ziming: get rid of new_wrapper for now - Instrumentor, - get_all_subclasses, -) +Explicit instrumentor imports retain their historical activation behavior. +Importing a lightweight observation helper must not patch framework functions. +""" + +from importlib import import_module +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from traincheck.developer.annotations import ( + annotate_answer_start_token_ids, + annotate_stage, + ) + from traincheck.instrumentor.tracer import Instrumentor, get_all_subclasses __all__ = [ "Instrumentor", @@ -14,3 +20,22 @@ "annotate_stage", "annotate_answer_start_token_ids", ] + +_EXPORT_MODULES: dict[str, str] = { + "Instrumentor": "traincheck.instrumentor.tracer", + "get_all_subclasses": "traincheck.instrumentor.tracer", + "annotate_stage": "traincheck.developer.annotations", + "annotate_answer_start_token_ids": "traincheck.developer.annotations", +} + + +def __getattr__(name: str) -> object: + if name not in _EXPORT_MODULES: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value: object = getattr(import_module(_EXPORT_MODULES[name]), name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(__all__)) diff --git a/traincheck/_model_state.py b/traincheck/_model_state.py new file mode 100644 index 00000000..63f054cf --- /dev/null +++ b/traincheck/_model_state.py @@ -0,0 +1,37 @@ +"""Shared, ordered state enumeration for model sampling and snapshots.""" + +from collections.abc import Iterator +from typing import Literal + +from torch import Tensor, nn + +TensorKind = Literal["parameter", "buffer"] + + +def iter_model_tensors( + model: nn.Module, + *, + include_buffers: bool = False, + remove_duplicate: bool = True, +) -> Iterator[tuple[TensorKind, str, Tensor]]: + """Yield named parameters, then optionally all buffers, in registration order. + + Defaults retain VarSampler's parameters-only, duplicate-removing traversal. + Snapshots request every registered name, including tied aliases and + nonpersistent buffers. None-valued registrations are not materialized state. + """ + parameters = ( + model.named_parameters() + if remove_duplicate + else model.named_parameters(remove_duplicate=False) + ) + for name, parameter in parameters: + yield "parameter", name, parameter + if include_buffers: + buffers = ( + model.named_buffers() + if remove_duplicate + else model.named_buffers(remove_duplicate=False) + ) + for name, buffer in buffers: + yield "buffer", name, buffer diff --git a/traincheck/_tensor_fingerprint.py b/traincheck/_tensor_fingerprint.py new file mode 100644 index 00000000..87679468 --- /dev/null +++ b/traincheck/_tensor_fingerprint.py @@ -0,0 +1,18 @@ +"""Exact tensor hashing shared by tracker dumping and explicit observations.""" + +import hashlib + +from torch import Tensor + + +def exact_tensor_hash(tensor: Tensor) -> int: + """Preserve the tracker's conventional SHA-256 implementation. + + Callers choose the input representation. The legacy caller passes a native + tensor, retaining its existing NumPy dtype limitations. Model snapshots pass + a contiguous uint8 view to cover all bytes of supported logical dtypes. + This module imports no instrumentor package and activates no instrumentation. + """ + if tensor.is_cuda: + tensor = tensor.cpu() + return int(hashlib.sha256(tensor.detach().numpy().tobytes()).hexdigest(), 16) diff --git a/traincheck/instrumentor/proxy_wrapper/hash.py b/traincheck/instrumentor/proxy_wrapper/hash.py index bb4a4cbb..d7d76f6c 100644 --- a/traincheck/instrumentor/proxy_wrapper/hash.py +++ b/traincheck/instrumentor/proxy_wrapper/hash.py @@ -3,6 +3,8 @@ from numba import cuda from torch import Tensor +from traincheck._tensor_fingerprint import exact_tensor_hash + MULTIPLIER = 6364136223846793005 INCREMENT = 1 MODULUS = 2**64 @@ -137,10 +139,4 @@ def tensor_hash(x: Tensor, with_parallel: bool = True, with_cuda: bool = True) - result = hash_tensor_cpu(x) return result else: # conventional approach using hashlib, use as the baseline to test the accuracy of the parallel approach - # if on cuda, move to cpu. using hashlib to hash the tensor - if x.is_cuda: - x = x.cpu() - import hashlib - - result = int(hashlib.sha256(x.detach().numpy().tobytes()).hexdigest(), 16) - return result + return exact_tensor_hash(x) diff --git a/traincheck/instrumentor/proxy_wrapper/proxy_observer.py b/traincheck/instrumentor/proxy_wrapper/proxy_observer.py index 0f0a6d28..7fcae6d7 100644 --- a/traincheck/instrumentor/proxy_wrapper/proxy_observer.py +++ b/traincheck/instrumentor/proxy_wrapper/proxy_observer.py @@ -19,10 +19,14 @@ def observe_proxy_var( var: typing.Union["Proxy", "ProxyParameter"], - phase, + phase: str, observe_api_name: str, -): - if config.DISABLE_WRAPPER: +) -> None: + # Imported at invocation to avoid the tracker's initialization dependency + # cycle. Recorder-owned tensor reads must not update observation timestamps. + from traincheck.instrumentor.tracer import tracing_suppressed + + if config.DISABLE_WRAPPER or tracing_suppressed(): return # update the proxy object's timestamp diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index 59b2658b..efe33d89 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -15,6 +15,7 @@ import torch import traincheck.config.config as config # needed to allow for change of values after import +from traincheck._model_state import iter_model_tensors from traincheck.config.config import ( INSTR_MODULES_TO_SKIP, SKIP_INSTR_APIS, @@ -1248,6 +1249,12 @@ def _instrument_module( return count_wrapped +class SampledParameter(TypedDict): + name: str + type: str + attributes: dict[str, object] + + class VarSampler: """ Tracker for the state of a variable. This variable itself cannot be reassigned, i.e. var.attr = new_value is allowed but not var = new_var. @@ -1278,7 +1285,7 @@ def __init__(self, var, var_name: str): **Many of the activations and intermediate tensors are not updated inplace, so this observer will not be able to detect the changes in those tensors.** """ - self.param_versions = {} # type: ignore + self.param_versions: dict[str, int] = {} timestamp = get_timestamp_ns() curr_meta_vars = get_meta_vars() @@ -1298,9 +1305,9 @@ def __init__(self, var, var_name: str): } ) - def _get_state_copy(self): - state_copy = [] - for name, param in self.var.named_parameters(): + def _get_state_copy(self) -> list[SampledParameter]: + state_copy: list[SampledParameter] = [] + for _, name, param in iter_model_tensors(self.var): param_global_name = self.var_name + "." + name self.param_versions[param_global_name] = param._version state_copy.append( diff --git a/traincheck/observation.py b/traincheck/observation.py new file mode 100644 index 00000000..f2c79622 --- /dev/null +++ b/traincheck/observation.py @@ -0,0 +1,525 @@ +"""Explicit model snapshots and offline comparison, without installing hooks. + +Snapshots reuse model sampling's enumeration and the tensor tracker's existing +SHA-256 implementation. They fingerprint *all* logical tensor bytes, not the +legacy quantized, first-row parallel hash. Equal digests are fingerprint +evidence, not a collision-free proof. No model or tensor is retained. + +Capture at a quiescent boundary: concurrent mutation is unsupported. For FSDP, +the caller must enter its framework's full-parameter context on every required +rank and pass the unwrapped module. There are no hidden collectives, automatic +shard gathering, implicit rank inference, or inference-engine requirements. +""" + +from __future__ import annotations + +import inspect +import json +import sys +import time +from collections.abc import Sequence +from contextlib import nullcontext +from dataclasses import asdict, dataclass +from pathlib import Path +from types import ModuleType +from typing import ContextManager, Literal + +import torch +from torch import nn + +from traincheck._model_state import TensorKind, iter_model_tensors +from traincheck._tensor_fingerprint import exact_tensor_hash + +_SUPPORTED_DTYPES: frozenset[str] = frozenset( + str(dtype) + for dtype in ( + torch.bool, + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, + torch.complex64, + torch.complex128, + ) +) + + +def _text(value: object, field: str) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"{field} must be a nonempty string") + return value + + +def _integer(value: object, field: str) -> int: + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise ValueError(f"{field} must be a nonnegative integer") + return value + + +@dataclass(frozen=True) +class TensorSnapshot: + """One registered tensor name; aliases reference the first name of the object.""" + + name: str + kind: TensorKind + shape: tuple[int, ...] + dtype: str + device: str + sha256: str + requires_grad: bool + object_alias_of: str | None + + def __post_init__(self) -> None: + _text(self.name, "tensor.name") + if self.kind not in ("parameter", "buffer"): + raise ValueError("tensor.kind must be parameter or buffer") + if not isinstance(self.shape, tuple): + raise ValueError("tensor.shape must be a tuple") + for dimension in self.shape: + _integer(dimension, "tensor.shape dimension") + if _text(self.dtype, "tensor.dtype") not in _SUPPORTED_DTYPES: + raise ValueError(f"unsupported tensor dtype: {self.dtype}") + _text(self.device, "tensor.device") + if ( + not isinstance(self.sha256, str) + or len(self.sha256) != 64 + or any(character not in "0123456789abcdef" for character in self.sha256) + ): + raise ValueError("tensor.sha256 must be a lowercase SHA-256 hex digest") + if not isinstance(self.requires_grad, bool): + raise ValueError("tensor.requires_grad must be bool") + if self.object_alias_of is not None: + _text(self.object_alias_of, "tensor.object_alias_of") + + +@dataclass(frozen=True) +class ModelSnapshot: + schema_version: Literal[1] + run_id: str + model_name: str + rank: int + step: int + phase: str + time_ns: int + tensors: tuple[TensorSnapshot, ...] + + def __post_init__(self) -> None: + if _integer(self.schema_version, "schema_version") != 1: + raise ValueError("unsupported snapshot schema_version") + _text(self.run_id, "run_id") + _text(self.model_name, "model_name") + _text(self.phase, "phase") + _integer(self.rank, "rank") + _integer(self.step, "step") + _integer(self.time_ns, "time_ns") + if not isinstance(self.tensors, tuple) or not self.tensors: + raise ValueError("snapshot must contain a nonempty tuple of tensors") + seen: dict[str, TensorSnapshot] = {} + for tensor in self.tensors: + if not isinstance(tensor, TensorSnapshot): + raise ValueError("snapshot.tensors must contain TensorSnapshot values") + if tensor.name in seen: + raise ValueError(f"duplicate tensor name: {tensor.name}") + alias = tensor.object_alias_of + if alias is not None: + if alias not in seen or seen[alias].object_alias_of is not None: + raise ValueError(f"invalid object alias for {tensor.name}: {alias}") + original = seen[alias] + if ( + tensor.shape, + tensor.dtype, + tensor.device, + tensor.sha256, + tensor.requires_grad, + ) != ( + original.shape, + original.dtype, + original.device, + original.sha256, + original.requires_grad, + ): + raise ValueError(f"inconsistent object alias for {tensor.name}") + seen[tensor.name] = tensor + + +@dataclass(frozen=True) +class TensorDifference: + name: str + fields: tuple[str, ...] + + +@dataclass(frozen=True) +class SnapshotComparison: + reference_rank: int + candidate_rank: int + tensors_compared: int + differences: tuple[TensorDifference, ...] + + @property + def equal(self) -> bool: + return not self.differences + + +@dataclass(frozen=True) +class RankComparison: + reference_rank: int + ranks: tuple[int, ...] + comparisons: tuple[SnapshotComparison, ...] + + @property + def equal(self) -> bool: + return all(comparison.equal for comparison in self.comparisons) + + +def _validate_tensor(tensor: torch.Tensor, name: str) -> None: + # Do not silently call DTensor.to_local(), accept a fake tensor, or execute + # arbitrary subclass dispatch. ProxyParameter has ordinary Tensor dispatch. + if type(tensor) not in (torch.Tensor, nn.Parameter) and not _is_tracked_parameter( + tensor + ): + raise TypeError(f"unsupported tensor subclass at {name}: {type(tensor)}") + if tensor.layout != torch.strided or tensor.is_quantized or tensor.is_meta: + raise ValueError(f"{name} requires a materialized, dense, unquantized tensor") + if tensor.device.type not in ("cpu", "cuda", "mps"): + raise ValueError(f"unsupported tensor device at {name}: {tensor.device}") + if str(tensor.dtype) not in _SUPPORTED_DTYPES: + raise ValueError(f"unsupported tensor dtype at {name}: {tensor.dtype}") + if tensor.__dict__.get("_is_flat_param", False): + raise ValueError( + f"{name} is a flat parameter; expose complete named parameters" + ) + + +def _is_tracked_parameter(tensor: torch.Tensor) -> bool: + """Recognize the exact existing tracker type without installing that tracker.""" + name = "traincheck.instrumentor.proxy_wrapper.subclass" + if name not in sys.modules: # Absence means the subclass tracker is inactive. + return False + module = sys.modules[name] + if not isinstance(module, ModuleType): + raise TypeError("the loaded subclass tracker must be a Python module") + parameter_type: object = vars(module).get("ProxyParameter") + if ( + not isinstance(parameter_type, type) + or parameter_type.__module__ != name + or not issubclass(parameter_type, nn.Parameter) + ): + raise TypeError("the loaded tracker has an invalid ProxyParameter class") + return type(tensor) is parameter_type + + +def _suppress_existing_tracer() -> ContextManager[object]: + """Use the installed recorder's guard, without importing an inactive recorder.""" + name = "traincheck.instrumentor.tracer" + if name not in sys.modules: + return nullcontext() + module = sys.modules[name] + if not isinstance(module, ModuleType): + raise TypeError("the loaded tracer must be a Python module") + # A missing guard in an already-loaded recorder must propagate as an error. + from traincheck.instrumentor.tracer import suppress_tracing + + return suppress_tracing() + + +def _validate_model_enumeration(model: nn.Module) -> None: + """Reject custom enumeration that could silently omit registered state.""" + visited: set[int] = set() + active: set[int] = set() + + def visit(module: nn.Module) -> None: + if id(module) in active: + raise ValueError("cyclic module registration is unsupported") + if id(module) in visited: + return + visited.add(id(module)) + active.add(id(module)) + for name in ( + "named_parameters", + "named_buffers", + "named_modules", + "_named_members", + ): + if inspect.getattr_static(module, name) is not inspect.getattr_static( + nn.Module, name + ): + raise TypeError( + f"custom {name} enumeration is unsupported for snapshots" + ) + children: object = vars(module)["_modules"] + if not isinstance(children, dict): + raise TypeError("module._modules must be a dictionary") + for child in children.values(): + if child is not None: + if not isinstance(child, nn.Module): + raise TypeError("module children must be nn.Module or None") + visit(child) + active.remove(id(module)) + + visit(model) + + +def snapshot_model( + model: nn.Module, + *, + run_id: str, + rank: int, + phase: str, + step: int = 0, + model_name: str = "model", +) -> ModelSnapshot: + """Fingerprint parameters and all buffers at an explicit observation boundary. + + Every registered alias is retained. Captured values do not share storage with + the model; requires_grad, gradients, optimizer references and RNG are not + changed. None registrations are not tensors. Device is retained as provenance + but excluded from cross-device comparison. Object aliasing is recorded; + storage aliasing between distinct tensor objects is not assessed. + + CPU transfer and hashing are synchronous. This is occasional observation, + not a zero-overhead hook or an atomic snapshot of concurrent training. + """ + if not isinstance(model, nn.Module): + raise TypeError("model must be a torch.nn.Module") + _text(run_id, "run_id") + _text(model_name, "model_name") + _text(phase, "phase") + _integer(rank, "rank") + _integer(step, "step") + with _suppress_existing_tracer(), torch.no_grad(): + started = time.monotonic_ns() + _validate_model_enumeration(model) + entries = tuple( + iter_model_tensors(model, include_buffers=True, remove_duplicate=False) + ) + for _, name, tensor in entries: + _validate_tensor(tensor, name) + aliases: dict[int, str] = {} + records: list[TensorSnapshot] = [] + for kind, name, tensor in entries: + # Resolve lazy conjugate/negative views and hash complete logical + # bytes. The fresh byte view cannot carry an application-cached hash. + value = tensor.detach().resolve_conj().resolve_neg().contiguous().cpu() + raw = value.reshape(-1).view(torch.uint8) + digest = exact_tensor_hash(raw) + alias = aliases.get(id(tensor)) # Absent means this object's first name. + records.append( + TensorSnapshot( + name=name, + kind=kind, + shape=tuple(tensor.shape), + dtype=str(tensor.dtype), + device=str(tensor.device), + sha256=f"{digest:064x}", + requires_grad=tensor.requires_grad, + object_alias_of=alias, + ) + ) + if alias is None: + aliases[id(tensor)] = name + return ModelSnapshot( + 1, run_id, model_name, rank, step, phase, started, tuple(records) + ) + + +def write_snapshot(snapshot: ModelSnapshot, path: Path) -> None: + """Write one immutable observation; an existing file is an error.""" + if not isinstance(snapshot, ModelSnapshot): + raise TypeError("snapshot must be ModelSnapshot") + with path.open("x", encoding="utf-8") as handle: + json.dump(asdict(snapshot), handle, allow_nan=False, separators=(",", ":")) + handle.write("\n") + + +def _object(value: object, fields: set[str], name: str) -> dict[str, object]: + if not isinstance(value, dict): + raise ValueError(f"{name} must be a JSON object") + result: dict[str, object] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise ValueError(f"{name} keys must be strings") + result[key] = item + if set(result) != fields: + raise ValueError(f"{name} fields must be exactly {sorted(fields)}") + return result + + +def _unique_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 _read_tensor(value: object) -> TensorSnapshot: + data = _object( + value, + { + "name", + "kind", + "shape", + "dtype", + "device", + "sha256", + "requires_grad", + "object_alias_of", + }, + "tensor", + ) + shape = data["shape"] + if not isinstance(shape, list): + raise ValueError("tensor.shape must be a JSON array") + dimensions = tuple(_integer(size, "tensor.shape dimension") for size in shape) + kind = data["kind"] + if kind not in ("parameter", "buffer"): + raise ValueError("tensor.kind must be parameter or buffer") + # The explicit branch narrows the untrusted JSON value to its literal contract. + tensor_kind: TensorKind = "parameter" if kind == "parameter" else "buffer" + requires_grad = data["requires_grad"] + if not isinstance(requires_grad, bool): + raise ValueError("tensor.requires_grad must be bool") + alias = data["object_alias_of"] + if alias is not None: + alias = _text(alias, "tensor.object_alias_of") + return TensorSnapshot( + _text(data["name"], "tensor.name"), + tensor_kind, + dimensions, + _text(data["dtype"], "tensor.dtype"), + _text(data["device"], "tensor.device"), + _text(data["sha256"], "tensor.sha256"), + requires_grad, + alias, + ) + + +def read_snapshot(path: Path) -> ModelSnapshot: + """Read and validate an observation, rejecting missing/duplicate fields.""" + with path.open(encoding="utf-8") as handle: + raw: object = json.load(handle, object_pairs_hook=_unique_object) + data = _object( + raw, + { + "schema_version", + "run_id", + "model_name", + "rank", + "step", + "phase", + "time_ns", + "tensors", + }, + "snapshot", + ) + if _integer(data["schema_version"], "schema_version") != 1: + raise ValueError("unsupported snapshot schema_version") + tensors = data["tensors"] + if not isinstance(tensors, list): + raise ValueError("snapshot.tensors must be a JSON array") + return ModelSnapshot( + 1, + _text(data["run_id"], "run_id"), + _text(data["model_name"], "model_name"), + _integer(data["rank"], "rank"), + _integer(data["step"], "step"), + _text(data["phase"], "phase"), + _integer(data["time_ns"], "time_ns"), + tuple(_read_tensor(tensor) for tensor in tensors), + ) + + +def compare_model_snapshots( + reference: ModelSnapshot, candidate: ModelSnapshot +) -> SnapshotComparison: + """Compare an explicitly selected pair, allowing different phases/steps/ranks. + + Both observations must identify the same run and model. All tensor names, + registration order, kinds, shapes, dtypes, contents and object aliases are + compared. Physical devices and observation times are provenance, not equality + conditions. A matching SHA-256 is fingerprint evidence, not a proof of equal + values. Missing tensors are differences, not a reduced comparison scope. + """ + if not isinstance(reference, ModelSnapshot) or not isinstance( + candidate, ModelSnapshot + ): + raise TypeError("comparison requires ModelSnapshot values") + if (reference.run_id, reference.model_name) != ( + candidate.run_id, + candidate.model_name, + ): + raise ValueError("snapshots must identify the same run and model") + left = {tensor.name: tensor for tensor in reference.tensors} + right = {tensor.name: tensor for tensor in candidate.tensors} + differences: list[TensorDifference] = [] + if tuple(left) != tuple(right): + differences.append(TensorDifference("", ("tensor_names_or_order",))) + compared = 0 + for name, tensor in left.items(): + if name not in right: + differences.append(TensorDifference(name, ("missing",))) + continue + other = right[name] + compared += 1 + fields: list[str] = [] + for field in ( + "kind", + "shape", + "dtype", + "sha256", + "requires_grad", + "object_alias_of", + ): + if getattr(tensor, field) != getattr(other, field): + fields.append(field) + if fields: + differences.append(TensorDifference(name, tuple(fields))) + for name in right: + if name not in left: + differences.append(TensorDifference(name, ("unexpected",))) + return SnapshotComparison( + reference.rank, candidate.rank, compared, tuple(differences) + ) + + +def compare_ranks( + snapshots: Sequence[ModelSnapshot], + *, + expected_ranks: Sequence[int], + reference_rank: int = 0, +) -> RankComparison: + """Compare complete rank coverage at one phase/step; missing ranks are errors.""" + ranks = tuple(expected_ranks) + if len(ranks) < 2 or len(set(ranks)) != len(ranks): + raise ValueError("expected_ranks must contain at least two distinct ranks") + for rank in ranks: + _integer(rank, "expected rank") + _integer(reference_rank, "reference_rank") + if reference_rank not in ranks: + raise ValueError("reference_rank must be in expected_ranks") + by_rank: dict[int, ModelSnapshot] = {} + for snapshot in snapshots: + if not isinstance(snapshot, ModelSnapshot): + raise TypeError("snapshots must contain ModelSnapshot values") + if snapshot.rank in by_rank: + raise ValueError(f"duplicate snapshot rank: {snapshot.rank}") + by_rank[snapshot.rank] = snapshot + if set(by_rank) != set(ranks): + raise ValueError( + f"rank coverage mismatch: expected {ranks}, observed {tuple(by_rank)}" + ) + reference = by_rank[reference_rank] + comparisons: list[SnapshotComparison] = [] + for rank in ranks: + snapshot = by_rank[rank] + if (snapshot.phase, snapshot.step) != (reference.phase, reference.step): + raise ValueError("rank snapshots must share phase and step") + if rank != reference_rank: + comparisons.append(compare_model_snapshots(reference, snapshot)) + return RankComparison(reference_rank, ranks, tuple(comparisons))