From 30e81b39e1e280c8fe51e13349ebec626a5353aa Mon Sep 17 00:00:00 2001 From: Void Agent Date: Sun, 2 Aug 2026 13:30:57 +0100 Subject: Add tests/test_loss_reweight.py: batch determinism + weighting semantics (runs in container) --- tests/test_loss_reweight.py | 53 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 tests/test_loss_reweight.py diff --git a/tests/test_loss_reweight.py b/tests/test_loss_reweight.py new file mode 100644 index 0000000..7df66c9 --- /dev/null +++ b/tests/test_loss_reweight.py @@ -0,0 +1,53 @@ +"""Minimal tests for the loss-reweighting experiment (no pytest needed). + +Run inside the meru container where torch exists: + python3 tests/test_loss_reweight.py +Skips gracefully when torch is unavailable (e.g. on voidlaptop bare metal). +""" +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src')) + +try: + import numpy as np + import torch + import torch.nn.functional as F + import loss_reweight as lr +except ImportError: + print("SKIP: torch not available (run inside the meru container)") + sys.exit(0) + +DEV = 'cuda' if torch.cuda.is_available() else 'cpu' + + +def test_batch_determinism(): + """Paired-model guarantee: identical minibatch order per seed, different across seeds.""" + data = np.random.RandomState(0).randint(0, 65, 50000).astype(np.uint16) + g1a = torch.Generator().manual_seed(20260731) + g1b = torch.Generator().manual_seed(20260731) + x1, y1 = lr._batch(data, 128, 16, g1a, DEV) + x1b, y1b = lr._batch(data, 128, 16, g1b, DEV) + assert torch.equal(x1, x1b) and torch.equal(y1, y1b), "same seed must give same batch" + g2 = torch.Generator().manual_seed(20260732) + x2, _ = lr._batch(data, 128, 16, g2, DEV) + assert not torch.equal(x1, x2), "different seeds must give different batches" + + +def test_weighting_semantics(): + """control == plain CE; ctrl_random picks a fixed set of non-q positions.""" + q_id = 3 + y = torch.randint(0, 65, (2, 4), device=DEV) + logits = torch.randn(2, 4, 65, device=DEV) + l_ctl = lr._weighted_loss(logits, y, 'control', q_id, 0, 65, DEV) + l_ce = F.cross_entropy(logits.view(-1, 65), y.view(-1)) + assert abs(l_ctl.item() - l_ce.item()) < 1e-6, "control must equal plain CE" + l_r1 = lr._weighted_loss(logits, y, 'ctrl_random', q_id, 7, 65, DEV) + l_r2 = lr._weighted_loss(logits, y, 'ctrl_random', q_id, 7, 65, DEV) + assert abs(l_r1.item() - l_r2.item()) < 1e-9, "ctrl_random must be deterministic" + + +if __name__ == '__main__': + test_batch_determinism() + test_weighting_semantics() + print(f"PASS ({DEV})") -- cgit v1.2.3