File size: 3,382 Bytes
c466f90 | 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 | """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):
@classmethod
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()
|