"""The comparison instrument, calibrated. `benchmarks/compare.py` is what will produce the CUDA-vs-MPS verdict, and `pretrain.py` is what produces the artifacts it reads. Both were untested, which is a bad property for an instrument: a wrong answer from a measuring tool looks exactly like a right one, and the CUDA run is a thing we get to do once before believing its output. The load-bearing assumption is the first test below. Everything the comparison claims rests on two machines starting from bit-identical weights on bit-identical data — if that quietly stopped being true, every number in the comparison would be measuring initialisation drift and still look plausible. It did quietly stop being true, and only on another machine: the seed gives identical weights on *one* platform, which is all a test here can check, but S4's `eigh` returns different eigenvectors on macOS and Linux and so different `B` and `P`. See tests/test_init_weights.py, which covers the mechanism that now carries the assumption instead of the seed. """ from __future__ import annotations import importlib.util import json import sys from pathlib import Path import pytest import torch ROOT = Path(__file__).resolve().parent.parent # The benchmark scripts import each other by plain name (`import init_weights`), # which works when they are run as scripts because Python puts the script's own # directory on sys.path. Loading them by file location here does not, so the # directory has to be added explicitly. sys.path.insert(0, str(ROOT / "benchmarks")) def _load(name: str, path: Path): """Import a script that is not part of the installed package.""" spec = importlib.util.spec_from_file_location(name, path) module = importlib.util.module_from_spec(spec) sys.modules[name] = module spec.loader.exec_module(module) return module pretrain = _load("_td_pretrain", ROOT / "benchmarks" / "pretrain.py") compare = _load("_td_compare", ROOT / "benchmarks" / "compare.py") # --- the assumption the whole comparison rests on ---------------------------- def test_initialisation_is_identical_across_calls(): """Models are built on CPU under a fixed seed and only then moved, so two machines start from the same weights *bitwise*. This is the property that makes a cross-device difference mean 'arithmetic' rather than 'different starting point', and nothing else in the suite checks it.""" lat = pretrain.sparse_2d() built = [] for _ in range(2): torch.manual_seed(pretrain.SEED) built.append(pretrain.ZOO["lstm_2d_sparse"]["build"](lat)) a, b = (dict(m.named_parameters()) for m in built) assert a.keys() == b.keys() for key in a: assert torch.equal(a[key], b[key]), f"{key} differs between two seeded builds" def test_the_data_stream_is_identical_across_calls(): """Same for the batches: drawn on CPU from a seeded generator, so both machines see the same numbers in the same order.""" def draw_all(): gen = torch.Generator().manual_seed(pretrain.SEED + 1) return [torch.randn(4, 6, 6, 8, 1, generator=gen) for _ in range(3)] for first, second in zip(draw_all(), draw_all(), strict=True): assert torch.equal(first, second) def test_the_lattice_is_the_same_every_time(): """The sparse mask is generated, not stored — if it drifted, the two runs would be training on different problems.""" a, b = pretrain.sparse_2d(), pretrain.sparse_2d() assert torch.equal(a.valid, b.valid) assert a.shape == b.shape and a.names == b.names @pytest.mark.parametrize("name", list(pretrain.ZOO)) def test_every_model_in_the_zoo_builds_and_runs(name): """A broken entry currently surfaces only part-way through a benchmark run — after minutes of training other models.""" cfg = pretrain.ZOO[name] lat = cfg["lat"]() torch.manual_seed(pretrain.SEED) model = cfg["build"](lat) x = torch.randn(1, 2, *lat.shape, 1) with torch.no_grad(): out = model(x) assert out.shape[:-1] == x.shape[:-1] assert torch.isfinite(out).all() # --- the comparison itself --------------------------------------------------- def _run_dir(tmp_path: Path, name: str, losses: list[float], weights: dict) -> Path: """A minimal `pretrain.py`-shaped output directory.""" root = tmp_path / name (root / "toy").mkdir(parents=True) torch.save({"model": weights, "head": {}}, root / "toy" / "weights.pt") record = { "name": "toy", "n_params": sum(v.numel() for v in weights.values()), "losses": losses, "loss_first": losses[0], "loss_final": losses[-1], "steps_per_second": 10.0, } (root / "toy" / "metrics.json").write_text(json.dumps(record)) (root / "manifest.json").write_text( json.dumps( { "device": "cpu", "device_name": name, "torch": torch.__version__, "torch_dimensions": "test", "seed": 1, "steps": len(losses), "batch": 1, "models": [record], } ) ) return root def test_two_identical_runs_report_no_divergence(tmp_path): w = {"a": torch.ones(4, 4)} left = _run_dir(tmp_path, "left", [1.0, 0.5, 0.25], w) right = _run_dir(tmp_path, "right", [1.0, 0.5, 0.25], dict(w)) a, b = compare.load(left), compare.load(right) assert compare.first_divergence(a["models"][0]["losses"], b["models"][0]["losses"]) is None rel, worst = compare.weight_delta( compare.weights_of(left, "toy"), compare.weights_of(right, "toy") ) assert rel == 0.0 and worst == 0.0 def test_divergence_is_found_at_the_step_it_happens(): """Off by one here would misattribute *when* two machines parted company, which is the number the comparison leans on most.""" base = [1.0, 0.5, 0.25, 0.125] moved = [1.0, 0.5, 0.2500004, 0.125] # 1.6e-6 relative, above the 1e-6 tolerance assert compare.first_divergence(base, moved) == 2 assert compare.first_divergence(base, list(base)) is None # A difference below tolerance is not a divergence. assert compare.first_divergence(base, [1.0, 0.5, 0.25000001, 0.125]) is None def test_weight_delta_is_the_arithmetic_it_claims(): """Relative Frobenius distance and the worst single element, checked against numbers computed by hand rather than against itself.""" a = {"model": {"w": torch.tensor([[3.0, 0.0], [0.0, 4.0]])}} # ‖w‖ = 5 b = {"model": {"w": torch.tensor([[3.0, 0.0], [0.0, 1.0]])}} # differs by 3 at one entry rel, worst = compare.weight_delta(a, b) assert worst == pytest.approx(3.0) assert rel == pytest.approx(3.0 / 5.0) def test_a_model_missing_from_one_run_is_reported_not_crashed(tmp_path, capsys): """One side failing to train a model must not take the whole comparison down — the other fifteen rows are still the point.""" left = _run_dir(tmp_path, "left", [1.0, 0.5], {"a": torch.ones(2, 2)}) right = _run_dir(tmp_path, "right", [1.0, 0.5], {"a": torch.ones(2, 2)}) manifest = json.loads((right / "manifest.json").read_text()) manifest["models"].append({"name": "only_here", "error": "boom"}) (right / "manifest.json").write_text(json.dumps(manifest)) argv = sys.argv sys.argv = ["compare.py", str(left), str(right)] try: assert compare.main() == 0 finally: sys.argv = argv assert "toy" in capsys.readouterr().out def test_the_cuda_harness_imports_and_skips_without_a_device(): """`scripts/cuda_check.py` runs its checks at import. Without CUDA every CUDA-only check must skip rather than fail, or the report from a machine that *has* one cannot be trusted either.""" if torch.cuda.is_available(): pytest.skip("this asserts the no-CUDA behaviour") module = _load("_td_cuda_check", ROOT / "scripts" / "cuda_check.py") statuses = {name: status for name, status, _ in module.RESULTS} assert statuses, "the harness recorded nothing at all" assert "fail" not in statuses.values(), f"failed without a device: {statuses}" assert any(s == "skip" for s in statuses.values())