# Source code for tensorplay.serialization.mega Source: https://www.tensorplay.cn/docs/_modules/tensorplay/serialization/mega.html ``` """Native MEGA artifact support for TensorPlay checkpoints.""" from __future__ import annotations import concurrent.futures import hashlib import json import os import struct import sys import tempfile import threading import zlib from collections import OrderedDict from collections.abc import Mapping from typing import Any import tensorplay as tp from .common import ( _apply_location, _contig_stride, resolve_restore_location, ) MEGA_EXTENSION = ".mega" DEFAULT_ALIGNMENT = 4096 _CHECKSUM_TYPES = {"none": 0, "crc32": 1, "sha256": 2} _CHECKSUM_NAMES = {value: key for key, value in _CHECKSUM_TYPES.items()} _MEGA_DTYPE_SIZES = { "BOOL": 1, "U8": 1, "I8": 1, "I16": 2, "U16": 2, "I32": 4, "U32": 4, "I64": 8, "U64": 8, "F16": 2, "BF16": 2, "F32": 4, "F64": 8, "F8_E5M2": 1, "F8_E4M3": 1, "F8_E8M0": 1, } _MEGA_TO_TP = { "BOOL": "bool", "U8": "uint8", "I8": "int8", "I16": "int16", "U16": "uint16", "I32": "int32", "U32": "uint32", "I64": "int64", "U64": "uint64", "F16": "float16", "BF16": "bfloat16", "F32": "float32", "F64": "float64", } _MEGA_TO_NUMPY = { "BOOL": "?", "U8": "u1", "I8": "i1", "I16": "i2", "U16": "u2", "I32": "i4", "U32": "u4", "I64": "i8", "U64": "u8", "F16": "f2", "BF16": "u2", "F32": "f4", "F64": "f8", } def _require_megatensors(): try: module = __import__("megatensors") except ImportError as error: raise ImportError( "TensorPlay MEGA serialization requires megatensors>=0.0.5" ) from error cpp = getattr(module, "cpp", None) if cpp is None or not callable(getattr(cpp, "write_file", None)): raise ImportError( "The installed megatensors package does not provide the MEGA file backend" ) if not callable(getattr(cpp, "parse_metadata_fd", None)): raise ImportError( "The installed megatensors package does not provide the MEGA parser" ) return module def _native_metadata(filename: str) -> dict: module = _require_megatensors() flags = os.O_RDONLY | int(getattr(os, "O_BINARY", 0)) fd = os.open(filename, flags) try: size = os.fstat(fd).st_size parsed = module.cpp.parse_metadata_fd(fd, filename, size) finally: os.close(fd) if not isinstance(parsed, dict): raise ValueError(f"{filename}: invalid MEGA metadata result") return parsed def _json_metadata_map(value, *, field: str) -> dict: if value is None or value == "": return {} if isinstance(value, str): try: value = json.loads(value) except json.JSONDecodeError as error: raise ValueError(f"invalid {field} metadata") from error if not isinstance(value, Mapping): raise ValueError(f"invalid {field} metadata") return {str(key): str(item) for key, item in value.items()} def _normalise_record( record, metadata: Mapping[str, Any], custom_checksums: Mapping[str, str] | None = None, checksum_mode: str | None = None, ) -> dict: if not isinstance(record, (tuple, list)) or len(record) < 14: raise ValueError("MEGA parser returned an invalid tensor record") ( name, _tensor_id, dtype, shape, payload_offset, logical_nbytes, stored_nbytes, tensor_flags, compression_codec, shuffle_elem_size, checksum_type, checksum, storage_format, chunks, ) = record[:14] name = str(name) dtype = str(dtype) shape = [int(dim) for dim in shape] if any(dim < 0 for dim in shape): raise ValueError(f"MEGA tensor {name!r} has a negative dimension") if dtype not in _MEGA_DTYPE_SIZES: raise NotImplementedError(f"unsupported MEGA dtype {dtype!r}") payload_offset = int(payload_offset) logical_nbytes = int(logical_nbytes) stored_nbytes = int(stored_nbytes) if min(payload_offset, logical_nbytes, stored_nbytes) < 0: raise ValueError(f"MEGA tensor {name!r} has a negative payload range") expected_nbytes = _MEGA_DTYPE_SIZES[dtype] for dim in shape: expected_nbytes *= dim if logical_nbytes != expected_nbytes: raise ValueError( f"MEGA tensor {name!r} has {logical_nbytes} logical bytes, " f"expected {expected_nbytes}" ) checksum_type = int(checksum_type) if checksum_type not in _CHECKSUM_NAMES: raise ValueError(f"unknown MEGA checksum type {checksum_type}") if isinstance(checksum, (bytes, bytearray)): checksum_bytes = bytes(checksum) else: checksum_bytes = b"" if checksum_bytes and len(checksum_bytes) != 32: raise ValueError(f"invalid checksum field for MEGA tensor {name!r}") if custom_checksums is None: custom_checksums = _json_metadata_map( metadata.get("tensorplay.tensor_checksums"), field="tensorplay.tensor_checksums", ) mode = checksum_mode or str(metadata.get("tensorplay.checksum", "none")).lower() if mode not in _CHECKSUM_TYPES: raise ValueError(f"unknown TensorPlay checksum mode {mode!r}") if checksum_type == 0 and mode != "none": checksum_type = _CHECKSUM_TYPES[mode] digest = custom_checksums.get(name) if digest: try: checksum_bytes = bytes.fromhex(digest) except ValueError as error: raise ValueError(f"invalid checksum for MEGA tensor {name!r}") from error if checksum_type == _CHECKSUM_TYPES["crc32"]: if len(checksum_bytes) != 4: raise ValueError(f"invalid CRC32 checksum for MEGA tensor {name!r}") checksum_bytes += b"\x00" * 28 elif len(checksum_bytes) != 32: raise ValueError(f"invalid SHA256 checksum for MEGA tensor {name!r}") return { "shape": shape, "dtype": dtype, "itemsize": _MEGA_DTYPE_SIZES[dtype], "payload_offset": payload_offset, "logical_nbytes": logical_nbytes, "stored_nbytes": stored_nbytes, "tensor_flags": int(tensor_flags), "compression_codec": int(compression_codec), "shuffle_elem_size": int(shuffle_elem_size), "checksum_type": checksum_type, "checksum": checksum_bytes, "storage_format": str(storage_format), "chunks": _normalise_chunks(chunks), "checksum_hex": custom_checksums.get(name), } def _normalise_chunks(chunks) -> list[dict]: if not chunks: return [] normalised = [] for chunk in chunks: if isinstance(chunk, Mapping): values = chunk elif isinstance(chunk, (tuple, list)) and len(chunk) >= 10: values = dict( zip( ( "tensor_id", "chunk_id", "logical_offset", "logical_size", "payload_offset", "stored_size", "codec", "flags", "checksum_type", "checksum", ), chunk[:10], ) ) else: raise ValueError("MEGA parser returned an invalid chunk record") item = { key: int(values[key]) for key in ( "tensor_id", "chunk_id", "logical_offset", "logical_size", "payload_offset", "stored_size", "codec", "flags", "checksum_type", ) } checksum = values.get("checksum", b"") if isinstance(checksum, bytearray): checksum = bytes(checksum) if not isinstance(checksum, bytes): raise ValueError("MEGA chunk checksum must be bytes") item["checksum"] = checksum normalised.append(item) return normalised [docs] def parse_mega_header(path: str | os.PathLike[str]) -> dict: """Read MEGA metadata and tensor ranges without materializing tensors.""" filename = os.fspath(path) file_size = os.path.getsize(filename) parsed = _native_metadata(filename) metadata = dict(parsed.get("metadata") or {}) custom_checksums = _json_metadata_map( metadata.get("tensorplay.tensor_checksums"), field="tensorplay.tensor_checksums", ) checksum_mode = str(metadata.get("tensorplay.checksum", "none")).lower() header_length = int(parsed.get("header_length", 0)) if header_length <= 0 or header_length > file_size: raise ValueError(f"{filename}: invalid MEGA header length") tensors = OrderedDict() for record in parsed.get("tensor_records", []): name = str(record[0]) if name in tensors: raise ValueError(f"{filename}: duplicate MEGA tensor name {name!r}") info = _normalise_record(record, metadata, custom_checksums, checksum_mode) end = header_length + info["payload_offset"] + info["stored_nbytes"] if end > file_size: raise ValueError(f"{filename}: tensor payload is outside the file") tensors[name] = info return { "version": int(parsed.get("version", 1)), "metadata": metadata, "tensors": tensors, "raw_header_size": header_length, "header_length": header_length, "file_size": file_size, } _PRIMITIVES = (bool, int, float, str) def _is_tensor(value: Any) -> bool: parameter_cls = getattr(tp.nn, "Parameter", None) return isinstance(value, tp.Tensor) or ( parameter_cls is not None and isinstance(value, parameter_cls) ) def _flatten_tree(obj: Any) -> tuple: if _is_tensor(obj): return OrderedDict([("0", obj)]), "tensor", None flat: OrderedDict[str, Any] = OrderedDict() names_used: set[str] = set() active: set[int] = set() tensor_names: dict[int, str] = {} def assign(value, path): identity = id(value) existing = tensor_names.get(identity) if existing is not None: return {"__tensor__": existing} name = ".".join(str(part) for part in path) or "0" base = name counter = 1 while name in names_used: name = f"{base}#{counter}" counter += 1 names_used.add(name) tensor_names[identity] = name flat[name] = value return {"__tensor__": name} def encode(node, path): if _is_tensor(node): return assign(node, path) if isinstance(node, _PRIMITIVES) or node is None: return node if isinstance(node, (dict, list, tuple)): node_id = id(node) if node_id in active: raise ValueError("cyclic containers cannot be saved in MEGA") active.add(node_id) try: if isinstance(node, dict): encoded = {} for key, value in node.items(): if not isinstance(key, str): location = ".".join(map(str, path)) or "" raise TypeError( f"dict keys must be str for MEGA serialization " f"(got {type(key).__name__} at {location})" ) encoded[key] = encode(value, path + [key]) if len(encoded) == 1 and next(iter(encoded)) in { "__tensor__", "__tuple__", "__dict__" }: return {"__dict__": [[key, value] for key, value in encoded.items()]} return encoded if isinstance(node, list): return [encode(value, path + [index]) for index, value in enumerate(node)] return { "__tuple__": [ encode(value, path + [index]) for index, value in enumerate(node) ] } finally: active.remove(node_id) location = ".".join(map(str, path)) or "" raise TypeError( f"unsupported leaf of type {type(node).__name__} at {location}: " "MEGA stores tensors and JSON primitives" ) layout = encode(obj, []) if isinstance(obj, Mapping) and all(_is_tensor(value) for value in obj.values()): if len(flat) == len(obj): return OrderedDict((key, obj[key]) for key in obj), "dict", None if isinstance(obj, tuple) and all(_is_tensor(value) for value in obj): if len(flat) == len(obj): return OrderedDict((str(index), value) for index, value in enumerate(obj)), "tuple", None if isinstance(obj, list) and all(_is_tensor(value) for value in obj): if len(flat) == len(obj): return OrderedDict((str(index), value) for index, value in enumerate(obj)), "list", None return flat, "tree", layout def _rebuild_tree(layout, tensors): if isinstance(layout, dict): if set(layout) == {"__tensor__"}: name = layout["__tensor__"] if name not in tensors: raise ValueError(f"MEGA layout references unknown tensor {name!r}") return tensors[name] if set(layout) == {"__tuple__"}: return tuple(_rebuild_tree(item, tensors) for item in layout["__tuple__"]) if set(layout) == {"__dict__"}: items = layout["__dict__"] if not isinstance(items, list): raise ValueError("invalid escaped MEGA dictionary layout") return { str(key): _rebuild_tree(value, tensors) for key, value in items } return {key: _rebuild_tree(value, tensors) for key, value in layout.items()} if isinstance(layout, list): return [_rebuild_tree(item, tensors) for item in layout] return layout def _metadata_value(value): if value is None: return "null" if isinstance(value, (str, int, float, bool)): return value try: return json.dumps(value, ensure_ascii=False, separators=(",", ":")) except (TypeError, ValueError) as error: raise TypeError(f"MEGA metadata value {value!r} is not serializable") from error def _metadata_for_writer(metadata, *, container, layout, devices, checksum, alignment): result = {str(key): _metadata_value(value) for key, value in (metadata or {}).items()} result.setdefault("general.architecture", "tensorplay") result.setdefault("general.alignment", int(alignment)) result.setdefault("mega.tensor_info.format", "self_describing") result["tensorplay.container"] = container result["tensorplay.byteorder"] = sys.byteorder result["tensorplay.tensor_devices"] = json.dumps( devices, ensure_ascii=False, separators=(",", ":") ) result["tensorplay.checksum"] = checksum result["tensorplay.tensor_checksums"] = "" if layout is not None: result["tensorplay.layout"] = json.dumps( layout, ensure_ascii=False, separators=(",", ":") ) return result def _append_u8(buffer: bytearray, value: int) -> None: buffer.extend(struct.pack(" None: buffer.extend(struct.pack(" None: buffer.extend(struct.pack(" None: encoded = str(value).encode("utf-8") _append_u64(buffer, len(encoded)) buffer.extend(encoded) def _append_compact_string(buffer: bytearray, value: str) -> None: encoded = str(value).encode("utf-8") if len(encoded) > 0xFFFFFFFF: raise ValueError("MEGA compact string is too large") _append_u32(buffer, len(encoded)) buffer.extend(encoded) def _append_metadata_value(buffer: bytearray, value) -> None: if isinstance(value, str): _append_u32(buffer, 11) _append_string(buffer, value) return if isinstance(value, bool): _append_u32(buffer, 10) _append_u8(buffer, int(value)) return if isinstance(value, int): if value < 0: raise ValueError("MEGA metadata integers must be non-negative") if value <= 0xFFFFFFFF: _append_u32(buffer, 4) _append_u32(buffer, value) else: _append_u32(buffer, 6) _append_u64(buffer, value) return if isinstance(value, (list, tuple)): _append_u32(buffer, 12) _append_u64(buffer, len(value)) if not value or isinstance(value[0], str): _append_u32(buffer, 11) for item in value: if not isinstance(item, str): raise TypeError("MEGA metadata string arrays must contain strings") _append_string(buffer, item) return if isinstance(value[0], int) and not isinstance(value[0], bool): _append_u32(buffer, 4) for item in value: if not isinstance(item, int) or isinstance(item, bool) or not 0 <= item <= 0xFFFFFFFF: raise TypeError("MEGA metadata integer arrays must contain uint32 values") _append_u32(buffer, item) return raise TypeError(f"unsupported MEGA metadata value type: {type(value).__name__}") def _build_raw_header(records, metadata: Mapping[str, Any], checksum: str) -> bytes: header = bytearray() _append_u32(header, 0x4147454D) _append_u32(header, 1) _append_u64(header, len(records)) _append_u64(header, len(metadata)) for key, value in metadata.items(): _append_string(header, str(key)) _append_metadata_value(header, value) checksum_type = _CHECKSUM_TYPES[checksum] for record in records: _append_compact_string(header, record["name"]) flags = 0 if record["stored_nbytes"] != record["logical_nbytes"]: flags |= 1 << 1 if checksum_type: flags |= 1 << 5 _append_u32(header, flags) shape = record["shape"] _append_u32(header, len(shape)) for dimension in reversed(shape): _append_u64(header, dimension) _append_compact_string(header, record["logical_dtype"]) _append_u64(header, record["payload_offset"]) if flags & (1 << 1): _append_u64(header, record["stored_nbytes"]) if checksum_type: _append_u32(header, checksum_type) digest = record["checksum_bytes"] if len(digest) != 32: raise ValueError("MEGA checksum must contain 32 bytes") header.extend(digest) return bytes(header) def _raw_tensor_bytes_view(tensor): from . import archive if archive._device_string_of(tensor) != "cpu": return None import numpy as np if archive._dtype_name_of(tensor) == "bfloat16": array = tensor.view(getattr(tp, "uint16")).numpy() else: array = tensor.numpy() if not array.flags.c_contiguous: array = np.ascontiguousarray(array) return memoryview(array).cast("B") def _write_raw_mega_file( output: str, records: list[dict], metadata: Mapping[str, Any], alignment: int, checksum: str, ) -> bool: from .policy import _skip_payload_data saver = getattr(tp.Tensor, "_save_file_segments", None) if not callable(saver) or _skip_payload_data(): return False writer_metadata = dict(metadata) checksums = {} for record in records: digest = b"" if checksum != "none": view = _raw_tensor_bytes_view(record["tensor"]) if view is None: return False if checksum == "crc32": value = zlib.crc32(view) & 0xFFFFFFFF digest = value.to_bytes(4, "little") + b"\x00" * 28 checksums[record["name"]] = f"{value:08x}" else: digest = hashlib.sha256(view).digest() checksums[record["name"]] = digest.hex() record["checksum_bytes"] = digest writer_metadata["tensorplay.tensor_checksums"] = json.dumps( checksums, ensure_ascii=False, separators=(",", ":") ) try: header = _build_raw_header(records, writer_metadata, checksum) except (TypeError, ValueError, struct.error): return False with open(output, "wb", buffering=0) as handle: handle.write(header) padding = (-len(header)) % int(alignment) if padding: handle.write(b"\x00" * padding) saver(output, [record["tensor"] for record in records]) return True def _worker_count(count: int) -> int: if count <= 1: return 1 raw = os.environ.get("TENSORPLAY_SERIALIZATION_WORKERS") try: requested = int(raw) if raw else min(4, os.cpu_count() or 1) except ValueError: requested = min(4, os.cpu_count() or 1) return max(1, min(requested, count)) def _write_at(fd: int, data: bytes, offset: int, lock: threading.Lock): if hasattr(os, "pwrite"): view = memoryview(data) position = 0 while position < len(view): written = os.pwrite(fd, view[position:], offset + position) if written <= 0: raise OSError("short write while staging MEGA payload") position += written return with lock: os.lseek(fd, offset, os.SEEK_SET) position = 0 while position < len(data): written = os.write(fd, data[position:]) if written <= 0: raise OSError("short write while staging MEGA payload") position += written def _stage_payloads(records, payload_path: str, total_nbytes: int, checksum: str): from . import archive from .policy import _skip_payload_data fd = os.open(payload_path, os.O_RDWR) try: os.ftruncate(fd, total_nbytes) lock = threading.Lock() def stage(record): data = ( b"\x00" * int(record["logical_nbytes"]) if _skip_payload_data() else archive._tensor_bytes(record["tensor"]) ) expected = int(record["logical_nbytes"]) if len(data) != expected: raise RuntimeError( f"MEGA tensor {record['name']!r} produced {len(data)} bytes, " f"expected {expected}" ) _write_at(fd, data, int(record["payload_offset"]), lock) digest = None if checksum == "crc32": digest = f"{zlib.crc32(data) & 0xFFFFFFFF:08x}" elif checksum == "sha256": digest = hashlib.sha256(data).hexdigest() return record["name"], digest workers = _worker_count(len(records)) if workers == 1: values = [stage(record) for record in records] else: with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool: values = list(pool.map(stage, records)) return {name: digest for name, digest in values if digest is not None} finally: os.close(fd) def save_mega(filename: str, obj: Any, *, metadata=None, alignment: int = DEFAULT_ALIGNMENT, checksum: str = "crc32"): from . import archive checksum = str(checksum).lower() if checksum not in _CHECKSUM_TYPES: raise ValueError( f"checksum must be one of {sorted(_CHECKSUM_TYPES)} (got {checksum!r})" ) alignment = int(alignment) if alignment <= 0: raise ValueError("MEGA alignment must be positive") flat, container, layout = _flatten_tree(obj) records = [] devices = {} payload_offset = 0 mega_dtypes = { "bool": "BOOL", "uint8": "U8", "int8": "I8", "int16": "I16", "uint16": "U16", "int32": "I32", "uint32": "U32", "int64": "I64", "uint64": "U64", "float16": "F16", "bfloat16": "BF16", "float32": "F32", "float64": "F64", } for name, tensor in flat.items(): dtype_name = archive._dtype_name_of(tensor) mega_dtype = mega_dtypes.get(dtype_name) if mega_dtype is None: raise NotImplementedError( f"TensorPlay dtype {dtype_name!r} is not representable in MEGA" ) shape = [int(dim) for dim in tensor.shape] logical_nbytes = int(tensor.numel()) * archive._ITEMSIZE[dtype_name] records.append( { "name": str(name), "shape": shape, "logical_dtype": mega_dtype, "storage_format": "raw_dense", "payload_offset": payload_offset, "logical_nbytes": logical_nbytes, "stored_nbytes": logical_nbytes, "tensor": tensor, } ) devices[str(name)] = archive._device_string_of(tensor) payload_offset += logical_nbytes output = os.fspath(filename) writer_metadata = _metadata_for_writer( metadata, container=container, layout=layout, devices=devices, checksum=checksum, alignment=alignment, ) if _write_raw_mega_file( output, records, writer_metadata, alignment, checksum ): return output_dir = os.path.dirname(os.path.abspath(output)) or "." temp_fd, payload_path = tempfile.mkstemp( prefix=".tensorplay-", suffix=".payload", dir=output_dir ) os.close(temp_fd) try: checksums = _stage_payloads(records, payload_path, payload_offset, checksum) writer_metadata = _metadata_for_writer( metadata, container=container, layout=layout, devices=devices, checksum=checksum, alignment=alignment, ) writer_metadata["tensorplay.tensor_checksums"] = json.dumps( checksums, ensure_ascii=False, separators=(",", ":") ) descriptors = [] for record in records: descriptor = { key: value for key, value in record.items() if key not in {"tensor", "name"} } descriptor["name"] = record["name"] descriptor["src_filename"] = payload_path descriptor["src_offset"] = record["payload_offset"] if checksum != "none": digest = bytes.fromhex(checksums[record["name"]]) if checksum == "crc32": digest += b"\x00" * 28 descriptor["checksum_type"] = _CHECKSUM_TYPES[checksum] descriptor["checksum"] = digest descriptors.append(descriptor) module = _require_megatensors() module.cpp.write_file(output, descriptors, writer_metadata, alignment) finally: try: os.unlink(payload_path) except FileNotFoundError: pass class _ArtifactAware(OrderedDict): pass class _ArtifactList(list): pass class _ArtifactTuple(tuple): pass def _attach_keepalive(result, owner): if isinstance(result, OrderedDict): holder = _ArtifactAware(result) holder._mega_keepalive = owner return holder if isinstance(result, dict): holder = _ArtifactAware(result) holder._mega_keepalive = owner return holder if isinstance(result, list): holder = _ArtifactList(result) holder._mega_keepalive = owner return holder if isinstance(result, tuple): holder = _ArtifactTuple(result) holder._mega_keepalive = owner return holder return result def _dtype_name(mega_dtype): name = _MEGA_TO_TP.get(str(mega_dtype)) if name is None or not hasattr(tp, name): raise NotImplementedError(f"unsupported MEGA dtype {mega_dtype!r}") return name def _read_at(fd: int, offset: int, size: int) -> bytes: if size == 0: return b"" if hasattr(os, "pread"): data = os.pread(fd, size, offset) else: os.lseek(fd, offset, os.SEEK_SET) data = os.read(fd, size) if len(data) != size: raise ValueError("truncated MEGA tensor payload") return data def _mmap_tensor(mapping, offset: int, info: dict): import numpy as np dtype_name = _dtype_name(info["dtype"]) if info["stored_nbytes"] != info["logical_nbytes"]: return None length = int(info["logical_nbytes"]) shape = list(info["shape"]) if length == 0: return tp.empty(shape, dtype=getattr(tp, dtype_name)) np_dtype = _MEGA_TO_NUMPY[info["dtype"]] array = np.frombuffer( memoryview(mapping)[offset:offset + length], dtype=np_dtype, count=length // info["itemsize"], ) try: flat = tp.from_dlpack(array) except (AttributeError, TypeError): flat = tp.from_dlpack(array.__dlpack__()) if dtype_name == "bfloat16": flat = flat.view(getattr(tp, "bfloat16")) try: flat._serialization_keepalive = mapping except (AttributeError, TypeError): pass result = flat.reshape(shape) if shape else flat.reshape([]) try: result._serialization_keepalive = mapping except (AttributeError, TypeError): pass return result def _verify_payload(fd, mapping, filename: str, header_length: int, name: str, info: dict): from .policy import _skip_payload_data if _skip_payload_data(): return if info.get("chunks"): return checksum_type = int(info["checksum_type"]) if checksum_type == 0: return offset = int(info["payload_offset"]) + int(header_length) size = int(info["stored_nbytes"]) if mapping is not None: payload = memoryview(mapping)[offset:offset + size] else: payload = _read_at(fd, offset, size) expected_hex = info.get("checksum_hex") if expected_hex: if checksum_type == _CHECKSUM_TYPES["crc32"]: actual = f"{zlib.crc32(payload) & 0xFFFFFFFF:08x}" else: actual = hashlib.sha256(payload).hexdigest() if actual.lower() != expected_hex.lower(): raise ValueError(f"MEGA tensor {name!r} checksum verification failed") return checksum = info.get("checksum", b"") if checksum_type == _CHECKSUM_TYPES["crc32"] and checksum and any(checksum): actual = (zlib.crc32(payload) & 0xFFFFFFFF).to_bytes(4, "little") if checksum[:4] != actual: raise ValueError(f"MEGA tensor {name!r} checksum verification failed") elif checksum_type == _CHECKSUM_TYPES["sha256"] and checksum and any(checksum): if hashlib.sha256(payload).digest() != checksum[:32]: raise ValueError(f"MEGA tensor {name!r} checksum verification failed") def load_mega(filename: str, map_location=None, *, mmap: bool = False): from . import archive from .policy import _skip_payload_data filename = os.fspath(filename) header = parse_mega_header(filename) metadata = header["metadata"] infos = header["tensors"] saved_byteorder = str(metadata.get("tensorplay.byteorder", "little")) if saved_byteorder not in {"little", "big"}: raise ValueError(f"unknown MEGA byteorder {saved_byteorder!r}") swap = saved_byteorder != sys.byteorder devices = _json_metadata_map( metadata.get("tensorplay.tensor_devices"), field="tensorplay.tensor_devices" ) layout = metadata.get("tensorplay.layout") if isinstance(layout, str): try: layout = json.loads(layout) except json.JSONDecodeError as error: raise ValueError("invalid tensorplay.layout metadata") from error flags = os.O_RDONLY | int(getattr(os, "O_BINARY", 0)) fd = os.open(filename, flags) mapping = None skip = _skip_payload_data() if mmap and not swap and not skip: try: from .policy import _mmap_file mapping = _mmap_file(fd, filename=filename) except (BufferError, OSError, TypeError, ValueError): mapping = None try: region_owner = {} loaded: dict[str, Any] = {} aliases: list[tuple[str, str, dict]] = [] cpp = None for name, info in infos.items(): _verify_payload(fd, mapping, filename, header["header_length"], name, info) raw_dense = ( info["storage_format"] == "raw_dense" and int(info["tensor_flags"]) == 0 and int(info["compression_codec"]) == 0 and not info["chunks"] ) region = ( info["payload_offset"], info["stored_nbytes"], info["dtype"], ) if raw_dense and region in region_owner: aliases.append((name, region_owner[region], info)) continue if raw_dense: region_owner[region] = name dtype_name = _dtype_name(info["dtype"]) saved_location = str(devices.get(name, "cpu")) target = resolve_restore_location( map_location, saved_location, info["logical_nbytes"] ) tensor = None absolute_offset = header["header_length"] + int(info["payload_offset"]) if ( mapping is not None and isinstance(target, str) and target == "cpu" and raw_dense ): try: tensor = _mmap_tensor(mapping, absolute_offset, info) except (BufferError, ValueError, RuntimeError): tensor = None if tensor is None and raw_dense: data = ( b"\x00" * int(info["stored_nbytes"]) if skip else _read_at(fd, absolute_offset, int(info["stored_nbytes"])) ) if swap: data = archive._byteswap_bytes(data, dtype_name) flat = archive._tensor_from_flat_bytes(data, dtype_name) if flat is None: tensor = tp.empty(info["shape"], dtype=getattr(tp, dtype_name)) else: tensor = ( flat.reshape(info["shape"]) if info["shape"] else flat.reshape([]) ) if tensor is None and not skip: if cpp is None: cpp = _require_megatensors().cpp decoder = getattr(cpp, "decode_payload_fd", None) if not callable(decoder): raise ImportError("MEGA payload decoder is unavailable") target_device = target if isinstance(target, str) else "cpu" decode_device = "cpu" if swap else target_device if decode_device == "cpu": tensor = tp.empty(info["shape"], dtype=getattr(tp, dtype_name)) else: tensor = tp.empty( info["shape"], dtype=getattr(tp, dtype_name), device=archive._parse_device(decode_device), ) if info["chunks"]: chunk_decoder = getattr(cpp, "decode_chunks_fd", None) if not callable(chunk_decoder): raise ImportError("MEGA chunk decoder is unavailable") chunk_decoder( fd, filename, str(name), int(header["header_length"]), info["chunks"], int(info["logical_nbytes"]), int(info["shuffle_elem_size"]), int(tensor.data_ptr()), decode_device != "cpu", ) else: decoder( fd, filename, str(name), absolute_offset, int(info["stored_nbytes"]), int(info["logical_nbytes"]), int(info["tensor_flags"]), int(info["compression_codec"]), int(info["shuffle_elem_size"]), int(info["checksum_type"]), bytes(info.get("checksum", b"")), int(tensor.data_ptr()), decode_device != "cpu", ) if swap: tensor = archive._byteswap_tensor(tensor) if tensor is None: tensor = tp.zeros(info["shape"], dtype=getattr(tp, dtype_name)) loaded[name] = _apply_location(tensor, target) for name, owner, info in aliases: base = loaded[owner] shape = list(info["shape"]) loaded[name] = base.as_strided(shape, _contig_stride(shape), 0) ordered = OrderedDict((name, loaded[name]) for name in infos) container = str(metadata.get("tensorplay.container", "dict")) if container == "tree" and layout is not None: result = _rebuild_tree(layout, ordered) elif container == "tuple": result = tuple(ordered.values()) elif container == "list": result = list(ordered.values()) elif container == "tensor": if len(ordered) != 1: raise ValueError("MEGA tensor container must contain exactly one tensor") result = next(iter(ordered.values())) else: result = ordered if mapping is not None: result = _attach_keepalive(result, mapping) mapping = None return result finally: os.close(fd) if mapping is not None: mapping.close() def _convert_safetensors_file( source: str, target: str, *, metadata=None, alignment: int = DEFAULT_ALIGNMENT, checksum: str = "none", ): from . import archive source = os.path.abspath(os.fspath(source)) target = os.fspath(target) alignment = int(alignment) if alignment <= 0: raise ValueError("MEGA alignment must be positive") with open(source, "rb") as handle: header, data_start = archive._read_safetensors_header(handle) file_size = os.path.getsize(source) source_metadata = header.get("__metadata__", {}) if not isinstance(source_metadata, Mapping): raise ValueError("safetensors metadata must be a mapping") merged_metadata = dict(source_metadata) merged_metadata.update(metadata or {}) records = [] devices = {} cursor = 0 payload_offset = 0 checksum_mode = str(checksum).lower() if checksum_mode not in _CHECKSUM_TYPES: raise ValueError( f"checksum must be one of {sorted(_CHECKSUM_TYPES)} " f"(got {checksum!r})" ) checksum_fd = None try: if checksum_mode != "none": checksum_fd = os.open(source, os.O_RDONLY | int(getattr(os, "O_BINARY", 0))) for name, spec in header.items(): if name == "__metadata__": continue if not isinstance(spec, Mapping): raise ValueError(f"safetensors tensor {name!r} spec must be an object") dtype = str(spec.get("dtype", "")) dtype_name = archive._SAFETENSORS_DTYPES.get(dtype) if dtype_name is None: raise NotImplementedError( f"unsupported safetensors dtype: {dtype!r}" ) raw_shape = spec.get("shape") if not isinstance(raw_shape, list): raise ValueError(f"invalid shape for safetensors tensor {name!r}") shape = [int(dim) for dim in raw_shape] if any(dim < 0 for dim in shape): raise ValueError(f"safetensors tensor {name!r} has a negative dimension") offsets = spec.get("data_offsets") if ( not isinstance(offsets, list) or len(offsets) != 2 or not all(isinstance(item, int) and not isinstance(item, bool) for item in offsets) ): raise ValueError(f"invalid data_offsets for tensor {name!r}") start, end = (int(offsets[0]), int(offsets[1])) if not (cursor <= start <= end): raise ValueError( f"safetensors tensor {name!r} has non-monotonic data_offsets" ) nbytes = end - start expected = archive._ITEMSIZE[dtype_name] for dim in shape: expected *= dim if nbytes != expected: raise ValueError( f"safetensors tensor {name!r} has {nbytes} payload bytes, " f"expected {expected}" ) if start < 0 or end > file_size - data_start: raise ValueError(f"safetensors tensor {name!r} payload is outside the file") descriptor = { "name": str(name), "shape": shape, "logical_dtype": dtype, "storage_format": "raw_dense", "payload_offset": payload_offset, "logical_nbytes": nbytes, "stored_nbytes": nbytes, "src_filename": source, "src_offset": data_start + start, } if checksum_fd is not None: if hasattr(os, "pread"): data = os.pread(checksum_fd, nbytes, data_start + start) else: os.lseek(checksum_fd, data_start + start, os.SEEK_SET) data = os.read(checksum_fd, nbytes) if len(data) != nbytes: raise ValueError(f"truncated safetensors tensor {name!r}") if checksum_mode == "crc32": digest = (zlib.crc32(data) & 0xFFFFFFFF).to_bytes(4, "little") descriptor["checksum_type"] = _CHECKSUM_TYPES[checksum_mode] descriptor["checksum"] = digest + b"\x00" * 28 else: descriptor["checksum_type"] = _CHECKSUM_TYPES[checksum_mode] descriptor["checksum"] = hashlib.sha256(data).digest() records.append(descriptor) devices[str(name)] = "cpu" cursor = end payload_offset += nbytes finally: if checksum_fd is not None: os.close(checksum_fd) writer_metadata = _metadata_for_writer( merged_metadata, container="dict", layout=None, devices=devices, checksum=checksum_mode, alignment=alignment, ) writer_metadata["tensorplay.byteorder"] = "little" writer_metadata["tensorplay.tensor_checksums"] = "" os.makedirs(os.path.dirname(os.path.abspath(target)) or ".", exist_ok=True) _require_megatensors().cpp.write_file(target, records, writer_metadata, int(alignment)) return target [docs] def convert_to_mega(model_dir, output_dir=None, **kwargs): """Convert a model directory or a supported checkpoint into MEGA.""" source = os.fspath(model_dir) if os.path.isfile(source): from pathlib import Path options = dict(kwargs) metadata = options.pop("metadata", None) alignment = options.pop("alignment", DEFAULT_ALIGNMENT) checksum = options.pop("checksum", "none") if options: names = ", ".join(sorted(options)) raise TypeError(f"unsupported single-file conversion options: {names}") source_path = Path(source) if output_dir is None: target = source_path.with_suffix(MEGA_EXTENSION) else: target_path = Path(output_dir) target = ( target_path / f"{source_path.stem}{MEGA_EXTENSION}" if target_path.is_dir() or target_path.suffix != MEGA_EXTENSION else target_path ) target.parent.mkdir(parents=True, exist_ok=True) if source_path.suffix.lower() == ".safetensors": return _convert_safetensors_file( source_path, target, metadata=metadata, alignment=alignment, checksum=checksum, ) value = tp.load(source_path, map_location="cpu", mmap=False) save_mega(target, value, metadata=metadata, alignment=alignment, checksum=checksum) return target return _require_megatensors().convert_model( model_dir, output_dir=output_dir, **kwargs ) convert_model = convert_to_mega __all__ = [ "DEFAULT_ALIGNMENT", "MEGA_EXTENSION", "_CHECKSUM_NAMES", "_CHECKSUM_TYPES", "_MEGA_DTYPE_SIZES", "_flatten_tree", "_rebuild_tree", "_require_megatensors", "convert_model", "convert_to_mega", "load_mega", "parse_mega_header", "save_mega", ] ```