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
|
import os, sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
import torch
import torch.nn as nn
from zbp_scaling.model import ScalingLM
from zbp_scaling.zbp import ZBPConfig, zbp_blocks
def test_exact_mode_matches_bp():
torch.manual_seed(0)
m = ScalingLM(101, 64, 2, 4, 32, cfg=ZBPConfig(mode="bp"))
x = torch.randint(0, 101, (2, 32)); y = torch.randint(0, 101, (2, 32))
g_bp = torch.autograd.grad(nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()), list(m.parameters()))
for b in zbp_blocks(m):
b.cfg = ZBPConfig(mode="exact")
g = torch.autograd.grad(nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()), list(m.parameters()))
err = max((a - b).norm().item() / (b.norm().item() + 1e-12) for a, b in zip(g, g_bp))
assert err < 1e-5, err
def test_zbp_cd_trains_shape():
torch.manual_seed(0)
m = ScalingLM(101, 64, 2, 4, 32, cfg=ZBPConfig(mode="cd", n_probes=4, eps=0.1, probe_chunk=4))
x = torch.randint(0, 101, (2, 32)); y = torch.randint(0, 101, (2, 32))
nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()).backward()
assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in m.parameters())
|
