| """Numerical contract only. These tests do not run Swift, Metal or an Apple app.""" |
| import importlib.util |
| import json |
| import math |
| from pathlib import Path |
| import unittest |
|
|
| import torch |
|
|
| ROOT = Path(__file__).resolve().parent |
| SPEC = importlib.util.spec_from_file_location('native_reference', ROOT / 'apple/generate_parity_fixture.py') |
| GENERATOR = importlib.util.module_from_spec(SPEC) |
| SPEC.loader.exec_module(GENERATOR) |
|
|
|
|
| class NativeNumericalReferenceTests(unittest.TestCase): |
| @classmethod |
| def setUpClass(cls): |
| cls.fixture = json.loads((ROOT / 'apple/parity.json').read_text()) |
| cls.actual = GENERATOR.generate(ROOT / 'dbmoe.py') |
|
|
| def test_fixture_contains_only_finite_json(self): |
| json.dumps(self.fixture, allow_nan=False) |
| self.assertEqual(self.fixture['format'], 'dbmoe.native.parity.v1') |
|
|
| def test_fixture_has_explicit_noise_and_both_experts(self): |
| self.assertEqual(self.fixture['batch'], self.actual['batch']) |
| self.assertEqual(self.fixture['sigma'], self.actual['sigma']) |
| self.assertEqual(self.fixture['epsilon'], self.actual['epsilon']) |
| self.assertEqual(self.fixture['cfg'], self.actual['cfg']) |
| self.assertEqual(self.fixture['expected']['expert_fractions'][0]['data'], [.5, .5]) |
|
|
| def test_losses_gradients_and_adamw_step_replay(self): |
| expected, actual = self.fixture['expected'], self.actual['expected'] |
| for name in ('loss', 'ce', 'auxiliary', 'gradient_norm_before_clip'): |
| self.assertTrue(math.isclose(expected[name], actual[name], rel_tol=1e-5, abs_tol=1e-6), name) |
| for section in ('initial',): |
| self.assertEqual(self.fixture[section], self.actual[section]) |
| for section in ('gradients', 'clipped_gradients', 'after_step', 'deltas'): |
| self.assertEqual(set(expected[section]), set(actual[section])) |
| for name, ref in expected[section].items(): |
| got = actual[section][name] |
| if ref is None: |
| self.assertIsNone(got, name) |
| continue |
| self.assertEqual(ref['shape'], got['shape'], name) |
| self.assertEqual(ref['dtype'], got['dtype'], name) |
| torch.testing.assert_close(torch.tensor(ref['data']), torch.tensor(got['data']), |
| rtol=1e-5, atol=1e-6, msg=section + ': ' + name) |
| for name in ('denoised', 'logits'): |
| torch.testing.assert_close(torch.tensor(expected[name]['data']), torch.tensor(actual[name]['data']), |
| rtol=1e-5, atol=1e-6, msg=name) |
|
|
| def test_attention_and_clipping_are_exercised(self): |
| expected = self.fixture['expected'] |
| self.assertGreater(expected['gradient_norm_before_clip'], 1.) |
| qkv = expected['gradients']['block.layers.0.qkv.weight']['data'] |
| self.assertGreater(max(abs(x) for x in qkv), 1e-8) |
| deltas = expected['deltas']['block.layers.0.qkv.weight']['data'] |
| self.assertGreater(max(abs(x) for x in deltas), 1e-8) |
|
|
| def test_optional_skewed_fixture_exercises_unequal_usage(self): |
| result = GENERATOR.generate(ROOT / 'dbmoe.py', skewed=True) |
| expected = result['expected'] |
| self.assertEqual(sorted(expected['expert_fractions'][0]['data']), [.25, .75]) |
| self.assertGreater(abs(expected['auxiliary'] - 1.), 1e-3) |
| self.assertAlmostEqual(expected['loss'], 3.6444308757781982, places=5) |
| json.dumps(result, allow_nan=False) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
