diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 3a1a888..5a811a9 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -5,6 +5,7 @@ Thanks for your interest. This page is the short version: how to build, test, li ## Prerequisites - [rustup](https://rustup.rs/) - the Rust toolchain is pinned by `rust-toolchain.toml` (stable + rustfmt + clippy). + Stable 1.85 or newer is required: `sha2` 0.11 and its dependency chain (`digest`, `crypto-common`, `hybrid-array`) are edition 2024. - [uv](https://docs.astral.sh/uv/) - manages the Python workspace and its single `uv.lock`. torch is pinned to the CPU-only PyTorch wheel index in the root `pyproject.toml`, so `uv sync` installs the ~100MB `+cpu` build instead of the default Linux wheel and its ~3GB of NVIDIA CUDA libraries. macOS wheels are unaffected (MPS still works); if you need CUDA locally, point the `torch` entry in `[tool.uv.sources]` at a CUDA index and re-lock, but don't commit that. diff --git a/Cargo.lock b/Cargo.lock index 1a144b0..850fdd7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -127,11 +127,11 @@ checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" [[package]] name = "block-buffer" -version = "0.10.4" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" dependencies = [ - "generic-array", + "hybrid-array", ] [[package]] @@ -215,6 +215,12 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "core-foundation-sys" version = "0.8.7" @@ -223,30 +229,30 @@ checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" [[package]] name = "cpufeatures" -version = "0.2.17" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" dependencies = [ "libc", ] [[package]] name = "crypto-common" -version = "0.1.7" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" dependencies = [ - "generic-array", - "typenum", + "hybrid-array", ] [[package]] name = "digest" -version = "0.10.7" +version = "0.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" dependencies = [ "block-buffer", + "const-oid", "crypto-common", ] @@ -302,16 +308,6 @@ dependencies = [ "slab", ] -[[package]] -name = "generic-array" -version = "0.14.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" -dependencies = [ - "typenum", - "version_check", -] - [[package]] name = "getrandom" version = "0.4.3" @@ -329,6 +325,15 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hybrid-array" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" +dependencies = [ + "typenum", +] + [[package]] name = "iana-time-zone" version = "0.1.65" @@ -580,9 +585,9 @@ dependencies = [ [[package]] name = "sha2" -version = "0.10.9" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" dependencies = [ "cfg-if", "cpufeatures", @@ -671,12 +676,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "version_check" -version = "0.9.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" - [[package]] name = "wasm-bindgen" version = "0.2.126" diff --git a/Cargo.toml b/Cargo.toml index b403068..df10f8c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,7 @@ anyhow = "1" uuid = { version = "1", features = ["v4"] } chrono = { version = "0.4", features = ["clock"] } inventory = "0.3" -sha2 = "0.10" +sha2 = "0.11" [workspace.lints.rust] unsafe_code = "forbid" diff --git a/crates/a2d-run/src/rundir.rs b/crates/a2d-run/src/rundir.rs index c1c7777..3115938 100644 --- a/crates/a2d-run/src/rundir.rs +++ b/crates/a2d-run/src/rundir.rs @@ -1,6 +1,7 @@ //! Run-directory layout and `manifest.json` lifecycle. use std::fs; +use std::io::Read; use std::path::Path; use a2d_contracts::{Manifest, RunStatus}; @@ -68,7 +69,14 @@ pub fn reopen_run_dir(run_dir: &Path) -> Result { Ok(manifest) } -/// sha256 (hex) of the primary source safetensors file, for manifest provenance. +/// sha256 (lowercase hex) of the primary source safetensors file, for manifest +/// provenance. Read in fixed-size chunks, so peak memory does not grow with the +/// weights. +/// +/// The digest string is a cross-language contract: eval re-derives it with +/// Python's `hashlib.sha256(...).hexdigest()` and refuses the AR baseline on +/// mismatch (`a2d_core/eval/likelihood.py`), so the encoding must stay +/// hexdigest-identical. /// /// ponytail: hash the primary `model.safetensors`; upgrade to a header + /// shard-manifest digest if big-model / sharded provenance gets expensive. @@ -76,6 +84,24 @@ pub fn source_hash(model_dir: &Path) -> Result { let path = model_dir.join("model.safetensors"); let mut file = fs::File::open(&path).with_context(|| format!("opening {}", path.display()))?; let mut hasher = Sha256::new(); - std::io::copy(&mut file, &mut hasher).with_context(|| format!("hashing {}", path.display()))?; - Ok(format!("{:x}", hasher.finalize())) + // digest 0.11 dropped the `io::Write` impl on hashers, so stream the file by + // hand rather than reading a multi-GB safetensors blob into memory. + let mut buf = vec![0u8; 64 * 1024]; + loop { + let n = match file.read(&mut buf) { + Ok(0) => break, + Ok(n) => n, + Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue, + Err(e) => { + return Err(e).with_context(|| format!("hashing {}", path.display())); + } + }; + hasher.update(&buf[..n]); + } + // digest 0.11 returns a `hybrid_array::Array`, which has no `LowerHex` impl. + Ok(hasher + .finalize() + .iter() + .map(|b| format!("{b:02x}")) + .collect()) } diff --git a/crates/a2d-run/src/worker.rs b/crates/a2d-run/src/worker.rs index db900f6..630ca42 100644 --- a/crates/a2d-run/src/worker.rs +++ b/crates/a2d-run/src/worker.rs @@ -567,4 +567,19 @@ printf '%s\n' '{"schema_version":"0.1.0","job_id":"test-job","seq":3,"ts":"2026- ); let _ = fs::remove_dir_all(&base); } + + #[test] + fn source_hash_spans_read_chunks() { + // 200_000 bytes forces several passes of the 64 KiB read loop, so a + // chunking or hex-encoding regression fails here rather than silently + // invalidating manifest provenance of already-converted models. + let base = std::env::temp_dir().join(format!("a2d-run-test-{}", Uuid::new_v4())); + fs::create_dir_all(&base).unwrap(); + fs::write(base.join("model.safetensors"), vec![b'a'; 200_000]).unwrap(); + assert_eq!( + rundir::source_hash(&base).unwrap(), + "2287d207f24a941ff3b56c04c8a25ad56b63e3023207b3bb5b4ac0c9869d74be" + ); + let _ = fs::remove_dir_all(&base); + } } diff --git a/packages/a2d-worker-hf/src/a2d_core/eval/likelihood.py b/packages/a2d-worker-hf/src/a2d_core/eval/likelihood.py index a18a7d0..8b3905e 100644 --- a/packages/a2d-worker-hf/src/a2d_core/eval/likelihood.py +++ b/packages/a2d-worker-hf/src/a2d_core/eval/likelihood.py @@ -51,9 +51,10 @@ def mdlm_bound( Corruption (the t-schedule + masking, which drives the RNG) always runs over the FULL chunk set, so the result is independent of ``eval_batch_size``; only the model forward is split into sub-batches of ``eval_batch_size`` chunks (``<= 0`` => one forward over all - chunks). Per-sequence rows do not attend to each other, so a sub-batched forward is - numerically identical to the single-batch one - it just caps peak memory at the sub-batch - size instead of at ``max_eval_tokens`` (avoids OOM on a single giant forward). + chunks). Per-sequence rows do not attend to each other, so a sub-batched forward matches the + single-batch one up to float32 round-off (the batch dim changes the forward's GEMM shapes, + so the last bits of the logits can move) - it just caps peak memory at the sub-batch size + instead of at ``max_eval_tokens`` (avoids OOM on a single giant forward). """ from a2d_core.objectives.mdlm import MDLM diff --git a/packages/a2d-worker-hf/src/a2d_core/transform/identity.py b/packages/a2d-worker-hf/src/a2d_core/transform/identity.py index 7711932..82614be 100644 --- a/packages/a2d-worker-hf/src/a2d_core/transform/identity.py +++ b/packages/a2d-worker-hf/src/a2d_core/transform/identity.py @@ -3,7 +3,9 @@ At ``alpha=0`` the annealed patch reproduces base causality to the bit, so a correct patch yields ``max_abs_diff == 0.0`` on CPU float32. The gate ALWAYS runs float32 on CPU regardless of ``--dtype`` (Risk 2). Grow adds a logit COLUMN, so patched logits -are sliced to ``base_vocab`` before comparing (Decision 7 / Risk 3). +are sliced to ``base_vocab`` before comparing (Decision 7 / Risk 3) - and because that +column also widens the logit GEMM, some CPU kernels round the shared columns a couple +of ulps apart, which is why the verdict is ``<= tolerance`` and not ``== 0.0``. This gate CANNOT prove the patch reaches the model's causality - a no-op seam that leaves the model fully causal passes it too - so bidirectionality is proven separately diff --git a/packages/a2d-worker-hf/tests/test_identity.py b/packages/a2d-worker-hf/tests/test_identity.py index 859dd7e..bee6fcb 100644 --- a/packages/a2d-worker-hf/tests/test_identity.py +++ b/packages/a2d-worker-hf/tests/test_identity.py @@ -24,7 +24,18 @@ def test_patched_at_alpha0_is_bit_identical_to_base(tiny_gpt2: Callable[..., Any probe = torch.randint(0, base_vocab, (2, 8)) result = check_identity(base, patched, state, probe, base_vocab) + # Grow widens the tied lm_head from base_vocab to base_vocab+1 columns, so the two + # models run DIFFERENTLY SHAPED logit GEMMs; some CPU kernels round the shared columns + # a couple of ulps apart (2.2e-8 seen on x86 CI). That is the resize, not the patch, so + # the grown pair is held to IDENTITY_TOLERANCE and bit-exactness is asserted below at + # equal lm_head width, where the patch is the only difference between the two models. assert result.passed assert result.tolerance == IDENTITY_TOLERANCE assert result.max_abs_diff <= IDENTITY_TOLERANCE - assert result.max_abs_diff == 0.0 # eager + fp32 is exact, not merely within tolerance + + ungrown_state = AnnealState() + ungrown = tiny_gpt2(0) + ungrown.load_state_dict(base.state_dict()) + install_anneal_patch(ungrown, ungrown_state) + exact = check_identity(base, ungrown, ungrown_state, probe, base_vocab) + assert exact.max_abs_diff == 0.0 # eager + fp32 is exact, not merely within tolerance diff --git a/packages/a2d-worker-hf/tests/test_likelihood.py b/packages/a2d-worker-hf/tests/test_likelihood.py index ec714ba..eee9b05 100644 --- a/packages/a2d-worker-hf/tests/test_likelihood.py +++ b/packages/a2d-worker-hf/tests/test_likelihood.py @@ -51,14 +51,20 @@ def test_mdlm_bound_sub_batch_matches_single_batch(convert_setup: ConvertSetup) single = mdlm_bound(model, tokenizer, mask_id, eval_batch_size=0, **kw) # one big forward subbed = mdlm_bound(model, tokenizer, mask_id, eval_batch_size=4, **kw) # 8 sub-batches tiny = mdlm_bound(model, tokenizer, mask_id, eval_batch_size=1, **kw) # one seq per forward + # Sub-batching changes the forward's batch dim, so fp32 kernels round the logits a few + # ulps apart; per-sequence nats are then summed in float64 in the same order, leaving only + # that round-off (1.7e-9 relative seen on x86 CI). A real splitting bug - wrong rows, wrong + # weights, resampled corruption - moves the bound by percent, so 1e-6 (~8 fp32 eps) still + # catches it while staying above what a shape-dependent GEMM kernel can do. + rel_tol = 1e-6 for other in (subbed, tiny): assert math.isclose( - single.nats_per_token, other.nats_per_token, rel_tol=1e-9, abs_tol=1e-12 + single.nats_per_token, other.nats_per_token, rel_tol=rel_tol, abs_tol=1e-12 ) assert math.isclose( - single.bits_per_token, other.bits_per_token, rel_tol=1e-9, abs_tol=1e-12 + single.bits_per_token, other.bits_per_token, rel_tol=rel_tol, abs_tol=1e-12 ) - assert math.isclose(single.std_error, other.std_error, rel_tol=1e-9, abs_tol=1e-12) + assert math.isclose(single.std_error, other.std_error, rel_tol=rel_tol, abs_tol=1e-12) assert single.n_tokens == subbed.n_tokens == tiny.n_tokens assert single.mc_samples == subbed.mc_samples == tiny.mc_samples