| |
| """Lossless, CPU-only safetensors sharding, including native Wan T5 .pth. |
| |
| Safetensors inputs need only the Python standard library. A PyTorch ZIP input |
| uses weights_only=True + FakeTensorMode for metadata, never unrestricted pickle |
| or model construction. Only whole, contiguous, unshared storages are accepted; |
| the actual tensor bytes are streamed directly from the read-only source file. |
| No tensor is cast, quantized, split, or passed through a model. No network calls. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| from contextlib import ExitStack |
| from datetime import datetime, timezone |
| import hashlib |
| import json |
| import math |
| import os |
| from pathlib import Path |
| import shutil |
| import struct |
| import sys |
| import time |
| import zipfile |
|
|
| BLOCK = 8 * 1024 * 1024 |
| SIZES = { |
| 'BOOL': 1, 'U8': 1, 'I8': 1, 'U16': 2, 'I16': 2, 'F16': 2, 'BF16': 2, |
| 'U32': 4, 'I32': 4, 'F32': 4, 'U64': 8, 'I64': 8, 'F64': 8, |
| 'F8_E4M3': 1, 'F8_E4M3FN': 1, 'F8_E5M2': 1, |
| } |
| SIDECARS = { |
| 'config.json', 'generation_config.json', 'tokenizer.json', |
| 'tokenizer_config.json', 'processor_config.json', 'preprocessor_config.json', |
| 'video_preprocessor_config.json', 'chat_template.jinja', 'chat_template.json', |
| 'special_tokens_map.json', 'added_tokens.json', 'merges.txt', 'vocab.json', |
| 'tokenizer.model', 'spiece.model', 'LICENSE', 'LICENSE.txt', 'LICENSE.md', |
| } |
|
|
|
|
| def now(): |
| return datetime.now(timezone.utc).isoformat() |
|
|
|
|
| def sha_file(path): |
| digest = hashlib.sha256() |
| with Path(path).open('rb') as handle: |
| for block in iter(lambda: handle.read(BLOCK), b''): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def no_duplicate_keys(pairs): |
| result = {} |
| for key, value in pairs: |
| if key in result: |
| raise ValueError('duplicate JSON key: ' + key) |
| result[key] = value |
| return result |
|
|
|
|
| def parse_json(value): |
| return json.loads(value, object_pairs_hook=no_duplicate_keys) |
|
|
|
|
| def read_header(path): |
| path = Path(path) |
| with path.open('rb') as handle: |
| prefix = handle.read(8) |
| if len(prefix) != 8: |
| raise ValueError('truncated safetensors prefix') |
| length = struct.unpack('<Q', prefix)[0] |
| if not 2 <= length <= 100 * 1024**2: |
| raise ValueError('invalid safetensors header length') |
| raw = handle.read(length) |
| if len(raw) != length: |
| raise ValueError('truncated safetensors header') |
| header = parse_json(raw) |
| if not isinstance(header, dict): |
| raise ValueError('safetensors header is not a dictionary') |
| metadata = header.get('__metadata__', {}) |
| if not isinstance(metadata, dict) or any( |
| not isinstance(k, str) or not isinstance(v, str) for k, v in metadata.items() |
| ): |
| raise ValueError('invalid safetensors metadata') |
| tensors = [] |
| for name, info in header.items(): |
| if name == '__metadata__': |
| continue |
| if not isinstance(info, dict) or set(info) != {'dtype', 'shape', 'data_offsets'}: |
| raise ValueError('invalid tensor header') |
| shape, offsets, dtype = info['shape'], info['data_offsets'], info['dtype'] |
| if dtype not in SIZES: |
| raise ValueError('unsupported dtype: ' + str(dtype)) |
| if not isinstance(shape, list) or any(type(x) is not int or x < 0 for x in shape): |
| raise ValueError('invalid tensor shape') |
| if not isinstance(offsets, list) or len(offsets) != 2 or any( |
| type(x) is not int or x < 0 for x in offsets |
| ): |
| raise ValueError('invalid tensor offsets') |
| size = offsets[1] - offsets[0] |
| if size < 0 or size != math.prod(shape) * SIZES[dtype]: |
| raise ValueError('shape/dtype/byte count differs') |
| tensors.append(dict( |
| name=name, dtype=dtype, shape=shape, nbytes=size, |
| source_offset=8 + length + offsets[0], relative_offset=offsets[0], |
| source_file=path.name, |
| )) |
| |
| tensors.sort(key=lambda item: (item['relative_offset'], item['nbytes'])) |
| cursor = 0 |
| for tensor in tensors: |
| if tensor['relative_offset'] != cursor: |
| raise ValueError('overlapping, gapped or unreferenced tensor bytes') |
| cursor += tensor['nbytes'] |
| if 8 + length + cursor != path.stat().st_size: |
| raise ValueError('file size/header coverage differs') |
| return tensors, metadata, prefix + raw |
|
|
|
|
| def source_record(path, header): |
| stat = path.stat() |
| return dict(name=path.name, bytes=stat.st_size, mtime_ns=stat.st_mtime_ns, |
| header_sha256=hashlib.sha256(header).hexdigest()) |
|
|
|
|
| def zip_payload_offset(handle, info): |
| handle.seek(info.header_offset) |
| raw = handle.read(30) |
| if len(raw) != 30: |
| raise ValueError('truncated PyTorch ZIP local header') |
| fields = struct.unpack('<IHHHHHIIIHH', raw) |
| if fields[0] != 0x04034B50 or fields[2] & 1 or fields[3] != zipfile.ZIP_STORED: |
| raise ValueError('encrypted or compressed PyTorch storage is not supported') |
| name_length, extra_length = fields[-2:] |
| return info.header_offset + 30 + name_length + extra_length |
|
|
|
|
| def load_torch_metadata(path): |
| if sys.byteorder != 'little' or not zipfile.is_zipfile(path): |
| raise ValueError('only little-endian, uncompressed PyTorch ZIP checkpoints are supported') |
| |
| os.environ['CUDA_VISIBLE_DEVICES'] = '' |
| import torch |
| from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode |
|
|
| with zipfile.ZipFile(path) as archive, path.open('rb') as raw: |
| infos = archive.infolist() |
| byteorders = [item for item in infos if item.filename.endswith('/byteorder')] |
| if len(byteorders) != 1 or archive.read(byteorders[0]) != b'little': |
| raise ValueError('checkpoint byteorder is not explicitly little-endian') |
| pickles = [item for item in infos if item.filename.endswith('/data.pkl')] |
| if len(pickles) != 1 or pickles[0].file_size > 100 * 1024**2: |
| raise ValueError('unexpected PyTorch checkpoint metadata') |
| metadata_bytes = archive.read(pickles[0]) |
| storage_ranges = {} |
| for info in infos: |
| if '/data/' not in info.filename: |
| continue |
| if info.compress_type != zipfile.ZIP_STORED or info.flag_bits & 1: |
| raise ValueError('compressed or encrypted storage') |
| offset = zip_payload_offset(raw, info) |
| if offset in storage_ranges: |
| raise ValueError('duplicate ZIP storage offset') |
| if offset + info.file_size > path.stat().st_size: |
| raise ValueError('out-of-bounds ZIP storage') |
| storage_ranges[offset] = info.file_size |
| with FakeTensorMode(): |
| state = torch.load(path, map_location='cpu', weights_only=True, mmap=True) |
| if isinstance(state, dict) and len(state) == 1: |
| for wrapper in ('state_dict', 'module', 'model_state'): |
| if wrapper in state: |
| state = state[wrapper] |
| break |
| if not isinstance(state, dict) or not state: |
| raise ValueError('only a nonempty flat tensor state_dict can be converted') |
| dtypes = { |
| torch.bool: 'BOOL', torch.uint8: 'U8', torch.int8: 'I8', |
| torch.int16: 'I16', torch.float16: 'F16', torch.bfloat16: 'BF16', |
| torch.int32: 'I32', torch.float32: 'F32', torch.int64: 'I64', |
| torch.float64: 'F64', |
| } |
| tensors, seen = [], set() |
| for name, tensor in state.items(): |
| if not isinstance(name, str) or name == '__metadata__' or not isinstance(tensor, FakeTensor): |
| raise ValueError('non-tensor entry or unsupported state_dict key') |
| if tensor.layout != torch.strided or not tensor.is_contiguous(): |
| raise ValueError('non-contiguous/sparse tensors are not supported') |
| if tensor.dtype not in dtypes: |
| raise ValueError('unsupported torch dtype: ' + str(tensor.dtype)) |
| storage = tensor.untyped_storage() |
| offset = getattr(storage, '_checkpoint_offset', None) |
| size = tensor.numel() * tensor.element_size() |
| if offset is None: |
| raise ValueError('PyTorch lacks FakeTensor checkpoint offsets; use a recent PyTorch (>=2.7)') |
| |
| if tensor.storage_offset() != 0 or storage.nbytes() != size or storage_ranges.get(offset) != size: |
| raise ValueError('only complete, zero-offset tensor storages are supported') |
| if offset in seen: |
| raise ValueError('shared storages are not supported') |
| seen.add(offset) |
| tensors.append(dict(name=name, dtype=dtypes[tensor.dtype], shape=list(tensor.shape), |
| nbytes=size, source_file=path.name, source_offset=offset)) |
| if seen != set(storage_ranges): |
| raise ValueError('unused or missing PyTorch storage') |
| tensors.sort(key=lambda item: item['source_offset']) |
| mapping = {name: path.name for name in state} |
| index = {'metadata': {'total_size': sum(t['nbytes'] for t in tensors)}, 'weight_map': mapping} |
| record = source_record(path, metadata_bytes) |
| return index, tensors, [record], {'format': 'pt'}, torch.__version__ |
|
|
|
|
| def load_source(source, source_index=None): |
| source = Path(source).resolve() |
| if source.is_dir(): |
| if source_index: |
| if Path(source_index).name != source_index: |
| raise ValueError('unsafe source index filename') |
| index_path = source / source_index |
| else: |
| indices = sorted(source.glob('*.safetensors.index.json')) |
| if len(indices) != 1: |
| raise ValueError('expected one safetensors index; select --source-index explicitly') |
| index_path = indices[0] |
| index = parse_json(index_path.read_text()) |
| mapping = index.get('weight_map') |
| if not isinstance(mapping, dict) or not mapping: |
| raise ValueError('empty or invalid safetensors index') |
| tensors, files, common, seen = [], [], None, set() |
| for name in sorted(set(mapping.values())): |
| if not isinstance(name, str) or Path(name).name != name or not name.endswith('.safetensors'): |
| raise ValueError('unsafe index filename') |
| path = source / name |
| parsed, metadata, prefix = read_header(path) |
| if common is None: |
| common = metadata |
| elif common != metadata: |
| raise ValueError('source shard metadata differs') |
| for tensor in parsed: |
| if tensor['name'] in seen or mapping.get(tensor['name']) != name: |
| raise ValueError('duplicate tensor or incorrect index mapping') |
| seen.add(tensor['name']) |
| tensors.append(tensor) |
| files.append(source_record(path, prefix)) |
| if seen != set(mapping): |
| raise ValueError('index tensor set differs from shard headers') |
| identity = {'type': 'indexed_safetensors', 'path': str(source), |
| 'index': index_path.name, 'index_sha256': sha_file(index_path)} |
| root = source |
| elif source.is_file() and source.suffix == '.safetensors': |
| tensors, common, prefix = read_header(source) |
| files = [source_record(source, prefix)] |
| index = {'weight_map': {tensor['name']: source.name for tensor in tensors}} |
| identity = {'type': 'single_safetensors', 'path': str(source)} |
| root = source.parent |
| elif source.is_file() and source.suffix in ('.pth', '.pt', '.bin'): |
| index, tensors, files, common, version = load_torch_metadata(source) |
| identity = {'type': 'torch_zip_state_dict', 'path': str(source), |
| 'metadata_loader': 'weights_only=True, mmap=True, FakeTensorMode', |
| 'torch_version': version} |
| root = source.parent |
| else: |
| raise ValueError('source must be indexed/single safetensors or a supported PyTorch ZIP') |
| if not tensors: |
| raise ValueError('empty checkpoint') |
| total = sum(tensor['nbytes'] for tensor in tensors) |
| if index.get('metadata', {}).get('total_size', total) != total: |
| raise ValueError('source index total_size differs') |
| return root, index, tensors, files, common, identity |
|
|
|
|
| def output_header(tensors, metadata): |
| result, offset = {'__metadata__': metadata}, 0 |
| for tensor in tensors: |
| result[tensor['name']] = dict(dtype=tensor['dtype'], shape=tensor['shape'], |
| data_offsets=[offset, offset + tensor['nbytes']]) |
| offset += tensor['nbytes'] |
| raw = json.dumps(result, ensure_ascii=False, separators=(',', ':')).encode() |
| raw += b' ' * ((-len(raw)) % 8) |
| return struct.pack('<Q', len(raw)) + raw |
|
|
|
|
| def plan_shards(tensors, metadata, max_bytes): |
| groups, current, size = [], [], 0 |
| for tensor in tensors: |
| candidate = current + [tensor] |
| if len(output_header(candidate, metadata)) + size + tensor['nbytes'] > max_bytes: |
| if not current: |
| raise ValueError('single tensor exceeds shard limit; a tensor cannot be split') |
| groups.append(current) |
| current, size = [tensor], tensor['nbytes'] |
| if len(output_header(current, metadata)) + size > max_bytes: |
| raise ValueError('single tensor exceeds shard limit') |
| else: |
| current, size = candidate, size + tensor['nbytes'] |
| if current: |
| groups.append(current) |
| return groups |
|
|
|
|
| def write_atomic_json(path, value): |
| path = Path(path) |
| temporary = path.with_name(path.name + '.partial') |
| temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2) + '\n') |
| temporary.replace(path) |
|
|
|
|
| def stable_sources(root, files): |
| for info in files: |
| stat = (root / info['name']).stat() |
| if (stat.st_size, stat.st_mtime_ns) != (info['bytes'], info['mtime_ns']): |
| raise RuntimeError('source checkpoint changed during sharding') |
|
|
|
|
| def verify_outputs(out, groups, metadata, tensor_hashes, records, max_bytes): |
| if len(records) != len(groups): |
| raise ValueError('shard count differs') |
| mapping, verified = {}, set() |
| for group, record in zip(groups, records): |
| path = out / record['name'] |
| parsed, actual_metadata, prefix = read_header(path) |
| expected = {tensor['name']: tensor for tensor in group} |
| if actual_metadata != metadata or path.stat().st_size > max_bytes: |
| raise ValueError('output metadata/size differs') |
| if {tensor['name'] for tensor in parsed} != set(expected): |
| raise ValueError('output tensor names differ') |
| full = hashlib.sha256(prefix) |
| with path.open('rb') as handle: |
| for tensor in parsed: |
| original = expected[tensor['name']] |
| if any(tensor[key] != original[key] for key in ('dtype', 'shape', 'nbytes')): |
| raise ValueError('tensor dtype/shape differs') |
| handle.seek(tensor['source_offset']) |
| digest, left = hashlib.sha256(), tensor['nbytes'] |
| while left: |
| block = handle.read(min(BLOCK, left)) |
| if not block: |
| raise ValueError('truncated output tensor') |
| digest.update(block) |
| full.update(block) |
| left -= len(block) |
| if digest.hexdigest() != tensor_hashes.get(tensor['name']): |
| raise ValueError('tensor bytes differ: ' + tensor['name']) |
| verified.add(tensor['name']) |
| mapping[tensor['name']] = path.name |
| if full.hexdigest() != record['sha256'] or path.stat().st_size != record['bytes']: |
| raise ValueError('output shard hash/size differs') |
| print(f'VERIFIED {out.name}/{path.name} ({len(group)} tensors)', flush=True) |
| if verified != set(tensor_hashes): |
| raise ValueError('tensor verification coverage differs') |
| return mapping |
|
|
|
|
| def model_card(spec): |
| return f'''--- |
| tags: |
| - safetensors |
| - lossless-resharding |
| --- |
| # {spec['kind']} |
| |
| {spec['shard_count']} standard safetensors shards, each at most |
| {spec['max_shard_bytes']:,} bytes **including the header**. Native tensor names, |
| shapes, dtypes and raw value bytes are unchanged. No model construction, |
| quantization, inference, training, GPU allocation or upload was performed. |
| |
| Source/base: `{spec.get('base_model') or 'see parent bundle README'}`. |
| Raw tensor data: {spec['raw_tensor_bytes']:,} bytes, {spec['tensor_count']:,} tensors. |
| `{spec['prefix']}.safetensors.index.json` maps all tensor names to their shards. |
| `reshard_manifest.json` records source/file hashes and independently verified |
| per-tensor SHA256. For a native DiffSynth component, pass the **list of shard |
| paths** to `ModelConfig(path=[...])`; this export is not an architecture/key |
| conversion into a Transformers or diffusers model. |
| |
| PyTorch ZIP conversion, when applicable, uses restricted `weights_only=True` |
| FakeTensor metadata loading and streams original complete storages without |
| executing unrestricted pickle. Obtain/review the original model license before |
| redistribution. This directory is prepared locally, not uploaded. |
| ''' |
|
|
|
|
| def reshard(source, out, max_bytes=5_000_000_000, kind='checkpoint', resume=False, |
| prefix='model', source_index=None, base_model=None): |
| if not prefix or Path(prefix).name != prefix or not all(c.isalnum() or c in '._-' for c in prefix): |
| raise ValueError('unsafe output prefix') |
| if type(max_bytes) is not int or max_bytes <= 0: |
| raise ValueError('positive integer shard size required') |
| root, index, tensors, files, metadata, identity = load_source(source, source_index) |
| out = Path(out).resolve() |
| source_path = Path(source).resolve() |
| if out == root or out in root.parents or root in out.parents: |
| raise ValueError('output must not overlap the original checkpoint directory') |
| groups = plan_shards(tensors, metadata, max_bytes) |
| spec = dict(source=identity, source_files=files, max_shard_bytes=max_bytes, |
| tensor_count=len(tensors), raw_tensor_bytes=sum(t['nbytes'] for t in tensors), |
| shard_count=len(groups), prefix=prefix, kind=kind, base_model=base_model) |
| final, state_path = out / 'reshard_manifest.json', out / '.reshard_state.json' |
| index_name = prefix + '.safetensors.index.json' |
| if final.is_file(): |
| old = parse_json(final.read_text()) |
| if old['spec'] != spec: |
| raise ValueError('existing completed export has different inputs') |
| verify_outputs(out, groups, metadata, old['tensor_sha256'], old['output_shards'], max_bytes) |
| if sha_file(out / index_name) != old['output_index_sha256']: |
| raise ValueError('completed output index changed') |
| for sidecar in old['small_files']: |
| if sha_file(out / sidecar['name']) != sidecar['sha256']: |
| raise ValueError('completed output sidecar changed') |
| for info in files: |
| if sha_file(root / info['name']) != old['source_file_sha256'][info['name']]: |
| raise ValueError('source bytes changed after completed export') |
| print('ALREADY COMPLETE + VERIFIED', out, flush=True) |
| return old |
| state = dict(spec=spec, completed=[], tensor_sha256={}) |
| if out.exists() and any(out.iterdir()): |
| if not resume or not state_path.is_file(): |
| raise ValueError('output is not empty; use fresh output or matching --resume') |
| state = parse_json(state_path.read_text()) |
| if state['spec'] != spec: |
| raise ValueError('resume inputs differ') |
| out.mkdir(parents=True, exist_ok=True) |
| write_atomic_json(state_path, state) |
| records, tensor_hashes = state['completed'], state['tensor_sha256'] |
| if len(records) > len(groups): |
| raise ValueError('too many completed shards in resume state') |
| print(json.dumps({'export': str(out), 'source_type': identity['type'], |
| 'tensors': len(tensors), 'raw_GB': round(spec['raw_tensor_bytes'] / 1e9, 3), |
| 'shards': len(groups), 'GPU_used': False}, indent=2), flush=True) |
| source_hashes = {info['name']: sha_file(root / info['name']) for info in files} |
| stable_sources(root, files) |
| start, last_report = time.monotonic(), 0 |
| done = sum(t['nbytes'] for group in groups[:len(records)] for t in group) |
| with ExitStack() as stack: |
| handles = {info['name']: stack.enter_context((root / info['name']).open('rb')) for info in files} |
| for ordinal, group in enumerate(groups, 1): |
| name = f'{prefix}-{ordinal:05d}-of-{len(groups):05d}.safetensors' |
| target = out / name |
| if ordinal <= len(records): |
| record = records[ordinal - 1] |
| if record['name'] != name or target.stat().st_size != record['bytes'] or sha_file(target) != record['sha256']: |
| raise ValueError('completed shard differs on resume') |
| continue |
| header = output_header(group, metadata) |
| full, temporary = hashlib.sha256(header), target.with_name(name + '.partial') |
| with temporary.open('wb') as handle: |
| handle.write(header) |
| for tensor in group: |
| source_handle = handles[tensor['source_file']] |
| source_handle.seek(tensor['source_offset']) |
| digest, left = hashlib.sha256(), tensor['nbytes'] |
| while left: |
| block = source_handle.read(min(BLOCK, left)) |
| if not block: |
| raise ValueError('truncated source tensor') |
| handle.write(block) |
| digest.update(block) |
| full.update(block) |
| left -= len(block) |
| done += len(block) |
| if done - last_report >= 1024**3: |
| print(f'COPY {out.name}: {done / 1e9:.1f}/{spec["raw_tensor_bytes"] / 1e9:.1f} GB; ' |
| f'shard {ordinal}/{len(groups)}; {time.monotonic() - start:.0f}s', flush=True) |
| last_report = done |
| tensor_hashes[tensor['name']] = digest.hexdigest() |
| handle.flush() |
| size = temporary.stat().st_size |
| if size != len(header) + sum(t['nbytes'] for t in group) or size > max_bytes: |
| raise ValueError('incorrect output shard size') |
| stable_sources(root, files) |
| temporary.replace(target) |
| records.append(dict(name=name, bytes=size, sha256=full.hexdigest(), tensor_count=len(group))) |
| write_atomic_json(state_path, state) |
| print('WROTE', target, f'{size / 1e9:.3f} GB', flush=True) |
| mapping = verify_outputs(out, groups, metadata, tensor_hashes, records, max_bytes) |
| if set(mapping) != set(index['weight_map']): |
| raise ValueError('output index tensor set differs') |
| result_index = {key: value for key, value in index.items() if key not in ('weight_map', 'metadata')} |
| result_index['metadata'] = {**index.get('metadata', {}), 'total_size': spec['raw_tensor_bytes']} |
| result_index['weight_map'] = {name: mapping[name] for name in index['weight_map']} |
| sidecars = [] |
| |
| if source_path.is_dir(): |
| for name in sorted(SIDECARS): |
| path = root / name |
| if not path.is_file(): |
| continue |
| if path.stat().st_size > 100 * 1024**2: |
| raise ValueError('unexpectedly large sidecar') |
| digest = sha_file(path) |
| shutil.copy2(path, out / name) |
| if sha_file(out / name) != digest: |
| raise ValueError('sidecar copy differs') |
| sidecars.append(dict(name=name, bytes=path.stat().st_size, sha256=digest)) |
| stable_sources(root, files) |
| for info in files: |
| if sha_file(root / info['name']) != source_hashes[info['name']]: |
| raise ValueError('source bytes changed during sharding') |
| write_atomic_json(out / index_name, result_index) |
| (out / 'README.md').write_text(model_card(spec)) |
| manifest = dict(spec=spec, created_at_UTC=now(), source_file_sha256=source_hashes, |
| output_shards=records, small_files=sidecars, tensor_sha256=tensor_hashes, |
| output_index_sha256=sha_file(out / index_name), |
| validation=dict(all_tensor_names_shapes_dtypes_equal=True, |
| all_tensor_bytes_SHA256_equal=True, |
| all_output_shards_read_back=True, |
| source_files_unchanged=True, max_file_limit_respected=True), |
| GPU_used=False, model_loaded=False, training_run=False, upload_run=False) |
| write_atomic_json(final, manifest) |
| state_path.unlink() |
| print('COMPLETE', out, flush=True) |
| return manifest |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--source', type=Path, required=True) |
| parser.add_argument('--output-dir', type=Path, required=True) |
| parser.add_argument('--source-index') |
| parser.add_argument('--prefix', default='model') |
| parser.add_argument('--max-shard-gb', type=float, default=5) |
| parser.add_argument('--kind', default='checkpoint') |
| parser.add_argument('--base-model') |
| parser.add_argument('--resume', action='store_true') |
| args = parser.parse_args() |
| if not math.isfinite(args.max_shard_gb) or args.max_shard_gb <= 0: |
| parser.error('--max-shard-gb must be finite and positive') |
| result = reshard(args.source, args.output_dir, int(args.max_shard_gb * 1e9), |
| args.kind, args.resume, args.prefix, args.source_index, args.base_model) |
| print(json.dumps({'output_dir': str(args.output_dir.resolve()), |
| 'shards': len(result['output_shards']), |
| 'tensors': result['spec']['tensor_count'], |
| 'tensor_bytes_verified': True, 'GPU_used': False, 'upload_run': False}, indent=2)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
