File size: 3,437 Bytes
93ebaf2 | 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 | """Read the explicit shared-bank aliases and architecture blocks of a v2 checkpoint."""
import json
import math
import struct
from pathlib import Path
import torch
from safetensors import safe_open
class SharedBankCheckpoint:
def __init__(self, directory):
self.directory = Path(directory).resolve()
index = json.loads((self.directory / "model.safetensors.index.json").read_text())
self.metadata = index["metadata"]
if self.metadata["format"] != "the-net-shared-aliases-v2":
raise ValueError("Expected the-net-shared-aliases-v2")
self.weight_map = index["weight_map"]
self.headers = {}
def _resolve(self, name):
filename = self.weight_map[name]
key = name
if name.startswith(self.metadata["alias_prefix"]):
filename = self.metadata["canonical_bank"]
key = self.metadata["alias_tensor"]
path = (self.directory / filename).resolve()
if not path.is_relative_to(self.directory):
raise ValueError("Checkpoint paths must stay inside the checkpoint directory")
if path not in self.headers:
with safe_open(path, framework="pt"):
pass
with path.open("rb") as stream:
length = struct.unpack("<Q", stream.read(8))[0]
self.headers[path] = (8 + length, json.loads(stream.read(length)))
base, header = self.headers[path]
return path, base, header[key]
def tensor_info(self, name):
_, _, entry = self._resolve(name)
return {"dtype": entry["dtype"], "shape": list(entry["shape"])}
def stored_slice(self, name, start, length):
path, base, entry = self._resolve(name)
size = math.prod(entry["shape"])
if type(start) is not int or type(length) is not int:
raise ValueError("Slice bounds must be integers")
if start < 0 or length < 0 or start + length > size or length > 1_000_000:
raise ValueError("Invalid slice or slice exceeds one million stored elements")
dtype, width = {"U8": (torch.uint8, 1), "F16": (torch.float16, 2)}[entry["dtype"]]
if length == 0:
return torch.empty(0, dtype=dtype)
with path.open("rb") as stream:
stream.seek(base + entry["data_offsets"][0] + start * width)
data = bytearray(stream.read(length * width))
if len(data) != length * width:
raise ValueError("Truncated checkpoint payload")
return torch.frombuffer(data, dtype=dtype).clone()
def bit_slice(self, name, start, length):
info = self.tensor_info(name)
if info["dtype"] != "U8":
raise ValueError("Bit slices require a U8 shared bank")
if type(start) is not int or type(length) is not int:
raise ValueError("Slice bounds must be integers")
if start < 0 or length < 0 or start + length > math.prod(info["shape"]) * 8:
raise ValueError("Invalid bit slice")
if length > 1_000_000:
raise ValueError("Read at most one million bits per slice")
if length == 0:
return torch.empty(0, dtype=torch.uint8)
offset = start % 8
data = self.stored_slice(name, start // 8, (offset + length + 7) // 8)
positions = torch.arange(offset, offset + length)
return ((data[positions // 8].long() >> (positions % 8)) & 1).to(torch.uint8)
|