| """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() |
|
|
