File size: 8,385 Bytes
e64c5d2 | 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 | """Narrow mutation/inheritance tests; no experimental datasets are read."""
from pathlib import Path
from tempfile import TemporaryDirectory
import pytest
import torch
from lightnest import rv32
from lightnest.morph_patch import (
MorphConfig, MorphPatchNetwork, MorphPredictor, compare_functions,
encode_batch, load_morph, route_statistics, state_fingerprint, supervised_loss,
)
from lightnest.relation_patch import RelationConfig, RelationPatchNetwork
@pytest.fixture(autouse=True)
def small_cpu_budget():
torch.set_num_threads(2)
def binary(words):
return b''.join(word.to_bytes(4, 'little') for word in words)
def fixtures():
first = binary([rv32.lw('x1', 0, 'x0'), rv32.lw('x2', 4, 'x0'),
rv32.add('x3', 'x1', 'x2'), rv32.addi('x3', 'x3', 0),
rv32.addi('x10', 'x3', 0), rv32.EBREAK])
short = binary([rv32.addi('x10', 'x0', 8), rv32.EBREAK])
return ['change add to sub', 'do not change sub to add'], [first, short]
def test_one_expert_exactly_inherits_relation_without_rng_or_storage_alias():
torch.manual_seed(1801)
parent = RelationPatchNetwork(RelationConfig(width=64)).eval()
state = torch.get_rng_state().clone()
child, report = MorphPatchNetwork.from_relation(parent, seed=81)
assert torch.equal(state, torch.get_rng_state())
assert child.parameter_count == 113138
assert report['all_copied_tensors_exact'] and report['newly_initialized_parameter_values'] == 0
assert report['parameter_name_map']['experts.0.weight_ih'] == 'update.weight_ih'
assert child.experts[0].weight_ih.data_ptr() != parent.update.weight_ih.data_ptr()
result = compare_functions(parent, child, *fixtures(), atol=0, rtol=0)
assert result['passed'] and result['scores_max_abs_diff'] == result['bits_max_abs_diff'] == 0
def test_selected_single_leaf_splits_preserve_function_and_old_tensor_names():
torch.manual_seed(13)
parent = MorphPatchNetwork().eval()
commands, codes = fixtures()
for index, expected in zip((0, 0, 1), (138227, 163316, 188405)):
before = state_fingerprint(parent)
rng = torch.get_rng_state().clone()
child, report = parent.spawn_child('split', expert_index=index, seed=271 + index)
assert child.parameter_count == expected
assert report['parameter_name_map'][f'experts.{parent.config.experts}.weight_hh'] == f'experts.{index}.weight_hh'
assert report['newly_initialized_parameter_values'] == 129
assert report['function_preserving_expected']
assert torch.equal(torch.get_rng_state(), rng)
assert state_fingerprint(parent) == before
assert compare_functions(parent, child, commands, codes)['passed']
assert child.experts[-1].weight_ih.data_ptr() != parent.experts[index].weight_ih.data_ptr()
parent = child
assert parent.config.router_children == ((1, 2), (-1, -3), (-2, -4))
with pytest.raises(ValueError, match='fewer than four'):
parent.spawn_child('split', expert_index=0)
def test_new_state_gates_break_gradient_symmetry_without_initial_output_noise():
torch.manual_seed(31)
one = MorphPatchNetwork(MorphConfig(width=32))
two, _ = one.spawn_child('split', expert_index=0, seed=79)
commands, codes = fixtures()
encoded = encode_batch(commands, codes, two.config)
scores, bits, diagnostics = two(*encoded, return_diagnostics=True)
initial = route_statistics(diagnostics, two)
assert initial['route_input_dependence_nats'] > 0
assert initial['expert_parameter_rms_distance'][0][1] == 0
assert initial['expert_output_rms_distance'][0][1] == 0
rows = [{'edits': [{'offset': 11, 'value': codes[0][11] ^ 64}]}, {'edits': []}]
optimizer = torch.optim.SGD(two.parameters(), lr=0.1)
loss, _, _ = supervised_loss(scores, bits, rows, two.none_index)
loss.backward()
difference = (two.experts[0].weight_ih.grad - two.experts[1].weight_ih.grad).abs().max()
assert float(difference) > 1e-8
optimizer.step()
_, _, after = two(*encoded, return_diagnostics=True)
measured = route_statistics(after, two)
assert measured['expert_parameter_rms_distance'][0][1] > 0
assert measured['expert_output_rms_distance'][0][1] > 0
optimizer.zero_grad(set_to_none=True)
scores, bits = two(*encoded)
supervised_loss(scores, bits, rows, two.none_index)[0].backward()
assert float(two.routers[0].weight.grad.abs().max()) > 1e-10
def test_padding_diagnostics_uniform_ablation_and_deeper_weight_inheritance():
torch.manual_seed(43)
model = MorphPatchNetwork(MorphConfig(width=32, experts=3)).eval()
commands, codes = fixtures()
encoded = encode_batch(commands, codes, model.config)
before = state_fingerprint(model)
with torch.inference_mode():
scores, bits, diag = model(*encoded, return_diagnostics=True)
uniform_scores, _, uniform = model(*encoded, return_diagnostics=True, routing='uniform')
solo = model(*encode_batch(commands[1:], codes[1:], model.config))
mask = diag['valid'][None].expand(diag['route_weights'].shape[:3])
torch.testing.assert_close(diag['route_weights'].sum(-1)[mask], torch.ones(int(mask.sum())))
assert int(torch.count_nonzero(diag['route_weights'][~mask])) == 0
torch.testing.assert_close(uniform['route_weights'][mask], torch.full((int(mask.sum()), 3), 1 / 3))
assert float((scores - uniform_scores).abs().max()) > 1e-7
assert state_fingerprint(model) == before
for together, single in zip((scores, bits), solo):
torch.testing.assert_close(together[1], single[0], atol=2e-6, rtol=2e-5)
deeper, report = model.spawn_child('deepen', message_steps=8, seed=2)
assert not report['function_preserving_expected']
assert state_fingerprint(deeper) == before and deeper.parameter_count == model.parameter_count
assert not compare_functions(model, deeper, commands, codes, atol=0, rtol=0)['passed']
assert compare_functions(model, deeper, commands, codes, message_steps=4, atol=0, rtol=0)['passed']
@pytest.mark.parametrize('tree', [((0, -1),), ((-1, -1),), ((-1, -3),), ((1, -2), (-1, -2))])
def test_invalid_router_topology_rejected(tree):
with pytest.raises(ValueError):
MorphConfig(experts=len(tree) + 1, router_children=tree)
def test_checkpoint_roundtrip_preserves_lineage_and_rejects_bad_tensors():
torch.manual_seed(59)
model = MorphPatchNetwork(MorphConfig(width=32))
model, _ = model.spawn_child('split', expert_index=0, seed=3)
model, _ = model.spawn_child('split', expert_index=1, seed=5)
with TemporaryDirectory(prefix='lightnest-morph-test-') as directory:
path = Path(directory) / 'model.pt'
model.save(path, purpose='mechanism test only')
restored = load_morph(path)
assert restored.config == model.config and restored.lineage == model.lineage
assert state_fingerprint(restored) == state_fingerprint(model)
assert compare_functions(model, restored, *fixtures(), atol=0, rtol=0)['passed']
saved = torch.load(path, weights_only=True)
saved['state_dict']['routers.0.weight'][0, 0] = float('nan')
torch.save(saved, path)
with pytest.raises(ValueError, match='tensor'):
load_morph(path)
def test_predictor_byte_precondition_and_stop_use_common_runner():
from lightnest.iterative_patch import run_iterative_patch
class FixedNetwork:
config = MorphConfig()
none_index = 1024
def eval(self):
return self
def __call__(self, command, code, relations, message_steps=None, routing='learned'):
scores = torch.full((1, 1025), -40.)
bits = torch.full((1, 1024, 8), -40.)
position = 1024 if int(code[0, 3]) - 1 == 64 else 3
scores[0, position] = 40.
bits[0, 3, 6] = 40.
return scores, bits
code = binary([rv32.add('x3', 'x1', 'x2'), rv32.addi('x10', 'x3', 0), rv32.EBREAK])
predictor = MorphPredictor(FixedNetwork(), routing='uniform')
decision = predictor('change add to sub', code)
assert decision.edits[0].expected == code[3:4]
result = run_iterative_patch('change add to sub', code, predictor, max_steps=2)
assert result.final[3] == 64 and result.final[:3] == code[:3] and result.final[4:] == code[4:]
assert result.stop_reason == 'predictor_stop'
|