File size: 7,738 Bytes
c28ffff | 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """Small synthetic checkpoints only, no models/data/network or GPU."""
import hashlib
import importlib.util
import json
from pathlib import Path
import sys
import tempfile
import unittest
from unittest.mock import patch
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import reshard_weights as R
def fixture(root):
src = root / 'source'
src.mkdir()
mapping, total = {}, 0
for part in range(2):
filename = f'model-{part+1:05d}-of-00002.safetensors'
data, info = [], []
for j in range(3):
name = f'layer{part}.{j}.weight'
blob = bytes([part*10+j]) * 64
data.append(blob)
info.append(dict(name=name, dtype='BF16', shape=[32], nbytes=64))
mapping[name] = filename
total += 64
(src / filename).write_bytes(R.output_header(info, {'format': 'pt'}) + b''.join(data))
(src / 'model.safetensors.index.json').write_text(json.dumps(
{'metadata': {'total_size': total}, 'weight_map': mapping}))
(src / 'config.json').write_text('{"model_type":"test"}')
(src / 'merges.txt').write_text('tokenizer fixture\n')
return src
class ShardingTests(unittest.TestCase):
def setUp(self):
self.temporary = tempfile.TemporaryDirectory()
self.addCleanup(self.temporary.cleanup)
self.root = Path(self.temporary.name)
self.src = fixture(self.root)
self.out = self.root / 'output'
def test_lossless_and_size_bounded(self):
before = {p.name: R.sha_file(p) for p in self.src.iterdir()}
result = R.reshard(self.src, self.out, 320, 'native test')
self.assertTrue(result['validation']['all_tensor_bytes_SHA256_equal'])
self.assertEqual(len(result['tensor_sha256']), 6)
self.assertGreater(len(result['output_shards']), 1)
self.assertTrue(all(row['bytes'] <= 320 for row in result['output_shards']))
self.assertEqual(before, {p.name: R.sha_file(p) for p in self.src.iterdir()})
self.assertEqual((self.src / 'merges.txt').read_bytes(), (self.out / 'merges.txt').read_bytes())
for p in self.src.glob('*.safetensors'):
self.assertEqual(result['source_file_sha256'][p.name], R.sha_file(p))
self.assertNotIn('Trick Prediction', (self.out / 'README.md').read_text())
def test_single_native_safetensors_and_prefix(self):
source = self.src / 'model-00001-of-00002.safetensors'
result = R.reshard(source, self.out, 320, 'Wan test', prefix='diffusion_pytorch_model')
self.assertEqual(result['spec']['source']['type'], 'single_safetensors')
self.assertTrue((self.out / 'diffusion_pytorch_model.safetensors.index.json').is_file())
self.assertFalse((self.out / 'config.json').exists())
def test_resumption_after_readback_interruption(self):
with patch.object(R, 'verify_outputs', side_effect=RuntimeError('interrupted')):
with self.assertRaisesRegex(RuntimeError, 'interrupted'):
R.reshard(self.src, self.out, 320)
self.assertTrue((self.out / '.reshard_state.json').exists())
result = R.reshard(self.src, self.out, 320, resume=True)
self.assertFalse((self.out / '.reshard_state.json').exists())
self.assertEqual(len(result['tensor_sha256']), 6)
def test_too_large_tensor_is_not_split(self):
with self.assertRaisesRegex(ValueError, 'single tensor exceeds'):
R.reshard(self.src, self.out, 64)
def test_corruption_rejected_on_completed_export(self):
result = R.reshard(self.src, self.out, 320)
shard = self.out / result['output_shards'][0]['name']
data = bytearray(shard.read_bytes())
data[-1] ^= 1
shard.write_bytes(data)
with self.assertRaisesRegex(ValueError, 'tensor bytes differ'):
R.reshard(self.src, self.out, 320, resume=True)
def test_unrelated_directory_not_overwritten(self):
self.out.mkdir()
(self.out / 'user_file').write_text('keep me')
with self.assertRaisesRegex(ValueError, 'not empty'):
R.reshard(self.src, self.out, 320)
self.assertEqual((self.out / 'user_file').read_text(), 'keep me')
def test_sidecar_corruption_rejected(self):
R.reshard(self.src, self.out, 320)
(self.out / 'config.json').write_text('changed')
with self.assertRaisesRegex(ValueError, 'sidecar changed'):
R.reshard(self.src, self.out, 320, resume=True)
def test_unsafe_index_path_rejected(self):
path = self.src / 'model.safetensors.index.json'
value = json.loads(path.read_text())
value['weight_map']['layer0.0.weight'] = '../outside.safetensors'
path.write_text(json.dumps(value))
with self.assertRaisesRegex(ValueError, 'unsafe index'):
R.load_source(self.src)
def test_source_output_overlap_rejected(self):
with self.assertRaisesRegex(ValueError, 'overlap'):
R.reshard(self.src, self.src / 'output', 320)
def test_duplicate_json_key_rejected(self):
with self.assertRaisesRegex(ValueError, 'duplicate JSON key'):
R.parse_json('{"a":1,"a":2}')
@unittest.skipUnless(importlib.util.find_spec('torch'), 'torch not installed; stdlib sharding tests still run')
class TorchConversionTests(unittest.TestCase):
def setUp(self):
import torch
self.torch = torch
self.temporary = tempfile.TemporaryDirectory()
self.addCleanup(self.temporary.cleanup)
self.root = Path(self.temporary.name)
(self.root / 'source').mkdir()
self.source = self.root / 'source/native.pth'
self.out = self.root / 'output'
def test_fake_metadata_direct_bytes_equal_native_reader(self):
torch = self.torch
from safetensors import safe_open
tensors = {'embedding.weight': torch.arange(96, dtype=torch.bfloat16).reshape(12, 8),
'norm.weight': torch.tensor([1., 2., 3., 4.]),
'counter': torch.tensor(5, dtype=torch.int64)}
torch.save(tensors, self.source)
original = R.sha_file(self.source)
result = R.reshard(self.source, self.out, 320, 'test T5')
self.assertEqual(result['spec']['source']['type'], 'torch_zip_state_dict')
self.assertEqual(result['source_file_sha256'][self.source.name], original)
self.assertEqual(R.sha_file(self.source), original)
seen = set()
for shard in result['output_shards']:
with safe_open(self.out / shard['name'], framework='pt', device='cpu') as handle:
for name in handle.keys():
self.assertTrue(torch.equal(tensors[name], handle.get_tensor(name)))
seen.add(name)
self.assertEqual(seen, set(tensors))
self.assertFalse(torch.cuda.is_initialized())
def test_noncontiguous_tensor_rejected(self):
torch = self.torch
torch.save({'tensor': torch.arange(12).reshape(3, 4).T}, self.source)
with self.assertRaisesRegex(ValueError, 'non-contiguous'):
R.reshard(self.source, self.out, 1024)
def test_shared_storage_rejected(self):
torch = self.torch
tensor = torch.arange(12)
torch.save({'a': tensor, 'b': tensor}, self.source)
with self.assertRaisesRegex(ValueError, 'shared storages'):
R.reshard(self.source, self.out, 1024)
def test_non_tensor_entries_rejected(self):
self.torch.save({'a': self.torch.arange(12), 'trainer': 'not a weight'}, self.source)
with self.assertRaisesRegex(ValueError, 'non-tensor entry'):
R.reshard(self.source, self.out, 1024)
if __name__ == '__main__':
unittest.main()
|