| """Lazy, read-only Release 188 NoNE page storage. |
| |
| The public release is one standard safetensors weight map. This module reads |
| only that small JSON index and the published direct-page index at startup. |
| Multi-gigabyte page shards are opened only when the model routes one of their |
| page IDs. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import mmap |
| import re |
| from pathlib import Path, PurePosixPath |
| from typing import Final, Iterator |
|
|
| import torch |
| from safetensors import safe_open |
|
|
| from resynthesis.none_paging import ( |
| NoNEPageRequestPacket, |
| NoNEPageWeights, |
| validate_page_weights, |
| ) |
|
|
| RELEASE_188_GENERATION: Final[int] = 188 |
| _DIRECT_INDEX_PATH: Final[str] = ( |
| "weights/safetensors/direct-page-index.safetensors" |
| ) |
| _PAGE_SHARD_PATTERN: Final[re.Pattern[str]] = re.compile( |
| r"weights/safetensors/pages-(\d{5})-of-(\d{5})\.safetensors" |
| ) |
| _PAGE_TENSOR_COMPONENTS: Final[tuple[str, ...]] = ( |
| "down_t", |
| "ffn_mode_t", |
| "format_revision_t", |
| "gate_t", |
| "glyph_down_t", |
| "glyph_up_t", |
| "outcome_memory_t", |
| "page_ids_t", |
| "repair_memory_t", |
| "transfer_memory_t", |
| "translation_gate_t", |
| "up_t", |
| ) |
| _PAGE_COMPONENT_BITS: Final[dict[str, int]] = { |
| name: 1 << index for index, name in enumerate(_PAGE_TENSOR_COMPONENTS) |
| } |
| _COMPLETE_PAGE_COMPONENT_MASK: Final[int] = ( |
| (1 << len(_PAGE_TENSOR_COMPONENTS)) - 1 |
| ) |
| _SELF_CONTAINED_PAGE_REVISIONS: Final[frozenset[int]] = frozenset((2, 3, 7)) |
| _RELEASE_SHARD_SCHEMA: Final[str] = ( |
| "nnf.resynthesis.release-page-shard.v1" |
| ) |
| _WHITESPACE: Final[bytes] = b" \t\r\n" |
|
|
|
|
| def _skip_whitespace(payload: mmap.mmap, position: int) -> int: |
| while position < len(payload) and payload[position] in _WHITESPACE: |
| position += 1 |
| return position |
|
|
|
|
| def _ascii_json_string( |
| payload: mmap.mmap, |
| position: int, |
| ) -> tuple[str, int]: |
| """Read one unescaped ASCII string from the generated public index.""" |
|
|
| if position >= len(payload) or payload[position] != ord('"'): |
| raise RuntimeError("Release 188 weight-map string is malformed") |
| start = position + 1 |
| position = start |
| while position < len(payload): |
| value = payload[position] |
| if value == ord('"'): |
| try: |
| return payload[start:position].decode("ascii"), position + 1 |
| except UnicodeDecodeError as error: |
| raise RuntimeError( |
| "Release 188 weight-map string is not ASCII" |
| ) from error |
| if value == ord("\\") or value < 0x20 or value > 0x7E: |
| raise RuntimeError( |
| "Release 188 weight-map contains an escaped or non-ASCII string" |
| ) |
| position += 1 |
| raise RuntimeError("Release 188 weight-map string is truncated") |
|
|
|
|
| def _weight_map_pairs(index_path: Path) -> Iterator[tuple[str, str]]: |
| """Stream the flat weight map without constructing a million-entry dict.""" |
|
|
| with index_path.open("rb") as handle: |
| with mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_READ) as payload: |
| marker = b'"weight_map"' |
| marker_position = payload.find(marker) |
| if marker_position < 0: |
| raise RuntimeError("Release 188 index omits its weight map") |
| position = _skip_whitespace( |
| payload, |
| marker_position + len(marker), |
| ) |
| if position >= len(payload) or payload[position] != ord(":"): |
| raise RuntimeError("Release 188 weight map is malformed") |
| position = _skip_whitespace(payload, position + 1) |
| if position >= len(payload) or payload[position] != ord("{"): |
| raise RuntimeError("Release 188 weight map is not an object") |
| position += 1 |
| while True: |
| position = _skip_whitespace(payload, position) |
| if position >= len(payload): |
| raise RuntimeError("Release 188 weight map is truncated") |
| if payload[position] == ord("}"): |
| break |
| name, position = _ascii_json_string(payload, position) |
| position = _skip_whitespace(payload, position) |
| if position >= len(payload) or payload[position] != ord(":"): |
| raise RuntimeError( |
| "Release 188 weight-map entry is malformed" |
| ) |
| position = _skip_whitespace(payload, position + 1) |
| shard, position = _ascii_json_string(payload, position) |
| yield name, shard |
| position = _skip_whitespace(payload, position) |
| if position >= len(payload): |
| raise RuntimeError("Release 188 weight map is truncated") |
| if payload[position] == ord(","): |
| position += 1 |
| continue |
| if payload[position] == ord("}"): |
| break |
| raise RuntimeError("Release 188 weight-map separator differs") |
|
|
|
|
| def _resolved_repo_file(repository_root: Path, relative_path: str) -> Path: |
| relative = PurePosixPath(relative_path) |
| if relative.is_absolute() or ".." in relative.parts: |
| raise RuntimeError("Release 188 weight path escapes its repository") |
| resolved = repository_root.joinpath(*relative.parts).resolve() |
| try: |
| resolved.relative_to(repository_root) |
| except ValueError as error: |
| raise RuntimeError( |
| "Release 188 weight path escapes its repository" |
| ) from error |
| return resolved |
|
|
|
|
| class Release188SafetensorsPageStore: |
| """Read-only page store compatible with ``NoNEPagedExpertRuntime``.""" |
|
|
| def __init__(self, index_path: Path) -> None: |
| resolved_index = index_path.expanduser().resolve() |
| if ( |
| resolved_index.name != "model.safetensors.index.json" |
| or not resolved_index.is_file() |
| ): |
| raise RuntimeError( |
| "Release 188 repository-root safetensors index is missing" |
| ) |
| self.index_path = resolved_index |
| self.root = resolved_index.parent |
| self._page_shard_by_id: dict[int, Path] = {} |
| component_masks: dict[int, int] = {} |
| direct_index_names: set[str] = set() |
| shard_ordinals: set[int] = set() |
| shard_count: int | None = None |
|
|
| for tensor_name, relative_shard in _weight_map_pairs(resolved_index): |
| if relative_shard == _DIRECT_INDEX_PATH: |
| direct_index_names.add(tensor_name) |
| page_stem, separator, component = tensor_name.partition(".") |
| page_digits = page_stem.removeprefix("page_") |
| if ( |
| not tensor_name.startswith("page_") |
| or not separator |
| or not page_digits.isdigit() |
| ): |
| continue |
| component_bit = _PAGE_COMPONENT_BITS.get(component) |
| shard_match = _PAGE_SHARD_PATTERN.fullmatch(relative_shard) |
| if ( |
| component_bit is None |
| or shard_match is None |
| ): |
| raise RuntimeError( |
| "Release 188 page tensor mapping is malformed" |
| ) |
| page_id = int(page_digits) |
| ordinal = int(shard_match.group(1)) |
| observed_count = int(shard_match.group(2)) |
| if ( |
| ordinal < 1 |
| or observed_count < 1 |
| or ordinal > observed_count |
| or ( |
| shard_count is not None |
| and observed_count != shard_count |
| ) |
| ): |
| raise RuntimeError("Release 188 page shard sequence differs") |
| shard_count = observed_count |
| shard_ordinals.add(ordinal) |
| shard_path = _resolved_repo_file(self.root, relative_shard) |
| previous_shard = self._page_shard_by_id.setdefault( |
| page_id, |
| shard_path, |
| ) |
| if previous_shard != shard_path: |
| raise RuntimeError( |
| "Release 188 page tensors cross physical shards" |
| ) |
| prior_mask = component_masks.get(page_id, 0) |
| if prior_mask & component_bit: |
| raise RuntimeError( |
| "Release 188 page tensor mapping contains a duplicate" |
| ) |
| component_masks[page_id] = prior_mask | component_bit |
|
|
| required_direct_names = { |
| "format_revisions_t", |
| "page_count_t", |
| "page_ids_t", |
| } |
| if not required_direct_names.issubset(direct_index_names): |
| raise RuntimeError( |
| "Release 188 weight map omits its direct page index" |
| ) |
| if ( |
| not component_masks |
| or any( |
| mask != _COMPLETE_PAGE_COMPONENT_MASK |
| for mask in component_masks.values() |
| ) |
| or shard_count is None |
| or shard_ordinals != set(range(1, shard_count + 1)) |
| ): |
| raise RuntimeError( |
| "Release 188 weight map has incomplete page tensors" |
| ) |
|
|
| direct_index_path = _resolved_repo_file( |
| self.root, |
| _DIRECT_INDEX_PATH, |
| ) |
| if not direct_index_path.is_file(): |
| raise RuntimeError("Release 188 direct page index is missing") |
| with safe_open( |
| str(direct_index_path), |
| framework="pt", |
| device="cpu", |
| ) as direct_index: |
| page_ids_t: torch.Tensor = direct_index.get_tensor( |
| "page_ids_t" |
| ).long() |
| revisions_t: torch.Tensor = direct_index.get_tensor( |
| "format_revisions_t" |
| ).long() |
| page_count_t: torch.Tensor = direct_index.get_tensor( |
| "page_count_t" |
| ).long() |
| if ( |
| page_ids_t.ndim != 1 |
| or page_ids_t.numel() < 1 |
| or revisions_t.shape != page_ids_t.shape |
| or page_count_t.shape != (1,) |
| or int(page_count_t[0]) != page_ids_t.shape[0] |
| or torch.unique(page_ids_t).numel() != page_ids_t.numel() |
| or page_ids_t.lt(0).any() |
| ): |
| raise RuntimeError("Release 188 direct page index is malformed") |
| page_ids = tuple(int(value) for value in page_ids_t.tolist()) |
| revisions = tuple(int(value) for value in revisions_t.tolist()) |
| if ( |
| set(page_ids) != set(self._page_shard_by_id) |
| or any( |
| revision not in _SELF_CONTAINED_PAGE_REVISIONS |
| for revision in revisions |
| ) |
| ): |
| raise RuntimeError( |
| "Release 188 direct index and page shards differ" |
| ) |
| self._page_catalog_ids_t = page_ids_t.detach().cpu().clone() |
| self._format_revision_by_id = dict( |
| zip(page_ids, revisions, strict=True) |
| ) |
| self._session_id_t: torch.Tensor | None = None |
| self._accepted_generation_t = torch.tensor( |
| RELEASE_188_GENERATION, |
| dtype=torch.long, |
| ) |
|
|
| @property |
| def page_catalog_ids_t(self) -> torch.Tensor: |
| """Return the complete public page catalog as a tensor.""" |
|
|
| return self._page_catalog_ids_t.clone() |
|
|
| def begin_session(self, session_id_t: torch.Tensor) -> torch.Tensor: |
| """Bind caller-owned inference session identity to immutable release 188.""" |
|
|
| resolved = session_id_t.detach().cpu().long().reshape(-1) |
| if ( |
| session_id_t.dtype != torch.long |
| or session_id_t.ndim != 1 |
| or resolved.numel() < 1 |
| ): |
| raise ValueError("Release 188 session identity is malformed") |
| self._session_id_t = resolved.clone() |
| return self._accepted_generation_t.clone() |
|
|
| def accepted_generation_t(self) -> torch.Tensor: |
| """Return the immutable accepted generation for the active session.""" |
|
|
| if self._session_id_t is None: |
| raise RuntimeError("Release 188 page store has no active session") |
| return self._accepted_generation_t.clone() |
|
|
| def accepted_page_ids_t_boundary(self) -> torch.Tensor: |
| """Return the complete accepted physical page identity in sorted order.""" |
|
|
| if self._session_id_t is None: |
| raise RuntimeError("Release 188 page store has no active session") |
| return torch.sort(self._page_catalog_ids_t).values.clone() |
|
|
| def _validated_request_page_ids_boundary( |
| self, |
| *, |
| session_id_t: torch.Tensor, |
| generation_t: torch.Tensor, |
| page_ids_t: torch.Tensor, |
| ) -> tuple[int, ...]: |
| active_session_t = self._session_id_t |
| requested_session_t = session_id_t.detach().cpu().long().reshape(-1) |
| requested_generation_t = ( |
| generation_t.detach().cpu().long().reshape(-1) |
| ) |
| requested_page_ids_t = ( |
| page_ids_t.detach().cpu().long().reshape(-1) |
| ) |
| if ( |
| active_session_t is None |
| or session_id_t.dtype != torch.long |
| or session_id_t.ndim != 1 |
| or not torch.equal(active_session_t, requested_session_t) |
| ): |
| raise RuntimeError("Release 188 page request crossed its session") |
| if ( |
| generation_t.dtype != torch.long |
| or requested_generation_t.shape != (1,) |
| or not torch.equal( |
| requested_generation_t, |
| self._accepted_generation_t.reshape(1), |
| ) |
| ): |
| raise RuntimeError("Release 188 page request generation differs") |
| if ( |
| page_ids_t.dtype != torch.long |
| or page_ids_t.ndim != 1 |
| or requested_page_ids_t.numel() < 1 |
| or torch.unique(requested_page_ids_t).numel() |
| != requested_page_ids_t.numel() |
| ): |
| raise ValueError("Release 188 page request is malformed") |
| requested_page_ids = tuple( |
| int(value) for value in requested_page_ids_t.tolist() |
| ) |
| if any( |
| page_id not in self._page_shard_by_id |
| for page_id in requested_page_ids |
| ): |
| raise KeyError("Release 188 page request is outside its catalog") |
| return requested_page_ids |
|
|
| def _load_page_row_boundary( |
| self, |
| handle: object, |
| *, |
| page_id: int, |
| ) -> NoNEPageWeights: |
| prefix = f"page_{page_id:06d}." |
| keys = set(handle.keys()) |
| expected_keys = { |
| f"{prefix}{component}" for component in _PAGE_TENSOR_COMPONENTS |
| } |
| if {name for name in keys if name.startswith(prefix)} != expected_keys: |
| raise RuntimeError("Release 188 routed page tensor schema differs") |
|
|
| def tensor(component: str) -> torch.Tensor: |
| return handle.get_tensor( |
| f"{prefix}{component}" |
| ) |
|
|
| stored_page_ids_t = tensor("page_ids_t").long() |
| stored_revision_t = tensor("format_revision_t").long() |
| if ( |
| stored_page_ids_t.shape != (1,) |
| or int(stored_page_ids_t[0]) != page_id |
| or stored_revision_t.shape != (1,) |
| or int(stored_revision_t[0]) |
| != self._format_revision_by_id[page_id] |
| ): |
| raise RuntimeError("Release 188 routed page identity differs") |
| weights = NoNEPageWeights( |
| page_ids_t=stored_page_ids_t, |
| ffn_mode_t=tensor("ffn_mode_t"), |
| gate_t=tensor("gate_t"), |
| up_t=tensor("up_t"), |
| down_t=tensor("down_t"), |
| glyph_down_t=tensor("glyph_down_t"), |
| glyph_up_t=tensor("glyph_up_t"), |
| translation_gate_t=tensor("translation_gate_t"), |
| outcome_memory_t=tensor("outcome_memory_t"), |
| repair_memory_t=tensor("repair_memory_t"), |
| transfer_memory_t=tensor("transfer_memory_t"), |
| ) |
| validate_page_weights(weights, synchronize_tensor_values=False) |
| return weights |
|
|
| def materialize_page_ids_boundary( |
| self, |
| *, |
| session_id_t: torch.Tensor, |
| generation_t: torch.Tensor, |
| page_ids_t: torch.Tensor, |
| device: torch.device, |
| dtype: torch.dtype, |
| trainable: bool, |
| ) -> NoNEPageWeights: |
| """Materialize only model-routed page rows from their mapped shards.""" |
|
|
| if trainable: |
| raise RuntimeError( |
| "Release 188 public safetensors page store is read-only" |
| ) |
| if not isinstance(device, torch.device): |
| raise TypeError("Release 188 page device must be torch.device") |
| if not dtype.is_floating_point: |
| raise TypeError("Release 188 page dtype must be floating point") |
| requested_page_ids = self._validated_request_page_ids_boundary( |
| session_id_t=session_id_t, |
| generation_t=generation_t, |
| page_ids_t=page_ids_t, |
| ) |
| positions_by_shard: dict[Path, list[tuple[int, int]]] = {} |
| for position, page_id in enumerate(requested_page_ids): |
| positions_by_shard.setdefault( |
| self._page_shard_by_id[page_id], |
| [], |
| ).append((position, page_id)) |
| rows: list[NoNEPageWeights | None] = [None] * len(requested_page_ids) |
| for shard_path, routed_rows in positions_by_shard.items(): |
| if not shard_path.is_file(): |
| raise RuntimeError( |
| f"Release 188 routed page shard is missing: {shard_path}" |
| ) |
| with safe_open( |
| str(shard_path), |
| framework="pt", |
| device="cpu", |
| ) as handle: |
| metadata = handle.metadata() |
| if ( |
| metadata is None |
| or metadata.get("schema") != _RELEASE_SHARD_SCHEMA |
| or metadata.get("release") != str(RELEASE_188_GENERATION) |
| ): |
| raise RuntimeError( |
| "Release 188 routed page shard metadata differs" |
| ) |
| for position, page_id in routed_rows: |
| rows[position] = self._load_page_row_boundary( |
| handle, |
| page_id=page_id, |
| ) |
| if any(row is None for row in rows): |
| raise RuntimeError("Release 188 page materialization is incomplete") |
| resolved_rows = tuple(row for row in rows if row is not None) |
| weights = ( |
| resolved_rows[0] |
| if len(resolved_rows) == 1 |
| else NoNEPageWeights( |
| page_ids_t=torch.cat( |
| tuple(row.page_ids_t for row in resolved_rows) |
| ), |
| ffn_mode_t=torch.cat( |
| tuple(row.ffn_mode_t for row in resolved_rows) |
| ), |
| gate_t=torch.cat(tuple(row.gate_t for row in resolved_rows)), |
| up_t=torch.cat(tuple(row.up_t for row in resolved_rows)), |
| down_t=torch.cat(tuple(row.down_t for row in resolved_rows)), |
| glyph_down_t=torch.cat( |
| tuple(row.glyph_down_t for row in resolved_rows) |
| ), |
| glyph_up_t=torch.cat( |
| tuple(row.glyph_up_t for row in resolved_rows) |
| ), |
| translation_gate_t=torch.cat( |
| tuple(row.translation_gate_t for row in resolved_rows) |
| ), |
| outcome_memory_t=torch.cat( |
| tuple(row.outcome_memory_t for row in resolved_rows) |
| ), |
| repair_memory_t=torch.cat( |
| tuple(row.repair_memory_t for row in resolved_rows) |
| ), |
| transfer_memory_t=torch.cat( |
| tuple(row.transfer_memory_t for row in resolved_rows) |
| ), |
| ) |
| ) |
| return weights.to( |
| device=device, |
| dtype=dtype, |
| trainable=False, |
| ) |
|
|
| def materialize_weights( |
| self, |
| request: NoNEPageRequestPacket, |
| *, |
| device: torch.device, |
| dtype: torch.dtype, |
| trainable: bool, |
| ) -> NoNEPageWeights: |
| """Compatibility adapter for the runtime's tensor request packet.""" |
|
|
| return self.materialize_page_ids_boundary( |
| session_id_t=request.session_id_t, |
| generation_t=request.generation_t, |
| page_ids_t=request.unique_page_ids_t, |
| device=device, |
| dtype=dtype, |
| trainable=trainable, |
| ) |
|
|
