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
69
|
"""Minimal tests for the faithful J-lens (jlens_v3) correctness.
Run inside the meru container where torch exists:
python3 tests/test_jlens_v3.py
Skips gracefully when torch is unavailable.
Key assertion: at the LAST layer the residual-to-residual Jacobian is the
identity, so the faithful J-lens vectors must equal the unembedding rows
(cosine similarity ~1.0). This is the ruler-check that validates the
W_U-probed VJP machinery.
"""
import os
import sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
try:
import numpy as np
import torch
import jlens_v3
except ImportError:
print("SKIP: torch not available (run inside the meru container)")
sys.exit(0)
from model import GPT, GPTConfig # noqa: E402
DEV = 'cuda' if torch.cuda.is_available() else 'cpu'
def _tiny_model(n_layer=2, d=16, block=32):
torch.manual_seed(0)
cfg = GPTConfig(n_layer=n_layer, n_head=2, n_embd=d, block_size=block,
bias=False, vocab_size=65, dropout=0.0)
return GPT(cfg).eval().to(DEV)
def _tiny_data(n=2000, block=32):
return np.random.RandomState(0).randint(0, 65, n).astype(np.uint16)
def test_shapes_and_finite():
model = _tiny_model()
data = _tiny_data()
batches = jlens_v3.make_batches(data, 32, 4, 3, DEV)
fv = jlens_v3.compute_faithful_jlens(model, 1, batches, DEV, chunk=8)
assert fv.shape == (65, 16), f"faithful vecs shape {fv.shape}"
assert torch.isfinite(fv).all(), "faithful vecs not finite"
pn = jlens_v3.compute_proxy_norms(model, 1, batches, DEV, 65, chunk=16)
assert all(np.isfinite(v) for v in pn.values()), "proxy norms not finite"
def test_last_layer_identity():
"""At the final layer J must be identity: faithful vectors == W_U rows."""
n_layer, d = 2, 16
model = _tiny_model(n_layer=n_layer, d=d)
data = _tiny_data()
batches = jlens_v3.make_batches(data, 32, 4, 3, DEV)
fv = jlens_v3.compute_faithful_jlens(model, n_layer - 1, batches, DEV,
chunk=8)
wu = model.lm_head.weight.detach().float()
sims = torch.nn.functional.cosine_similarity(fv.float(), wu.cpu().float(),
dim=1)
assert sims.mean().item() > 0.99, f"identity check failed: {sims.mean():.4f}"
if __name__ == '__main__':
test_shapes_and_finite()
test_last_layer_identity()
print(f"PASS ({DEV})")
|
