1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
|
"""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"
def test_forward_full_logits():
"""Regression: model(X, Y)[0] must be full (B, T, V) logits.
Karpathy nanoGPT returns LAST-position logits only when targets=None."""
from model import GPT, GPTConfig
torch.manual_seed(0)
cfg = GPTConfig(n_layer=2, n_head=2, n_embd=16, block_size=32, bias=False,
vocab_size=65, dropout=0.0)
m = GPT(cfg).to(DEV)
x = torch.randint(0, 65, (2, 32), device=DEV)
y = torch.randint(0, 65, (2, 32), device=DEV)
logits = m(x, y)[0]
assert logits.shape == (2, 32, 65), f"full logits expected, got {logits.shape}"
if __name__ == '__main__':
test_batch_determinism()
test_weighting_semantics()
test_forward_full_logits()
print(f"PASS ({DEV})")
|
