diff --git a/python/freetoken/checkpoint/__init__.py b/python/freetoken/checkpoint/__init__.py
index f3ce70fcd..406fdf997 100644
--- a/python/freetoken/checkpoint/__init__.py
+++ b/python/freetoken/checkpoint/__init__.py
@@ -5,6 +5,7 @@
"""
from .ftw import (
+ FTWFormatError,
FTWReader,
FTWWriter,
is_ftw_checkpoint,
@@ -14,6 +15,6 @@
from .convert import convert_checkpoint
__all__ = [
- "FTWReader", "FTWWriter", "is_ftw_checkpoint",
+ "FTWFormatError", "FTWReader", "FTWWriter", "is_ftw_checkpoint",
"iter_ftw_weights", "load_ftw_banks", "convert_checkpoint",
]
diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py
index e2bd0de24..76729b79e 100644
--- a/python/freetoken/checkpoint/ftw.py
+++ b/python/freetoken/checkpoint/ftw.py
@@ -37,6 +37,7 @@
import mmap
import os
import re
+import stat
import threading
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
@@ -59,6 +60,14 @@
# Per-layer expert-bank entry name (converter streaming path, see checkpoint/convert.py):
# each layer of a bank is its own FTW tensor instead of one flat [num_layers*E, ...] region.
_LAYER_ENTRY_RE = re.compile(r"^(?P.+)#L(?P\d{5})$")
+_MAX_INDEX_INT = (1 << 63) - 1
+_INDEX_FIELDS = {
+ "format", "version", "align", "shard_limit", "total_bytes", "tensors", "shards",
+}
+
+
+class FTWFormatError(ValueError):
+ """The FTW v1 index or one of its declared shards is structurally invalid."""
def layer_bank_entry_name(bank_name: str, layer_id: int) -> str:
@@ -102,6 +111,321 @@ def _elsize(dt: torch.dtype) -> int:
return torch.empty((), dtype=dt).element_size()
+def _index_int(value, field: str, *, minimum: int = 0) -> int:
+ """Return a JSON integer that is safe to pass to file/tensor APIs.
+
+ ``bool`` is deliberately excluded even though it subclasses ``int`` in Python.
+ File offsets, byte counts, and dimensions ultimately cross signed 64-bit APIs, so
+ accepting larger arbitrary-precision JSON numbers only postpones a less useful
+ overflow failure until allocation or I/O.
+ """
+ if type(value) is not int or not minimum <= value <= _MAX_INDEX_INT:
+ raise FTWFormatError(
+ f"{field} must be an integer in [{minimum}, {_MAX_INDEX_INT}], got {value!r}"
+ )
+ return value
+
+
+def _required(obj: dict, key: str, owner: str):
+ if key not in obj:
+ raise FTWFormatError(f"{owner} is missing required field {key!r}")
+ return obj[key]
+
+
+def _json_object(pairs: list[tuple[str, object]]) -> dict:
+ """Build one JSON object while rejecting ambiguous duplicate keys."""
+ obj = {}
+ for key, value in pairs:
+ if key in obj:
+ raise FTWFormatError(f"FTW index contains duplicate JSON key {key!r}")
+ obj[key] = value
+ return obj
+
+
+def _invalid_json_constant(value: str):
+ raise FTWFormatError(f"FTW index contains non-standard JSON constant {value}")
+
+
+def _json_float(value: str) -> float:
+ parsed = float(value)
+ if not math.isfinite(parsed):
+ raise FTWFormatError(f"FTW index floating-point value is out of range: {value}")
+ return parsed
+
+
+def _validate_metadata(value, field: str = "metadata") -> None:
+ """Reject metadata that cannot round-trip through strict JSON unambiguously."""
+ if isinstance(value, dict):
+ for key, item in value.items():
+ if not isinstance(key, str):
+ raise ValueError(f"{field} keys must be strings, got {key!r}")
+ _validate_metadata(item, f"{field}.{key}")
+ elif isinstance(value, (list, tuple)):
+ for i, item in enumerate(value):
+ _validate_metadata(item, f"{field}[{i}]")
+ elif isinstance(value, float):
+ if not math.isfinite(value):
+ raise ValueError(f"{field} must not contain NaN or infinity")
+ elif value is not None and not isinstance(value, (str, int, bool)):
+ raise ValueError(
+ f"{field} must contain only JSON-compatible values, got {type(value).__name__}"
+ )
+
+
+def _safe_shard_basename(value, field: str) -> str:
+ if (not isinstance(value, str) or not value or value in {".", ".."}
+ or os.path.basename(value) != value or "/" in value or "\\" in value
+ or "\x00" in value):
+ raise FTWFormatError(f"{field} must be a safe shard basename, got {value!r}")
+ return value
+
+
+def _validate_ftw_v1_index(path: str, index) -> tuple[list[dict], list[dict]]:
+ """Eagerly validate all v1 structure used by :class:`FTWReader`.
+
+ This is intentionally a structural check, not a content-authenticity check. It
+ prevents malformed JSON from turning into path traversal, integer overflow,
+ overlapping tensor views, or delayed mmap/O_DIRECT failures. A future format can
+ add hashes without weakening the valid v1 contract checked here.
+ """
+ if not isinstance(index, dict):
+ raise FTWFormatError(f"FTW index root must be an object, got {type(index).__name__}")
+
+ fmt = _required(index, "format", "FTW index")
+ if fmt != FORMAT_TAG:
+ raise FTWFormatError(f"not a {FORMAT_TAG} checkpoint: {path} (format={fmt!r})")
+
+ version = _index_int(_required(index, "version", "FTW index"), "version")
+ if version != FORMAT_VERSION:
+ raise FTWFormatError(
+ f"unsupported FTW version {version}; this reader supports {FORMAT_VERSION}"
+ )
+
+ align = _index_int(_required(index, "align", "FTW index"), "align", minimum=1)
+ if align != ALIGN:
+ raise FTWFormatError(f"unsupported FTW alignment {align}; expected {ALIGN}")
+
+ shard_limit = _index_int(
+ _required(index, "shard_limit", "FTW index"), "shard_limit", minimum=ALIGN
+ )
+ if shard_limit % ALIGN:
+ raise FTWFormatError(f"shard_limit must be a multiple of {ALIGN}, got {shard_limit}")
+
+ total_bytes = _index_int(
+ _required(index, "total_bytes", "FTW index"), "total_bytes"
+ )
+ if total_bytes % ALIGN:
+ raise FTWFormatError(f"total_bytes must be {ALIGN}-aligned, got {total_bytes}")
+
+ raw_shards = _required(index, "shards", "FTW index")
+ if not isinstance(raw_shards, list):
+ raise FTWFormatError("shards must be an array")
+
+ shard_names: set[str] = set()
+ shards: list[dict] = []
+ for i, shard in enumerate(raw_shards):
+ owner = f"shards[{i}]"
+ if not isinstance(shard, dict):
+ raise FTWFormatError(f"{owner} must be an object")
+ file = _safe_shard_basename(_required(shard, "file", owner), f"{owner}.file")
+ if file in shard_names:
+ raise FTWFormatError(f"duplicate shard file {file!r}")
+ shard_names.add(file)
+
+ global_off = _index_int(
+ _required(shard, "global_off", owner), f"{owner}.global_off"
+ )
+ nbytes = _index_int(_required(shard, "nbytes", owner), f"{owner}.nbytes")
+ if global_off % ALIGN:
+ raise FTWFormatError(
+ f"{owner}.global_off must be {ALIGN}-aligned, got {global_off}"
+ )
+ if nbytes % ALIGN:
+ raise FTWFormatError(f"{owner}.nbytes must be {ALIGN}-aligned, got {nbytes}")
+ if nbytes > shard_limit:
+ raise FTWFormatError(
+ f"{owner}.nbytes {nbytes} exceeds shard_limit {shard_limit}"
+ )
+ if global_off > _MAX_INDEX_INT - nbytes:
+ raise FTWFormatError(f"{owner} byte range exceeds signed 64-bit range")
+
+ shard_path = os.path.join(path, file)
+ try:
+ # Follow symlinks: Hugging Face hub snapshots are symlink farms, so a shard
+ # entry routinely resolves into blobs/. The regular-file and size checks
+ # below therefore apply to the link target.
+ shard_stat = os.stat(shard_path)
+ except (OSError, UnicodeError) as exc:
+ # UnicodeError: a JSON-escaped lone surrogate is a valid str but not a
+ # filesystem name; keep it inside the format-error contract.
+ raise FTWFormatError(f"cannot stat FTW shard {file!r}: {exc}") from exc
+ if not stat.S_ISREG(shard_stat.st_mode):
+ raise FTWFormatError(f"FTW shard {file!r} is not a regular file")
+ if shard_stat.st_size != nbytes:
+ raise FTWFormatError(
+ f"FTW shard {file!r} size mismatch: index declares {nbytes}, "
+ f"file has {shard_stat.st_size}"
+ )
+ shards.append(shard)
+
+ shards.sort(key=lambda shard: shard["global_off"])
+ if any(shard["nbytes"] == 0 for shard in shards):
+ # FTWWriter can produce one empty shard for a checkpoint containing only
+ # zero-sized tensors. It cannot produce empty shards in a non-empty stream.
+ if total_bytes != 0 or len(shards) != 1:
+ raise FTWFormatError("zero-length shards are only valid as one empty FTW shard")
+
+ cursor = 0
+ for shard in shards:
+ if shard["global_off"] != cursor:
+ relation = "an overlap" if shard["global_off"] < cursor else "a gap"
+ raise FTWFormatError(
+ f"FTW shard coverage has {relation}: expected offset {cursor}, "
+ f"got {shard['global_off']} for {shard['file']!r}"
+ )
+ cursor += shard["nbytes"]
+ if cursor != total_bytes:
+ raise FTWFormatError(
+ f"shards cover {cursor} logical bytes but total_bytes is {total_bytes}"
+ )
+
+ raw_tensors = _required(index, "tensors", "FTW index")
+ if not isinstance(raw_tensors, list):
+ raise FTWFormatError("tensors must be an array")
+
+ tensor_names: set[str] = set()
+ tensors: list[dict] = []
+ allocations: list[tuple[int, int, str]] = []
+ zero_allocations: list[tuple[int, str]] = []
+ dtype_sizes: dict[torch.dtype, int] = {}
+ for i, tensor in enumerate(raw_tensors):
+ owner = f"tensors[{i}]"
+ if not isinstance(tensor, dict):
+ raise FTWFormatError(f"{owner} must be an object")
+
+ name = _required(tensor, "name", owner)
+ if not isinstance(name, str):
+ raise FTWFormatError(f"{owner}.name must be a string")
+ if name in tensor_names:
+ raise FTWFormatError(f"duplicate tensor name {name!r}")
+ tensor_names.add(name)
+
+ kind = _required(tensor, "kind", owner)
+ if not isinstance(kind, str):
+ raise FTWFormatError(f"{owner}.kind must be a string")
+
+ dtype_name = _required(tensor, "dtype", owner)
+ if not isinstance(dtype_name, str):
+ raise FTWFormatError(f"{owner}.dtype must be a string")
+ dtype = getattr(torch, dtype_name, None)
+ if not isinstance(dtype, torch.dtype):
+ raise FTWFormatError(f"{owner}.dtype is not a known torch dtype: {dtype_name!r}")
+ element_size = dtype_sizes.get(dtype)
+ if element_size is None:
+ try:
+ element_size = _elsize(dtype)
+ # Match the actual reconstruction primitive in iter_ftw_weights rather
+ # than accepting every object torch happens to classify as a dtype.
+ torch.frombuffer(bytearray(element_size), dtype=dtype, count=1)
+ except (RuntimeError, TypeError, ValueError) as exc:
+ raise FTWFormatError(
+ f"{owner}.dtype is not readable FTW storage: {dtype_name!r}"
+ ) from exc
+ dtype_sizes[dtype] = element_size
+
+ shape = _required(tensor, "shape", owner)
+ if not isinstance(shape, list):
+ raise FTWFormatError(f"{owner}.shape must be an array")
+ dims = [_index_int(dim, f"{owner}.shape[{j}]") for j, dim in enumerate(shape)]
+ numel = 1
+ for dim in dims:
+ if dim == 0:
+ numel = 0
+ elif numel and numel > _MAX_INDEX_INT // dim:
+ raise FTWFormatError(f"{owner}.shape element count exceeds signed 64-bit range")
+ else:
+ numel *= dim
+ if numel == 0:
+ try:
+ # Zero-numel shapes allocate no storage, so checking PyTorch's stride
+ # arithmetic here is cheap and catches dimensions that cannot be rebuilt.
+ torch.empty(tuple(dims), dtype=dtype)
+ except (RuntimeError, TypeError, ValueError) as exc:
+ raise FTWFormatError(
+ f"{owner}.shape cannot be reconstructed by torch: {shape!r}"
+ ) from exc
+ if numel and numel > _MAX_INDEX_INT // element_size:
+ raise FTWFormatError(f"{owner} byte size exceeds signed 64-bit range")
+ expected_nbytes = numel * element_size
+
+ global_off = _index_int(
+ _required(tensor, "global_off", owner), f"{owner}.global_off"
+ )
+ nbytes = _index_int(_required(tensor, "nbytes", owner), f"{owner}.nbytes")
+ if nbytes != expected_nbytes:
+ raise FTWFormatError(
+ f"{owner}.nbytes is {nbytes}, but shape {shape!r} and dtype "
+ f"{dtype_name!r} require {expected_nbytes}"
+ )
+ if global_off % ALIGN:
+ raise FTWFormatError(
+ f"{owner}.global_off must be {ALIGN}-aligned, got {global_off}"
+ )
+ if global_off > _MAX_INDEX_INT - nbytes:
+ raise FTWFormatError(f"{owner} byte range exceeds signed 64-bit range")
+ end = global_off + nbytes
+ padded_end = _align_up(end)
+ if end > total_bytes or padded_end > total_bytes:
+ raise FTWFormatError(
+ f"{owner} range [{global_off}, {end}) (padded to {padded_end}) "
+ f"exceeds total_bytes {total_bytes}"
+ )
+ if nbytes:
+ allocations.append((global_off, padded_end, name))
+ else:
+ zero_allocations.append((global_off, name))
+ tensors.append(tensor)
+
+ allocations.sort(key=lambda item: (item[0], item[1], item[2]))
+ allocation_boundaries = {0}
+ cursor = 0
+ previous_name: str | None = None
+ for start, end, name in allocations:
+ if start < cursor:
+ raise FTWFormatError(
+ f"tensor {name!r} overlaps padded allocation for {previous_name!r}"
+ )
+ if start > cursor:
+ raise FTWFormatError(
+ f"FTW tensor coverage has a gap: expected offset {cursor}, "
+ f"got {start} for {name!r}"
+ )
+ allocation_boundaries.add(start)
+ allocation_boundaries.add(end)
+ cursor = end
+ previous_name = name
+ if cursor != total_bytes:
+ raise FTWFormatError(
+ f"tensors cover {cursor} padded bytes but total_bytes is {total_bytes}"
+ )
+
+ for offset, name in zero_allocations:
+ if offset not in allocation_boundaries:
+ raise FTWFormatError(
+ f"zero-sized tensor {name!r} at offset {offset} is not on a tensor boundary"
+ )
+
+ # Preserve both zero-byte forms emitted by FTWWriter: finalize-without-add_tensor
+ # has no shards, while one or more zero-sized tensors produce one empty shard.
+ if total_bytes == 0:
+ if tensors and not shards:
+ raise FTWFormatError("a zero-sized tensor checkpoint must contain one empty shard")
+ if not tensors and shards:
+ raise FTWFormatError("an empty checkpoint must not contain shard entries")
+
+ return shards, tensors
+
+
def is_ftw_checkpoint(path: str) -> bool:
"""True if ``path`` is a directory holding a FreeToken Weight (FTW) index."""
return os.path.isfile(os.path.join(path, INDEX_NAME))
@@ -133,11 +457,17 @@ class FTWWriter:
"""
def __init__(self, out_dir: str, *, shard_limit: int = DEFAULT_SHARD_LIMIT):
- assert shard_limit % ALIGN == 0, "shard_limit must be a multiple of ALIGN"
+ if (type(shard_limit) is not int or shard_limit < ALIGN
+ or shard_limit > _MAX_INDEX_INT or shard_limit % ALIGN):
+ raise ValueError(
+ f"shard_limit must be an integer multiple of {ALIGN} in "
+ f"[{ALIGN}, {_MAX_INDEX_INT}], got {shard_limit!r}"
+ )
os.makedirs(out_dir, exist_ok=True)
self.out_dir = out_dir
self.shard_limit = shard_limit
self._tensors: list[dict] = []
+ self._tensor_names: set[str] = set()
self._shards: list[dict] = []
self._global = 0 # running FTW offset (incl. padding)
self._f = None # current shard file handle
@@ -171,6 +501,10 @@ def _write_raw(self, data: memoryview) -> None:
self._global += take
def add_tensor(self, name: str, tensor: torch.Tensor, kind: str = "weight") -> None:
+ if not isinstance(name, str) or not isinstance(kind, str):
+ raise ValueError("FTW tensor name and kind must be strings")
+ if name in self._tensor_names:
+ raise ValueError(f"duplicate FTW tensor name {name!r}")
t = tensor.detach().cpu().contiguous()
raw = t.reshape(-1).view(torch.uint8)
nbytes = int(raw.numel())
@@ -179,16 +513,24 @@ def add_tensor(self, name: str, tensor: torch.Tensor, kind: str = "weight") -> N
and self._cur + nbytes > self.shard_limit):
self._roll()
global_off = self._global
- assert global_off % ALIGN == 0, "tensor start must be aligned (invariant)"
+ if global_off % ALIGN:
+ raise RuntimeError("FTW writer invariant failed: tensor start is not aligned")
self._write_raw(memoryview(raw.numpy()))
self._tensors.append({"name": name, "kind": kind, "dtype": _dtype_str(t.dtype),
"shape": list(t.shape), "global_off": global_off, "nbytes": nbytes})
+ self._tensor_names.add(name)
# pad to ALIGN so the next tensor starts aligned
pad = _align_up(self._global) - self._global
if pad:
self._write_raw(memoryview(bytes(pad)))
def finalize(self, meta: dict) -> dict:
+ if not isinstance(meta, dict):
+ raise ValueError(f"FTW metadata must be an object, got {type(meta).__name__}")
+ collisions = _INDEX_FIELDS & meta.keys()
+ if collisions:
+ raise ValueError(f"FTW metadata cannot override index fields: {sorted(collisions)}")
+ _validate_metadata(meta)
if self._f is not None:
self._shards.append({"file": _SHARD_FMT.format(self._shard_idx),
"global_off": self._shard_start, "nbytes": self._cur})
@@ -199,7 +541,7 @@ def finalize(self, meta: dict) -> dict:
"tensors": self._tensors, "shards": self._shards, **meta}
tmp = os.path.join(self.out_dir, INDEX_NAME + ".tmp")
with open(tmp, "w") as f:
- json.dump(index, f)
+ json.dump(index, f, allow_nan=False)
os.replace(tmp, os.path.join(self.out_dir, INDEX_NAME))
return index
@@ -215,12 +557,22 @@ class FTWReader:
(the rounding reads into the region's padding, which is discarded by the tensor view)."""
def __init__(self, path: str):
- with open(os.path.join(path, INDEX_NAME)) as f:
- self.index = json.load(f)
- assert self.index.get("format") == FORMAT_TAG, f"not a {FORMAT_TAG}: {path}"
+ index_path = os.path.join(path, INDEX_NAME)
+ try:
+ with open(index_path, encoding="utf-8") as f:
+ self.index = json.load(
+ f,
+ object_pairs_hook=_json_object,
+ parse_constant=_invalid_json_constant,
+ parse_float=_json_float,
+ )
+ except FTWFormatError:
+ raise
+ except (ValueError, RecursionError) as exc:
+ raise FTWFormatError(f"malformed FTW index {index_path!r}: {exc}") from exc
self.dir = path
- self.shards = sorted(self.index["shards"], key=lambda s: s["global_off"])
- self.tensors = {t["name"]: t for t in self.index["tensors"]}
+ self.shards, tensors = _validate_ftw_v1_index(path, self.index)
+ self.tensors = {tensor["name"]: tensor for tensor in tensors}
self._fds: dict[str, int] = {}
self._maps: dict[str, tuple[mmap.mmap, memoryview]] = {}
# O_DIRECT (DMA straight from disk, bypassing the page cache) is the fast path but a
@@ -359,7 +711,9 @@ def rd(job):
def _transient_buffer(nbytes: int) -> mmap.mmap:
- return mmap.mmap(-1, _align_up(nbytes))
+ # mmap rejects a zero-length mapping, but FTWWriter legitimately emits zero-sized
+ # tensors. Give those tensors one aligned backing page while exposing count=0 below.
+ return mmap.mmap(-1, _align_up(nbytes) or ALIGN)
def iter_ftw_weights(path: str, *, kinds=("weight",), keep: Callable[[str], bool] | None = None,
@@ -402,9 +756,14 @@ def _producer():
buf = _transient_buffer(e["nbytes"])
reader.read_into(memoryview(buf), e, workers=workers, chunk=chunk)
dt = _dtype_of(e["dtype"])
- t = torch.frombuffer(buf, dtype=dt, count=e["nbytes"] // _elsize(dt))
- # a 0-d entry (a per-tensor scale) comes back 0-d, not [1]
- if not _put((e["name"], t.view(*e["shape"]) if e["shape"] else t.view(()), buf, e["nbytes"])):
+ if e["nbytes"]:
+ t = torch.frombuffer(buf, dtype=dt, count=e["nbytes"] // _elsize(dt))
+ # a 0-d entry (a per-tensor scale) comes back 0-d, not [1]
+ t = t.view(tuple(e["shape"])) if e["shape"] else t.view(())
+ else:
+ # torch.frombuffer rejects count=0 even with a non-empty backing map.
+ t = torch.empty(tuple(e["shape"]), dtype=dt)
+ if not _put((e["name"], t, buf, e["nbytes"])):
return
except BaseException as ex: # surface to consumer
err.append(ex)
@@ -465,13 +824,20 @@ def load_ftw_banks(
vectors, unaffected by the row split (fixed GPU residency; see
``cache_budget.expert_bytes_per_slot``).
"""
+ if type(num_layers) is not int or num_layers < 1:
+ raise ValueError(f"num_layers must be a positive integer, got {num_layers!r}")
+ if layer_residency is not None and len(layer_residency) != num_layers:
+ raise ValueError(
+ f"layer_residency has {len(layer_residency)} entries; expected {num_layers}"
+ )
+
from freetoken.moe.host_banks import (
HostBank, HostResidency, PinPipeline, alloc_banks, born_pinned_default,
)
from freetoken.utils.progress import byte_bar
- residency = layer_residency or [HostResidency.PINNED.value] * num_layers
- assert len(residency) == num_layers, (len(residency), num_layers)
+ residency = ([HostResidency.PINNED.value] * num_layers
+ if layer_residency is None else layer_residency)
# PINNED layers are born-pinned (cudaHostAlloc) where that wins (see born_pinned_default); LOCKED/PAGEABLE layers stay lazy mmaps
born = born_pinned_default()
@@ -514,7 +880,11 @@ def _backing(layer_id: int) -> str:
per_layer_groups.setdefault(m.group("base"), {})[int(m.group("layer"))] = e
mixed = {e["name"] for e in flat_entries} & per_layer_groups.keys()
- assert not mixed, f"FTW bank(s) mix flat and per-layer row layouts: {sorted(mixed)}"
+ if mixed:
+ reader.close()
+ raise FTWFormatError(
+ f"FTW bank(s) mix flat and per-layer row layouts: {sorted(mixed)}"
+ )
# Row banks: one padded-window HostBank per (name, layer_id) for the flat layout, plus
# how to carve the real [num_experts, *row_shape] tensor out of its head; ``None`` marks
@@ -526,13 +896,25 @@ def _backing(layer_id: int) -> str:
for e in flat_entries:
name = e["name"]
+ if not e["shape"]:
+ reader.close()
+ raise FTWFormatError(f"FTW bank {name!r} must have at least one dimension")
total, *row_shape = e["shape"]
- assert total % num_layers == 0, (name, total, num_layers)
+ if total % num_layers:
+ reader.close()
+ raise FTWFormatError(
+ f"FTW bank {name!r} has {total} rows, not divisible by "
+ f"num_layers={num_layers}"
+ )
num_experts = total // num_layers
dtype = _dtype_of(e["dtype"])
row_bytes = (math.prod(row_shape) if row_shape else 1) * _elsize(dtype)
layer_bytes = num_experts * row_bytes
- assert layer_bytes * num_layers == e["nbytes"], (name, layer_bytes, num_layers, e["nbytes"])
+ if layer_bytes * num_layers != e["nbytes"]:
+ reader.close()
+ raise FTWFormatError(
+ f"FTW bank {name!r} byte geometry does not match its layer layout"
+ )
row_hb[name] = []
row_view_args[name] = []
for layer_id in range(num_layers):
@@ -546,15 +928,21 @@ def _backing(layer_id: int) -> str:
row_jobs.append((name, bank, win_off, win_end - win_off, layer_bytes, layer_id))
for base, by_layer in per_layer_groups.items():
- assert sorted(by_layer) == list(range(num_layers)), (
- f"FTW bank {base!r} has per-layer entries for layers {sorted(by_layer)}, "
- f"expected exactly range({num_layers})"
- )
+ if sorted(by_layer) != list(range(num_layers)):
+ reader.close()
+ raise FTWFormatError(
+ f"FTW bank {base!r} has per-layer entries for layers {sorted(by_layer)}, "
+ f"expected exactly range({num_layers})"
+ )
row_hb[base] = []
row_view_args[base] = []
for layer_id in range(num_layers):
e = by_layer[layer_id]
- assert e["global_off"] % ALIGN == 0, (base, layer_id, e["global_off"]) # writer invariant
+ if e["global_off"] % ALIGN:
+ reader.close()
+ raise FTWFormatError(
+ f"FTW bank {base!r} layer {layer_id} is not {ALIGN}-aligned"
+ )
bank = HostBank(tuple(e["shape"]), _dtype_of(e["dtype"]), backing=_backing(layer_id))
row_hb[base].append(bank)
row_view_args[base].append(None)
@@ -661,6 +1049,6 @@ def _read_layer(job):
__all__ = [
"INDEX_NAME", "FORMAT_TAG", "FORMAT_VERSION", "ALIGN", "DEFAULT_SHARD_LIMIT",
- "is_ftw_checkpoint", "ftw_tensor_names", "FTWWriter", "FTWReader",
+ "FTWFormatError", "is_ftw_checkpoint", "ftw_tensor_names", "FTWWriter", "FTWReader",
"iter_ftw_weights", "load_ftw_banks", "layer_bank_entry_name",
]
diff --git a/tests/checkpoint/__init__.py b/tests/checkpoint/__init__.py
new file mode 100644
index 000000000..5b8db6daf
--- /dev/null
+++ b/tests/checkpoint/__init__.py
@@ -0,0 +1 @@
+"""Checkpoint-format tests."""
diff --git a/tests/checkpoint/test_ftw_validation.py b/tests/checkpoint/test_ftw_validation.py
new file mode 100644
index 000000000..0f9acb446
--- /dev/null
+++ b/tests/checkpoint/test_ftw_validation.py
@@ -0,0 +1,640 @@
+from __future__ import annotations
+
+import copy
+import json
+import mmap
+import os
+from pathlib import Path
+import subprocess
+import sys
+import types
+
+import pytest
+import torch
+
+from freetoken.checkpoint.ftw import (
+ ALIGN,
+ FORMAT_TAG,
+ FORMAT_VERSION,
+ INDEX_NAME,
+ FTWFormatError,
+ FTWReader,
+ FTWWriter,
+ iter_ftw_weights,
+)
+
+MAX_INDEX_INT = (1 << 63) - 1
+
+
+def _write_index(path, index) -> None:
+ with open(path / INDEX_NAME, "w") as f:
+ json.dump(index, f)
+
+
+def _write_raw_index(path, data: bytes) -> None:
+ with open(path / INDEX_NAME, "wb") as f:
+ f.write(data)
+
+
+def _two_tensor_checkpoint(tmp_path):
+ path = tmp_path / "checkpoint"
+ writer = FTWWriter(str(path), shard_limit=ALIGN)
+ writer.add_tensor("first", torch.arange(1000, dtype=torch.float32)) # 4000 + padding
+ writer.add_tensor("second", torch.arange(4, dtype=torch.int16)) # 8 + padding
+ return path, writer.finalize({"model_type": "test"})
+
+
+def _three_tensor_checkpoint(tmp_path):
+ path = tmp_path / "checkpoint"
+ writer = FTWWriter(str(path), shard_limit=ALIGN)
+ for name in ("first", "middle", "last"):
+ writer.add_tensor(name, torch.arange(8, dtype=torch.uint8))
+ return path, writer.finalize({})
+
+
+def _expect_invalid(path, index, match: str) -> None:
+ _write_index(path, index)
+ with pytest.raises(FTWFormatError, match=match):
+ FTWReader(str(path))
+
+
+def test_valid_writer_output_reads_tensor_spanning_shards(tmp_path):
+ path = tmp_path / "checkpoint"
+ expected = (torch.arange(5000, dtype=torch.int64) % 251).to(torch.uint8)
+ writer = FTWWriter(str(path), shard_limit=ALIGN)
+ writer.add_tensor("spanning", expected)
+ index = writer.finalize({})
+
+ assert len(index["shards"]) == 2
+ reader = FTWReader(str(path))
+ reader._direct = 0 # make this portable across tmpfs/overlayfs and CI kernels
+ entry = reader.tensors["spanning"]
+ pieces = list(reader._pieces(entry["global_off"], entry["nbytes"]))
+ assert [piece[3] for piece in pieces] == [ALIGN, 5000 - ALIGN]
+
+ buf = mmap.mmap(-1, 2 * ALIGN)
+ dest = memoryview(buf)
+ try:
+ reader.read_into(dest, entry, workers=1, chunk=ALIGN)
+ actual = torch.frombuffer(buf, dtype=torch.uint8, count=5000).clone()
+ assert torch.equal(actual, expected)
+ finally:
+ dest.release()
+ buf.close()
+ reader.close()
+
+
+def test_valid_empty_and_zero_tensor_writer_outputs(tmp_path):
+ empty_path = tmp_path / "empty"
+ empty_index = FTWWriter(str(empty_path), shard_limit=ALIGN).finalize({})
+ assert empty_index["shards"] == []
+ empty_reader = FTWReader(str(empty_path))
+ assert empty_reader.shards == []
+ empty_reader.close()
+
+ zero_path = tmp_path / "zero"
+ zero_writer = FTWWriter(str(zero_path), shard_limit=ALIGN)
+ zero_writer.add_tensor("zero", torch.empty(0, dtype=torch.float32))
+ zero_index = zero_writer.finalize({})
+ assert zero_index["shards"][0]["nbytes"] == 0
+ zero_reader = FTWReader(str(zero_path))
+ assert zero_reader.tensors["zero"]["nbytes"] == 0
+ zero_reader.close()
+
+
+def test_valid_zero_tensors_at_writer_allocation_boundaries(tmp_path):
+ path = tmp_path / "mixed-zero"
+ writer = FTWWriter(str(path), shard_limit=ALIGN)
+ writer.add_tensor("zero-before", torch.empty(0, dtype=torch.float32))
+ writer.add_tensor("first", torch.arange(1000, dtype=torch.float32))
+ writer.add_tensor("zero-between", torch.empty((2, 0, 3), dtype=torch.int16))
+ writer.add_tensor("second", torch.arange(4, dtype=torch.int16))
+ writer.add_tensor("zero-after", torch.empty(0, dtype=torch.uint8))
+ index = writer.finalize({})
+
+ reader = FTWReader(str(path))
+ assert [reader.tensors[name]["global_off"] for name in (
+ "zero-before", "zero-between", "zero-after",
+ )] == [0, ALIGN, 2 * ALIGN]
+ assert index["total_bytes"] == 2 * ALIGN
+ reader.close()
+
+
+def test_zero_tensor_can_be_iterated(tmp_path, monkeypatch):
+ path = tmp_path / "zero-iteration"
+ writer = FTWWriter(str(path), shard_limit=ALIGN)
+ writer.add_tensor("zero", torch.empty((2, 0, 3), dtype=torch.float32))
+ writer.finalize({})
+
+ class _Progress:
+ def update(self, _nbytes):
+ pass
+
+ def close(self):
+ pass
+
+ progress = types.ModuleType("freetoken.utils.progress")
+ progress.byte_bar = lambda *_args, **_kwargs: _Progress()
+ monkeypatch.setitem(sys.modules, "freetoken.utils.progress", progress)
+
+ loaded = list(iter_ftw_weights(str(path), workers=1, chunk=ALIGN, prefetch=1))
+ assert len(loaded) == 1
+ assert loaded[0][0] == "zero"
+ assert loaded[0][1].shape == (2, 0, 3)
+
+
+def test_scalar_tensor_round_trips_with_scalar_shape(tmp_path, monkeypatch):
+ path = tmp_path / "scalar-iteration"
+ writer = FTWWriter(str(path), shard_limit=ALIGN)
+ writer.add_tensor("scalar", torch.tensor(17, dtype=torch.int32))
+ writer.finalize({})
+
+ class _Progress:
+ def update(self, _nbytes):
+ pass
+
+ def close(self):
+ pass
+
+ progress = types.ModuleType("freetoken.utils.progress")
+ progress.byte_bar = lambda *_args, **_kwargs: _Progress()
+ monkeypatch.setitem(sys.modules, "freetoken.utils.progress", progress)
+
+ loaded = list(iter_ftw_weights(str(path), workers=1, chunk=ALIGN, prefetch=1))
+ assert loaded[0][0] == "scalar"
+ assert loaded[0][1].shape == ()
+ assert loaded[0][1].item() == 17
+
+
+def test_shards_may_be_declared_out_of_order(tmp_path):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ index["shards"].reverse()
+ _write_index(path, index)
+
+ reader = FTWReader(str(path))
+ assert [shard["global_off"] for shard in reader.shards] == [0, ALIGN]
+ reader.close()
+
+
+@pytest.mark.parametrize(
+ ("field", "value", "match"),
+ [
+ ("format", "other", "not a freetoken_weight"),
+ ("version", FORMAT_VERSION + 1, "unsupported FTW version"),
+ ("version", True, "version must be an integer"),
+ ("align", ALIGN * 2, "unsupported FTW alignment"),
+ ("shard_limit", ALIGN + 1, "shard_limit must be a multiple"),
+ ("shard_limit", 0, "shard_limit must be an integer"),
+ ("total_bytes", ALIGN + 1, "total_bytes must be 4096-aligned"),
+ ("total_bytes", -1, "total_bytes must be an integer"),
+ ("total_bytes", MAX_INDEX_INT + 1, "total_bytes must be an integer"),
+ ],
+)
+def test_rejects_invalid_header_fields(tmp_path, field, value, match):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ index[field] = value
+ _expect_invalid(path, index, match)
+
+
+def test_rejects_non_object_root_and_missing_fields(tmp_path):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ _expect_invalid(path, [], "index root must be an object")
+
+ index.pop("tensors")
+ _expect_invalid(path, index, "missing required field 'tensors'")
+
+
+@pytest.mark.parametrize("raw", [b"{", b"\xff"])
+def test_rejects_malformed_json_and_encoding(tmp_path, raw):
+ path, _index = _two_tensor_checkpoint(tmp_path)
+ _write_raw_index(path, raw)
+ with pytest.raises(FTWFormatError, match="malformed FTW index"):
+ FTWReader(str(path))
+
+
+def test_rejects_duplicate_json_object_keys(tmp_path):
+ path, _index = _two_tensor_checkpoint(tmp_path)
+ _write_raw_index(
+ path,
+ b'{"format":"freetoken_weight","version":1,"version":1}',
+ )
+ with pytest.raises(FTWFormatError, match="duplicate JSON key 'version'"):
+ FTWReader(str(path))
+
+
+@pytest.mark.parametrize("constant", ["NaN", "Infinity", "-Infinity"])
+def test_rejects_non_standard_json_constants_even_in_metadata(tmp_path, constant):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ raw = json.dumps(index).removesuffix("}") + f', "metadata_value": {constant}}}'
+ _write_raw_index(path, raw.encode())
+ with pytest.raises(FTWFormatError, match="non-standard JSON constant"):
+ FTWReader(str(path))
+
+
+def test_rejects_json_float_that_overflows_python_float(tmp_path):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ raw = json.dumps(index).removesuffix("}") + ', "metadata_value": 1e999}'
+ _write_raw_index(path, raw.encode())
+ with pytest.raises(FTWFormatError, match="floating-point value is out of range"):
+ FTWReader(str(path))
+
+
+@pytest.mark.parametrize(
+ ("field", "value", "match"),
+ [
+ ("shards", {}, "shards must be an array"),
+ ("tensors", {}, "tensors must be an array"),
+ ],
+)
+def test_rejects_non_array_index_collections(tmp_path, field, value, match):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ index[field] = value
+ _expect_invalid(path, index, match)
+
+
+@pytest.mark.parametrize(
+ "file",
+ ["../outside.ftw", "/tmp/outside.ftw", "sub/file.ftw", "sub\\file.ftw", "..", "bad\x00.ftw"],
+)
+def test_rejects_unsafe_shard_names_before_filesystem_access(tmp_path, file):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ index["shards"][0]["file"] = file
+ _expect_invalid(path, index, "safe shard basename")
+
+
+def test_rejects_duplicate_shard_names(tmp_path):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ index["shards"][1]["file"] = index["shards"][0]["file"]
+ _expect_invalid(path, index, "duplicate shard file")
+
+
+@pytest.mark.parametrize(
+ ("entry", "match"),
+ [
+ (1, r"shards\[0\] must be an object"),
+ ({"global_off": 0, "nbytes": ALIGN}, "missing required field 'file'"),
+ ({"file": 1, "global_off": 0, "nbytes": ALIGN}, "safe shard basename"),
+ ],
+)
+def test_rejects_malformed_shard_entries(tmp_path, entry, match):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ index["shards"][0] = entry
+ _expect_invalid(path, index, match)
+
+
+def test_rejects_missing_truncated_extended_and_non_regular_shards_eagerly(tmp_path):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ first = path / index["shards"][0]["file"]
+ first.unlink()
+ _expect_invalid(path, index, "cannot stat FTW shard")
+
+ path, index = _two_tensor_checkpoint(tmp_path / "truncated")
+ first = path / index["shards"][0]["file"]
+ with open(first, "r+b") as f:
+ f.truncate(ALIGN - 1)
+ _expect_invalid(path, index, "size mismatch")
+
+ path, index = _two_tensor_checkpoint(tmp_path / "extended")
+ first = path / index["shards"][0]["file"]
+ with open(first, "ab") as f:
+ f.write(b"\x00")
+ _expect_invalid(path, index, "size mismatch")
+
+ path, index = _two_tensor_checkpoint(tmp_path / "directory")
+ first = path / index["shards"][0]["file"]
+ first.unlink()
+ first.mkdir()
+ _expect_invalid(path, index, "not a regular file")
+
+
+def test_accepts_symlinked_shard_whose_target_matches(tmp_path):
+ # Hugging Face hub snapshots are symlink farms (every file links into blobs/), so a
+ # shard entry that resolves through a symlink to a matching regular file must load.
+ path, index = _two_tensor_checkpoint(tmp_path)
+ blobs = tmp_path / "blobs"
+ blobs.mkdir()
+ first = path / index["shards"][0]["file"]
+ target = blobs / "first-blob"
+ target.write_bytes(first.read_bytes())
+ first.unlink()
+ first.symlink_to(target)
+
+ reader = FTWReader(str(path))
+ reader._direct = 0
+ buf = mmap.mmap(-1, ALIGN)
+ dest = memoryview(buf)
+ try:
+ reader.read_into(dest, reader.tensors["first"], workers=1, chunk=ALIGN)
+ actual = torch.frombuffer(buf, dtype=torch.float32, count=1000).clone()
+ assert torch.equal(actual, torch.arange(1000, dtype=torch.float32))
+ finally:
+ dest.release()
+ buf.close()
+ reader.close()
+
+
+def test_symlinked_shard_is_checked_against_its_target(tmp_path):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ external = tmp_path / "outside.ftw"
+ external.write_bytes(b"\x00" * (ALIGN + 1))
+ first = path / index["shards"][0]["file"]
+ first.unlink()
+ first.symlink_to(external)
+
+ _expect_invalid(path, index, "size mismatch")
+
+
+def test_rejects_shard_name_that_is_not_a_filesystem_name(tmp_path):
+ # A JSON-escaped lone surrogate decodes to a valid str that os.stat cannot encode;
+ # it must surface as FTWFormatError, not UnicodeEncodeError.
+ path, index = _two_tensor_checkpoint(tmp_path)
+ index["shards"][0]["file"] = "\ud800.ftw"
+
+ _expect_invalid(path, index, "cannot stat FTW shard")
+
+
+@pytest.mark.parametrize(
+ ("mutation", "match"),
+ [
+ (lambda index: index["shards"][1].update(global_off=3 * ALIGN), "coverage has a gap"),
+ (lambda index: index["shards"][1].update(global_off=0), "coverage has an overlap"),
+ (lambda index: index.update(total_bytes=3 * ALIGN), "shards cover .* total_bytes"),
+ (lambda index: index["shards"][0].update(global_off=1), "global_off must be 4096-aligned"),
+ (lambda index: index["shards"][0].update(global_off=MAX_INDEX_INT + 1), "global_off must be an integer"),
+ (lambda index: index["shards"][0].update(
+ global_off=(MAX_INDEX_INT // ALIGN) * ALIGN,
+ ), "byte range exceeds signed 64-bit"),
+ (lambda index: index["shards"][0].update(nbytes=2 * ALIGN), "exceeds shard_limit"),
+ (lambda index: index["shards"][0].update(nbytes=ALIGN - 1), "nbytes must be 4096-aligned"),
+ (lambda index: index["shards"][0].update(nbytes=True), "nbytes must be an integer"),
+ ],
+)
+def test_rejects_invalid_shard_geometry(tmp_path, mutation, match):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ mutation(index)
+ _expect_invalid(path, index, match)
+
+
+def test_rejects_zero_length_shard_in_nonempty_stream(tmp_path):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ empty_name = "empty.ftw"
+ (path / empty_name).touch()
+ index["shards"].insert(0, {"file": empty_name, "global_off": 0, "nbytes": 0})
+ _expect_invalid(path, index, "zero-length shards")
+
+
+def test_rejects_non_writer_zero_byte_shard_forms(tmp_path):
+ empty_path = tmp_path / "empty"
+ empty_index = FTWWriter(str(empty_path), shard_limit=ALIGN).finalize({})
+ (empty_path / "empty.ftw").touch()
+ empty_index["shards"] = [{"file": "empty.ftw", "global_off": 0, "nbytes": 0}]
+ _expect_invalid(empty_path, empty_index, "empty checkpoint must not contain shard")
+
+ zero_path = tmp_path / "zero"
+ zero_writer = FTWWriter(str(zero_path), shard_limit=ALIGN)
+ zero_writer.add_tensor("zero", torch.empty(0, dtype=torch.float32))
+ zero_index = zero_writer.finalize({})
+ zero_index["shards"] = []
+ _expect_invalid(zero_path, zero_index, "zero-sized tensor checkpoint must contain")
+
+
+@pytest.mark.parametrize(
+ ("mutation", "match"),
+ [
+ (lambda index: index["tensors"][1].update(name="first"), "duplicate tensor name"),
+ (lambda index: index["tensors"][0].update(name=1), "name must be a string"),
+ (lambda index: index["tensors"][0].update(kind=1), "kind must be a string"),
+ (lambda index: index["tensors"][0].update(dtype="not_a_dtype"), "not a known torch dtype"),
+ (lambda index: index["tensors"][0].update(dtype="strided"), "not a known torch dtype"),
+ (lambda index: index["tensors"][0].update(dtype=1), "dtype must be a string"),
+ (lambda index: index["tensors"][0].update(shape="1000"), "shape must be an array"),
+ (lambda index: index["tensors"][0].update(shape=[True]), r"shape\[0\] must be an integer"),
+ (lambda index: index["tensors"][0].update(shape=[1.0]), r"shape\[0\] must be an integer"),
+ (lambda index: index["tensors"][0].update(shape=[-1]), r"shape\[0\] must be an integer"),
+ (lambda index: index["tensors"][0].update(shape=[(1 << 63) - 1, 2]), "element count exceeds"),
+ (lambda index: index["tensors"][0].update(
+ shape=[0, MAX_INDEX_INT, MAX_INDEX_INT], nbytes=0,
+ ), "shape cannot be reconstructed by torch"),
+ (lambda index: index["tensors"][0].update(shape=[MAX_INDEX_INT], dtype="float32"), "byte size exceeds"),
+ (lambda index: index["tensors"][0].update(nbytes=4), "shape .* require 4000"),
+ (lambda index: index["tensors"][0].update(global_off=1), "global_off must be 4096-aligned"),
+ (lambda index: index["tensors"][0].update(
+ dtype="uint8", shape=[ALIGN], nbytes=ALIGN,
+ global_off=(MAX_INDEX_INT // ALIGN) * ALIGN,
+ ), "byte range exceeds signed 64-bit"),
+ (lambda index: index["tensors"][1].update(global_off=2 * ALIGN), "exceeds total_bytes"),
+ (lambda index: index["tensors"][0].update(nbytes=True), "nbytes must be an integer"),
+ ],
+)
+def test_rejects_invalid_tensor_entries(tmp_path, mutation, match):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ mutation(index)
+ _expect_invalid(path, index, match)
+
+
+@pytest.mark.parametrize(
+ ("entry", "match"),
+ [
+ (1, r"tensors\[0\] must be an object"),
+ ({"kind": "weight", "dtype": "float32", "shape": [1],
+ "global_off": 0, "nbytes": 4}, "missing required field 'name'"),
+ ],
+)
+def test_rejects_malformed_tensor_entries(tmp_path, entry, match):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ index["tensors"][0] = entry
+ _expect_invalid(path, index, match)
+
+
+def test_rejects_dtype_that_torch_cannot_reconstruct(tmp_path, monkeypatch):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ original = torch.frombuffer
+
+ def reject_float32(buffer, *, dtype, count):
+ if dtype == torch.float32:
+ raise RuntimeError("unsupported storage dtype")
+ return original(buffer, dtype=dtype, count=count)
+
+ monkeypatch.setattr(torch, "frombuffer", reject_float32)
+ _expect_invalid(path, index, "dtype is not readable FTW storage")
+
+
+def test_rejects_tensor_overlap_including_padded_allocation(tmp_path):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ index["tensors"][1]["global_off"] = 0
+ _expect_invalid(path, index, "overlaps padded allocation")
+
+
+@pytest.mark.parametrize(
+ ("remove", "match"),
+ [
+ (0, "tensor coverage has a gap"),
+ (1, "tensor coverage has a gap"),
+ (2, "tensors cover .* total_bytes"),
+ ],
+)
+def test_rejects_leading_internal_and_trailing_tensor_gaps(tmp_path, remove, match):
+ path, index = _three_tensor_checkpoint(tmp_path)
+ index["tensors"].pop(remove)
+ _expect_invalid(path, index, match)
+
+
+def test_rejects_zero_sized_tensor_inside_nonzero_allocation(tmp_path):
+ path = tmp_path / "checkpoint"
+ writer = FTWWriter(str(path), shard_limit=ALIGN)
+ writer.add_tensor("spanning", torch.arange(5000, dtype=torch.uint8))
+ index = writer.finalize({})
+ index["tensors"].append({
+ "name": "zero-inside",
+ "kind": "weight",
+ "dtype": "float32",
+ "shape": [0],
+ "global_off": ALIGN,
+ "nbytes": 0,
+ })
+ _expect_invalid(path, index, "zero-sized tensor .* is not on a tensor boundary")
+
+
+def test_index_validation_does_not_mutate_tensor_order(tmp_path):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ index["tensors"].reverse()
+ expected = copy.deepcopy(index["tensors"])
+ _write_index(path, index)
+
+ reader = FTWReader(str(path))
+ assert reader.index["tensors"] == expected
+ assert [entry["name"] for entry in reader.entries()] == ["second", "first"]
+ reader.close()
+
+
+def test_writer_emits_declared_v1_header(tmp_path):
+ path, index = _two_tensor_checkpoint(tmp_path)
+ assert index["format"] == FORMAT_TAG
+ assert index["version"] == FORMAT_VERSION
+ assert index["align"] == ALIGN
+ assert sum(os.path.getsize(path / shard["file"]) for shard in index["shards"]) == index["total_bytes"]
+
+
+@pytest.mark.parametrize("shard_limit", [0, ALIGN - 1, ALIGN + 1, True, MAX_INDEX_INT + 1])
+def test_writer_rejects_invalid_shard_limit_without_assertions(tmp_path, shard_limit):
+ with pytest.raises(ValueError, match="shard_limit must be an integer multiple"):
+ FTWWriter(str(tmp_path / "checkpoint"), shard_limit=shard_limit)
+
+
+def test_writer_does_not_emit_indexes_its_reader_would_reject(tmp_path):
+ path = tmp_path / "checkpoint"
+ writer = FTWWriter(str(path), shard_limit=ALIGN)
+ writer.add_tensor("tensor", torch.arange(4, dtype=torch.float32))
+
+ with pytest.raises(ValueError, match="duplicate FTW tensor name"):
+ writer.add_tensor("tensor", torch.arange(4, dtype=torch.float32))
+ with pytest.raises(ValueError, match="cannot override index fields"):
+ writer.finalize({"version": FORMAT_VERSION})
+ with pytest.raises(ValueError, match="must not contain NaN or infinity"):
+ writer.finalize({"metadata_value": float("nan")})
+ with pytest.raises(ValueError, match="keys must be strings"):
+ writer.finalize({"nested": {1: "ambiguous with a JSON string key"}})
+ with pytest.raises(ValueError, match="JSON-compatible values"):
+ writer.finalize({"metadata_value": object()})
+
+ writer.finalize({"metadata_value": 1.5})
+ FTWReader(str(path)).close()
+
+
+def test_validation_remains_active_under_optimized_python(tmp_path):
+ repo_root = Path(__file__).resolve().parents[2]
+ script = r'''
+import json
+import logging
+from pathlib import Path
+import sys
+import types
+
+import torch
+
+utils = types.ModuleType("freetoken.utils")
+utils.init_logger = logging.getLogger
+sys.modules["freetoken.utils"] = utils
+
+from freetoken.checkpoint.ftw import FTWFormatError, FTWReader, FTWWriter, INDEX_NAME, load_ftw_banks
+
+moe = types.ModuleType("freetoken.moe")
+moe.__path__ = []
+host_banks = types.ModuleType("freetoken.moe.host_banks")
+host_banks.HostBank = object
+host_banks.HostResidency = types.SimpleNamespace(PINNED=types.SimpleNamespace(value="pinned"))
+host_banks.PinPipeline = object
+host_banks.alloc_banks = lambda _specs: {}
+host_banks.born_pinned_default = lambda: False
+progress = types.ModuleType("freetoken.utils.progress")
+progress.byte_bar = lambda *_args, **_kwargs: None
+sys.modules["freetoken.moe"] = moe
+sys.modules["freetoken.moe.host_banks"] = host_banks
+sys.modules["freetoken.utils.progress"] = progress
+
+path = Path(sys.argv[1]) / "optimized"
+path.mkdir()
+(path / INDEX_NAME).write_text(json.dumps({"format": "wrong"}))
+
+try:
+ FTWReader(str(path))
+except FTWFormatError:
+ pass
+else:
+ raise SystemExit("malformed index was accepted under -O")
+
+for action in (
+ lambda: FTWWriter(str(path / "writer"), shard_limit=0),
+ lambda: load_ftw_banks(str(path), num_layers=0),
+ lambda: load_ftw_banks(str(path), num_layers=1, layer_residency=[]),
+):
+ try:
+ action()
+ except ValueError:
+ pass
+ else:
+ raise SystemExit("invalid writer/reader argument was accepted under -O")
+
+schema_cases = []
+
+mixed = path / "mixed"
+writer = FTWWriter(str(mixed), shard_limit=4096)
+writer.add_tensor("bank", torch.ones(1), kind="experts_bank")
+writer.add_tensor("bank#L00000", torch.ones(1), kind="experts_bank")
+writer.finalize({})
+schema_cases.append((mixed, 1))
+
+nondivisible = path / "nondivisible"
+writer = FTWWriter(str(nondivisible), shard_limit=4096)
+writer.add_tensor("bank", torch.ones(3), kind="experts_bank")
+writer.finalize({})
+schema_cases.append((nondivisible, 2))
+
+missing_layer = path / "missing-layer"
+writer = FTWWriter(str(missing_layer), shard_limit=4096)
+writer.add_tensor("bank#L00000", torch.ones(1), kind="experts_bank")
+writer.finalize({})
+schema_cases.append((missing_layer, 2))
+
+for checkpoint, layers in schema_cases:
+ try:
+ load_ftw_banks(str(checkpoint), num_layers=layers)
+ except FTWFormatError:
+ pass
+ else:
+ raise SystemExit("malformed expert-bank schema was accepted under -O")
+'''
+ env = os.environ.copy()
+ env["PYTHONPATH"] = os.pathsep.join(filter(None, [
+ str(repo_root / "python"), env.get("PYTHONPATH", ""),
+ ]))
+ result = subprocess.run(
+ [sys.executable, "-O", "-c", script, str(tmp_path)],
+ cwd=repo_root,
+ env=env,
+ capture_output=True,
+ text=True,
+ timeout=30,
+ check=False,
+ )
+ assert result.returncode == 0, result.stdout + result.stderr