| """Scientific controls for CPU budgeting, paired supervision and saved state. |
| |
| Fixtures are tiny synthetic RV32 programs; no experiment or sealed data is read. |
| """ |
| from copy import deepcopy |
| import gzip |
| import hashlib |
| import io |
|
|
| import pytest |
| import torch |
|
|
| from lightnest import rv32 |
| from lightnest.relation_patch import RelationConfig, RelationPatchNetwork |
| from experiments.dependency_train import ( |
| ARMS, DEFAULT_SELECTION, PairedBatchStream, build_predictor, json_bytes, |
| load_checkpoint, make_initial_model, promotion_decision, public_path, |
| relation_objectives, run_cpu_budget, save_checkpoint, tensor_inventory, |
| training_step, |
| ) |
|
|
|
|
| @pytest.fixture(autouse=True) |
| def bounded_threads(): |
| torch.set_num_threads(2) |
|
|
|
|
| def assemble(words): |
| return b''.join(word.to_bytes(4, 'little') for word in words) |
|
|
|
|
| def synthetic_rows(): |
| rows = [] |
| for origin in range(3): |
| for nuisance in (0, 1): |
| for variant in (0, 1): |
| for kind in ('edit', 'noop'): |
| binary = assemble([rv32.lw('x1', 0, 'x0'), rv32.lw('x2', 4, 'x0'), |
| rv32.add('x3', 'x1', 'x2'), rv32.add('x4', 'x1', 'x2'), |
| rv32.addi('x0', 'x0', 10 + nuisance), |
| rv32.addi('x10', 'x3' if variant == 0 else 'x4', 0), rv32.EBREAK]) |
| edits = [{'offset': 11 + 4 * variant, 'value': 64}] if kind == 'edit' else [] |
| target = bytearray(binary) |
| for edit in edits: |
| target[edit['offset']] = edit['value'] |
| row = dict(id=f'g{origin}/n{nuisance}/v{variant}/{kind}', |
| episode_id=f'g{origin}/n{nuisance}/v{variant}/{kind}', |
| source='new', origin_id=f'g{origin}', variant=variant, |
| nuisance_variant=nuisance, state_mask=0, |
| pair_id=f'g{origin}/n{nuisance}/return' if kind == 'edit' else None, |
| pair_state_key=f'g{origin}/n{nuisance}/return#0' if kind == 'edit' else None, |
| invariance_state_key=f'g{origin}/v{variant}/{kind}#0', |
| command='change returned add to sub' if kind == 'edit' else 'keep every byte', |
| binary=binary.hex(), binary_bytes=binary, target=bytes(target).hex(), edits=edits) |
| rows.append(row) |
| for index in range(8): |
| binary = assemble([rv32.addi('x10', 'x0', index), rv32.EBREAK]) |
| rows.append(dict(id=f'replay{index}', episode_id=f'replay{index}', source='replay', |
| command='keep every byte', binary=binary.hex(), binary_bytes=binary, |
| target=binary.hex(), edits=[], pair_state_key=None)) |
| return rows |
|
|
|
|
| def tiny_model(): |
| return RelationPatchNetwork(RelationConfig(width=32, binary_capacity=64, command_capacity=128)) |
|
|
|
|
| def test_alternating_streams_preserve_every_new_row_noops_and_identical_prefix(): |
| rows = synthetic_rows() |
| first, control = PairedBatchStream(rows, 104729), PairedBatchStream(rows, 104729) |
| exposure = {'return': [], 'invariance': []} |
| for _ in range(4): |
| indices = first.next_batch() |
| assert indices == control.next_batch() |
| assert first.last_units == control.last_units |
| assert len(indices) == 16 |
| assert all(rows[i]['source'] == 'new' for i in indices[:12]) |
| assert all(rows[i]['source'] == 'replay' for i in indices[12:]) |
| exposure[first.last_stream].extend(indices[:12]) |
| for unit in first.last_units: |
| a, b = [rows[indices[pos]] for pos in unit['positions']] |
| if unit['kind'] == 'return': |
| assert a['pair_state_key'] == b['pair_state_key'] == unit['key'] |
| assert {a['variant'], b['variant']} == {0, 1} |
| elif unit['kind'] == 'invariance': |
| assert a['invariance_state_key'] == b['invariance_state_key'] == unit['key'] |
| assert {a['nuisance_variant'], b['nuisance_variant']} == {0, 1} |
| else: |
| assert not a['edits'] and not b['edits'] |
| expected = [i for i, row in enumerate(rows) if row['source'] == 'new'] |
| for stream in exposure.values(): |
| assert sorted(stream) == expected |
| a, b = first.state_dict(), control.state_dict() |
| assert a == b |
|
|
|
|
| def test_sampler_resume_reconstructs_next_batches_and_explicit_units(): |
| rows = synthetic_rows() |
| stream = PairedBatchStream(rows, 29) |
| for _ in range(3): |
| stream.next_batch() |
| state = deepcopy(stream.state_dict()) |
| expected = [] |
| for _ in range(9): |
| expected.append((stream.next_batch(), deepcopy(stream.last_units), stream.last_stream)) |
| restored = PairedBatchStream(rows, 29) |
| restored.load_state_dict(state) |
| actual = [] |
| for _ in range(9): |
| actual.append((restored.next_batch(), deepcopy(restored.last_units), restored.last_stream)) |
| assert actual == expected |
| with pytest.raises(ValueError, match='different stream'): |
| PairedBatchStream(rows, 30).load_state_dict(state) |
|
|
|
|
| def test_pair_metadata_is_validated_not_guessed_from_names(): |
| rows = synthetic_rows() |
| rows[0].pop('invariance_state_key') |
| with pytest.raises(ValueError, match='invariance_state_key'): |
| PairedBatchStream(rows, 1) |
| rows = synthetic_rows() |
| rows[0]['state_mask'] = 7 |
| with pytest.raises(ValueError, match='return-choice'): |
| PairedBatchStream(rows, 1) |
|
|
|
|
| def test_return_margin_uses_exclusive_sets_two_directions_and_skips_stops(): |
| scores = torch.zeros(4, 9, requires_grad=True) |
| with torch.no_grad(): |
| scores[0, 1] = scores[1, 3] = 4. |
| scores[:2, 2] = 90. |
| group = [dict(binary_bytes=b'12345678', pair_state_key='same', edits=edits) for edits in ( |
| [{'offset': 1}, {'offset': 2}], [{'offset': 2}, {'offset': 3}], [], [])] |
| units = [dict(kind='return', positions=[0, 1], key='same'), |
| dict(kind='return', positions=[2, 3], key='same')] |
| margin, js, counts = relation_objectives(scores, group, units, 8) |
| torch.testing.assert_close(margin, torch.nn.functional.softplus(torch.tensor(-3.))) |
| assert js.item() == 0. |
| assert counts['return_candidate_pairs'] == 2 |
| assert counts['return_effective_pairs'] == 1 |
| assert counts['return_effective_directions'] == 2 |
| margin.backward() |
| assert scores.grad[0, 1] < 0 and scores.grad[0, 3] > 0 |
| assert scores.grad[1, 3] < 0 and scores.grad[1, 1] > 0 |
| assert torch.count_nonzero(scores.grad[:, 2]) == 0 |
| assert torch.count_nonzero(scores.grad[2:]) == 0 |
|
|
|
|
| def test_js_includes_stop_excludes_padding_and_is_symmetric(): |
| group = [dict(binary_bytes=b'1234', invariance_state_key='inv', edits=[]) for _ in range(2)] |
| unit = [dict(kind='invariance', positions=[0, 1], key='inv')] |
| scores = torch.zeros(2, 9, requires_grad=True) |
| with torch.no_grad(): |
| scores[0, 4:8] = 1000. |
| scores[1, 4:8] = -1000. |
| _, js, counts = relation_objectives(scores, group, unit, 8) |
| torch.testing.assert_close(js, torch.tensor(0.), atol=1e-7, rtol=0.) |
| assert counts['invariance_effective_pairs'] == 1 |
| changed = scores.detach().clone().requires_grad_(True) |
| with torch.no_grad(): |
| changed[0, 8] = 4. |
| _, unequal, _ = relation_objectives(changed, group, unit, 8) |
| _, reversed_js, _ = relation_objectives(changed.flip(0), group, unit, 8) |
| assert unequal.item() > 0. |
| torch.testing.assert_close(unequal, reversed_js) |
| unequal.backward() |
| assert changed.grad[0, 8] != 0 |
| assert torch.count_nonzero(changed.grad[:, 4:8]) == 0 |
|
|
|
|
| def test_cpu_budget_excludes_evaluation_and_logs_without_equating_updates(): |
| clock = {'cpu': 0., 'wall': 0.} |
| boundaries = [] |
|
|
| def step(): |
| clock['cpu'] += 3. |
| clock['wall'] += 2. |
| return {'loss': 1.} |
|
|
| def log(_): |
| clock['cpu'] += 11. |
| clock['wall'] += 11. |
|
|
| def checkpoint(boundary): |
| boundaries.append(boundary) |
| clock['cpu'] += 100. |
| clock['wall'] += 500. |
| return {'stage': boundary['stage']} |
|
|
| result = run_cpu_budget(step, checkpoint, cpu_seconds=10., on_step=log, |
| cpu_clock=lambda: clock['cpu'], wall_clock=lambda: clock['wall']) |
| assert result['actual_updates'] == 4 |
| assert result['training_cpu_seconds'] == 12. |
| assert result['training_wall_seconds'] == 8. |
| assert result['examples_seen'] == 64 |
| assert result['budget_overshoot_seconds'] == 2. |
| assert [row['fraction'] for row in boundaries] == [.25, .5, .75, 1.] |
| assert [row['step'] for row in boundaries] == [1, 2, 3, 4] |
| assert all(row['cpu_overshoot_seconds'] < row['crossing_step_cpu_seconds'] for row in boundaries) |
|
|
|
|
| def metric_fixture(): |
| return dict(edit_total=64, edit_strict=20, pair_total=32, pair_both_strict=5, |
| noop_total=64, noop_strict=60, false_edit_on_correct=2, |
| invariance_edit_total=32, invariance_edit_both_strict=8, |
| retention_total=128, retention_strict=100, development_loss=1.) |
|
|
|
|
| def test_promotion_requires_task_gain_and_keeps_original_retention_anchor(): |
| original = metric_fixture() |
| better = dict(original, edit_strict=21, development_loss=2., retention_strict=98) |
| assert promotion_decision(better, original, original, DEFAULT_SELECTION)['accept'] |
| |
| for candidate in (dict(original, development_loss=.01), |
| dict(original, invariance_edit_both_strict=9), |
| dict(better, invariance_edit_both_strict=7), |
| dict(better, pair_both_strict=4), |
| dict(better, false_edit_on_correct=3)): |
| assert not promotion_decision(candidate, original, original, DEFAULT_SELECTION)['accept'] |
| incumbent = dict(original, retention_strict=99) |
| candidate = dict(better, retention_strict=97) |
| decision = promotion_decision(candidate, incumbent, original, DEFAULT_SELECTION) |
| assert not decision['accept'] |
| assert decision['retention_floor'] == 98 |
| with pytest.raises(ValueError, match='denominators'): |
| promotion_decision(dict(better, edit_total=63), original, original, DEFAULT_SELECTION) |
|
|
|
|
| def test_gzip_checkpoint_replays_optimizer_sampler_and_exact_next_update(tmp_path): |
| rows = synthetic_rows() |
| torch.manual_seed(137) |
| model = tiny_model() |
| sampler = PairedBatchStream(rows, 104729) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=.001, weight_decay=.01) |
| first = training_step(model, 'relation', rows, sampler, optimizer) |
| assert first['loss'] == pytest.approx(first['base_loss'] + .2 * first['return_margin_loss'] |
| + .1 * first['invariance_js_loss'], rel=1e-6) |
| path = tmp_path / 'checkpoint.pt.gz' |
| saved = save_checkpoint(path, model, optimizer, sampler, kind='relation', metadata={'role': 'synthetic-test'}) |
| assert path.read_bytes()[4:8] == b'\0\0\0\0' |
| assert saved['model_sha256'] == tensor_inventory(model.state_dict())['sha256'] |
| restored, payload = load_checkpoint(path) |
| assert tensor_inventory(restored.state_dict()) == tensor_inventory(model.state_dict()) |
| restored_optimizer = torch.optim.AdamW(restored.parameters(), lr=.001, weight_decay=.01) |
| restored_optimizer.load_state_dict(payload['optimizer']) |
| restored_sampler = PairedBatchStream(rows, 104729) |
| restored_sampler.load_state_dict(payload['sampler']) |
| expected = training_step(model, 'relation', rows, sampler, optimizer) |
| torch.set_rng_state(payload['torch_rng_state']) |
| actual = training_step(restored, 'relation', rows, restored_sampler, restored_optimizer) |
| assert expected == actual |
| for name, value in model.state_dict().items(): |
| torch.testing.assert_close(value, restored.state_dict()[name], atol=0., rtol=0.) |
| for first_state, second_state in zip(optimizer.state.values(), restored_optimizer.state.values()): |
| for key in first_state: |
| torch.testing.assert_close(first_state[key], second_state[key], atol=0., rtol=0.) |
|
|
|
|
| def test_initial_relation_import_and_portable_predictor_are_explicit(tmp_path): |
| model = tiny_model() |
| parent = tmp_path / 'parent.pt' |
| model.save(parent) |
| for name, depth in (('relation-d4', 4), ('relation-d8', 8)): |
| child, report = make_initial_model(ARMS[name], parent, 104729) |
| assert child.config.message_steps == depth |
| assert report['new_numel'] == 0 |
| assert report['function_preserving'] == (depth == 4) |
| assert report['parameter_name_map'] == {key: key for key in model.state_dict()} |
| assert tensor_inventory(child.state_dict()) == tensor_inventory(model.state_dict()) |
| path = tmp_path / 'portable.pt.gz' |
| saved = save_checkpoint(path, model, None, None, kind='relation', metadata={}) |
| predictor = build_predictor(dict(id='tiny', kind='relation', path=str(path), sha256=saved['sha256'], message_steps=8)) |
| assert predictor.message_steps == 8 |
| assert tensor_inventory(predictor.model.state_dict()) == tensor_inventory(model.state_dict()) |
| with pytest.raises(ValueError, match='digest'): |
| build_predictor(dict(kind='relation', path=str(path), sha256='0' * 64)) |
|
|
|
|
| def test_sealed_paths_are_rejected_before_opening(): |
| with pytest.raises(ValueError, match='prohibited'): |
| public_path('work/unopened-round/sealed/final.json') |
|
|
|
|
| def test_checkpoint_internal_tensor_commitment_rejects_tampering(tmp_path): |
| path = tmp_path / 'tamper.pt.gz' |
| save_checkpoint(path, tiny_model(), None, None, kind='relation', metadata={}) |
| payload = torch.load(io.BytesIO(gzip.decompress(path.read_bytes())), weights_only=True) |
| payload['state_dict']['pointer.bias'][0] += 1. |
| buffer = io.BytesIO() |
| torch.save(payload, buffer) |
| path.write_bytes(gzip.compress(buffer.getvalue(), mtime=0)) |
| with pytest.raises(ValueError, match='digest'): |
| load_checkpoint(path) |
|
|
