| """The update test: honest training passes, empty uploads and noise don't, and what is tested is what goes in.""" |
| import random |
| import threading |
| import time |
| import unittest |
| from types import SimpleNamespace |
| from unittest import mock |
|
|
| import torch |
| import dbmoe |
|
|
|
|
| CFG = dbmoe.make_cfg(dict(d=32, heads=2, layers_per_block=1, blocks=2, experts=2, ffn=64, ctx=32)) |
| TEXT = torch.tensor(list(b'the quick brown fox jumps over the lazy dog. ' * 400), dtype=torch.uint8) |
|
|
|
|
| def windows(n, seed): |
| start = torch.randint(0, TEXT.numel() - CFG['ctx'] - 1, (n,), generator=torch.Generator().manual_seed(seed)) |
| return torch.stack([TEXT[i:i + CFG['ctx'] + 1] for i in start.tolist()]) |
|
|
|
|
| def lease(model, block=1, steps=40, batch=4): |
| return dict(version=dbmoe.VERSION, lease='t', block=block, block_version=0, cfg=model.cfg, sigmas=model.sigmas, |
| shared={k: v.clone() for k, v in model.shared.state_dict().items()}, |
| block_state={k: v.clone() for k, v in model.blocks[block].state_dict().items()}, |
| steps=steps, batch=batch, lr=3e-3, ttl=60, data=windows(steps * batch, 1)) |
|
|
|
|
| def measure(model, L, d_shared, d_block, margin=0.002): |
| b = L['block'] |
| return dbmoe.gate_measure(model.cfg, model.sigmas, windows(8, 2), b, 0, L['shared'], L['block_state'], L['shared'] and model.blocks[0].state_dict(), |
| d_shared, d_block, 1.0, 1.0, L['data'][:6], margin) |
|
|
|
|
| class UpdateTestTests(unittest.TestCase): |
| def setUp(self): |
| torch.manual_seed(0) |
| self.model = dbmoe.Model(CFG) |
| self.L = lease(self.model) |
|
|
| def test_honest_training_learns_from_its_own_text(self): |
| up = dbmoe.train_lease(self.L, 'cpu') |
| m = measure(self.model, self.L, up['d_shared'], up['d_block']) |
| self.assertLess(m['trained_after'], m['trained_before']) |
| self.assertIsNotNone(m['before']) |
| self.assertIsNotNone(m['other_after']) |
|
|
| def test_empty_upload_stops_at_the_first_stage(self): |
| zero = lambda sd: {k: torch.zeros_like(v) for k, v in sd.items()} |
| m = measure(self.model, self.L, zero(self.L['shared']), zero(self.L['block_state'])) |
| self.assertEqual(m['trained_after'], m['trained_before']) |
| self.assertIsNone(m['before']) |
|
|
| def test_noise_stops_at_the_first_stage(self): |
| noise = lambda sd: {k: torch.randint(0, 2, v.shape).float() * 0.1 - 0.05 for k, v in sd.items()} |
| m = measure(self.model, self.L, noise(self.L['shared']), noise(self.L['block_state'])) |
| self.assertGreater(m['trained_after'], m['trained_before']) |
| self.assertIsNone(m['before']) |
|
|
| def test_not_a_number_stops_at_the_first_stage(self): |
| nan = lambda sd: {k: torch.full_like(v, float('nan')) for k, v in sd.items()} |
| m = measure(self.model, self.L, nan(self.L['shared']), nan(self.L['block_state'])) |
| self.assertIsNone(m['before']) |
|
|
| def test_a_harmful_round_measures_worse_and_a_real_one_does_not(self): |
| noise = lambda sd: {k: torch.randint(0, 2, v.shape).float() * 0.1 - 0.05 for k, v in sd.items()} |
| zero = lambda sd: {k: torch.zeros_like(v) for k, v in sd.items()} |
| x, blocks = windows(8, 2), {1: self.L['block_state']} |
| bad = dbmoe.round_measure(self.model.cfg, self.model.sigmas, x, self.L['shared'], blocks, noise(self.L['shared'])) |
| self.assertGreater(bad[1][1], bad[1][0]) |
| same = dbmoe.round_measure(self.model.cfg, self.model.sigmas, x, self.L['shared'], blocks, zero(self.L['shared'])) |
| self.assertEqual(same[1][0], same[1][1]) |
|
|
|
|
| class Source: |
| def __init__(self): |
| self.seed = 100 |
|
|
| def windows(self, n, T): |
| self.seed += 1 |
| return windows(n, self.seed) |
|
|
|
|
| def coordinator(seed=0): |
| random.seed(seed) |
| torch.manual_seed(seed) |
| C = dbmoe.Coordinator.__new__(dbmoe.Coordinator) |
| C.model = dbmoe.Model(CFG) |
| C.lr, C.steps, C.batch, C.ttl = 3e-3, 40, 4, 1800 |
| C.max_steps, C.max_batch, C.min_steps = 400, 64, 1 |
| C.lock = threading.RLock() |
| C.leases, C.data = {}, Source() |
| C.pending_workers = set() |
| C.gate_all, C.gate_tol, C.learn_margin, C.gate_strict = True, 0.03, 0.002, False |
| C.gate_x, C.gate_slots, C.gate_recent, C.gate_drift, C.check_windows = windows(12, 999), threading.BoundedSemaphore(2), [], {}, 6 |
| C.shared_bank, C.shared_round, C.round_busy = {}, 12, False |
| C.block_gates = dbmoe.collections.defaultdict(threading.Lock) |
| C.meta = dict(updates=0, tokens=0, block_updates=[0] * CFG['blocks'], block_ce=[None] * CFG['blocks'], workers={}, |
| evals=[], rejected=0, gate=dict(passed=0, refused=0), worker_alias_salt='test') |
| C.events, C.strikes = [], [] |
| C.log = lambda **kw: C.events.append(kw) |
| C.is_open = lambda owner: owner.startswith('open') |
| C.enrol = SimpleNamespace(strike=lambda owner, weight=1: C.strikes.append((owner, weight)), passed=lambda owner: None) |
| C.save = lambda *a, **kw: None |
| C.snap_every, C.eval_every, C.last_save, C.eval_busy, C.eval_started = 0, 0, time.time(), False, 0 |
| return C |
|
|
|
|
| def grant(C, worker='trusted/node'): |
| return C._lease(worker, 40, 4, 's' + str(random.random()), owner=worker.split('/')[0]) |
|
|
|
|
| def state(C): |
| return {k: v.clone() for k, v in C.model.state_dict().items()} |
|
|
|
|
| def zero(L): |
| return dict(lease=L['lease'], block=L['block'], losses=[1.0] * L['steps'], seconds=1.0, d_shared={k: torch.zeros_like(v) for k, v in L['shared'].items()}, |
| d_block={k: torch.zeros_like(v) for k, v in L['block_state'].items()}) |
|
|
|
|
| class CoordinatorTests(unittest.TestCase): |
| def assertSame(self, a, b): |
| for k in a: |
| self.assertTrue(torch.equal(a[k], b[k]), k) |
|
|
| def test_nothing_is_taken_on_trust_whoever_sends_it(self): |
| for worker in ('trusted/node', 'open-1/node'): |
| C = coordinator() |
| before = state(C) |
| code, body = C.update(worker, zero(grant(C, worker))) |
| self.assertEqual(code, 422, body) |
| self.assertEqual(C.meta['updates'], 0) |
| self.assertSame(before, state(C)) |
| code, body = C.update(worker, dbmoe.train_lease(grant(C, worker), 'cpu')) |
| self.assertEqual(code, 200, body) |
| self.assertEqual(C.meta['updates'], 1) |
|
|
| def test_a_block_takes_its_updates_one_at_a_time(self): |
| C = coordinator() |
| C.block_weights = lambda: [0.0001, 1.0] |
| ups = [(w, dbmoe.train_lease(grant(C, w), 'cpu')) for w in ('trusted/a', 'trusted/b')] |
| self.assertEqual({u['block'] for _, u in ups}, {1}) |
| seen, real, inside = [], dbmoe.gate_measure, threading.Semaphore(1) |
|
|
| def measure(*a, **kw): |
| self.assertTrue(inside.acquire(blocking=False), 'two tests at once on one block') |
| seen.append(C.meta['updates']) |
| time.sleep(0.2) |
| inside.release() |
| return real(*a, **kw) |
|
|
| with mock.patch.object(dbmoe, 'gate_measure', side_effect=measure): |
| ts = [threading.Thread(target=C.update, args=u) for u in ups] |
| [t.start() for t in ts] |
| [t.join() for t in ts] |
| self.assertEqual(seen, [0, 1]) |
|
|
| def test_an_update_merges_at_the_weight_it_was_tested_at(self): |
| C = coordinator() |
| L = grant(C) |
| up = dbmoe.train_lease(L, 'cpu') |
| before, real = state(C), dbmoe.gate_measure |
|
|
| def measure(*a, **kw): |
| grant(C, 'trusted/other') if L['block'] == 0 else None |
| C.leases[L['lease']]['same'] = 5 |
| return real(*a, **kw) |
|
|
| with mock.patch.object(dbmoe, 'gate_measure', side_effect=measure): |
| self.assertEqual(C.update('trusted/node', up)[0], 200) |
| k = next(iter(up['d_block'])) |
| got = C.model.blocks[L['block']].state_dict()[k] - before[f"blocks.{L['block']}.{k}"] |
| self.assertTrue(torch.allclose(got, up['d_block'][k], atol=1e-6)) |
|
|
| def test_duplicate_upload_counts_once(self): |
| C = coordinator() |
| up = dbmoe.train_lease(grant(C), 'cpu') |
| self.assertEqual(C.update('trusted/node', up)[0], 200) |
| self.assertEqual(C.update('trusted/node', up)[0], 409) |
| self.assertEqual(C.meta['updates'], 1) |
|
|
| def test_strikes_are_for_the_worker_not_the_network(self): |
| C = coordinator() |
| C.gate_strict = True |
| self.assertEqual(C.update('open-1/node', zero(grant(C, 'open-1/node')))[0], 422) |
| self.assertEqual(C.strikes, []) |
| C.gate_strict = False |
| self.assertEqual(C.update('open-1/node', zero(grant(C, 'open-1/node')))[0], 422) |
| self.assertEqual(C.strikes, [('open-1', 0.5)]) |
| L = grant(C, 'open-1/node') |
| with mock.patch.object(dbmoe, 'gate_measure', return_value=dict(trained_before=1.0, trained_after=0.9995, before=None, after=None)): |
| self.assertEqual(C.update('open-1/node', zero(L))[0], 422) |
| self.assertEqual(len(C.strikes), 1) |
| L = grant(C, 'open-1/node') |
| junk = {k: torch.randint(0, 2, v.shape).float() * 0.8 - 0.4 for k, v in L['block_state'].items()} |
| self.assertEqual(C.update('open-1/node', dict(zero(L), d_block=junk))[0], 422) |
| self.assertEqual(C.strikes[-1], ('open-1', 1)) |
|
|
| def test_a_round_goes_in_only_if_it_passes_its_own_test(self): |
| C = coordinator() |
| L = grant(C) |
| noise = {k: torch.randint(0, 2, v.shape).float() * 8 - 4 for k, v in L['shared'].items()} |
| before = state(C) |
| C._commit_round({L['block']: dict(n=1, sum=noise)}) |
| self.assertEqual(C.meta['gate']['rounds_refused'], 1) |
| self.assertEqual(C.meta.get('shared_rounds', 0), 0) |
| self.assertSame(before, state(C)) |
| self.assertFalse(C.round_busy) |
| up = dbmoe.train_lease(L, 'cpu') |
| C._commit_round({L['block']: dict(n=1, sum=up['d_shared'])}) |
| self.assertEqual(C.meta['shared_rounds'], 1) |
| k = next(iter(up['d_shared'])) |
| got = C.model.shared.state_dict()[k] - before['shared.' + k] |
| self.assertTrue(torch.allclose(got, up['d_shared'][k] * 0.5, atol=1e-6)) |
| self.assertFalse(any(g.locked() for g in C.block_gates.values())) |
|
|
| def test_helpful_shared_part_cannot_hide_a_harmful_block_part(self): |
| C = coordinator() |
| real = dbmoe.train_lease(grant(C), 'cpu') |
| L = C.leases[real['lease']] |
| x = C.gate_x |
| before = {k: v.clone() for k, v in C.model.blocks[L['block']].state_dict().items()} |
| bad = {k: torch.randint(0, 2, v.shape).float() * 0.2 - 0.1 for k, v in before.items()} |
| sh = {k: v.clone() for k, v in C.model.shared.state_dict().items()} |
| zero_shared = {k: torch.zeros_like(v) for k, v in sh.items()} |
| m = dbmoe.gate_measure(C.model.cfg, C.model.sigmas, x, L['block'], None, sh, before, None, real['d_shared'], real['d_block'], 0.5, 1.0, L['check'], -1.0) |
| joint = dbmoe.gate_measure(C.model.cfg, C.model.sigmas, x, L['block'], None, sh, before, None, zero_shared, real['d_block'], 0.0, 1.0, L['check'], -1.0) |
| self.assertEqual(m['after'], joint['after']) |
| harm = dbmoe.gate_measure(C.model.cfg, C.model.sigmas, x, L['block'], None, sh, before, None, real['d_shared'], bad, 0.5, 1.0, L['check'], -1.0) |
| self.assertGreater(harm['after'], harm['before']) |
| self.assertLessEqual(harm['shared_after'], harm['after']) |
|
|
| def test_a_full_identity_gets_told_to_retry_and_nobody_loses_a_job(self): |
| C = coordinator() |
| C.is_open = lambda owner: True |
| C.pending_leases, C.max_open_leases, C.max_identity_leases, C.max_active_leases, C.reserved_leases = {}, 2, 2, 64, 8 |
| C.open_lease_interval, C.last_open_lease = 0.0, {} |
| a = C.lease('open-1/a', 2, 1, 's', owner='open-1')['lease'] |
| b = C.lease('open-1/b', 2, 1, 's', owner='open-1')['lease'] |
| with self.assertRaises(dbmoe.LeaseCapacityError): |
| C.lease('open-1/c', 2, 1, 's', owner='open-1') |
| self.assertEqual(set(C.leases), {a, b}) |
| again = C.lease('open-1/a', 2, 1, 's', owner='open-1')['lease'] |
| self.assertEqual(set(C.leases), {again, b}) |
|
|
| def test_a_restarted_machine_can_reclaim_its_old_job_at_the_cap(self): |
| C = coordinator() |
| C.is_open = lambda owner: True |
| C.pending_leases, C.max_open_leases, C.max_identity_leases, C.max_active_leases, C.reserved_leases = {}, 2, 2, 64, 8 |
| C.open_lease_interval, C.last_open_lease = 0.0, {} |
| a = C.lease('open-1/a', 2, 1, 'before', owner='open-1')['lease'] |
| b = C.lease('open-1/b', 2, 1, 'before', owner='open-1')['lease'] |
| again = C.lease('open-1/a', 2, 1, 'after', owner='open-1')['lease'] |
| self.assertEqual(set(C.leases), {again, b}) |
| self.assertNotIn(a, C.leases) |
|
|
| def test_replacing_a_job_does_not_need_room_for_a_new_one(self): |
| C = coordinator() |
| C.is_open = lambda owner: owner.startswith('open') |
| C.pending_leases, C.max_open_leases, C.max_identity_leases, C.max_active_leases, C.reserved_leases = {}, 8, 8, 4, 2 |
| C.open_lease_interval, C.last_open_lease = 0.0, {} |
| C.lease('open-1/a', 2, 1, 's', owner='open-1') |
| C.lease('open-1/b', 2, 1, 's', owner='open-1') |
| with self.assertRaises(dbmoe.LeaseCapacityError): |
| C.lease('open-1/c', 2, 1, 's', owner='open-1') |
| C.lease('trusted/x', 2, 1, 's', owner='trusted') |
| C.lease('trusted/y', 2, 1, 's', owner='trusted') |
| self.assertEqual(len(C.leases), 4) |
| C.lease('open-1/a', 2, 1, 's', owner='open-1') |
| self.assertEqual(len(C.leases), 4) |
|
|
| def test_a_restart_at_the_global_cap_keeps_its_old_job_unless_it_gets_a_new_one(self): |
| C = coordinator() |
| C.is_open = lambda owner: owner.startswith('open') |
| C.pending_leases, C.max_open_leases, C.max_identity_leases, C.max_active_leases, C.reserved_leases = {}, 8, 8, 3, 1 |
| C.open_lease_interval, C.last_open_lease = 0.0, {} |
| a = C.lease('open-1/a', 2, 1, 'before', owner='open-1')['lease'] |
| C.lease('trusted/x', 2, 1, 's', owner='trusted') |
| C.lease('trusted/y', 2, 1, 's', owner='trusted') |
| with self.assertRaises(dbmoe.LeaseCapacityError): |
| C.lease('open-2/z', 2, 1, 's', owner='open-2') |
| self.assertIn(a, C.leases) |
| again = C.lease('open-1/a', 2, 1, 'after', owner='open-1')['lease'] |
| self.assertNotIn(a, C.leases) |
| self.assertIn(again, C.leases) |
| self.assertEqual(len(C.leases), 3) |
|
|
| def test_an_upload_must_account_for_every_step_of_its_job(self): |
| C = coordinator() |
| for change in (dict(losses=[1.0]), dict(tokens=1), dict(tokens=True)): |
| L = grant(C) |
| up = dbmoe.train_lease(L, 'cpu') |
| before = state(C) |
| code, body = C.update('trusted/node', dict(up, **change)) |
| self.assertEqual(code, 422, body) |
| self.assertEqual(C.meta['updates'], 0) |
| self.assertSame(before, state(C)) |
|
|
| def test_compute_follows_what_is_left_to_learn(self): |
| C = coordinator() |
| C.meta['evals'] = [dict(per_block_ce=[0.02, 2.0])] |
| got = [0, 0] |
| for i in range(120): |
| L = C._lease(f'trusted/w{i}', 2, 1, 's', owner='trusted') |
| got[L['block']] += 1 |
| C.leases.pop(L['lease']) |
| self.assertGreater(got[1], got[0] * 3) |
| self.assertGreater(got[0], 0) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
