| import copy |
| import json |
| from pathlib import Path |
| import random |
| import tempfile |
| import unittest |
|
|
| import torch |
|
|
| from scripts import verify_expanded_startup as audit |
|
|
|
|
| def predictions(): |
| return [{'id': f'{family}-{i}', 'group': f'{family}-{i}', 'family': family, |
| 'target': 0, 'choices': ['yes', 'no'], 'logits': [0.0, 0.0], |
| 'probabilities': [0.5, 0.5], 'log_probabilities': [-0.693147, -0.693147]} |
| for family in ('arc', 'banking', 'boolq', 'snli') for i in range(128)] |
|
|
|
|
| class StartupVerificationTests(unittest.TestCase): |
| def test_initial_checkpoint_proves_empty_optimizer_and_parent_weights(self): |
| parent = {'trainable_state': {'head': torch.tensor([1.0])}} |
| pilot = {'config': {'seed': 433}, 'initialization': {'restores_rng': False}, |
| 'data_signature': 'expanded', 'source_commit': 'source', 'model_provenance': 'base'} |
| initial = {**copy.deepcopy(pilot), 'step': 0, 'optimizer': {'state': {}}, |
| 'trainable_state': copy.deepcopy(parent['trainable_state']), |
| 'random_state': random.Random(433).getstate(), |
| 'torch_rng': torch.tensor([1]), 'cuda_rng': [torch.tensor([2])]} |
| self.assertTrue(audit.verify_initial_state(initial, pilot, parent, torch)['adam_state_empty']) |
| for kind in ('weights', 'optimizer', 'rng'): |
| changed = copy.deepcopy(initial) |
| if kind == 'weights': changed['trainable_state']['head'][0] += 1 |
| elif kind == 'optimizer': changed['optimizer']['state'][0] = {'step': 1500} |
| else: changed['random_state'] = random.Random(432).getstate() |
| with self.subTest(kind=kind), self.assertRaises(ValueError): |
| audit.verify_initial_state(changed, pilot, parent, torch) |
|
|
| def test_prediction_comparison_uses_identity_not_order(self): |
| rows = predictions() |
| proof = audit.prediction_parity(rows, list(reversed(copy.deepcopy(rows)))) |
| self.assertEqual(proof['rows'], 512) |
| self.assertTrue(all(v == 0 for v in proof['max_abs_difference'].values())) |
|
|
| def test_inexact_predictions_changed_identity_and_duplicates_rejected(self): |
| rows = predictions() |
| for key, value in [('logits', [1e-15, 0.0]), ('target', 1), |
| ('choices', ['no', 'yes']), ('group', 'other'), |
| ('id', rows[1]['id']), ('probabilities', [0.5])]: |
| with self.subTest(key=key): |
| changed = copy.deepcopy(rows) |
| changed[0][key] = value |
| with self.assertRaises(ValueError): |
| audit.prediction_parity(rows, changed) |
|
|
| def test_prediction_nan_and_partial_cut_rejected(self): |
| rows = predictions() |
| with self.assertRaisesRegex(ValueError, '512'): |
| audit.prediction_parity(rows, rows[:-1]) |
| rows[0]['logits'][0] = float('nan') |
| with self.assertRaisesRegex(ValueError, 'Non-finite'): |
| audit.prediction_parity(rows, rows) |
|
|
| def test_recursive_state_catches_adam_rng_and_dtype_changes(self): |
| state = {'optimizer': {0: {'exp_avg': torch.tensor([1.0]), 'step': torch.tensor(8.0)}}, |
| 'rng': [torch.tensor([1, 2], dtype=torch.uint8)], 'python': (3, (1, 2), None)} |
| audit.exact_state(state, copy.deepcopy(state), torch) |
| for key in ('optimizer', 'rng', 'dtype'): |
| changed = copy.deepcopy(state) |
| if key == 'optimizer': changed[key][0]['exp_avg'][0] += 1 |
| elif key == 'rng': changed[key][0][0] += 1 |
| else: changed['rng'][0] = changed['rng'][0].long() |
| with self.subTest(key=key), self.assertRaises(ValueError): |
| audit.exact_state(state, changed, torch) |
|
|
| def test_adam_requires_all_parameter_states_at_pilot_step(self): |
| artifact = {'optimizer': {'param_groups': [{'params': [0, 1]}], 'state': { |
| i: {'step': torch.tensor(8.0), 'exp_avg': torch.zeros(2), 'exp_avg_sq': torch.ones(2)} |
| for i in (0, 1)}}} |
| self.assertEqual(audit.optimizer_steps(artifact, 8, torch)['parameter_states'], 2) |
| for kind in ('old_step', 'missing', 'nonfinite'): |
| changed = copy.deepcopy(artifact) |
| if kind == 'old_step': changed['optimizer']['state'][0]['step'] = 1508 |
| elif kind == 'missing': del changed['optimizer']['state'][1] |
| else: changed['optimizer']['state'][1]['exp_avg'][0] = float('inf') |
| with self.subTest(kind=kind), self.assertRaises(ValueError): |
| audit.optimizer_steps(changed, 8, torch) |
|
|
| def test_finite_updates_need_contiguous_post_resume_rows_and_cap(self): |
| with tempfile.TemporaryDirectory() as directory: |
| path = Path(directory)/'training.jsonl' |
| rows = [{'step': n, 'loss': 0.2, 'gradient_norm': 0.1, |
| 'peak_cuda_allocated_bytes': audit.CAP} for n in range(9, 13)] |
| def write(values): path.write_text(''.join(json.dumps(r)+'\n' for r in values)) |
| write(rows) |
| self.assertEqual(audit.finite_updates(path, 8, 4)['last_step'], 12) |
| with path.open('a') as f: f.write('{"step":') |
| self.assertEqual(audit.finite_updates(path, 8, 4)['count'], 4) |
| for kind in ('nan', 'cap', 'gap', 'insufficient'): |
| changed = copy.deepcopy(rows) |
| if kind == 'nan': changed[0]['gradient_norm'] = float('nan') |
| elif kind == 'cap': changed[0]['peak_cuda_allocated_bytes'] += 1 |
| elif kind == 'gap': changed[0]['step'] = 7 |
| else: changed.pop() |
| write(changed) |
| with self.subTest(kind=kind), self.assertRaises(ValueError): |
| audit.finite_updates(path, 8, 4) |
|
|
| def test_checkpoint_reader_is_cpu_only_and_rejects_nonfinite_weights(self): |
| with tempfile.TemporaryDirectory() as directory: |
| path = Path(directory)/'checkpoint.pt' |
| artifact = {'format': 'opensysone-adapter-v1', 'trainable_state': {'head': torch.ones(2)}} |
| torch.save(artifact, path) |
| loaded, digest = audit.load_checkpoint(path, torch) |
| self.assertEqual(loaded['trainable_state']['head'].device.type, 'cpu') |
| self.assertEqual(digest, audit.sha256(path)) |
| artifact['trainable_state']['head'][0] = float('nan') |
| torch.save(artifact, path) |
| with self.assertRaisesRegex(ValueError, 'tensor'): |
| audit.load_checkpoint(path, torch) |
| self.assertFalse(torch.cuda.is_initialized()) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
