File size: 6,654 Bytes
1a0a7fb | 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 | 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()
|