import argparse import copy import importlib import json import math from pathlib import Path import unittest from run_lab import experiment parser = argparse.ArgumentParser() parser.add_argument('--reference', action='store_true') parser.add_argument('--topic', choices=['rlhf', 'lora', 'grpo']) args, rest = parser.parse_known_args() m = importlib.import_module('reference' if args.reference else 'starter') fixtures = json.loads(Path(__file__).with_name('fixtures.json').read_text()) class RLHF(unittest.TestCase): def test_margin_and_reference(self): self.assertAlmostEqual(m.dpo_loss(math.log(4), 0, 1), math.log(1.25)) self.assertAlmostEqual(m.dpo_loss(2, 2, 0.5), math.log(2)) def test_stability(self): self.assertAlmostEqual(m.dpo_loss(-1000, 0, 1), 1000) self.assertAlmostEqual(m.dpo_loss(1000, 0, 1), 0) def test_invalid(self): for values in [(0, 0, 0), (0, 0, -1), (float('nan'), 0, 1), (0, float('inf'), 1)]: with self.assertRaises(ValueError): m.dpo_loss(*values) def test_gradient(self): for theta in [-4, 0, 3]: delta = 1e-5 numerical = (m.dpo_loss(theta + delta, 0, 2) - m.dpo_loss(theta - delta, 0, 2)) / (2 * delta) self.assertAlmostEqual(m.dpo_gradient(theta, 2), numerical, places=8) self.assertAlmostEqual(m.dpo_gradient(-1000, 1), -1) self.assertAlmostEqual(m.dpo_gradient(1000, 1), 0) def test_update_and_counterexample(self): result = experiment('rlhf', m, fixtures) before, after = result['history'][0], result['history'][-1] self.assertLess(after['pair_loss_nats'], before['pair_loss_nats']) self.assertLess(after['independent_correct_probability'], before['independent_correct_probability']) self.assertAlmostEqual(result['history'][1]['theta'], 0.1) self.assertFalse(result['transfer_used_for_selection']) class LoRA(unittest.TestCase): def setUp(self): self.W, self.A, self.B, self.x = [[1, 0], [0, 1]], [[1, -1]], [[0.5], [-0.5]], [2, 1] def test_forward(self): self.assertEqual(m.lora_forward(self.W, self.A, self.B, self.x, 2), [3, 0]) def test_zero_adapter(self): self.assertEqual(m.lora_forward(self.W, self.A, [[0], [0]], self.x, 2), [2, 1]) def test_rectangular_rank_two(self): self.assertEqual(m.lora_forward([[1, 2, 3]], [[1, 0, 0], [0, 1, 0]], [[2, 4]], [1, 2, 3], 1), [19]) def test_merge_parity_and_no_mutation(self): before = copy.deepcopy((self.W, self.A, self.B)) merged = m.merge(self.W, self.A, self.B, 2) self.assertEqual(merged, [[2, -1], [-1, 2]]) self.assertEqual((self.W, self.A, self.B), before) self.assertIsNot(merged, self.W) self.assertEqual(merged, m.merge(self.W, self.A, self.B, 2)) for x in ([2, 1], [0.3, -0.7], [-2, 3]): expected = [sum(w * v for w, v in zip(row, x)) for row in merged] for a, b in zip(m.lora_forward(self.W, self.A, self.B, x, 2), expected): self.assertAlmostEqual(a, b) def test_rounding_and_zeros(self): codes, scale = m.quantize([-1, -0.3, 0, 0.3, 1], 4) self.assertEqual(codes, [-7, -2, 0, 2, 7]) self.assertAlmostEqual(scale, 1 / 7) self.assertEqual(m.quantize([0, 0], 4), ([0, 0], 1)) self.assertEqual(m.quantize([7, 0.5, 1.5, -0.5, -1.5], 4)[0], [7, 0, 2, 0, -2]) def test_outlier_and_invalid(self): self.assertEqual(m.quantize([-1, -0.3, 0, 0.3, 10], 4)[0][3], 0) for values, bits in [([], 4), ([float('nan')], 4), ([1], 3)]: with self.assertRaises(ValueError): m.quantize(values, bits) def test_four_way_report(self): result = experiment('lora', m, fixtures) self.assertEqual(len(result['four_way_outputs']), 4) self.assertGreater(result['weight_mse'], 0) self.assertLess(experiment('lora', m, fixtures, bits=8)['weight_mse'], result['weight_mse']) class GRPO(unittest.TestCase): def test_population_normalization(self): self.assertEqual(m.advantages([0, 0, 1, 1]), [-1, -1, 1, 1]) values = m.advantages([1, 2, 4]) self.assertAlmostEqual(sum(values), 0) self.assertAlmostEqual(sum(v*v for v in values) / 3, 1) def test_equal_and_separate_groups(self): self.assertEqual(m.advantages([5, 5]), [0, 0]) for a, b in zip(m.advantages([1, 2, 4]), m.advantages([11, 12, 14])): self.assertAlmostEqual(a, b) def test_invalid_group(self): for values in ([], [1], [0, float('nan')], [1, float('inf')]): with self.assertRaises(ValueError): m.advantages(values) def test_both_clipping_signs(self): for ratio, advantage, expected in [(1.4, 1, 1.2), (1.4, -1, -1.4), (0.6, 1, 0.6), (0.6, -1, -0.8)]: self.assertAlmostEqual(m.surrogate([ratio], advantage, [1]), expected) def test_mask_invariance(self): self.assertAlmostEqual(m.surrogate([1, 1.4], 1, [1, 1]), 1.1) self.assertAlmostEqual(m.surrogate([1, 1.4, 99], 1, [1, 1, 0]), 1.1) self.assertEqual(m.surrogate([1.4], 0, [1]), 0) def test_invalid_masks_and_ratios(self): for ratios, mask in [([1], [0]), ([1], [1, 0]), ([0], [1]), ([float('inf')], [1]), ([1], [2])]: with self.assertRaises(ValueError): m.surrogate(ratios, 1, mask) def test_verifier_counterexample(self): base = experiment('grpo', m, fixtures, format_weight=0) exploited = experiment('grpo', m, fixtures, format_weight=3) self.assertEqual(base['correctness'], exploited['correctness']) self.assertLess(base['advantages'][1], 0) self.assertGreater(exploited['advantages'][1], 0) self.assertEqual(exploited['transfer_equal_group'], [0, 0, 0, 0]) if __name__ == '__main__': classes = {'rlhf': RLHF, 'lora': LoRA, 'grpo': GRPO} suite = unittest.TestSuite(unittest.defaultTestLoader.loadTestsFromTestCase(cls) for key, cls in classes.items() if args.topic is None or args.topic == key) result = unittest.TextTestRunner(verbosity=2).run(suite) raise SystemExit(0 if result.wasSuccessful() else 1)