tiny-cube-value-iter / code /test_policy_distill.py

code: detached probe, so every version can be measured

e61febb verified about 1 month ago

7.52 kB

"""
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")

	"""
	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")