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