| """Mechanism tests use hand-written programs only, never experiment/final data.""" |
| from dataclasses import replace |
| from pathlib import Path |
| from tempfile import TemporaryDirectory |
|
|
| import pytest |
| import torch |
|
|
| from lightnest import rv32 |
| from lightnest.relation_patch import RelationConfig, RelationPatchNetwork |
| from lightnest.transformer_patch import ( |
| VARIANTS, TransformerConfig, TransformerPatchNetwork, TransformerPredictor, |
| attention_statistics, encode_batch, load_transformer, state_fingerprint, |
| supervised_loss, |
| ) |
|
|
|
|
| @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 inherited(variant='standard', seed=43, *, width=32): |
| torch.manual_seed(7711) |
| parent = RelationPatchNetwork(RelationConfig(width=width)) |
| return TransformerPatchNetwork.from_relation( |
| parent, config=TransformerConfig(variant=variant, width=width, ff_width=2 * width), seed=seed) |
|
|
|
|
| def test_exact_common_inheritance_parameter_counts_roles_named_initialization_and_rng(): |
| torch.manual_seed(7711) |
| parent = RelationPatchNetwork(RelationConfig()).eval() |
| rng = torch.get_rng_state().clone() |
| children = [TransformerPatchNetwork.from_relation(parent, config=TransformerConfig(variant=v), seed=43) |
| for v in VARIANTS] |
| assert torch.equal(torch.get_rng_state(), rng) |
| counts = dict(standard=146930, relation_bias=146946, reverse_query=146930, |
| shared_recurrent=113458, differential=147122, centered_attention=146938) |
| reference = dict(children[0][0].named_parameters()) |
| for child, report in children: |
| assert child.parameter_count == report['parameters'] == counts[child.config.variant] |
| assert report['copied_parameter_values'] == 79858 |
| assert report['newly_initialized_parameter_values'] == child.parameter_count - 79858 |
| assert report['all_copied_tensors_exact'] and not report['function_preserving'] |
| parameters = dict(child.named_parameters()) |
| assert set(report['parameter_name_map']) == set(parameters) == set(child.parameter_roles()) |
| assert sum(child.parameter_counts_by_role().values()) == child.parameter_count |
| for name, source in report['parameter_name_map'].items(): |
| if source is not None: |
| old = dict(parent.named_parameters())[source] |
| assert torch.equal(parameters[name], old) |
| assert parameters[name].data_ptr() != old.data_ptr() |
| if name in reference: |
| assert torch.equal(parameters[name], reference[name]), name |
| roles = child.parameter_roles() |
| assert roles['command_conv1.weight'] == roles['command_conv2.weight'] == 'auxiliary' |
| assert roles['command_embedding.weight'] == 'embedding' |
| assert roles['pointer.weight'] == 'output' |
| assert roles['blocks.0.attention.q_proj.weight'] == 'hidden_matrix' |
| assert all(parameters[name].ndim == 2 for name, role in roles.items() if role == 'hidden_matrix') |
| repeated, _ = TransformerPatchNetwork.from_relation(parent, config=child.config, seed=43) |
| assert state_fingerprint(repeated) == state_fingerprint(child) |
| different, _ = TransformerPatchNetwork.from_relation(parent, config=child.config, seed=44) |
| assert not torch.equal(different.blocks[0].attention.q_proj.weight, child.blocks[0].attention.q_proj.weight) |
| assert torch.equal(different.command_fc.weight, child.command_fc.weight) |
|
|
|
|
| def test_standard_attention_is_scaled_qk_softmax_and_separate_qkv(): |
| model, _ = inherited() |
| encoded = encode_batch(*fixtures(), model.config) |
| _, _, diag = model(*encoded, return_diagnostics=True) |
| block = model.blocks[0] |
| h = block.norm1(diag['initial_states']) |
| batch, length, width = h.shape |
| heads, dim = model.config.heads, width // model.config.heads |
| q = block.attention.q_proj(h).reshape(batch, length, heads, dim).transpose(1, 2) |
| k = block.attention.k_proj(h).reshape(batch, length, heads, dim).transpose(1, 2) |
| logits = q @ k.transpose(-1, -2) / dim**.5 |
| expected = logits.masked_fill(~diag['valid'][:, None, None, :], -torch.inf).softmax(-1) |
| expected = expected.masked_fill(~diag['valid'][:, None, :, None], 0.) |
| actual = diag['attention'][0] |
| torch.testing.assert_close(actual['qk_logits'], logits, rtol=0, atol=0) |
| torch.testing.assert_close(actual['probabilities'], expected, rtol=0, atol=0) |
| pointers = [getattr(block.attention, name).weight.data_ptr() for name in ('q_proj', 'k_proj', 'v_proj')] |
| assert len(set(pointers)) == 3 |
| assert not any(isinstance(module, torch.nn.GRUCell) for module in model.modules()) |
|
|
|
|
| @pytest.mark.parametrize('variant', VARIANTS) |
| def test_all_variants_have_finite_supervised_gradients_and_masked_probabilities(variant): |
| model, _ = inherited(variant) |
| commands, codes = fixtures() |
| scores, bits, diag = model(*encode_batch(commands, codes, model.config), return_diagnostics=True) |
| assert scores.shape == (2, 1025) and bits.shape == (2, 1024, 8) |
| assert bool(torch.isfinite(scores).all() and torch.isfinite(bits).all()) |
| assert bool((scores[1, len(codes[1]):model.none_index] == -1e9).all()) |
| assert int(torch.count_nonzero(bits[1, len(codes[1]):])) == 0 |
| for stage in diag['attention']: |
| for field in ('probabilities', 'secondary_probabilities'): |
| if field not in stage: |
| continue |
| weights = stage[field] |
| expected = stage['query_valid'][:, None].expand(-1, weights.shape[1], -1).to(weights.dtype) |
| torch.testing.assert_close(weights.sum(-1), expected, atol=1e-6, rtol=1e-6) |
| assert int(torch.count_nonzero(weights.masked_select(~stage['key_valid'][:, None, None, :]))) == 0 |
| rows = [{'edits': [{'offset': 11, 'value': codes[0][11] ^ 64}]}, {'edits': []}] |
| loss = supervised_loss(scores, bits, rows, model.none_index)[0] |
| loss.backward() |
| for block in model.blocks: |
| for layer in (block.attention.q_proj, block.attention.k_proj, block.attention.v_proj, |
| block.attention.out_proj, block.ff_in, block.ff_out): |
| grad = layer.weight.grad |
| assert grad is not None and bool(torch.isfinite(grad).all()) |
| assert float(grad.abs().max()) > 1e-10 |
| assert all(value.grad is None or bool(torch.isfinite(value.grad).all()) for value in model.parameters()) |
| extra = dict(relation_bias='relation_bias', differential='lambda_q1', |
| centered_attention='suppression_logits').get(variant) |
| if extra is not None: |
| assert float(getattr(model.blocks[0].attention, extra).grad.abs().max()) > 1e-10 |
| stats = attention_statistics(diag) |
| assert len(stats['stages']) == model.config.steps |
|
|
|
|
| @pytest.mark.parametrize('variant', VARIANTS) |
| def test_padding_batch_invariance_and_no_edges_outside_explicit_bias_arm(variant): |
| model, _ = inherited(variant) |
| model.eval() |
| commands, codes = fixtures() |
| encoded = encode_batch(commands, codes, model.config) |
| before = state_fingerprint(model) |
| with torch.inference_mode(): |
| together = model(*encoded) |
| solo = model(*encode_batch(commands[1:], codes[1:], model.config)) |
| n = encoded[2].shape[-1] |
| extended_command = torch.nn.functional.pad(encoded[0], (0, 3)) |
| extended_binary = torch.nn.functional.pad(encoded[1], (0, 8)) |
| extended_edges = torch.nn.functional.pad(encoded[2], (0, 2, 0, 2)) |
| extended = model(extended_command, extended_binary, extended_edges) |
| if variant != 'relation_bias': |
| altered = model(encoded[0], encoded[1], torch.rand_like(encoded[2])) |
| for original, changed in zip(together, altered): |
| torch.testing.assert_close(original, changed, atol=0, rtol=0) |
| for combined, single, padded in zip(together, solo, extended): |
| torch.testing.assert_close(combined[1], single[0], atol=2e-6, rtol=2e-5) |
| torch.testing.assert_close(combined, padded, atol=2e-6, rtol=2e-5) |
| assert state_fingerprint(model) == before |
|
|
|
|
| def test_direct_relation_bias_zero_is_standard_and_nonzero_adds_only_direct_edges(): |
| standard, _ = inherited('standard') |
| biased, _ = inherited('relation_bias') |
| encoded = encode_batch(*fixtures(), standard.config) |
| with torch.inference_mode(): |
| for original, zero_bias in zip(standard(*encoded), biased(*encoded)): |
| torch.testing.assert_close(original, zero_bias, atol=0, rtol=0) |
| biased.blocks[0].attention.relation_bias[0].fill_(2.5) |
| _, _, diag = biased(*encoded, return_diagnostics=True) |
| stage = diag['attention'][0] |
| edges = stage['typed_direct_edges'] |
| assert int(torch.count_nonzero(edges.diagonal(dim1=-2, dim2=-1))) == 0 |
| assert float(edges[0, 0, 4, 3]) == 1. |
| assert float(edges[0, 0, 4, 2]) == 0. |
| expected = 2.5 * edges[:, 0, None].expand(-1, biased.config.heads, -1, -1) |
| torch.testing.assert_close(stage['biased_logits'] - stage['qk_logits'], expected, atol=1e-6, rtol=1e-6) |
|
|
|
|
| def test_shared_block_really_reuses_weights_and_reverse_query_is_last_valid_token(): |
| commands, codes = fixtures() |
| shared, _ = inherited('shared_recurrent') |
| calls = [] |
| hook = shared.blocks[0].register_forward_hook(lambda module, args, result: calls.append(id(module))) |
| with torch.inference_mode(): |
| default = shared(*encode_batch(commands, codes, shared.config), return_diagnostics=True) |
| single = shared(*encode_batch(commands, codes, shared.config), steps=1) |
| hook.remove() |
| assert len(shared.blocks) == 1 and len(calls) == 5 and len(set(calls)) == 1 |
| assert default[2]['compute']['executed_blocks'] == 4 |
| assert not torch.equal(default[0], single[0]) |
| reverse, _ = inherited('reverse_query') |
| with torch.inference_mode(): |
| _, _, diag = reverse(*encode_batch(commands, codes, reverse.config), return_diagnostics=True) |
| assert diag['return_query_index'].tolist() == [5, 1] |
| torch.testing.assert_close(diag['return_query_initial'], diag['stage_states'][0][torch.arange(2), [5, 1]]) |
| assert diag['attention'][0]['probabilities'].shape == (2, 4, 6, 6) |
| assert diag['attention'][1]['probabilities'].shape == (2, 4, 1, 6) |
| assert diag['compute']['dense_attention_score_elements'] == 2 * 4 * (6 * 6 + 6) |
| assert default[2]['compute']['dense_attention_score_elements'] == 4 * 2 * 4 * 6 * 6 |
|
|
|
|
| def test_differential_and_centered_subtractions_report_signed_weights_truthfully(): |
| differential, _ = inherited('differential') |
| with torch.no_grad(): |
| differential.blocks[0].attention.lambda_q1.fill_(.5) |
| differential.blocks[0].attention.lambda_k1.fill_(.5) |
| differential.blocks[0].attention.lambda_q2.zero_() |
| differential.blocks[0].attention.lambda_k2.zero_() |
| with torch.inference_mode(): |
| _, _, diag = differential(*encode_batch(*fixtures(), differential.config), return_diagnostics=True) |
| stage = diag['attention'][0] |
| assert stage['effective_heads'] == 2 and stage['softmax_streams'] == 2 |
| assert stage['probabilities'].shape == (2, 2, 6, 6) |
| torch.testing.assert_close(stage['effective_weights'], stage['probabilities'] - |
| stage['coefficient'] * stage['secondary_probabilities'], atol=0, rtol=0) |
| assert float(stage['coefficient']) > 1 and bool((stage['effective_weights'] < 0).any()) |
| stats = attention_statistics(diag) |
| assert min(stats['stages'][0]['effective_negative_mass_per_head']) > 0 |
| assert max(stats['stages'][0]['effective_row_sum_per_head']) < 0 |
| centered, _ = inherited('centered_attention') |
| with torch.inference_mode(): |
| _, _, diag = centered(*encode_batch(*fixtures(), centered.config), return_diagnostics=True) |
| stage = diag['attention'][0] |
| uniform = stage['key_valid'].float() / stage['key_valid'].sum(1, keepdim=True) |
| expected = stage['probabilities'] - .5 * uniform[:, None, None] |
| expected = expected.masked_fill(~stage['query_valid'][:, None, :, None], 0.) |
| torch.testing.assert_close(stage['effective_weights'], expected, atol=0, rtol=0) |
|
|
|
|
| @pytest.mark.parametrize('variant', VARIANTS) |
| def test_checkpoint_roundtrip_keeps_tensor_inventory_config_roles_and_rng(variant): |
| model, _ = inherited(variant) |
| with TemporaryDirectory(prefix='lightnest-transformer-test-') as directory: |
| for suffix in ('.pt', '.pt.gz'): |
| path = Path(directory) / ('model' + suffix) |
| model.save(path, purpose='mechanism test only') |
| rng = torch.get_rng_state().clone() |
| loaded = load_transformer(path) |
| assert torch.equal(torch.get_rng_state(), rng) |
| assert loaded.config == model.config |
| assert loaded.initialization_report == model.initialization_report |
| assert loaded.parameter_roles() == model.parameter_roles() |
| assert state_fingerprint(loaded) == state_fingerprint(model) |
| with torch.inference_mode(): |
| encoded = encode_batch(*fixtures(), model.config) |
| for actual, expected in zip(loaded(*encoded), model(*encoded)): |
| torch.testing.assert_close(actual, expected, atol=0, rtol=0) |
|
|
|
|
| def test_loader_rejects_nonfinite_or_wrong_shape_and_predictor_uses_expected_byte(): |
| from lightnest.iterative_patch import run_iterative_patch |
| model, _ = inherited() |
| with TemporaryDirectory(prefix='lightnest-transformer-invalid-') as directory: |
| path = Path(directory) / 'invalid.pt' |
| model.save(path) |
| saved = torch.load(path, weights_only=True) |
| saved['state_dict']['blocks.0.attention.q_proj.weight'][0, 0] = float('nan') |
| torch.save(saved, path) |
| with pytest.raises(ValueError, match='tensor'): |
| load_transformer(path) |
| model.save(path) |
| saved = torch.load(path, weights_only=True) |
| saved['state_dict']['blocks.0.attention.q_proj.weight'] = torch.zeros(1) |
| torch.save(saved, path) |
| with pytest.raises(ValueError, match='tensor'): |
| load_transformer(path) |
|
|
| class FixedNetwork: |
| config = TransformerConfig() |
| none_index = 1024 |
|
|
| def eval(self): |
| return self |
|
|
| def __call__(self, command, code, relations): |
| 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 = TransformerPredictor(FixedNetwork()) |
| 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' |
|
|
|
|
| def test_declared_capacity_covers_512_byte_program_and_invalid_config_rejected(): |
| model, _ = inherited() |
| code = binary([rv32.addi('x3', 'x3', 0)] * 126 + [rv32.addi('x10', 'x3', 0), rv32.EBREAK]) |
| assert len(code) == 512 |
| with torch.inference_mode(): |
| scores, bits = model(*encode_batch(['change add to sub'], [code], model.config)) |
| assert scores.shape == (1, 1025) and bits.shape == (1, 1024, 8) |
| assert bool(torch.isfinite(scores).all() and torch.isfinite(bits).all()) |
| with pytest.raises(ValueError, match='even'): |
| TransformerConfig(variant='differential', heads=1) |
| with pytest.raises(ValueError, match='divide'): |
| TransformerConfig(heads=3) |
| with pytest.raises(ValueError, match='query'): |
| TransformerConfig(variant='reverse_query', steps=3) |
| with pytest.raises(ValueError, match='runtime'): |
| model(*encode_batch(*fixtures(), model.config), steps=4) |
|
|
