| """No base weights, network, CUDA or neural inference used by these tests.""" | |
| import importlib.util | |
| import json | |
| from pathlib import Path | |
| import struct | |
| import sys | |
| import unittest | |
| ROOT = Path(__file__).resolve().parents[1] | |
| sys.path.insert(0, str(ROOT)) | |
| sys.path.insert(0, str(ROOT / 'training')) | |
| import infer | |
| from adapters import attach, LoRALinear | |
| from objectives import group_advantages, learning_rate, policy_sum | |
| import torch | |
| from torch import nn | |
| class ReleaseTests(unittest.TestCase): | |
| def setUpClass(cls): | |
| cls.cfg = json.loads((ROOT / 'adapter_config.json').read_text()) | |
| with (ROOT / 'adapter.safetensors').open('rb') as stream: | |
| size = struct.unpack('<Q', stream.read(8))[0] | |
| cls.header = json.loads(stream.read(size)) | |
| def test_integrity(self): | |
| self.assertEqual(infer.sha256(ROOT / 'adapter.safetensors'), self.cfg['adapter_sha256']) | |
| def test_all_pairs_fp32_and_count(self): | |
| h = self.header | |
| self.assertEqual(len(h), 292) | |
| self.assertEqual(sum(entry['shape'][0] * entry['shape'][1] for entry in h.values()), 96337920) | |
| for key, entry in h.items(): | |
| self.assertEqual(entry['dtype'], 'F32') | |
| if key.endswith('.a'): | |
| pair = h[key[:-1] + 'b'] | |
| self.assertEqual(entry['shape'][0], 128) | |
| self.assertEqual(pair['shape'][1], 128) | |
| def test_custom_loader_linear_shapes(self): | |
| model = nn.Module() | |
| for key, entry in self.header.items(): | |
| if not key.endswith('.a'): | |
| continue | |
| parts = key[:-2].split('.') | |
| node = model | |
| for part in parts[:-1]: | |
| if part not in node._modules: | |
| node.add_module(part, nn.Module()) | |
| node = node._modules[part] | |
| out_features = self.header[key[:-1] + 'b']['shape'][0] | |
| node.add_module(parts[-1], nn.Linear(entry['shape'][1], out_features, bias=False, device='meta')) | |
| targets = attach(model, 128) | |
| self.assertEqual(len(targets), 146) | |
| shapes = {name: list(p.shape) for name, p in model.named_parameters() if p.requires_grad} | |
| self.assertEqual(shapes, {name: value['shape'] for name, value in self.header.items()}) | |
| def test_residual_zero_and_fp32(self): | |
| layer = LoRALinear(nn.Linear(3, 4), rank=2) | |
| x = torch.randn(5, 3) | |
| self.assertTrue(torch.equal(layer(x), layer.base(x))) | |
| self.assertEqual(layer.a.dtype, torch.float32) | |
| self.assertEqual(layer.b.dtype, torch.float32) | |
| def test_advantages_are_finite_and_masked(self): | |
| values = torch.tensor([[.6, .9, .7], [.7, .8, .8], [float('nan')]*3]) | |
| advantage = group_advantages(values, torch.tensor([True, True, False]), [.33, .33, .34]) | |
| self.assertTrue(torch.isfinite(advantage).all()) | |
| self.assertEqual(float(advantage[-1]), 0.) | |
| def test_schedule_endpoints(self): | |
| self.assertLess(learning_rate(0, 1500, 5e-6, 5e-7), 5e-6) | |
| self.assertAlmostEqual(learning_rate(1499, 1500, 5e-6, 5e-7), 5e-7) | |
| def test_vespo_gradient(self): | |
| x = torch.tensor([-.3, -.7], requires_grad=True) | |
| loss = policy_sum(x, x.detach(), torch.tensor(1.)) | |
| loss.backward() | |
| self.assertTrue(torch.isfinite(x.grad).all()) | |
| if __name__ == '__main__': | |
| unittest.main() | |