File size: 7,521 Bytes
0bcae1d e61febb 0bcae1d e61febb 0bcae1d | 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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | """
Smoke tests for train_policy_distill.py, runnable on CPU with no checkpoint.
The point is not coverage. It is that a rented GPU box must not be the place
where a shape error or a misaligned target index shows up: that costs metered
time, and a *silently* misaligned target costs a whole run that looks healthy.
So the alignment between the language-model sequence and the rollout is checked
against the string engine in gen_data, which is a completely separate
implementation from the numpy one in cube_ops.
python test_policy_distill.py
"""
import sys, tempfile
from pathlib import Path
import numpy as np
import torch
sys.path.insert(0, str(Path(__file__).parent))
import cube_ops as C
import cube_tokenizer as T
import train_policy_distill as D
from gen_data import SOLVED, apply_sequence
from train_value_iteration import ValueNet
def make_teacher(tmp):
"""A randomly initialised value net, saved in the value-iteration format."""
net = ValueNet(32, 2, 2)
path = Path(tmp) / "teacher.pt"
torch.save({"model": net.state_dict(), "hidden": 32, "layers": 2, "heads": 2,
"step": 0}, path)
return D.load_teacher(str(path), "cpu")
def test_token_tables():
assert len(D.ANSWER_IDS) == 19
assert D.MOVE_IDS == [T.STOI[m] for m in T.MOVES]
assert D.ANSWER_IDS[D.EOS_SLOT] == T.EOS
# The move a target index names must be the move cube_ops applies for it.
s = C.encode(SOLVED)
for i, mv in enumerate(T.MOVES):
arr = C.decode(C.apply_moves(s, [i]))[0]
assert arr == apply_sequence(SOLVED, [mv]), mv
print("ok token tables agree with both engines")
def test_pack_alignment(net):
"""Every masked position must carry the state its prefix actually reaches."""
rng = np.random.default_rng(0)
states0, _ = C.random_walk(6, 12, rng)
batch = D.rollout_batch(net, states0, "cpu", 8, 0.35, 0.0, rng, amp=False)
ids = batch["ids"].numpy()
faces = batch["faces"].numpy()
mask = batch["posmask"].numpy()
order = batch["order"].numpy()
for b in range(len(states0)):
start = C.decode(states0[b:b + 1])[0]
assert list(ids[b, :T.STATE_LEN + 2]) == T.encode_state(start)
idx = np.where(mask[b])[0]
assert list(order[b, idx]) == list(range(len(idx)))
for t, pos in enumerate(idx):
# The prefix that reaches this position is the prompt plus t moves.
prefix = ids[b, D.PROMPT_LEN:D.PROMPT_LEN + t]
moves = [T.ITOS[int(i)] for i in prefix]
assert apply_sequence(start, moves) == C.decode(faces[b, pos:pos + 1])[0], (b, t)
print(f"ok pack alignment on {len(states0)} rollouts, string engine agrees")
def test_targets(net):
rng = np.random.default_rng(1)
states0, _ = C.random_walk(32, 6, rng)
batch = D.rollout_batch(net, states0, "cpu", 10, 0.35, 0.0, rng, amp=False)
soft = batch["soft"].numpy()
mask = batch["posmask"].numpy()
faces = batch["faces"].numpy()
# float32 before summing: the packed target is float16, so a float16
# accumulation over 19 entries has more error than the assertion allows.
rows = soft[mask].astype(np.float32)
assert np.allclose(rows.sum(axis=1), 1.0, atol=1e-3), rows.sum(axis=1)[:5]
# EOS carries all the mass exactly where the state is solved, and none anywhere else.
solved = C.is_solved(faces[mask])
assert np.allclose(rows[solved][:, D.EOS_SLOT], 1.0)
assert np.allclose(rows[~solved][:, D.EOS_SLOT], 0.0)
assert (soft[~mask] == 0).all()
print(f"ok targets normalised; EOS mass matches solved on {mask.sum()} positions")
def test_hard_targets(net):
rng = np.random.default_rng(2)
states0, _ = C.random_walk(8, 6, rng)
batch = D.rollout_batch(net, states0, "cpu", 6, 0.35, 0.0, rng, hard=True, amp=False)
rows = batch["soft"].numpy()[batch["posmask"].numpy()]
assert np.allclose(rows.max(axis=1), 1.0)
print("ok --hard-targets produces one-hot rows")
def test_train_step(net):
class Args:
hidden, layers, heads = 32, 2, 2
model = D.build_policy(Args())
head = D.StateHead(32)
rng = np.random.default_rng(3)
states0, _ = C.random_walk(8, 10, rng)
batch = D.rollout_batch(net, states0, "cpu", 8, 0.35, 0.2, rng, amp=False)
loss, mloss, floss, facc = D.losses(model, head, batch, "cpu", 1.0,
torch.bfloat16, False)
assert torch.isfinite(loss), loss
loss.backward()
grads = [q.grad for q in head.parameters() if q.grad is not None]
assert grads and any(g.abs().sum() > 0 for g in grads)
# The facelet probe must reach the trunk, not just its own linear layer.
trunk = model.model.layers[0].self_attn.q_proj.weight.grad
assert trunk is not None and trunk.abs().sum() > 0
print(f"ok one training step: move {mloss.item():.3f} face {floss.item():.3f} "
f"acc {facc.item():.1%}, gradients reach the trunk")
return model, head
def test_probe_is_detached(net):
"""The probe must train itself and leave the model it measures alone."""
class Args:
hidden, layers, heads = 32, 2, 2
model = D.build_policy(Args())
head, probe = D.StateHead(32), D.StateHead(32)
rng = np.random.default_rng(11)
states0, _ = C.random_walk(8, 8, rng)
batch = D.rollout_batch(net, states0, "cpu", 6, 0.35, 0.0, rng, amp=False)
# state_weight 0 removes the attached head, so anything reaching the trunk
# after that comes from the move loss alone -- never from the probe.
loss, _, _, _ = D.losses(model, head, batch, "cpu", 0.0, torch.bfloat16, False,
probe)
loss.backward()
assert all(q.grad is not None and q.grad.abs().sum() > 0 for q in probe.parameters())
assert all(q.grad is None or q.grad.abs().sum() == 0 for q in head.parameters())
print("ok the probe trains itself and does not reach the attached head")
def test_eval(net, model, head):
rates = D.student_greedy(model, "cpu", [3], 8, np.random.default_rng(4))
assert set(rates) == {3} and 0.0 <= rates[3] <= 1.0
ceil = D.teacher_greedy(net, "cpu", [3], 8, np.random.default_rng(5), max_steps=4,
amp=False)
assert 0.0 <= ceil[3] <= 1.0
per_pos, exact = D.state_accuracy(model, head, net, "cpu", [3], 8,
np.random.default_rng(6), 6, 0.35, False)
assert "0:" in per_pos and 0.0 <= exact <= 1.0
print(f"ok eval paths run: student {rates[3]:.0%}, teacher {ceil[3]:.0%}, "
f"state {per_pos}")
def test_solved_start(net):
"""A start state that is already solved must produce EOS and nothing else."""
states0 = C.SOLVED_ARR[None, :].copy()
batch = D.rollout_batch(net, states0, "cpu", 8, 0.35, 0.0,
np.random.default_rng(7), amp=False)
mask = batch["posmask"].numpy()[0]
assert mask.sum() == 1
assert batch["soft"].numpy()[0, mask][0, D.EOS_SLOT] == 1.0
print("ok a solved start state yields exactly one EOS position")
if __name__ == "__main__":
torch.manual_seed(0)
with tempfile.TemporaryDirectory() as tmp:
teacher = make_teacher(tmp)
test_token_tables()
test_pack_alignment(teacher)
test_targets(teacher)
test_hard_targets(teacher)
test_solved_start(teacher)
test_probe_is_detached(teacher)
model, head = test_train_step(teacher)
test_eval(teacher, model, head)
print("\nall smoke tests passed")
|