114 lines
5.7 KiB
Python
114 lines
5.7 KiB
Python
import json
|
|
import io
|
|
import pathlib
|
|
import struct
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
|
|
from pilot.analyze import (build_blocks, exact_pair, object_record, read_object,
|
|
sample_object, sharing_groups)
|
|
from pilot.tensors import decode, load_tensors, metrics
|
|
|
|
|
|
class PilotTests(unittest.TestCase):
|
|
def test_directional_replacement_error_and_scaling(self):
|
|
x = np.array([1., 2., 3.])
|
|
m = metrics(x, 2 * x)
|
|
self.assertAlmostEqual(m['cosine'], 1.)
|
|
self.assertAlmostEqual(m['error_member_to_rep'], 1.)
|
|
self.assertAlmostEqual(m['error_rep_to_member'], .5)
|
|
self.assertAlmostEqual(m['scaled_error'], 0.)
|
|
self.assertAlmostEqual(m['scale'], .5)
|
|
|
|
def test_bfloat16(self):
|
|
raw = np.array([0x3f80, 0xc000, 0x0000], dtype='<u2')
|
|
np.testing.assert_array_equal(decode(raw, 'BF16'), [1., -2., 0.])
|
|
|
|
def test_ranged_download_preserves_prefix(self):
|
|
from pilot.download import ranged_download
|
|
data = b'the entire known checkpoint contents'
|
|
class Response(io.BytesIO):
|
|
status = 206
|
|
def request(req, timeout):
|
|
span = req.headers['Range'].split('=')[1]
|
|
a, b = map(int, span.split('-'))
|
|
response = Response(data[a:b + 1])
|
|
response.headers = {'Content-Range': f'bytes {a}-{b}/{len(data)}'}
|
|
return response
|
|
with tempfile.TemporaryDirectory() as td:
|
|
part = pathlib.Path(td) / 'file.partial'
|
|
part.write_bytes(data[:7])
|
|
with patch('pilot.download.urllib.request.urlopen', side_effect=request):
|
|
ranged_download('https://example.invalid/file', part, len(data), workers=2)
|
|
self.assertEqual(part.read_bytes(), data)
|
|
self.assertFalse(part.with_suffix('.partial.ranges.json').exists())
|
|
|
|
def test_no_transitive_chaining(self):
|
|
objects = [{'id': str(i)} for i in range(3)]
|
|
pairs = [{'a': 0, 'b': 1, 'error_member_to_rep': .01, 'error_rep_to_member': .5},
|
|
{'a': 1, 'b': 2, 'error_member_to_rep': .01, 'error_rep_to_member': .5}]
|
|
groups = sharing_groups(objects, pairs, .02)
|
|
self.assertFalse(any(set(g['members']) == {0, 1, 2} for g in groups))
|
|
self.assertTrue(all(g['max_error'] <= .02 for g in groups))
|
|
|
|
def test_3d_reader_and_exact_block(self):
|
|
x = np.arange(8 * 4 * 4, dtype='<f4').reshape(8, 4, 4)
|
|
names = ['model.layers.0.conv.weight', 'model.layers.1.conv.weight']
|
|
header = {name: {'dtype': 'F32', 'shape': list(x.shape),
|
|
'data_offsets': [i * x.nbytes, (i + 1) * x.nbytes]}
|
|
for i, name in enumerate(names)}
|
|
encoded = json.dumps(header).encode()
|
|
with tempfile.TemporaryDirectory() as td:
|
|
path = pathlib.Path(td)
|
|
(path / 'model.safetensors').write_bytes(struct.pack('<Q', len(encoded)) + encoded + x.tobytes() * 2)
|
|
tensors = load_tensors(path)
|
|
a = object_record(tensors[names[0]], 'block', (4, 2, 2), (4, 2, 2), '3d')
|
|
b = object_record(tensors[names[1]], 'block', (4, 2, 2), (4, 2, 2), '3d')
|
|
np.testing.assert_array_equal(read_object(a, tensors), x[4:8, 2:4, 2:4].ravel())
|
|
np.testing.assert_array_equal(sample_object(a, tensors), read_object(a, tensors))
|
|
self.assertEqual(exact_pair(a, b, tensors, float('inf'))['error_member_to_rep'], 0.)
|
|
blocks, coverage, totals = build_blocks(tensors, 1, float('inf'), 2)
|
|
self.assertTrue(blocks)
|
|
self.assertTrue(all(r['sampled_blocks'] <= 2 for r in coverage))
|
|
self.assertTrue(all(r['bytes'] + r['edge_bytes'] == x.nbytes * 2 for r in totals.values()))
|
|
|
|
def test_end_to_end_synthetic_groups_and_savings(self):
|
|
from pilot.analyze import main
|
|
arrays = {}
|
|
for layer in (0, 1):
|
|
factor = 1 + layer * .005
|
|
arrays[f'model.layers.{layer}.proj.weight'] = (np.arange(64 * 64, dtype='<f4').reshape(64, 64) + 1) * factor
|
|
arrays[f'model.layers.{layer}.conv.weight'] = (np.arange(8 * 4 * 4, dtype='<f4').reshape(8, 4, 4) + 1) * factor
|
|
header, payload, offset = {}, [], 0
|
|
for name, x in arrays.items():
|
|
header[name] = {'dtype': 'F32', 'shape': list(x.shape), 'data_offsets': [offset, offset + x.nbytes]}
|
|
payload.append(x.tobytes())
|
|
offset += x.nbytes
|
|
encoded = json.dumps(header).encode()
|
|
with tempfile.TemporaryDirectory() as td:
|
|
root = pathlib.Path(td)
|
|
(root / 'weights').mkdir()
|
|
(root / 'weights/model.safetensors').write_bytes(struct.pack('<Q', len(encoded)) + encoded + b''.join(payload))
|
|
(root / 'manifest.json').write_text(json.dumps({'repo': 'synthetic', 'revision': 'test'}))
|
|
with patch('sys.argv', ['analyze', '--root', str(root), '--seconds', '60']):
|
|
main()
|
|
summary = json.loads((root / 'analysis/summary.json').read_text())
|
|
self.assertTrue(summary['completed'])
|
|
groups = json.loads((root / 'analysis/groups_tensor_whole_0.01.json').read_text())
|
|
self.assertEqual(len(groups), 2)
|
|
self.assertTrue(all(len(g['members']) == 2 for g in groups))
|
|
self.assertTrue(all(g['max_error'] < .01 for g in groups))
|
|
from pilot.plot import main as plot_main
|
|
from pilot.validate import validate
|
|
with patch('sys.argv', ['plot', '--root', str(root)]):
|
|
plot_main()
|
|
self.assertTrue(validate(root)['passed'])
|
|
self.assertEqual(len(list((root / 'figures').glob('*.pdf'))), 6)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|