File size: 6,835 Bytes
92edcfa | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | """Versioned reference packing for strict groupwise binary matrices."""
from __future__ import annotations
from dataclasses import dataclass
import os
from pathlib import Path
import struct
import torch
from .binary import pack_binary_codes, unpack_binary_codes
MAGIC = b"WALB1G1\0"
VERSION = 1
# magic, version, rows, cols, group_size, padding, total_groups,
# code_bytes, scale_bytes
HEADER = struct.Struct("<8sIIIHHQQQ")
def _raw_bytes(tensor: torch.Tensor) -> bytes:
value = tensor.detach().contiguous().cpu()
return value.view(torch.uint8).numpy().tobytes()
def _tensor_from_bytes(data: bytes, dtype: torch.dtype) -> torch.Tensor:
if not data:
return torch.empty(0, dtype=dtype)
return torch.frombuffer(bytearray(data), dtype=dtype).clone()
@dataclass(frozen=True)
class PackedB1Matrix:
"""A fully binary matrix: one sign bit per padded weight plus FP16 scales."""
shape: tuple[int, int]
group_size: int
padding: int
total_groups: int
codes_packed: torch.Tensor
scales_fp16: torch.Tensor
def __post_init__(self) -> None:
rows, cols = self.shape
if rows <= 0 or cols <= 0:
raise ValueError("matrix shape must be positive")
if self.group_size <= 0 or self.group_size % 8:
raise ValueError("group_size must be positive and divisible by eight")
if self.padding != (-cols) % self.group_size:
raise ValueError("padding does not match shape/group_size")
expected_groups = rows * ((cols + self.padding) // self.group_size)
if self.total_groups != expected_groups:
raise ValueError("total group count does not match shape")
expected_codes = (self.total_groups * self.group_size + 7) // 8
if (
self.codes_packed.dtype != torch.uint8
or self.codes_packed.numel() != expected_codes
):
raise ValueError("invalid packed binary codes")
if (
self.scales_fp16.dtype != torch.float16
or self.scales_fp16.numel() != self.total_groups
):
raise ValueError("invalid FP16 scales")
if not torch.isfinite(self.scales_fp16).all() or torch.any(
self.scales_fp16 <= 0
):
raise ValueError("binary scales must be finite and positive")
@property
def payload_nbytes(self) -> int:
return self.codes_packed.numel() + self.scales_fp16.numel() * 2
@property
def serialized_nbytes(self) -> int:
return HEADER.size + self.payload_nbytes
def true_bpw(self, *, include_header: bool = True) -> float:
size = self.serialized_nbytes if include_header else self.payload_nbytes
return size * 8 / (self.shape[0] * self.shape[1])
def pack_binary_matrix(
codes: torch.Tensor,
scales: torch.Tensor,
*,
shape: tuple[int, int],
group_size: int = 128,
) -> PackedB1Matrix:
"""Pack a fully committed ``{-1,+1}`` matrix without an FP fallback."""
rows, cols = shape
if rows <= 0 or cols <= 0:
raise ValueError("matrix shape must be positive")
if group_size <= 0 or group_size % 8:
raise ValueError("group_size must be positive and divisible by eight")
padding = (-cols) % group_size
groups_per_row = (cols + padding) // group_size
expected_codes = (rows, groups_per_row, group_size)
if tuple(codes.shape) != expected_codes:
raise ValueError(
f"grouped code shape is {tuple(codes.shape)}, expected {expected_codes}"
)
if tuple(scales.shape) != (rows, groups_per_row):
raise ValueError("scale shape does not match grouped binary codes")
code_values = codes.detach().to(torch.int8).contiguous().cpu()
if not torch.all((code_values == -1) | (code_values == 1)):
raise ValueError("binary codes must be in {-1, +1}")
scale_values = scales.detach().half().contiguous().cpu()
if not torch.isfinite(scale_values).all() or torch.any(scale_values <= 0):
raise ValueError("binary scales must be finite and positive")
return PackedB1Matrix(
shape=shape,
group_size=group_size,
padding=padding,
total_groups=rows * groups_per_row,
codes_packed=pack_binary_codes(code_values),
scales_fp16=scale_values,
)
def unpack_binary_matrix(
packed: PackedB1Matrix, *, dtype: torch.dtype = torch.float32
) -> torch.Tensor:
count = packed.total_groups * packed.group_size
codes = unpack_binary_codes(packed.codes_packed, count).view(
packed.total_groups, packed.group_size
)
groups = codes.float() * packed.scales_fp16.float().view(-1, 1)
rows, cols = packed.shape
return groups.view(rows, -1)[:, :cols].to(dtype)
def write_packed_binary_matrix(path: str | Path, packed: PackedB1Matrix) -> int:
output = Path(path)
header = HEADER.pack(
MAGIC,
VERSION,
packed.shape[0],
packed.shape[1],
packed.group_size,
packed.padding,
packed.total_groups,
packed.codes_packed.numel(),
packed.scales_fp16.numel() * 2,
)
with output.open("xb") as handle:
handle.write(header)
handle.write(_raw_bytes(packed.codes_packed))
handle.write(_raw_bytes(packed.scales_fp16))
handle.flush()
os.fsync(handle.fileno())
size = output.stat().st_size
if size != packed.serialized_nbytes:
raise RuntimeError(f"serialized size is {size}, expected {packed.serialized_nbytes}")
return size
def read_packed_binary_matrix(path: str | Path) -> PackedB1Matrix:
data = Path(path).read_bytes()
if len(data) < HEADER.size:
raise ValueError("packed binary file is shorter than the header")
(
magic,
version,
rows,
cols,
group_size,
padding,
total_groups,
code_bytes,
scale_bytes,
) = HEADER.unpack_from(data)
if magic != MAGIC or version != VERSION:
raise ValueError("unsupported packed binary matrix format")
if scale_bytes % 2:
raise ValueError("FP16 scale payload length is not aligned")
if HEADER.size + code_bytes + scale_bytes != len(data):
raise ValueError("packed binary file length does not match its header")
split = HEADER.size + code_bytes
packed = PackedB1Matrix(
shape=(rows, cols),
group_size=group_size,
padding=padding,
total_groups=total_groups,
codes_packed=_tensor_from_bytes(data[HEADER.size:split], torch.uint8),
scales_fp16=_tensor_from_bytes(data[split:], torch.float16),
)
# Force decoding now so malformed bit payloads cannot be deferred to inference.
unpack_binary_codes(
packed.codes_packed, packed.total_groups * packed.group_size
)
return packed
|