diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..e3a13d1 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,41 @@ +name: tests + +on: + push: + branches: [main] + pull_request: + workflow_dispatch: + +jobs: + test: + name: CPU test suite (py${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.12"] + + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + + - name: Install (CPU torch only) + # The CPU wheel index keeps this a ~200 MB install instead of pulling the CUDA stack. + # The suite intentionally needs nothing else -- no transformers, no video codecs, no data. + run: | + python -m pip install --upgrade pip + pip install --index-url https://download.pytorch.org/whl/cpu torch + pip install pytest + + - name: Run tests + run: pytest -v + + - name: Distributed SIGReg gate (gloo, world=2) + # Called out as its own step because it is the repo's load-bearing correctness claim: + # sharding the batch across ranks must not change the objective. The GPU/NCCL twin of + # this runs as `torchrun --nproc_per_node=2 src/sigreg.py --verify`. + run: pytest tests/test_distributed_sigreg.py -v diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..4a01756 --- /dev/null +++ b/LICENSE @@ -0,0 +1,202 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2026 Koushik Aditya V I + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + diff --git a/README.md b/README.md index a072364..aef9379 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,8 @@ # vjepa2-engine +[![tests](https://github.com/KoushikVGitHub/vjepa2-engine/actions/workflows/ci.yml/badge.svg)](https://github.com/KoushikVGitHub/vjepa2-engine/actions/workflows/ci.yml) +[![license](https://img.shields.io/badge/license-Apache--2.0-blue.svg)](LICENSE) + **A from-scratch, dataset-agnostic engineering pipeline for self-supervised world-models:** data curation at scale → distributed training with principled anti-collapse (LeJEPA / SIGReg) → optimized inference. The three things a world-model lab runs every day, built end to end and benchmarked with logged numbers. > Self-supervised models live or die on infrastructure: how fast you can curate the data, how efficiently you shard a model across GPUs without it collapsing, and how cheaply you can serve it. This repo builds and benchmarks that engineering surface. The engine is **dataset-agnostic** — the loss, the distributed SIGReg, and the FSDP trainer don't change with the data; only the loader and input dims do. @@ -61,6 +64,29 @@ Built as a focused engineering intensive. Method credibility (from-scratch JEPA - **Curation is a *policy*, not a constant** — rejection thresholds read off the empirical distribution's tail; disk-cached manifests make reruns instant and turn N re-scans into 1. - **Collapse is monitored, not hoped for** — target-std *and* effective rank are logged every run. That caught a *dimensional* collapse hiding behind a perfectly healthy target-std, and fixing it lifted the probe 2.2× ([the debugging story](study/notes/collapse_resolution.md), including why the loss going **up 150×** meant the model got healthier). +## Correctness gates (CI) + +Every claim above rests on properties that fail *silently* — a broken all-reduce still trains, +a non-circular pad still converges, a blind regularizer still logs a healthy loss. So the ones +that can be checked without a GPU are checked on every push (`pytest`, ~11 s, CPU-only): + +| Gate | What would otherwise break unnoticed | +|---|---| +| `tests/test_distributed_sigreg.py` | **world=2 × batch-B ≡ world=1 × batch-2B**, over gloo on CPU — the same invariant `src/sigreg.py --verify` gates on 2 GPUs over NCCL. Reduce the ECF sums *after* the nonlinearity instead of before, or lose differentiability through the collective, and training silently optimizes a different objective. | +| `tests/test_anticollapse.py` | The repo's central finding, as an executable claim: on a rank-2-of-64 batch whose per-dim std is a healthy 1.00, SIGReg's gradient measures **2.3e-4** while the covariance penalty's is **0.25 (~1000×)**. This is *why* `--var-coef/--cov-coef` exist and why `eff_rank` — not `tgt_std` — is the detector that matters. | +| `tests/test_conv_stem.py` | Phase-2 tokenizer stays a drop-in: identical token grid, bit-identical linear path, legacy checkpoint keys intact, and padding genuinely circular (periodic-shift equivariance to 8e-7) — CAMELS boxes are periodic, so a zero/reflect pad would fabricate an edge in exactly the high-k signal the conv stem exists to preserve. | +| `tests/test_sigreg.py` | The objective is ~0 for N(0,I), large under collapse, penalizes mis-scaling, and is bit-reproducible under a seeded generator — the property distributed correctness depends on. | +| `tests/test_masking.py` | Context/target partition is disjoint and exhaustive, blocks are contiguous, context is never empty at the trained default (n-blocks=4, block=8, 32×32). | +| `tests/test_jepa.py` | Mode wiring: the EMA teacher is built **only** where the mode needs one (allocating it in lejepa/visreg wastes a full encoder of dead memory per GPU), the teacher is frozen, `stop_grad` really does sever the target, and collapse detectors land in range. | + +```bash +pip install -r requirements-dev.txt && pytest # 41 tests, CPU, no data required +``` + +The suite needs only `torch` — no `transformers`, no video codecs, no CAMELS download — so it +stays runnable on a bare runner. One test (`gloo`, world=2) skips on Windows, where gloo cannot +create a device; it runs on the Linux CI matrix (3.10 / 3.12). + ## Repository layout ``` @@ -74,6 +100,7 @@ scripts/ # production drivers run_probe.py train + evaluate the cosmology probe on a frozen checkpoint rank_report.py representation-geometry report: token vs pooled effective rank, PCA, ridge probe analyze_all.py per-field statistics -> curation thresholds (see study/notes/camels_field_stats.md) +tests/ # CPU correctness gates, run in CI on every push (see Correctness gates) study/ # the from-scratch fundamentals — imports the library from src/, nothing here is imported back collapse_study.py synthetic study: stop-grad vs EMA vs SIGReg (what actually stops collapse) sigreg_demo.py SIGReg sanity demo (~0 for N(0,I), large for collapsed) @@ -138,6 +165,9 @@ The engine is dataset-agnostic, so each stage swaps only the loader + input dims 2. **Stage 2 — CAMELS 3D grids across redshift.** The only registered temporal axis in CAMELS (z = 0, 0.5, 1, 1.5, 2) → genuine spatiotemporal prediction of structure formation over cosmic time. 3. **Stage 3 — SDOML solar observations.** NASA Solar Dynamics Observatory ML dataset — real *observed* multi-waveband video at scale; forecast-next-frame = world dynamics. +## License +[Apache-2.0](LICENSE) — free to use, modify, and build on, with an explicit patent grant. + ## Status Active engineering build. Not affiliated with Meta or AMI Labs. Demonstrates production-engineering skills for self-supervised world-models: large-scale data curation, distributed training with principled anti-collapse, and inference optimization. diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..aa6b182 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,5 @@ +[pytest] +testpaths = tests +addopts = -q --strict-markers +filterwarnings = + ignore::DeprecationWarning diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..a91e6c5 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,5 @@ +# CI/test dependencies only. The suite is CPU-only and deliberately does NOT need the heavy +# runtime stack (transformers / decord / av / CAMELS data) -- it gates the engine's math and +# wiring, so it must stay installable and fast on a bare runner. +pytest>=8.0 +torch>=2.3 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..aa6fcfc --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,29 @@ +"""Shared test fixtures. + +Every test here is CPU-only and CI-sized: the point is to gate the ENGINE's invariants +(distribution math, anti-collapse geometry, tokenizer shape-compatibility, mask structure) +on every push, not to reproduce a training run. Transformer DEPTH is kept at <=2 layers +throughout -- torch's CPU build is unstable on the real 24-layer attention, and depth is +irrelevant to every property under test. +""" +import os +import sys + +import pytest +import torch + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.join(REPO_ROOT, "src")) + + +@pytest.fixture(autouse=True) +def _deterministic(): + """Seed every test. These assertions are statistical; an unseeded run would flake.""" + torch.manual_seed(0) + yield + + +@pytest.fixture +def gen(): + """A seeded generator, matching how the training loop feeds SIGReg (base + step).""" + return torch.Generator().manual_seed(1234) diff --git a/tests/test_anticollapse.py b/tests/test_anticollapse.py new file mode 100644 index 0000000..bffa7c2 --- /dev/null +++ b/tests/test_anticollapse.py @@ -0,0 +1,108 @@ +"""The anisotropic-collapse result, as an executable claim. + +This is the repo's central finding (study/notes/collapse_resolution.md): SIGReg's Cramer-Wold +sketch tests each 1-D projection MARGINALLY, so a low-rank blob whose marginals still look +like N(0,1) sails through it with a vanishing gradient -- while the VICReg off-diagonal +covariance term sees it immediately. That is why `--var-coef/--cov-coef` exist, and why +effective rank (not target-std) is the collapse detector that matters. + +The README quotes ||grad sigreg|| ~ 2e-4 vs ~1.25 for the covariance penalty on the real +ViT-L shape. `test_sigreg_is_blind_where_covariance_is_not` rebuilds that comparison at CI +scale so the claim is re-checked on every push instead of resting on a screenshot. +""" +import torch +import torch.nn.functional as F + +from sigreg import sigreg_loss, variance_covariance_reg + + +def rank_r_batch(n=4096, d=64, r=2): + """A rank-r embedding whose MARGINALS are exactly N(0,1). + + Rows of A are unit-norm, so z_j = has variance 1 for every dim j: per-dim std + looks perfectly healthy, and target-std -- the naive collapse detector -- sees nothing + wrong. But the batch lives in an r-dim subspace of d, which is precisely the dimensional + collapse that silently caps probe R^2. + """ + A = F.normalize(torch.randn(d, r), dim=1) + return torch.randn(n, r) @ A.t() + + +def effective_rank(z): + """Participation ratio PR = tr(C)^2 / ||C||_F^2, in [1, d] -- the trainer's `eff_rank`.""" + z = z - z.mean(dim=0, keepdim=True) + C = (z.t() @ z) / z.size(0) + return (torch.diagonal(C).sum() ** 2 / C.pow(2).sum()).item() + + +def test_effective_rank_recovers_the_true_subspace_dimension(): + """The detector itself must be correct before anything it reports can be trusted.""" + for r in (1, 2, 8): + assert abs(effective_rank(rank_r_batch(r=r)) - r) < 0.25 * r + # ...and a genuinely isotropic batch should read out near full rank. + assert effective_rank(torch.randn(8192, 64)) > 55 + + +def test_collapsed_batch_has_healthy_per_dimension_std(): + """The trap, stated as a test: std is ~1.0 while the batch is rank-2 of 64. + + Anyone monitoring only tgt_std would call this run healthy. + """ + z = rank_r_batch() + assert abs(z.std(dim=0).mean().item() - 1.0) < 0.05 + assert effective_rank(z) < 3 + + +def test_sigreg_is_blind_where_covariance_is_not(): + """The load-bearing claim: at a dimensionally-collapsed point the covariance penalty's + gradient is ORDERS of magnitude stronger than SIGReg's, so it is the term that actually + escapes the low-rank basin.""" + z = rank_r_batch() + + za = z.clone().requires_grad_() + sigreg_loss(za, generator=torch.Generator().manual_seed(3)).backward() + + zb = z.clone().requires_grad_() + _, cov = variance_covariance_reg(zb) + cov.backward() + + g_sigreg, g_cov = za.grad.norm().item(), zb.grad.norm().item() + assert g_sigreg < 1e-2, f"SIGReg gradient unexpectedly large at collapse: {g_sigreg:.2e}" + assert g_cov > 100 * g_sigreg, ( + f"covariance penalty lost its advantage: |grad cov| {g_cov:.3e} " + f"vs |grad sigreg| {g_sigreg:.3e} (expected >100x)" + ) + + +def test_covariance_term_separates_decorrelated_from_low_rank(): + """cov is the RANK knob: ~0 for an isotropic batch, large for a collapsed one.""" + _, cov_ok = variance_covariance_reg(torch.randn(4096, 64)) + _, cov_bad = variance_covariance_reg(rank_r_batch()) + assert cov_ok.item() < 0.05 + assert cov_bad.item() > 10 * cov_ok.item() + + +def test_variance_hinge_is_the_scale_knob(): + """var is the SCALE knob: silent at std >= gamma, active below it, and its gradient does + NOT vanish as std -> 0 (which is exactly what SIGReg's does).""" + # Comfortably above the hinge. (At std == gamma exactly the hinge is not silent: finite-batch + # std fluctuates either side of 1.0, so a real N(0,1) sample still scores ~5e-3. That is the + # term working as specified, not a defect -- so the "silent" claim is tested with margin.) + healthy = torch.randn(2048, 32) * 1.5 + var_ok, _ = variance_covariance_reg(healthy, gamma=1.0) + assert var_ok.item() == 0.0 + + shrunk = torch.randn(2048, 32) * 0.01 + var_bad, _ = variance_covariance_reg(shrunk, gamma=1.0) + assert var_bad.item() > 0.9, "hinge should read ~gamma - std when the batch is shrinking" + + z = (torch.randn(2048, 32) * 1e-3).requires_grad_() + variance_covariance_reg(z, gamma=1.0)[0].backward() + assert z.grad.norm().item() > 0, "variance hinge must keep pushing on a near-dead direction" + + +def test_regularizers_are_finite_on_a_degenerate_batch(): + """Guard the eps floors: an exactly-constant batch must not produce NaN/Inf.""" + z = torch.full((512, 32), 0.7) + var_l, cov_l = variance_covariance_reg(z) + assert torch.isfinite(var_l) and torch.isfinite(cov_l) diff --git a/tests/test_conv_stem.py b/tests/test_conv_stem.py new file mode 100644 index 0000000..4afceac --- /dev/null +++ b/tests/test_conv_stem.py @@ -0,0 +1,100 @@ +"""Phase-2 conv-stem tokenizer: a drop-in that must not disturb anything already trained. + +Pytest port of `scripts/test_conv_stem.py` (kept as the standalone runnable demo) so the +guarantees run in CI. The claims that matter for the Phase-2 A/B: + - the conv stem emits the SAME token grid, so pos-embed / predictor / probe are untouched; + - the LINEAR path is bit-identical and keeps legacy checkpoint keys -- every checkpoint + trained before Phase 2 must still load; + - padding is genuinely circular, because CAMELS boxes are periodic and zero/reflect padding + would inject a spurious edge into exactly the high-k signal the stem exists to preserve. + +Real Phase-2 geometry (patch-8, d=1024, 32x32=1024 tokens). Depth is 2 layers only: torch's +CPU build is unstable on the true 24-layer attention, and depth cannot affect tokenizer shape +compatibility, which is all this file claims. +""" +import pytest +import torch + +from jepa_loss import JEPA, ConvStem, ViTEncoder, ViTPredictor, random_block_mask + +IMG, PATCH, D = 256, 8, 1024 +HEADS, LAYERS = 8, 2 +GRID = IMG // PATCH # 32 +N = GRID * GRID # 1024 tokens -- what the conv stem must reproduce +B = 2 + + +def build(stem, seed=0): + torch.manual_seed(seed) + return ViTEncoder(img=IMG, patch=PATCH, d=D, heads=HEADS, layers=LAYERS, stem=stem) + + +@pytest.fixture(scope="module") +def x(): + torch.manual_seed(0) + return torch.randn(B, 1, IMG, IMG) + + +def test_both_stems_emit_the_same_token_grid(x): + with torch.no_grad(): + y_lin, y_conv = build("linear")(x), build("conv")(x) + assert y_lin.shape == (B, N, D) + assert y_conv.shape == y_lin.shape, "conv stem must match the linear token layout" + + +def test_linear_path_is_bit_identical_and_keeps_legacy_keys(): + """Old checkpoints must load unchanged: same seed -> same params, same key namespace.""" + a, b = build("linear").state_dict(), build("linear").state_dict() + assert list(a) == list(b) + for k in a: + assert torch.equal(a[k], b[k]), f"linear path not reproducible at {k}" + assert {k.split(".")[0] for k in a} == {"proj", "pos", "blocks"} + + +def test_conv_key_namespace_is_disjoint_from_linear(): + """conv_stem.* vs proj.*: a conv checkpoint can never silently half-load into a linear model.""" + top = {k.split(".")[0] for k in build("conv").state_dict()} + assert "conv_stem" in top and "proj" not in top + assert not hasattr(build("linear"), "conv_stem"), "linear encoder must not build conv_stem" + + +def test_conv_stem_is_deterministic(x): + with torch.no_grad(): + assert torch.equal(build("conv")(x), build("conv")(x)) + + +def test_padding_is_circular(x): + """Periodic-shift equivariance: rolling the input by one patch must roll the token grid by + one column. Zero or reflect padding breaks this -- and would fabricate an edge in a + periodic simulation box.""" + stem = ConvStem(PATCH, D).eval() + with torch.no_grad(): + t0 = stem(x).view(B, GRID, GRID, D) + shifted = stem(torch.roll(x, shifts=PATCH, dims=3)).view(B, GRID, GRID, D) + expected = torch.roll(t0, shifts=1, dims=2) + err = (shifted - expected).abs().max().item() + assert err < 1e-4, f"circular equivariance broken (max err {err:.2e}) -- padding not circular?" + + +def test_conv_stem_rejects_non_power_of_two_patch(): + with pytest.raises(ValueError, match="power-of-2"): + ConvStem(patch=12, d=64) + + +def test_encoder_rejects_an_unknown_stem(): + with pytest.raises(ValueError, match="unknown stem"): + ViTEncoder(img=IMG, patch=PATCH, d=64, heads=4, layers=1, stem="wavelet") + + +@pytest.mark.parametrize("stem", ["linear", "conv"]) +def test_masked_forward_and_full_jepa_step_run_under_both_stems(x, stem): + enc = build(stem) + pred = ViTPredictor(N, d=D, pred_d=384, heads=6, layers=6) + jepa = JEPA(enc, pred, loss_mode="lejepa", var_coef=5.0, cov_coef=4e-2, target_norm=True) + ctx_idx, tgt_idx = random_block_mask(GRID, block=8, n_blocks=4, device=x.device) + + ctx = enc(x, keep=ctx_idx) + assert ctx.shape == (B, len(ctx_idx), D), "masked forward must drop exactly the target tokens" + + loss, _ = jepa(x, ctx_idx, tgt_idx) + assert torch.isfinite(loss).all(), f"{stem}: non-finite loss" diff --git a/tests/test_distributed_sigreg.py b/tests/test_distributed_sigreg.py new file mode 100644 index 0000000..779a757 --- /dev/null +++ b/tests/test_distributed_sigreg.py @@ -0,0 +1,120 @@ +"""Distributed SIGReg correctness: world=2 x batch-B must equal world=1 x batch-2B. + +This is the property that makes SIGReg worth using at scale -- its regularizer is an +expectation over the batch, so ranks all-reduce partial ECF sums instead of gathering +negative pairs across the wire (contrast SimCLR). If the all-reduce were non-differentiable, +or the sums were reduced AFTER the nonlinearity instead of before, training would still run +and still look plausible -- it would just be silently optimizing the wrong objective. + +`src/sigreg.py --verify` gates this on 2 GPUs over NCCL. That cannot run in CI, so this is +the same invariant over GLOO on CPU: same claim, no hardware, every push. +""" +import os +import sys + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +REPO_SRC = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src") + +B, D, SEED = 256, 32, 0 +WORLD = 2 + + +def _worker(rank, world, port, out): + """One rank: run distributed SIGReg on its shard, report loss + local grad.""" + sys.path.insert(0, REPO_SRC) + from sigreg import sigreg_loss + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group("gloo", rank=rank, world_size=world) + try: + torch.manual_seed(SEED) + full = torch.randn(world * B, D) # identical on every rank + shard = full[rank * B:(rank + 1) * B].detach().requires_grad_() + + # Rank-SYNCED generator: V and t must match across ranks or the all-reduced + # statistic mixes different projections and means nothing. + gen = torch.Generator().manual_seed(SEED) + loss = sigreg_loss(shard, generator=gen, distributed=True) + loss.backward() + out[rank] = (loss.item(), shard.grad.clone()) + finally: + dist.destroy_process_group() + + +def _reference(): + """The single-process answer on the FULL batch -- what sharding must reproduce.""" + sys.path.insert(0, REPO_SRC) + from sigreg import sigreg_loss + + torch.manual_seed(SEED) + full = torch.randn(WORLD * B, D).detach().requires_grad_() + gen = torch.Generator().manual_seed(SEED) + loss = sigreg_loss(full, generator=gen, distributed=False) + loss.backward() + return loss.item(), full.grad + + +def test_ecf_sums_are_additive_across_shards(gen): + """The reason the all-reduce is legal, checked WITHOUT a process group. + + SIGReg reduces SUMS and divides by the global N; the nonlinearity is applied after. So + concatenating shards must give the same ECF as summing their partial sums. This is the + algebra the collective relies on -- gated here so it also runs where gloo cannot (Windows), + and so a regression points at the math rather than at the networking. + """ + sys.path.insert(0, REPO_SRC) + from sigreg import random_directions + + torch.manual_seed(SEED) + full = torch.randn(WORLD * B, D) + V = random_directions(D, 16, full.device, generator=gen) + t = torch.randn(8, generator=gen) + + def ecf_sums(z): + tp = (z @ V).unsqueeze(-1) * t + return tp.cos().sum(dim=0), tp.sin().sum(dim=0) + + cos_full, sin_full = ecf_sums(full) + shards = [ecf_sums(full[r * B:(r + 1) * B]) for r in range(WORLD)] + cos_parts = sum(s[0] for s in shards) + sin_parts = sum(s[1] for s in shards) + + assert torch.allclose(cos_full, cos_parts, atol=1e-4) + assert torch.allclose(sin_full, sin_parts, atol=1e-4) + + +@pytest.mark.skipif(sys.platform == "win32", + reason="gloo cannot create a device on Windows (makeDeviceForHostname)") +@pytest.mark.skipif(not dist.is_gloo_available(), reason="gloo backend unavailable") +def test_sharded_sigreg_matches_the_full_batch(): + ctx = mp.get_context("spawn") + out = ctx.Manager().dict() + mp.start_processes(_worker, args=(WORLD, 29511, out), nprocs=WORLD, + start_method="spawn", join=True) + assert len(out) == WORLD, "a rank died before reporting" + + ref_loss, ref_grad = _reference() + + # (1) Every rank computes the SAME global-batch loss, equal to the unsharded reference. + for rank in range(WORLD): + got = out[rank][0] + assert abs(got - ref_loss) < 1e-5, ( + f"rank {rank} loss {got:.8f} != single-process reference {ref_loss:.8f}" + ) + + # (2) Each rank's gradient matches its slice of the reference gradient. The differentiable + # all-reduce sums grads on the way back, so each rank sees the reference slice scaled by + # world -- the same factor DDP/FSDP then averages out. Assert the exact relationship + # rather than "close enough either way", so a change in it is a test failure, not a shrug. + for rank in range(WORLD): + grad = out[rank][1] + expect = ref_grad[rank * B:(rank + 1) * B] * WORLD + assert torch.allclose(grad, expect, rtol=1e-4, atol=1e-6), ( + f"rank {rank}: sharded gradient != world x reference slice " + f"(max diff {(grad - expect).abs().max().item():.3e})" + ) diff --git a/tests/test_jepa.py b/tests/test_jepa.py new file mode 100644 index 0000000..dd7c53c --- /dev/null +++ b/tests/test_jepa.py @@ -0,0 +1,106 @@ +"""JEPA wiring: the three anti-collapse modes and the structural switches between them. + +These are cheap tests of expensive mistakes. Each mode is a different answer to "what stops +the representation collapsing", and they differ in what gets built (a frozen EMA teacher or +not), what carries gradient, and what gets logged. Getting that wrong does not raise -- it +produces a training run that looks fine and learns nothing. + +Tiny geometry (32px, patch 4, d=64, 1 layer): this file tests wiring, not capacity. +""" +import pytest +import torch + +from jepa_loss import JEPA, LOSS_MODES, ViTEncoder, ViTPredictor, random_block_mask + +IMG, PATCH, D, GRID = 32, 4, 64, 8 +N = GRID * GRID +DEV = torch.device("cpu") + + +def make(loss_mode="lejepa", **kw): + enc = ViTEncoder(img=IMG, patch=PATCH, d=D, heads=4, layers=1) + pred = ViTPredictor(N, d=D, pred_d=32, heads=4, layers=1) + return JEPA(enc, pred, loss_mode=loss_mode, **kw) + + +@pytest.fixture +def batch(): + x = torch.randn(4, 1, IMG, IMG) + ctx_idx, tgt_idx = random_block_mask(GRID, block=2, n_blocks=2, device=DEV) + return x, ctx_idx, tgt_idx + + +def test_encoder_token_shape_and_masked_subsetting(): + enc = ViTEncoder(img=IMG, patch=PATCH, d=D, heads=4, layers=1) + x = torch.randn(2, 1, IMG, IMG) + assert enc(x).shape == (2, N, D) + keep = torch.arange(10) + assert enc(x, keep=keep).shape == (2, 10, D), "keep= must encode only the context tokens" + + +@pytest.mark.parametrize("mode", list(LOSS_MODES)) +def test_every_mode_produces_a_finite_loss_and_trains(batch, mode): + x, ctx_idx, tgt_idx = batch + jepa = make(mode, var_coef=5.0, cov_coef=4e-2, target_norm=True) + loss, tgt = jepa(x, ctx_idx, tgt_idx) + assert torch.isfinite(loss).all() + assert tgt.shape[1] == len(tgt_idx) + + loss.backward() + grads = [p.grad for p in jepa.context_encoder.parameters() if p.grad is not None] + assert grads, f"{mode}: no gradient reached the encoder" + assert all(torch.isfinite(g).all() for g in grads) + + +@pytest.mark.parametrize("mode,teacher", [(m, LOSS_MODES[m]["needs_teacher"]) for m in LOSS_MODES]) +def test_teacher_is_built_only_where_the_mode_needs_one(mode, teacher): + """lejepa/visreg replace the teacher with a distributional regularizer -- so allocating one + would waste a whole encoder's parameters of dead memory per GPU (it matters at ViT-L).""" + jepa = make(mode) + assert (jepa.target_encoder is not None) == teacher + + +def test_ema_teacher_is_frozen(): + jepa = make("ema") + assert all(not p.requires_grad for p in jepa.target_encoder.parameters()) + + +def test_step_ema_moves_the_teacher_only_in_ema_mode(): + jepa = make("ema", ema_decay=0.9) + before = [p.clone() for p in jepa.target_encoder.parameters()] + with torch.no_grad(): + for p in jepa.context_encoder.parameters(): + p.add_(1.0) # force a divergence for the EMA to chase + jepa.step_ema() + assert any(not torch.equal(a, b) for a, b in zip(before, jepa.target_encoder.parameters())) + + lejepa = make("lejepa") + lejepa.step_ema() # must be a no-op, not an AttributeError + + +def test_stop_grad_is_the_switch_that_severs_the_target(batch): + """The collapse control from the study: with stop_grad the target is detached; without it, + gradient flows into BOTH sides and a constant vector becomes the global minimum.""" + x, ctx_idx, tgt_idx = batch + _, tgt_on = make("ema", stop_grad=True)(x, ctx_idx, tgt_idx) + assert not tgt_on.requires_grad + + _, tgt_off = make("ema", stop_grad=False)(x, ctx_idx, tgt_idx) + assert tgt_off.requires_grad + + +def test_collapse_detectors_are_stashed_in_range(batch): + """The trainer's abort guard reads these every step; they must exist and be sane.""" + x, ctx_idx, tgt_idx = batch + jepa = make("lejepa", var_coef=5.0, cov_coef=4e-2) + jepa(x, ctx_idx, tgt_idx) + assert 1.0 <= jepa.last_eff_rank <= D, f"eff_rank {jepa.last_eff_rank} outside [1, {D}]" + assert jepa.last_tgt_std > 0 + for name in ("last_pred", "last_reg", "last_var", "last_cov"): + assert torch.isfinite(torch.tensor(getattr(jepa, name))) + + +def test_unknown_loss_mode_is_rejected_at_construction(): + """Fail at build time, not 4000 steps into a pod run.""" + with pytest.raises(ValueError, match="unknown loss_mode"): + make("contrastive") diff --git a/tests/test_masking.py b/tests/test_masking.py new file mode 100644 index 0000000..a8f539b --- /dev/null +++ b/tests/test_masking.py @@ -0,0 +1,52 @@ +"""I-JEPA block masking. + +The context/target split is the supervision signal itself. A silent bug here (overlap, an +empty context, a target ratio far off the intended ~15-25%) would not crash -- it would just +make the prediction task trivially easy, which is exactly the low-rank cheat that feeds +dimensional collapse. Phase 1 swept this ratio, so its structure is worth pinning down. +""" +import torch + +from jepa_loss import random_block_mask + +GRID, BLOCK = 32, 8 +DEV = torch.device("cpu") + + +def test_context_and_target_partition_the_grid(): + """Disjoint and exhaustive: every token is either seen or predicted, never both.""" + ctx, tgt = random_block_mask(GRID, BLOCK, DEV, n_blocks=4) + assert len(set(ctx.tolist()) & set(tgt.tolist())) == 0 + assert sorted(ctx.tolist() + tgt.tolist()) == list(range(GRID * GRID)) + + +def test_target_ratio_respects_the_block_budget(): + """Union of n blocks, so the ratio is AT MOST n*block^2/grid^2 -- less when they overlap.""" + for n_blocks in (1, 4, 8): + _, tgt = random_block_mask(GRID, BLOCK, DEV, n_blocks=n_blocks) + ratio = len(tgt) / (GRID * GRID) + assert 0 < ratio <= n_blocks * BLOCK ** 2 / GRID ** 2 + 1e-9 + assert ratio >= BLOCK ** 2 / GRID ** 2 - 1e-9, "at least one full block must be masked" + + +def test_a_single_block_is_contiguous_and_square(): + """Masked regions are spatial blocks, not scattered tokens -- the whole point of I-JEPA.""" + _, tgt = random_block_mask(GRID, BLOCK, DEV, n_blocks=1) + assert len(tgt) == BLOCK ** 2 + rows, cols = tgt // GRID, tgt % GRID + assert rows.max() - rows.min() == BLOCK - 1 + assert cols.max() - cols.min() == BLOCK - 1 + + +def test_context_is_never_empty_at_production_settings(): + """The encoder needs something to read: n-blocks=4, block=8 on a 32x32 grid (the trained + default) must always leave context tokens behind.""" + for _ in range(25): + ctx, _ = random_block_mask(GRID, BLOCK, DEV, n_blocks=4) + assert len(ctx) > 0 + + +def test_mask_is_random_across_calls(): + """A fixed mask would let the encoder memorize one hole instead of learning structure.""" + masks = {tuple(random_block_mask(GRID, BLOCK, DEV, n_blocks=2)[1].tolist()) for _ in range(8)} + assert len(masks) > 1 diff --git a/tests/test_sigreg.py b/tests/test_sigreg.py new file mode 100644 index 0000000..1f9b278 --- /dev/null +++ b/tests/test_sigreg.py @@ -0,0 +1,68 @@ +"""SIGReg (Cramer-Wold / characteristic-function) objective. + +The contract: sigreg_loss(z) is ~0 iff z is isotropic Gaussian, and grows as z departs from +it. If this breaks, LeJEPA training silently loses its only anti-collapse signal, so it is +gated here rather than discovered 4000 steps into a pod run. +""" +import pytest +import torch + +from sigreg import random_directions, sigreg_loss + + +def test_random_directions_are_unit_vectors(): + V = random_directions(dim=64, n_proj=32, device=torch.device("cpu")) + assert V.shape == (64, 32) + norms = V.norm(dim=0) + assert torch.allclose(norms, torch.ones_like(norms), atol=1e-6) + + +def test_near_zero_for_isotropic_gaussian(gen): + """z ~ N(0, I) is the objective's global minimum: every 1-D projection is N(0,1).""" + z = torch.randn(4096, 64) + loss = sigreg_loss(z, generator=gen) + assert loss.item() < 1e-3, f"N(0,I) should score ~0, got {loss.item():.2e}" + + +def test_large_for_complete_collapse(gen): + """Every sample mapped to the same vector -- the failure SIGReg exists to catch.""" + z = torch.ones(4096, 64) * 0.3 + loss = sigreg_loss(z, generator=gen) + assert loss.item() > 0.1, f"collapsed batch should score high, got {loss.item():.2e}" + + +def test_penalizes_wrong_scale(gen): + """Correct shape, wrong variance still fails: the target is N(0,1), not 'any Gaussian'.""" + ok = sigreg_loss(torch.randn(4096, 64), generator=gen).item() + scaled = sigreg_loss(torch.randn(4096, 64) * 5.0, generator=gen).item() + assert scaled > ok * 10, f"mis-scaled batch scored {scaled:.2e} vs isotropic {ok:.2e}" + + +def test_deterministic_under_a_seeded_generator(): + """Distributed correctness DEPENDS on this: ranks must draw identical V and t.""" + z = torch.randn(512, 32) + a = sigreg_loss(z, generator=torch.Generator().manual_seed(7)) + b = sigreg_loss(z, generator=torch.Generator().manual_seed(7)) + assert torch.equal(a, b) + + +def test_differs_across_generator_seeds(): + """Sanity on the above: the sketch really is random, not a constant projection.""" + z = torch.randn(512, 32) + a = sigreg_loss(z, generator=torch.Generator().manual_seed(7)) + b = sigreg_loss(z, generator=torch.Generator().manual_seed(8)) + assert not torch.equal(a, b) + + +def test_gradient_flows_to_embeddings(gen): + z = torch.randn(256, 32, requires_grad=True) + sigreg_loss(z, generator=gen).backward() + assert z.grad is not None and torch.isfinite(z.grad).all() + assert z.grad.abs().sum() > 0 + + +def test_distributed_mode_requires_a_synced_generator(): + """Without a rank-synced generator each rank would project onto DIFFERENT directions and + the all-reduced statistic would be meaningless. Fail loudly instead.""" + with pytest.raises(AssertionError): + sigreg_loss(torch.randn(64, 16), generator=None, distributed=True)