Skip to content

Add model snapshots and cross-rank diagnostics for agents - #172

Open
Essoz wants to merge 7 commits into
codex/real-rank-runtime-pilotfrom
codex/model-snapshot-agent-workflow
Open

Essoz wants to merge 7 commits into
codex/real-rank-runtime-pilotfrom
codex/model-snapshot-agent-workflow

Conversation

@Essoz

@Essoz Essoz commented Nov 2, 2026 •

Copy link
Copy Markdown
Collaborator

Agents previously had to implement model fingerprints and rank comparisons themselves. This adds a typed model snapshot API, strict readers and pair/rank comparisons, plus runnable CPU/DDP examples. The caller selects observation points and predicates; missing ranks, malformed snapshots and unsupported tensors fail explicitly.

Approach and contracts

The implementation reuses existing model-state enumeration and exact tensor SHA-256, covering materialized registered parameters and buffers, nonpersistent buffers and repeated object aliases. It preserves application state and does not implicitly gather shards or invoke the invariant engine. Lightweight observation imports leave PyTorch hooks unchanged through lazy exports that preserve existing public object identities. The existing parallel hash remains unchanged.

A regression also fixes the legacy proxy observer ignoring tracing suppression: snapshot collection changes from four extra variable records and modified parameter timestamps to zero extra records and unchanged timestamps. Ordinary observations still execute. This is a correctness result, not a tracing-overhead speedup claim.

Snapshots synchronously copy complete tensor bytes to CPU. Callers must choose a quiescent boundary and expose full parameters for sharded models. This does not provide atomic snapshots during concurrent mutation, automatic collectives, negligible overhead, or storage-alias/checkpoint-completeness proof.

Validation

  • 102 focused tests passed, five CUDA-only skips; configured Ruff, mypy, Black and isort hooks plus scoped strict typing passed.
  • Full suite: 1,074 passed, nine skipped, two failures reproduced with the same assertions at parent 5530554: test_traincheck_internal_attr_hidden and test_inject_eval_loop. No new failure observed; the full suite is not completely green.
  • Actual two-process CPU example passed.
  • Native Ring18 parity: 60 complete model/buffer/gradient/optimizer/input/RNG state-file comparisons and 665 independently checked tensor fingerprints across seven snapshots passed. These single paired runs include private observation I/O and are not an overhead benchmark.
  • Final report passed diff/link checks and independent accounting review. All 203 exported grading files, exact source/mode bindings, phase outcomes and costs were independently verified.

Fresh optional-tool comparison

Two fresh native agent conversations received the same neutral readiness task and existing 14.1M-parameter native FSDP case. Neither was told there was a bug. TrainCheck was optional in one arm. The implementation was frozen at signed commit 498bf5f; the preceding case, source/corpus/oracles and six-scenario criterion were unchanged. Requested model/settings, exact source hashes and environment are in the report.

The optional-tool agent voluntarily used snapshot capture and comparison: nine runtime snapshots covering 95 tensors each, plus executed read_snapshot, pair comparisons and complete-rank comparisons with assertions. No TrainCheck integration failure was observed. It still supplied its own FSDP boundary adapter and observation predicates.

Independent grading retained all nine submitted occurrences, replaying six exact arm/content/permission maps under six scenarios: 36 actual replays, 32 passes, four initial mixed/swapped mismatches, zero unknowns. Both agents first submitted a valid repair in turn two; both final versions passed all six scenarios. Both fixed the observed inconsistent-rank initialization; neither established its underlying cache-version origin.

At first independently valid submitted repair Control Optional TrainCheck Reduction vs control
Agent CLI execution, seconds 188.260 141.211 24.99%
Input plus output tokens, including cached input once 981,070 715,584 27.06%
Allocated solver GPU-seconds 35.031 25.289 27.81%
Actual elapsed including interruptions, seconds 2,113.871 2,075.498 1.82%

These costs stop at the frozen turn-two submission boundary, before the second GPU request and later visible verification. CLI time excludes remote execution and queue/operator pauses. Token volumes are not billed dollars; uncached input plus output is 79,694 versus 74,560.

This is one interrupted pair with different agent-authored diagnostics and checks, not a causal or general speedup. Full-trial control/treatment totals are 1,852,226/1,269,560 tokens, 362.937/294.554 CLI seconds and 118.342/72.822 GPU-seconds. Control made an extra ordinary-entry validation that explains 78.27% of the terminal GPU difference. Actual terminal wall time is longer for treatment because it required another transport recovery. All failed attempts, original conversations and cumulative budgets remain accounted for; raw and adjusted clocks are separate.

New experiment GPU costs: snapshot compatibility preflight 95.853, both solvers 191.164, independent grading 995.453, total 1,282.470 GPU-seconds. GPUs are released. Large tensor states and transcripts are outside git.

See the full measurements and evidence report and agent-facing guide.

Review dependency

Stacked on #171 (codex/real-rank-runtime-pilot), itself stacked on #170. Signed implementation and checkpoint commits preserve the experiment history. Final result commit 2fb022a and review correction f5a212c complete the report. Ready for review; do not merge automatically.

@Essoz
Essoz marked this pull request as ready for review November 2, 2026 07:39

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant