From 4e9f9c8306028c5eb31986d141ffe19e7c757eee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Mon, 28 Sep 2026 19:02:52 -0400 Subject: [PATCH 1/4] fix(ple): check the shard table's aggregate schema and exact geometry before either backend allocates `source_from_safetensors` (disk) and `load_ple_table` (pinned) each parsed the `shard_` headers on their own and accepted a table that merely covered the hash range: dtype and equal shapes were checked, but not the byte extent against `rows x cols`, extents against the file, the scale's dtype/shape/value, the layer the tensors belong to, or the total row count against the config. Both now share `ple_table_layout()`, which refuses by name (tensor, file, expected vs actual): a part that is not F8_E4M3 or not a positive 2-D shape, a shape that differs from the first part's, `data_offsets` that are not `rows x cols x 1` byte or run past the payload, a `weight_scale` that is not exactly one finite positive BF16 scalar, tensors from two layers, duplicate or non-contiguous part indices. `check_ple_geometry()` then requires the layout to be exactly the table the config addresses: `split_ngram_parts` parts, `ngram_head_dim` wide, on a configured PLE layer, and a total row count EQUAL to the padded sum of the per-head n-gram primes (`expected_ple_rows`, derived with the same `derive_ngram_hash_constants` the dummy-weight path uses; Qwen3.8-Flash-Next: 320,001,536 = 128 x 2,500,012). The model passes its args to `resolve_row_source`, so the disk backend is gated the same way as pinned. The scale is read as its two raw bytes instead of through `safe_open`, so a malformed header can no longer surface as a `SafetensorError`. The two existing callers that passed a bare `SimpleNamespace` now carry the geometry fields; the 28-row toy table in test_weight derives from primes 3+5+7+11 padded to 4. --- python/freetoken/models/qwen4_exp/model.py | 2 +- python/freetoken/models/qwen4_exp/ple_disk.py | 72 ++---- python/freetoken/models/qwen4_exp/weight.py | 223 +++++++++++++---- tests/models/qwen4_exp/test_ple.py | 5 + tests/models/qwen4_exp/test_ple_integrity.py | 227 ++++++++++++++++++ tests/models/qwen4_exp/test_weight.py | 12 +- 6 files changed, 434 insertions(+), 107 deletions(-) create mode 100644 tests/models/qwen4_exp/test_ple_integrity.py diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 4eb69165d5..06a752810c 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -186,7 +186,7 @@ def load_host_tables(self, engine_config) -> int: "image_token_id": args.image_token_id, } disk_table = DiskRowTable( - resolve_row_source(engine_config.model_path), + resolve_row_source(engine_config.model_path, args), constants, max_graph_rows=max(256, engine_config.cuda_graph_max_bs or 0), max_extend_tokens=engine_config.max_extend_tokens, diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index 536853be3b..78b2a8e218 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -10,7 +10,6 @@ from dataclasses import dataclass from typing import Sequence -import safetensors import torch from freetoken.mm import MM_PAD_SHIFT_VALUE, restore_placeholder @@ -19,13 +18,7 @@ from freetoken.kernel.pinned import alloc_pinned_tensor from freetoken.utils import init_logger -from .weight import ( - _PLE_SCALE_SUFFIX, - _PLE_SHARD_RE, - _PLE_ST_DTYPES, - _ple_table_files, - _safetensors_header, -) +from .weight import check_ple_geometry, ple_table_layout _IO_URING_ENV = "FREETOKEN_PLE_IO_URING" _SYNC_ENV = "FREETOKEN_PLE_SYNC" # auto | wait | gate @@ -60,50 +53,25 @@ def total_rows(self) -> int: return len(self.extent_base) * self.rows_per_extent -def source_from_safetensors(folder: str) -> PleRowSource: - """Map the checkpoint's ``ngram_embedding.shard_`` tensors in place: one extent per shard, no copy.""" - rows = cols = 0 - dtype: torch.dtype | None = None - scale: torch.Tensor | None = None +def source_from_safetensors(folder: str, qwen4_args=None) -> PleRowSource: + """Map the checkpoint's ``ngram_embedding.shard_`` tensors in place: one extent per shard, no copy. + + The shard headers must describe one consistent table (:func:`ple_table_layout`) and, given the config, exactly the table it addresses (:func:`check_ple_geometry`).""" + layout = ple_table_layout(folder) + if qwen4_args is not None: + check_ple_geometry(layout, qwen4_args) paths: list[str] = [] path_idx: dict[str, int] = {} - shards: dict[int, tuple[int, int]] = {} - for path in _ple_table_files(folder): - header, base = _safetensors_header(path) - for key, meta in header.items(): - if key == "__metadata__": - continue - if key.endswith(_PLE_SCALE_SUFFIX): - with safetensors.safe_open(path, framework="pt", device="cpu") as f: - scale = f.get_tensor(key).reshape(()) - continue - match = _PLE_SHARD_RE.search(key) - if match is None: - continue - shard_dtype = _PLE_ST_DTYPES.get(meta["dtype"]) - if shard_dtype is None or dtype not in (None, shard_dtype): - raise ValueError(f"PLE shard {key} has dtype {meta['dtype']}, expected one of {sorted(_PLE_ST_DTYPES)} for every shard") - dtype = shard_dtype - if rows and tuple(meta["shape"]) != (rows, cols): - raise ValueError(f"PLE shard {key} is {meta['shape']}, expected {[rows, cols]}") - rows, cols = meta["shape"] - if path not in path_idx: - path_idx[path] = len(paths) - paths.append(path) - idx = int(match.group("shard")) - if idx in shards: - raise ValueError(f"duplicate PLE shard {idx} in {path}") - shards[idx] = (path_idx[path], base + meta["data_offsets"][0]) - if sorted(shards) != list(range(len(shards))) or not shards: - raise ValueError(f"PLE shard indices are not contiguous 0..N-1: {sorted(shards)[:8]}") - if scale is None: - if dtype is torch.float8_e4m3fn: - raise ValueError("PLE table has no weight_scale") - scale = torch.tensor(1.0) - order = [shards[i] for i in range(len(shards))] - row_bytes = cols * dtype.itemsize - return PleRowSource(paths, [f for f, _ in order], [b for _, b in order], rows, row_bytes, row_bytes, float(scale), - row_values=cols, row_dtype=dtype) + for part in layout.parts: + if part.path not in path_idx: + path_idx[part.path] = len(paths) + paths.append(part.path) + row_bytes = layout.cols * layout.dtype.itemsize + return PleRowSource( + paths, [path_idx[p.path] for p in layout.parts], [p.file_offset for p in layout.parts], + layout.rows_per_part, row_bytes, row_bytes, float(layout.scale), + row_values=layout.cols, row_dtype=layout.dtype, + ) def source_from_gguf(model_path: str) -> PleRowSource: @@ -118,7 +86,7 @@ def source_from_gguf(model_path: str) -> PleRowSource: return PleRowSource([info.path], [0], [info.offset], rows, nbytes, nbytes, 1.0, info.ggml_type, values) -def resolve_row_source(model_path: str) -> PleRowSource: +def resolve_row_source(model_path: str, qwen4_args=None) -> PleRowSource: """Pick the row source for a checkpoint; the seam where a repacked format would plug in.""" from freetoken.models.gguf.reader import gguf_config_source @@ -126,7 +94,7 @@ def resolve_row_source(model_path: str) -> PleRowSource: return source_from_gguf(gguf_src) from freetoken.utils import download_hf_weight - return source_from_safetensors(download_hf_weight(model_path)) + return source_from_safetensors(download_hf_weight(model_path), qwen4_args) class DiskRowTable: diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 6c62d141cb..a93fe15fc0 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -61,9 +61,9 @@ # The n-gram table itself: too big for the dense state dict, loaded by load_ple_table. _PLE_TABLE_INFIX = ".ple.ple_embedding.ngram_embedding." _PLE_SHARD_RE = re.compile( - r"\.ple\.ple_embedding\.ngram_embedding\.shard_(?P\d+)\.weight$" + r"\.layers\.(?P\d+)\.ple\.ple_embedding\.ngram_embedding\.shard_(?P\d+)\.weight$" ) -_PLE_SCALE_SUFFIX = ".ple.ple_embedding.ngram_embedding.weight_scale" +_PLE_SCALE_RE = re.compile(r"\.layers\.(?P\d+)\.ple\.ple_embedding\.ngram_embedding\.weight_scale$") _PLE_FILE_BYTES = 4 << 30 # ple-table-*.safetensors written by ftw_side_files # Zero-centered Qwen4ExpTextRMSNorm weights, loaded RAW: GroupedPlusOneRMSNorm / GemmaPlusOneRMSNorm @@ -325,6 +325,38 @@ def tensor(self) -> torch.Tensor: _PLE_ST_DTYPES = {"F8_E4M3": torch.float8_e4m3fn, "BF16": torch.bfloat16} +_PLE_SCALE_ST_DTYPE = "BF16" + + +@dataclass(frozen=True) +class PleShardPart: + """One ``shard_`` tensor of the table: where its ``rows_per_part x cols`` bytes sit on disk.""" + + index: int + name: str + path: str + file_offset: int + nbytes: int + + +@dataclass(frozen=True) +class PleTableLayout: + """The table as the shard headers describe it, after the aggregate schema check: parts ``0..N-1`` of equal ``[rows_per_part, cols]`` blocks of one dtype (F8_E4M3 or BF16), one BF16 scale (1 for a bf16 table that stores none), one layer.""" + + parts: tuple[PleShardPart, ...] + scale: torch.Tensor # bf16 scalar + rows_per_part: int + cols: int + layer: int + dtype: torch.dtype = torch.float8_e4m3fn + + @property + def total_rows(self) -> int: + return len(self.parts) * self.rows_per_part + + @property + def total_bytes(self) -> int: + return self.total_rows * self.cols * self.dtype.itemsize def _safetensors_header(path: str) -> tuple[dict, int]: @@ -333,6 +365,130 @@ def _safetensors_header(path: str) -> tuple[dict, int]: return json.loads(fh.read(n)), 8 + n +def _read_bytes(path: str, offset: int, nbytes: int) -> bytes: + with open(path, "rb") as fh: + fh.seek(offset) + return fh.read(nbytes) + + +def _data_offsets(where: str, meta: dict, payload_bytes: int) -> tuple[int, int]: + """``meta["data_offsets"]`` as ``(begin, end)``, both inside the file's payload.""" + offsets = meta.get("data_offsets") + ok = isinstance(offsets, list) and len(offsets) == 2 and all(type(o) is int for o in offsets) + if not ok or not 0 <= offsets[0] <= offsets[1] <= payload_bytes: + raise ValueError(f"PLE tensor {where} data_offsets {offsets} run past the {payload_bytes}-byte payload") + return offsets[0], offsets[1] + + +def ple_table_layout(folder: str) -> PleTableLayout: + """Parse every shard header holding a piece of the table and check that they describe one consistent table. + + Refuses, naming the tensor and file: a part that is not F8_E4M3 or BF16, or not the first part's dtype, not a positive 2-D shape, or a different shape from the first part; a byte extent that is not ``rows x cols x itemsize``; an extent outside its file; a scale that is not exactly one finite positive BF16 scalar (an fp8 table must store one; a bf16 table may omit it and is unscaled); tensors from two layers; duplicate or non-contiguous part indices. No table bytes are read. + """ + parts: dict[int, PleShardPart] = {} + scale: tuple[str, bytes] | None = None # (where, raw bf16) + layer: tuple[int, str] | None = None # (layer id, the tensor it was taken from) + rows = cols = 0 + dtype: torch.dtype | None = None # the first part's dtype, which every other part must share + first = "" # the part whose shape every other part must match + for path in _ple_table_files(folder): + header, base = _safetensors_header(path) + payload_bytes = os.path.getsize(path) - base + for key, meta in header.items(): + if key == "__metadata__": + continue + part = _PLE_SHARD_RE.search(key) + scale_match = None if part is not None else _PLE_SCALE_RE.search(key) + if part is None and scale_match is None: + continue + where = f"{key} in {os.path.basename(path)}" + key_layer = int((part or scale_match).group("layer")) + if layer is None: + layer = (key_layer, where) + elif key_layer != layer[0]: + raise ValueError(f"PLE tensor {where} is for layer {key_layer}; {layer[1]} is for layer {layer[0]}") + begin, end = _data_offsets(where, meta, payload_bytes) + if scale_match is not None: + if meta.get("dtype") != _PLE_SCALE_ST_DTYPE or meta.get("shape") not in ([], [1]) or end - begin != 2: + raise ValueError( + f"PLE weight_scale {where} must be one BF16 scalar, got dtype {meta.get('dtype')}, " + f"shape {meta.get('shape')}, {end - begin} bytes" + ) + if scale is not None: + raise ValueError(f"PLE table has two weight_scale tensors: {scale[0]} and {where}") + scale = (where, _read_bytes(path, base + begin, 2)) + continue + shard_dtype = _PLE_ST_DTYPES.get(meta.get("dtype")) + if shard_dtype is None or dtype not in (None, shard_dtype): + raise ValueError(f"PLE shard {where} has dtype {meta.get('dtype')}, expected one of {sorted(_PLE_ST_DTYPES)} for every shard") + dtype = shard_dtype + shape = meta.get("shape") + if not (isinstance(shape, list) and len(shape) == 2 and all(type(d) is int and d > 0 for d in shape)): + raise ValueError(f"PLE shard {where} has shape {shape}, expected a positive 2-D [rows, cols]") + if first and shape != [rows, cols]: + raise ValueError(f"PLE shard {where} is {shape}, expected {[rows, cols]} like {first}") + rows, cols = shape + first = first or where + itemsize = dtype.itemsize + if end - begin != rows * cols * itemsize: + raise ValueError( + f"PLE shard {where} spans {end - begin} bytes; shape {shape} x {itemsize} byte{'s' if itemsize != 1 else ''} " + f"needs {rows * cols * itemsize}" + ) + index = int(part.group("shard")) + if index in parts: + seen = parts[index] + raise ValueError(f"duplicate PLE shard {index}: {seen.name} in {os.path.basename(seen.path)} and {where}") + parts[index] = PleShardPart(index, key, path, base + begin, end - begin) + if not parts or sorted(parts) != list(range(len(parts))): + raise ValueError(f"PLE shard indices are not contiguous 0..N-1: {sorted(parts)[:8]}") + if scale is None: + if dtype is torch.float8_e4m3fn: + raise ValueError("PLE table has no weight_scale") + value = torch.tensor(1.0, dtype=torch.bfloat16) # a bf16 table without one is unscaled + else: + value = torch.frombuffer(bytearray(scale[1]), dtype=torch.bfloat16).clone().reshape(()) + if not bool(torch.isfinite(value)) or float(value) <= 0: + raise ValueError(f"PLE weight_scale {scale[0]} must be finite and positive, got {float(value)}") + return PleTableLayout(tuple(parts[i] for i in range(len(parts))), value, rows, cols, layer[0], dtype) + + +def expected_ple_rows(qwen4_args, *, ple_index: int = 0) -> int: + """The row count the checkpoint table must have for this config. + + HF sizes the table at init as the sum of the per-head n-gram vocabulary primes (the same + :func:`derive_ngram_hash_constants` the dummy-weight path uses), padded up to + ``make_ngram_vocab_size_divisible_by``; Qwen3.8-Flash-Next: 16 primes after 19,999,999 -> 320,001,536. + """ + from .ple import derive_ngram_hash_constants + + heads = (int(qwen4_args.ngram_size) - 1) * int(qwen4_args.heads_per_ngram) + _, sizes, _ = derive_ngram_hash_constants( # vocab_size only shapes the multipliers, unused here + vocab_size=1, ngram_size=int(qwen4_args.ngram_size), num_ngram_heads=heads, + ngram_vocab_size_base=int(qwen4_args.ngram_vocab_size_base), ple_layer_index=ple_index, + ) + divisible = int(qwen4_args.make_ngram_vocab_size_divisible_by) + return -(-sum(sizes) // divisible) * divisible + + +def check_ple_geometry(layout: PleTableLayout, qwen4_args) -> None: + """The layout must be exactly the table this config addresses: part count, row width, layer, and the padded total row count (equal, not merely covering).""" + parts = int(qwen4_args.split_ngram_parts) + if len(layout.parts) != parts: + raise ValueError(f"PLE table needs shards 0..{parts - 1}, found {len(layout.parts)}") + if layout.cols != qwen4_args.ngram_head_dim: + raise ValueError(f"PLE table row is {layout.cols} wide, config says {qwen4_args.ngram_head_dim}") + layer_ids = tuple(getattr(qwen4_args, "ple_layer_ids", None) or ()) + if layer_ids and layout.layer not in layer_ids: + raise ValueError(f"PLE table tensors are for layer {layout.layer}, config ple_layer_ids is {list(layer_ids)}") + expected = expected_ple_rows(qwen4_args, ple_index=layer_ids.index(layout.layer) if layer_ids else 0) + if layout.total_rows != expected: + raise ValueError( + f"PLE table has {layout.total_rows} rows ({len(layout.parts)} x {layout.rows_per_part}); " + f"config geometry requires {expected}" + ) + + def _ple_table_files(folder: str) -> list[str]: """Shards holding a piece of the n-gram table, from the index when there is one.""" index = os.path.join(folder, "model.safetensors.index.json") @@ -388,61 +544,23 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, An fp8 table needs the checkpoint's scalar ``weight_scale``; a bf16 table without one is unscaled. """ folder = download_hf_weight(model_path) - parts: dict[int, tuple[str, int, int]] = {} # shard index -> (path, file offset, bytes) - scale: torch.Tensor | None = None - rows = cols = 0 - dtype: torch.dtype | None = None - for path in _ple_table_files(folder): - header, base = _safetensors_header(path) - for key, meta in header.items(): - if key == "__metadata__": - continue - if key.endswith(_PLE_SCALE_SUFFIX): - with safetensors.safe_open(path, framework="pt", device="cpu") as f: - scale = f.get_tensor(key).reshape(()) - continue - match = _PLE_SHARD_RE.search(key) - if match is None: - continue - shard_dtype = _PLE_ST_DTYPES.get(meta["dtype"]) - if shard_dtype is None or dtype not in (None, shard_dtype): - raise ValueError(f"PLE table shard {key} has dtype {meta['dtype']}, expected one of {sorted(_PLE_ST_DTYPES)} for every shard") - dtype = shard_dtype - shape = meta["shape"] - if rows and tuple(shape) != (rows, cols): - raise ValueError(f"PLE table shard {key} is {shape}, expected {[rows, cols]}") - rows, cols = shape - begin, end = meta["data_offsets"] - parts[int(match.group("shard"))] = (path, base + begin, end - begin) - - expected = int(qwen4_args.split_ngram_parts) - if sorted(parts) != list(range(expected)): - raise ValueError( - f"PLE table needs shards 0..{expected - 1}, found {len(parts)}: {sorted(parts)[:8]}" - ) - if cols != qwen4_args.ngram_head_dim: - raise ValueError(f"PLE table row is {cols} wide, config says {qwen4_args.ngram_head_dim}") - if scale is None: - if dtype is torch.float8_e4m3fn: - raise ValueError("PLE table has no weight_scale") - scale = torch.tensor(1.0, dtype=torch.bfloat16) + layout = ple_table_layout(folder) + check_ple_geometry(layout, qwen4_args) - bank = HostBank((expected * rows, cols), dtype) - shard_bytes = rows * cols * dtype.itemsize - bar = byte_bar(expected * shard_bytes, "Loading PLE table") + bank = HostBank((layout.total_rows, layout.cols), layout.dtype) + part_bytes = layout.rows_per_part * layout.cols * layout.dtype.itemsize + bar = byte_bar(layout.total_bytes, "Loading PLE table") try: buf = bank.memoryview() - for shard in range(expected): - path, offset, nbytes = parts[shard] - assert nbytes == shard_bytes, f"PLE shard {shard} is {nbytes} B, expected {shard_bytes}" - read_range_into(buf, path, file_offset=offset, nbytes=nbytes, - dest_offset=shard * shard_bytes, workers=workers, chunk=chunk) - bar.update(nbytes) + for part in layout.parts: + read_range_into(buf, part.path, file_offset=part.file_offset, nbytes=part.nbytes, + dest_offset=part.index * part_bytes, workers=workers, chunk=chunk) + bar.update(part.nbytes) finally: bar.close() if pin and torch.cuda.is_available(): bank.pin() - return PleTable(bank=bank, weight_scale=scale) + return PleTable(bank=bank, weight_scale=layout.scale) # ====================================================================================== @@ -467,7 +585,12 @@ def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: __all__ = [ "nvfp4_expert_spec", + "PleShardPart", "PleTable", + "PleTableLayout", + "check_ple_geometry", + "expected_ple_rows", "iter_weights", "load_ple_table", + "ple_table_layout", ] diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 0d2461427e..fbe8f4dd5b 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -240,6 +240,11 @@ def test_pinned_uva_real_table(): args = SimpleNamespace( split_ngram_parts=text["split_ngram_parts"], ngram_head_dim=text["ple_embed_dim"] // heads, + ngram_size=text["ngram_size"], + heads_per_ngram=text["heads_per_ngram"], + ngram_vocab_size_base=text["ngram_vocab_size_base"], + make_ngram_vocab_size_divisible_by=text["make_ngram_vocab_size_divisible_by"], + ple_layer_ids=tuple(i - 1 for i in text["ple_layer_ids"]), # HF stores them 1-based ) table = load_ple_table(path, args) scale = float(table.weight_scale) diff --git a/tests/models/qwen4_exp/test_ple_integrity.py b/tests/models/qwen4_exp/test_ple_integrity.py new file mode 100644 index 0000000000..e0a02ed1e3 --- /dev/null +++ b/tests/models/qwen4_exp/test_ple_integrity.py @@ -0,0 +1,227 @@ +"""PLE n-gram table integrity: the aggregate shard schema and the config-derived geometry are checked, by name, before either backend touches a row. + +The fixture models the shipping checkpoint: ``shard_`` F8_E4M3 blocks of equal shape spread over several ``model-*.safetensors`` files, one BF16 ``weight_scale`` in one of them, and a ``model.safetensors.index.json`` naming them. The geometry is the toy config's own (``ngram_vocab_size_base`` 1000 -> primes 1009+1013+1019+1021 = 4062, padded to 8 -> 4064 rows over 4 parts of 1016), so the exact-row-count gate is exercised for real, not with a table sized to whatever the loader accepts. +""" + +from __future__ import annotations + +import json +import os +import random +import struct +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.models.qwen4_exp import weight +from freetoken.models.qwen4_exp.config import parse_config +from freetoken.models.qwen4_exp.ple_disk import source_from_safetensors +from freetoken.models.qwen4_exp.weight import ( + check_ple_geometry, + expected_ple_rows, + load_ple_table, + ple_table_layout, +) + +from .common import toy_hf_config + +LAYER = 1 # toy config ple_layer_ids [2] is 1-based -> the checkpoint keys say layers.1 +PREFIX = f"model.language_model.layers.{LAYER}.ple.ple_embedding.ngram_embedding" +ROWS, COLS, PARTS = 1016, 16, 4 # toy geometry: 4064 rows = 4 x 1016, ple_embed_dim 64 / 4 heads + + +def _args(): + return parse_config(toy_hf_config()).qwen4_args + + +def _st_bytes(entries: list[tuple[str, str, list[int], bytes]]) -> bytes: + """A safetensors file from ``(name, dtype, shape, payload)`` entries; the header says whatever it is told.""" + header: dict = {} + payload, offset = [], 0 + for name, dtype, shape, raw in entries: + header[name] = {"dtype": dtype, "shape": shape, "data_offsets": [offset, offset + len(raw)]} + payload.append(raw) + offset += len(raw) + encoded = json.dumps(header).encode() + encoded += b" " * (-len(encoded) % 8) + return struct.pack(" bytes: + return struct.pack(" dict[str, bytes]: + """Write the toy table; returns ``{tensor name: payload bytes}``.""" + prefix = f"model.language_model.layers.{layer}.ple.ple_embedding.ngram_embedding" + raw = {f"{prefix}.shard_{i}.weight": random.Random(i).randbytes(rows * cols) for i in range(parts)} + per_file: list[list[tuple[str, str, list[int], bytes]]] = [[] for _ in range(files)] + for i, (name, data) in enumerate(raw.items()): + per_file[i % files].append((name, "F8_E4M3", [rows, cols], data)) + per_file[-1].append((f"{prefix}.weight_scale", "BF16", [], _bf16(scale))) + weight_map = {} + for n, entries in enumerate(per_file): + filename = f"model-{n:05d}-of-{files:05d}.safetensors" + (folder / filename).write_bytes(_st_bytes(entries)) + weight_map.update({name: filename for name, *_ in entries}) + if index: + (folder / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + return raw + + +def _file_of(folder: Path, name: str) -> Path: + weight_map = json.loads((folder / "model.safetensors.index.json").read_text())["weight_map"] + return folder / weight_map[name] + + +def _rewrite_header(path: Path, edit) -> None: + """Re-serialise ``path``'s header after ``edit(header_dict)``; the payload bytes stay where they are.""" + data = path.read_bytes() + n = struct.unpack(" 28 +NGRAM_ARGS = SimpleNamespace( + split_ngram_parts=NGRAM_SHARDS, ngram_head_dim=NGRAM_DIM, ngram_size=3, heads_per_ngram=2, + ngram_vocab_size_base=3, make_ngram_vocab_size_divisible_by=4, ple_layer_ids=(0,), +) @pytest.fixture(scope="session", autouse=True) @@ -397,8 +402,7 @@ def test_hash_constants_stay_int64(loaded): def test_load_ple_table_concatenates_shards_in_index_order(checkpoint): folder, raw = checkpoint - args = SimpleNamespace(split_ngram_parts=NGRAM_SHARDS, ngram_head_dim=NGRAM_DIM) - table = load_ple_table(folder, args, pin=False) + table = load_ple_table(folder, NGRAM_ARGS, pin=False) assert table.tensor.shape == (NGRAM_SHARDS * NGRAM_ROWS, NGRAM_DIM) assert table.tensor.dtype is torch.float8_e4m3fn prefix = "model.language_model.layers.0.ple.ple_embedding.ngram_embedding" @@ -414,14 +418,14 @@ def test_load_ple_table_reads_an_unscaled_bf16_table(tmp_path): prefix = "model.language_model.layers.0.ple.ple_embedding.ngram_embedding" shards = {f"{prefix}.shard_{i}.weight": _bf16(NGRAM_ROWS, NGRAM_DIM) for i in range(NGRAM_SHARDS)} save_file(shards, str(tmp_path / "model.safetensors")) - table = load_ple_table(str(tmp_path), SimpleNamespace(split_ngram_parts=NGRAM_SHARDS, ngram_head_dim=NGRAM_DIM), pin=False) + table = load_ple_table(str(tmp_path), NGRAM_ARGS, pin=False) assert torch.equal(table.tensor, torch.cat([shards[f"{prefix}.shard_{i}.weight"] for i in range(NGRAM_SHARDS)])) assert float(table.weight_scale) == 1.0 def test_load_ple_table_rejects_a_shard_count_mismatch(checkpoint): folder, _raw = checkpoint - args = SimpleNamespace(split_ngram_parts=NGRAM_SHARDS + 1, ngram_head_dim=NGRAM_DIM) + args = SimpleNamespace(**{**vars(NGRAM_ARGS), "split_ngram_parts": NGRAM_SHARDS + 1}) with pytest.raises(ValueError, match="shards 0"): load_ple_table(folder, args, pin=False) From b5a40a3049b862f61aca4829ad35b2c11b421697 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Mon, 28 Sep 2026 19:05:54 -0400 Subject: [PATCH 2/4] fix(ple): open shards nofollow and bound the header read at discovery `_safetensors_header` opened whatever path the index named with a plain `open()` and read as many bytes as the first eight said: a shard swapped for a symlink was followed, a FIFO would park the loader, and a corrupt length field was handed straight to `fh.read(n)` (u64 max -> OverflowError, a length one past the JSON -> an uncaught UnicodeDecodeError from the payload bytes, a six-byte file -> struct.error). Discovery now canonicalises each shard path exactly once (`realpath(strict)`, because the HF cache is symlinks into blobs/; a dangling link is refused by name) and every later open is `O_RDONLY|O_NOFOLLOW|O_CLOEXEC|O_NONBLOCK` on that target, accepting only a regular file: a symlink that has since appeared at the canonical path is refused as "is a symlink; refusing to follow it". The header length is checked against a 64 MiB budget and the file size before a byte of it is read, and the refusal carries both numbers; the decoded header must be a UTF-8 JSON object or the refusal is a typed ValueError naming the file. The scale's two bytes are read through the same open. --- python/freetoken/models/qwen4_exp/weight.py | 90 ++++++++++++++++--- tests/models/qwen4_exp/test_ple_integrity.py | 94 ++++++++++++++++++++ 2 files changed, 172 insertions(+), 12 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index a93fe15fc0..e6f53b2edf 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -11,9 +11,11 @@ from __future__ import annotations +import errno import json import os import re +import stat import struct from dataclasses import dataclass from typing import Iterator @@ -359,16 +361,79 @@ def total_bytes(self) -> int: return self.total_rows * self.cols * self.dtype.itemsize +_PLE_HEADER_MAX_BYTES = 64 << 20 # a real shard header is ~150 KB; a length past this is refused unread + + +def _canonical_shard_path(path: str) -> str: + """Resolve a discovered shard path once (the HF cache is symlinks into ``blobs/``); every later open is O_NOFOLLOW on the target.""" + try: + return os.path.realpath(path, strict=True) + except OSError as exc: + raise ValueError(f"cannot resolve PLE shard {path}: {exc.strerror}") from exc + + +def _open_ple_shard(path: str) -> int: + """Open a canonical shard path without following a symlink that has since appeared there; only a regular file is accepted.""" + flags = os.O_RDONLY | os.O_NOFOLLOW | os.O_CLOEXEC | getattr(os, "O_NONBLOCK", 0) # NONBLOCK: a FIFO must not park us + try: + fd = os.open(path, flags) + except OSError as exc: + if exc.errno == errno.ELOOP: + raise ValueError(f"PLE shard {path} is a symlink; refusing to follow it") from exc + raise + if not stat.S_ISREG(os.fstat(fd).st_mode): + os.close(fd) + raise ValueError(f"PLE shard {path} is not a regular file") + return fd + + +def _pread_exact(fd: int, nbytes: int, offset: int, path: str) -> bytes: + chunks = [] + while nbytes: + chunk = os.pread(fd, min(nbytes, 1 << 20), offset) + if not chunk: + raise ValueError(f"PLE shard {path} ended after {offset} bytes; a further {nbytes} were declared") + chunks.append(chunk) + nbytes -= len(chunk) + offset += len(chunk) + return b"".join(chunks) + + def _safetensors_header(path: str) -> tuple[dict, int]: - with open(path, "rb") as fh: - n = struct.unpack(" _PLE_HEADER_MAX_BYTES: + raise ValueError( + f"PLE shard {path} declares a {n}-byte safetensors header; the budget is {_PLE_HEADER_MAX_BYTES} bytes" + ) + if 8 + n > size: + raise ValueError( + f"PLE shard {path} declares a {n}-byte safetensors header that runs past the end of the {size}-byte file" + ) + raw = _pread_exact(fd, n, 8, path) + try: + header = json.loads(raw.decode("utf-8")) + except (UnicodeDecodeError, ValueError) as exc: # json.JSONDecodeError is a ValueError + raise ValueError(f"PLE shard {path} safetensors header is not UTF-8 JSON: {exc}") from exc + if not isinstance(header, dict): + raise ValueError(f"PLE shard {path} safetensors header is a JSON {type(header).__name__}, expected an object") + return header, 8 + n + finally: + os.close(fd) def _read_bytes(path: str, offset: int, nbytes: int) -> bytes: - with open(path, "rb") as fh: - fh.seek(offset) - return fh.read(nbytes) + fd = _open_ple_shard(path) + try: + return _pread_exact(fd, nbytes, offset, path) + finally: + os.close(fd) def _data_offsets(where: str, meta: dict, payload_bytes: int) -> tuple[int, int]: @@ -490,14 +555,15 @@ def check_ple_geometry(layout: PleTableLayout, qwen4_args) -> None: def _ple_table_files(folder: str) -> list[str]: - """Shards holding a piece of the n-gram table, from the index when there is one.""" + """Shards holding a piece of the n-gram table, from the index when there is one; canonical paths (see :func:`_canonical_shard_path`).""" index = os.path.join(folder, "model.safetensors.index.json") if not os.path.exists(index): - return sorted(iter_weight_files(folder)) - with open(index, encoding="utf-8") as fh: - weight_map = json.load(fh)["weight_map"] - files = {shard for name, shard in weight_map.items() if _PLE_TABLE_INFIX in name} - return sorted(os.path.join(folder, shard) for shard in files) + files = iter_weight_files(folder) + else: + with open(index, encoding="utf-8") as fh: + weight_map = json.load(fh)["weight_map"] + files = [os.path.join(folder, shard) for shard in {shard for name, shard in weight_map.items() if _PLE_TABLE_INFIX in name}] + return sorted(_canonical_shard_path(path) for path in files) def ftw_side_files(model_path: str, out_dir: str) -> list[str]: diff --git a/tests/models/qwen4_exp/test_ple_integrity.py b/tests/models/qwen4_exp/test_ple_integrity.py index e0a02ed1e3..fdbf1e2486 100644 --- a/tests/models/qwen4_exp/test_ple_integrity.py +++ b/tests/models/qwen4_exp/test_ple_integrity.py @@ -225,3 +225,97 @@ def test_indexless_folder_is_discovered_from_the_shard_headers(tmp_path): assert layout.total_rows == 4064 assert {os.path.basename(p.path) for p in layout.parts} == {"model-00000-of-00002.safetensors", "model-00001-of-00002.safetensors"} assert len(raw) == PARTS + + +# ====================================================================================== +# 2. discovery: nofollow shard opens, bounded header reads +# ====================================================================================== + + +def test_hf_cache_symlink_layout_resolves_once_and_then_opens_nofollow(tmp_path): + """The HF cache is symlinks into blobs/: discovery canonicalises each shard path exactly once; every later open is O_NOFOLLOW on the target.""" + snapshot, blobs = tmp_path / "snapshot", tmp_path / "blobs" + snapshot.mkdir() + blobs.mkdir() + raw = _write_table(snapshot) + for shard in sorted(snapshot.glob("*.safetensors")): + target = blobs / f"blob-{shard.name}" + shard.rename(target) + shard.symlink_to(target) + layout = ple_table_layout(str(snapshot)) + assert layout.total_rows == 4064 + assert {p.path for p in layout.parts} == {str(blobs / f"blob-model-{n:05d}-of-00002.safetensors") for n in range(2)} + table = load_ple_table(str(snapshot), _args(), pin=False) + assert table.tensor[:ROWS].view(torch.uint8).numpy().tobytes() == raw[f"{PREFIX}.shard_0.weight"] + + +def test_symlink_and_non_regular_shard_paths_are_refused_by_name(tmp_path): + _write_table(tmp_path) + shard = tmp_path / "model-00000-of-00002.safetensors" + link = tmp_path / "link.safetensors" + link.symlink_to(shard) + with pytest.raises(ValueError, match=rf"PLE shard {link} is a symlink; refusing to follow it"): + weight._safetensors_header(str(link)) + fifo = tmp_path / "fifo.safetensors" + os.mkfifo(fifo) + with pytest.raises(ValueError, match=rf"PLE shard {fifo} is not a regular file"): + weight._safetensors_header(str(fifo)) + with pytest.raises(ValueError, match=rf"PLE shard {tmp_path} is not a regular file"): + weight._safetensors_header(str(tmp_path)) + dangling = tmp_path / "model-00001-of-00002.safetensors" + dangling.unlink() + dangling.symlink_to(tmp_path / "gone") + with pytest.raises(ValueError, match=rf"cannot resolve PLE shard {dangling}: No such file or directory"): + ple_table_layout(str(tmp_path)) + + +@pytest.mark.parametrize( + "case, header_len, message", + [ + ("one byte over budget", (64 << 20) + 1, + r"declares a 67108865-byte safetensors header; the budget is 67108864 bytes"), + ("u64 max", (1 << 64) - 1, + r"declares a 18446744073709551615-byte safetensors header; the budget is 67108864 bytes"), + ("past the end of the file", None, + r"declares a {n}-byte safetensors header that runs past the end of the {size}-byte file"), + ], +) +def test_header_length_is_bounded_before_it_is_read(tmp_path, no_bank, case, header_len, message): + _write_table(tmp_path) + shard = tmp_path / "model-00000-of-00002.safetensors" + size = shard.stat().st_size + if header_len is None: + header_len = size - 7 # in budget, one byte past the file + message = message.format(n=header_len, size=size) + with shard.open("r+b") as fh: + fh.write(struct.pack(" Date: Mon, 28 Sep 2026 19:09:08 -0400 Subject: [PATCH 3/4] fix(ple): revalidate each shard's identity and header digest before its rows are read Both backends parsed the shard headers once and then read payload bytes from whatever sat at those paths later: the pinned loader per O_DIRECT read, the disk store when it pinned its descriptors. Nothing tied the read to the file preflight had admitted. Preflight now records, per shard file, what it saw through the nofollow open (`PleShardIdentity`: device, inode, size, mtime, ctime, header base and the SHA-256 of the header bytes it parsed). `revalidate_ple_shard` re-opens the canonical path nofollow and refuses with the field that moved ("changed after preflight: st_ctime_ns A -> B", "... st_ino ...", "is a symlink") and, for a same-inode edit whose timestamps were put back, with the digest that moved ("header changed after preflight: sha256 A -> B"). The pinned loader runs it right before each file's first read and re-stats every file once the bank is filled ("changed while reading"); `PleRowSource` carries the identities and `DiskRowTable` runs it before `PleStore` opens its descriptors. Best-effort concurrent-change evidence, as the reviewed T16 branch framed it, not authentication of an untrusted checkpoint. --- python/freetoken/models/qwen4_exp/ple_disk.py | 7 +- python/freetoken/models/qwen4_exp/weight.py | 110 ++++++++++++--- tests/models/qwen4_exp/test_ple_integrity.py | 130 ++++++++++++++++++ 3 files changed, 224 insertions(+), 23 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index 78b2a8e218..ba96d43075 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -18,7 +18,7 @@ from freetoken.kernel.pinned import alloc_pinned_tensor from freetoken.utils import init_logger -from .weight import check_ple_geometry, ple_table_layout +from .weight import PleShardIdentity, check_ple_geometry, ple_table_layout, revalidate_ple_shard _IO_URING_ENV = "FREETOKEN_PLE_IO_URING" _SYNC_ENV = "FREETOKEN_PLE_SYNC" # auto | wait | gate @@ -47,6 +47,7 @@ class PleRowSource: ggml_type: int | None = None row_values: int | None = None row_dtype: torch.dtype = torch.float8_e4m3fn + files: tuple[PleShardIdentity, ...] = () # what preflight saw of each file, re-checked before the store opens them @property def total_rows(self) -> int: @@ -70,7 +71,7 @@ def source_from_safetensors(folder: str, qwen4_args=None) -> PleRowSource: return PleRowSource( paths, [path_idx[p.path] for p in layout.parts], [p.file_offset for p in layout.parts], layout.rows_per_part, row_bytes, row_bytes, float(layout.scale), - row_values=layout.cols, row_dtype=layout.dtype, + row_values=layout.cols, row_dtype=layout.dtype, files=layout.files, ) @@ -109,6 +110,8 @@ def __init__( max_extend_tokens: int = 8192, dtype: torch.dtype = torch.bfloat16, ) -> None: + for shard in source.files: # the files the store is about to pin descriptors to are the ones preflight parsed + revalidate_ple_shard(shard) from freetoken.kernel.row_store import PleStore from freetoken.layers.quantization.ggml_types import float_dtype diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index e6f53b2edf..268e0aa697 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -12,6 +12,7 @@ from __future__ import annotations import errno +import hashlib import json import os import re @@ -341,15 +342,36 @@ class PleShardPart: nbytes: int +@dataclass(frozen=True) +class PleShardIdentity: + """What preflight saw of one shard file: its inode identity, size, timestamps and the SHA-256 of the header bytes it parsed. + + Best-effort evidence that the file a payload read opens is the file the header came from; it does not authenticate the checkpoint.""" + + path: str + st_dev: int + st_ino: int + st_size: int + st_mtime_ns: int + st_ctime_ns: int + header_base: int + header_sha256: str + + @classmethod + def of(cls, path: str, st: os.stat_result, header_base: int, header_sha256: str) -> "PleShardIdentity": + return cls(path, st.st_dev, st.st_ino, st.st_size, st.st_mtime_ns, st.st_ctime_ns, header_base, header_sha256) + + @dataclass(frozen=True) class PleTableLayout: - """The table as the shard headers describe it, after the aggregate schema check: parts ``0..N-1`` of equal ``[rows_per_part, cols]`` blocks of one dtype (F8_E4M3 or BF16), one BF16 scale (1 for a bf16 table that stores none), one layer.""" + """The table as the shard headers describe it, after the aggregate schema check: parts ``0..N-1`` of equal ``[rows_per_part, cols]`` blocks of one dtype (F8_E4M3 or BF16), one BF16 scale (1 for a bf16 table that stores none), one layer; ``files`` is what preflight saw of each shard file.""" parts: tuple[PleShardPart, ...] scale: torch.Tensor # bf16 scalar rows_per_part: int cols: int layer: int + files: tuple[PleShardIdentity, ...] dtype: torch.dtype = torch.float8_e4m3fn @property @@ -399,31 +421,65 @@ def _pread_exact(fd: int, nbytes: int, offset: int, path: str) -> bytes: return b"".join(chunks) -def _safetensors_header(path: str) -> tuple[dict, int]: - """The shard's JSON header and its payload base, through a nofollow open and a bounded read: a header length past the budget or past the file is refused with the number, not read.""" +def _header_bytes(fd: int, path: str, size: int) -> bytes: + """The raw header of an open shard, bounded: a header length past the budget or past the file is refused with the number, not read.""" + prefix = os.pread(fd, 8, 0) + if len(prefix) != 8: + raise ValueError(f"PLE shard {path} is {size} bytes; not a safetensors file") + n = struct.unpack(" _PLE_HEADER_MAX_BYTES: + raise ValueError( + f"PLE shard {path} declares a {n}-byte safetensors header; the budget is {_PLE_HEADER_MAX_BYTES} bytes" + ) + if 8 + n > size: + raise ValueError( + f"PLE shard {path} declares a {n}-byte safetensors header that runs past the end of the {size}-byte file" + ) + return _pread_exact(fd, n, 8, path) + + +def _safetensors_header(path: str) -> tuple[dict, PleShardIdentity]: + """The shard's JSON header plus what was seen of the file, through a nofollow open and a bounded read.""" fd = _open_ple_shard(path) try: - size = os.fstat(fd).st_size - prefix = os.pread(fd, 8, 0) - if len(prefix) != 8: - raise ValueError(f"PLE shard {path} is {size} bytes; not a safetensors file") - n = struct.unpack(" _PLE_HEADER_MAX_BYTES: - raise ValueError( - f"PLE shard {path} declares a {n}-byte safetensors header; the budget is {_PLE_HEADER_MAX_BYTES} bytes" - ) - if 8 + n > size: - raise ValueError( - f"PLE shard {path} declares a {n}-byte safetensors header that runs past the end of the {size}-byte file" - ) - raw = _pread_exact(fd, n, 8, path) + st = os.fstat(fd) + raw = _header_bytes(fd, path, st.st_size) try: header = json.loads(raw.decode("utf-8")) except (UnicodeDecodeError, ValueError) as exc: # json.JSONDecodeError is a ValueError raise ValueError(f"PLE shard {path} safetensors header is not UTF-8 JSON: {exc}") from exc if not isinstance(header, dict): raise ValueError(f"PLE shard {path} safetensors header is a JSON {type(header).__name__}, expected an object") - return header, 8 + n + return header, PleShardIdentity.of(path, st, 8 + len(raw), hashlib.sha256(raw).hexdigest()) + finally: + os.close(fd) + + +_IDENTITY_FIELDS = ("st_dev", "st_ino", "st_size", "st_mtime_ns", "st_ctime_ns") + + +def _check_same_file(identity: PleShardIdentity, st: os.stat_result, when: str) -> None: + for field in _IDENTITY_FIELDS: + seen, now = getattr(identity, field), getattr(st, field) + if seen != now: + raise ValueError(f"PLE shard {identity.path} changed {when}: {field} {seen} -> {now}") + + +def revalidate_ple_shard(identity: PleShardIdentity) -> None: + """Immediately before a payload read: the canonical path must still open nofollow to the same regular file (device, inode, size, mtime, ctime) and its header bytes must still hash to what preflight parsed. Refuses with the field or digest that moved.""" + fd = _open_ple_shard(identity.path) + try: + _check_same_file(identity, os.fstat(fd), "after preflight") + raw = _header_bytes(fd, identity.path, identity.st_size) + if 8 + len(raw) != identity.header_base: + raise ValueError( + f"PLE shard {identity.path} header changed after preflight: length {identity.header_base - 8} -> {len(raw)}" + ) + digest = hashlib.sha256(raw).hexdigest() + if digest != identity.header_sha256: + raise ValueError( + f"PLE shard {identity.path} header changed after preflight: sha256 {identity.header_sha256} -> {digest}" + ) finally: os.close(fd) @@ -456,9 +512,12 @@ def ple_table_layout(folder: str) -> PleTableLayout: rows = cols = 0 dtype: torch.dtype | None = None # the first part's dtype, which every other part must share first = "" # the part whose shape every other part must match + files: list[PleShardIdentity] = [] for path in _ple_table_files(folder): - header, base = _safetensors_header(path) - payload_bytes = os.path.getsize(path) - base + header, identity = _safetensors_header(path) + files.append(identity) + base = identity.header_base + payload_bytes = identity.st_size - base for key, meta in header.items(): if key == "__metadata__": continue @@ -515,7 +574,7 @@ def ple_table_layout(folder: str) -> PleTableLayout: value = torch.frombuffer(bytearray(scale[1]), dtype=torch.bfloat16).clone().reshape(()) if not bool(torch.isfinite(value)) or float(value) <= 0: raise ValueError(f"PLE weight_scale {scale[0]} must be finite and positive, got {float(value)}") - return PleTableLayout(tuple(parts[i] for i in range(len(parts))), value, rows, cols, layer[0], dtype) + return PleTableLayout(tuple(parts[i] for i in range(len(parts))), value, rows, cols, layer[0], tuple(files), dtype) def expected_ple_rows(qwen4_args, *, ple_index: int = 0) -> int: @@ -615,13 +674,20 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, bank = HostBank((layout.total_rows, layout.cols), layout.dtype) part_bytes = layout.rows_per_part * layout.cols * layout.dtype.itemsize + identity = {f.path: f for f in layout.files} bar = byte_bar(layout.total_bytes, "Loading PLE table") try: buf = bank.memoryview() + verified: set[str] = set() for part in layout.parts: + if part.path not in verified: # the file is what preflight parsed, right before its first read + revalidate_ple_shard(identity[part.path]) + verified.add(part.path) read_range_into(buf, part.path, file_offset=part.file_offset, nbytes=part.nbytes, dest_offset=part.index * part_bytes, workers=workers, chunk=chunk) bar.update(part.nbytes) + for shard in layout.files: # ... and still is once its bytes are in the bank + _check_same_file(shard, os.stat(shard.path), "while reading") finally: bar.close() if pin and torch.cuda.is_available(): @@ -651,6 +717,7 @@ def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: __all__ = [ "nvfp4_expert_spec", + "PleShardIdentity", "PleShardPart", "PleTable", "PleTableLayout", @@ -659,4 +726,5 @@ def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: "iter_weights", "load_ple_table", "ple_table_layout", + "revalidate_ple_shard", ] diff --git a/tests/models/qwen4_exp/test_ple_integrity.py b/tests/models/qwen4_exp/test_ple_integrity.py index fdbf1e2486..a853b4c9d1 100644 --- a/tests/models/qwen4_exp/test_ple_integrity.py +++ b/tests/models/qwen4_exp/test_ple_integrity.py @@ -319,3 +319,133 @@ def test_header_must_be_a_utf8_json_object(tmp_path, no_bank, case, edit, messag shard.write_bytes(struct.pack(" None: + """Same inode, same size, mtime restored: only ctime and the header bytes betray the edit.""" + before = shard.stat() + data = shard.read_bytes() + n = struct.unpack(" \d+"), + ("one byte appended", lambda p: p.open("ab").write(b"\0"), r"changed after preflight: st_size (\d+) -> \d+"), + ("replaced by a copy", lambda p: (p.with_suffix(".new").write_bytes(p.read_bytes()), os.replace(p.with_suffix(".new"), p)), + r"changed after preflight: st_ino \d+ -> \d+"), + ("swapped for a symlink", lambda p: (p.rename(p.with_suffix(".blob")), p.symlink_to(p.with_suffix(".blob"))), + r"is a symlink; refusing to follow it"), + ], +) +def test_shard_changed_between_preflight_and_read_is_refused_with_the_reason(tmp_path, case, mutate, message): + _write_table(tmp_path) + layout = ple_table_layout(str(tmp_path)) + shard = Path(layout.files[1].path) + mutate(shard) + with pytest.raises(ValueError, match=rf"PLE shard {shard} {message}"): + weight.revalidate_ple_shard(layout.files[1]) + weight.revalidate_ple_shard(layout.files[0]) # the other shard is still what preflight saw + + +def test_header_digest_catches_an_edit_the_stat_gate_cannot_see(tmp_path): + """A writer that also restores ctime (privileged, or a coarse filesystem) still changes the bytes preflight hashed.""" + import dataclasses + + _write_table(tmp_path) + layout = ple_table_layout(str(tmp_path)) + identity = layout.files[0] + _flip_dtype_in_place(Path(identity.path)) + forged = dataclasses.replace(identity, st_ctime_ns=os.stat(identity.path).st_ctime_ns) + with pytest.raises(ValueError, match=rf"PLE shard {identity.path} header changed after preflight: sha256 {identity.header_sha256[:16]}\S* -> [0-9a-f]{{16}}"): + weight.revalidate_ple_shard(forged) + + +def test_pinned_load_revalidates_each_shard_before_its_first_read(tmp_path, monkeypatch): + _write_table(tmp_path) + layout = ple_table_layout(str(tmp_path)) + other = Path(layout.files[1].path) # parts 1 and 3 + reads = [] + real = weight.read_range_into + + def read_then_mutate(buf, path, **kwargs): + reads.append(os.path.basename(path)) + if len(reads) == 1: + _flip_dtype_in_place(other) # while part 0 is being read, model-00001 changes under us + return real(buf, path, **kwargs) + + monkeypatch.setattr(weight, "read_range_into", read_then_mutate) + with pytest.raises(ValueError, match=rf"PLE shard {other} changed after preflight: st_ctime_ns"): + load_ple_table(str(tmp_path), _args(), pin=False) + assert reads == ["model-00000-of-00002.safetensors"] # part 1's read never happened + + +def test_pinned_load_refuses_a_shard_that_changed_while_it_was_being_read(tmp_path, monkeypatch): + _write_table(tmp_path) + layout = ple_table_layout(str(tmp_path)) + first = Path(layout.files[0].path) + reads = [] + real = weight.read_range_into + + def read_then_mutate(buf, path, **kwargs): + reads.append(os.path.basename(path)) + out = real(buf, path, **kwargs) + if len(reads) == 4: + first.open("ab").write(b"\0") # after the last read, model-00000 grows + return out + + monkeypatch.setattr(weight, "read_range_into", read_then_mutate) + with pytest.raises(ValueError, match=rf"PLE shard {first} changed while reading: st_size \d+ -> \d+"): + load_ple_table(str(tmp_path), _args(), pin=False) + assert len(reads) == 4 + + +def test_disk_source_is_revalidated_before_the_store_opens_its_descriptors(tmp_path): + from freetoken.models.qwen4_exp.ple_disk import DiskRowTable + + from .common import EOS, hash_constants + + _write_table(tmp_path) + args = _args() + source = source_from_safetensors(str(tmp_path), args) + multipliers, sizes, offsets = hash_constants(args) + constants = { + "num_ngram_heads": args.num_ngram_heads, "layer_multipliers": multipliers.tolist(), + "per_head_vocab_sizes": sizes.tolist(), "per_head_offsets": offsets.tolist(), "eos_token_id": EOS, + } + shard = Path(source.files[0].path) + _flip_dtype_in_place(shard) + with pytest.raises(ValueError, match=rf"PLE shard {shard} changed after preflight: st_ctime_ns"): + DiskRowTable(source, constants) From 918a626eacf22ed960aab812913b236d455c5712 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Mon, 28 Sep 2026 19:11:48 -0400 Subject: [PATCH 4/4] fix(ple): admit the pinned bank by effective memory and discard it on any later failure `--ple-backend pinned` allocated the 47.7 GiB `HostBank` with no look at the host, and a failure after that (a read error, a shard that changed, a refused cudaHostRegister) closed the progress bar and leaked the bank: the lazy mmap stayed in `_LIVE_BUFFERS` with whatever had been read into it. Admission now happens after the schema/geometry checks and before the bank exists: `effective_memory_available()` is the tighter of `MemAvailable` and the cgroup v2 headroom (`memory.max - memory.current`, a known zero when over the limit; `None` when neither is readable, which admits best-effort like the rest of the loader). The table must fit under 88% of it, and the refusal carries the numbers ("needs N bytes ...; effective memory available is A bytes, C after the 12% headroom"). `HostBank.discard()` unmaps an unpublished bank and forgets it (a registered mmap cannot be unmapped without cudaHostUnregister, which the extension does not expose, so such a bank is only dropped by the object); the loader calls it from a `BaseException` handler around fill and pin, releasing its memoryview first and never letting the cleanup mask the failure. The `disk` backend allocates ~21 MiB of fixed staging and needs neither. --- python/freetoken/models/qwen4_exp/weight.py | 53 +++++++- python/freetoken/moe/host_banks.py | 18 +++ tests/models/qwen4_exp/test_ple_integrity.py | 122 +++++++++++++++++++ 3 files changed, 190 insertions(+), 3 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 268e0aa697..ec5fb2224a 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -11,6 +11,7 @@ from __future__ import annotations +import contextlib import errno import hashlib import json @@ -658,6 +659,31 @@ def flush(): return written +_PLE_HEADROOM_PCT = 12 # of effective memory kept free of the pinned bank: Python heaps, reader bounces, page-alignment slack + + +def effective_memory_available(meminfo: str = "/proc/meminfo", cgroup: str = "/sys/fs/cgroup") -> int | None: + """Host bytes a startup allocation may take: the tighter of ``MemAvailable`` and the cgroup v2 headroom (``memory.max - memory.current``, a known zero when over); ``None`` when neither is readable.""" + bounds = [] + try: + with open(meminfo, encoding="utf-8") as fh: + for line in fh: + if line.startswith("MemAvailable:"): + bounds.append(int(line.split()[1]) * 1024) + break + except (OSError, ValueError): + pass + try: + with open(os.path.join(cgroup, "memory.max"), encoding="utf-8") as fh: + limit = fh.read().strip() + if limit != "max": + with open(os.path.join(cgroup, "memory.current"), encoding="utf-8") as fh: + bounds.append(max(0, int(limit) - int(fh.read().strip()))) + except (OSError, ValueError): + pass + return min(bounds) if bounds else None + + def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, workers: int = 8, chunk: int = 8 << 20) -> PleTable: """Concatenate the checkpoint's ``ngram_embedding.shard_`` tensors into one pinned host bank. @@ -667,15 +693,27 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, bank is filled shard by shard at ``shard_index * rows_per_shard``. Each read is O_DIRECT: the table is ~47.7 GiB (fp8) and must not also sit in the page cache while the bank holds the same bytes. An fp8 table needs the checkpoint's scalar ``weight_scale``; a bf16 table without one is unscaled. + + Schema, geometry and the effective-memory admission (12% headroom) all finish before the bank + exists; any failure after that (a read, a shard that changed, the pin) discards the bank and + publishes nothing. """ folder = download_hf_weight(model_path) layout = ple_table_layout(folder) check_ple_geometry(layout, qwen4_args) + available = effective_memory_available() + if available is not None and layout.total_bytes > available * (100 - _PLE_HEADROOM_PCT) // 100: + raise ValueError( + f"PLE table needs {layout.total_bytes} bytes of host memory for the pinned bank; effective memory " + f"available is {available} bytes, {available * (100 - _PLE_HEADROOM_PCT) // 100} after the " + f"{_PLE_HEADROOM_PCT}% headroom" + ) bank = HostBank((layout.total_rows, layout.cols), layout.dtype) part_bytes = layout.rows_per_part * layout.cols * layout.dtype.itemsize identity = {f.path: f for f in layout.files} bar = byte_bar(layout.total_bytes, "Loading PLE table") + buf: memoryview | None = None try: buf = bank.memoryview() verified: set[str] = set() @@ -688,10 +726,18 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, bar.update(part.nbytes) for shard in layout.files: # ... and still is once its bytes are in the bank _check_same_file(shard, os.stat(shard.path), "while reading") - finally: + buf.release() + buf = None bar.close() - if pin and torch.cuda.is_available(): - bank.pin() + if pin and torch.cuda.is_available(): + bank.pin() + except BaseException: + with contextlib.suppress(Exception): # never mask the failure with its own cleanup + if buf is not None: + buf.release() + bar.close() + bank.discard() + raise return PleTable(bank=bank, weight_scale=layout.scale) @@ -722,6 +768,7 @@ def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: "PleTable", "PleTableLayout", "check_ple_geometry", + "effective_memory_available", "expected_ple_rows", "iter_weights", "load_ple_table", diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index 3cdc8d7796..25208f93a4 100644 --- a/python/freetoken/moe/host_banks.py +++ b/python/freetoken/moe/host_banks.py @@ -184,6 +184,24 @@ def unpin(self) -> None: refund_pin(len(self._buf)) self._pinned = False + def discard(self) -> None: + """Free a bank that was never published: drop the tensor view, unmap the buffer and forget it. + + For a loader rolling back after ``HostBank(...)``: nothing else may hold the bank's ``memoryview`` (release it first, or ``mmap.close`` refuses). A registered (``pin()``-ed) mmap is not unmapped here: it is only forgotten by this object and stays mapped and registered (``unpin()`` is the explicit way back). Idempotent.""" + if self.tensor is None: + return + buf, self._buf = self._buf, None + self.tensor = None + self.nbytes = 0 + if not isinstance(buf, mmap.mmap): + return # cuda backing: dropping the numpy slice (and its .base) frees the pinned allocation + if self._pinned: + logger.warning(f"discarding a registered {len(buf) / 2**30:.1f} GiB bank: the mapping stays registered until exit") + return + buf.close() + with contextlib.suppress(ValueError): + _LIVE_BUFFERS.remove(buf) + def release(self) -> None: """Drop the resident pages; the address space stays valid, the contents become undefined. diff --git a/tests/models/qwen4_exp/test_ple_integrity.py b/tests/models/qwen4_exp/test_ple_integrity.py index a853b4c9d1..ae9735f6cb 100644 --- a/tests/models/qwen4_exp/test_ple_integrity.py +++ b/tests/models/qwen4_exp/test_ple_integrity.py @@ -449,3 +449,125 @@ def test_disk_source_is_revalidated_before_the_store_opens_its_descriptors(tmp_p _flip_dtype_in_place(shard) with pytest.raises(ValueError, match=rf"PLE shard {shard} changed after preflight: st_ctime_ns"): DiskRowTable(source, constants) + + +# ====================================================================================== +# 4. pinned backend: effective-memory admission before the bank exists; rollback after +# ====================================================================================== + + +def _live_mmaps(): + from freetoken.moe import host_banks + + return [m for m in host_banks._LIVE_BUFFERS if not m.closed] + + +def test_effective_memory_is_the_tighter_of_meminfo_and_the_cgroup(tmp_path): + meminfo = tmp_path / "meminfo" + cgroup = tmp_path / "cgroup" + cgroup.mkdir() + meminfo.write_text("MemTotal: 65536000 kB\nMemFree: 1000 kB\nMemAvailable: 60000000 kB\n") + (cgroup / "memory.max").write_text("max\n") + (cgroup / "memory.current").write_text("1000\n") + assert weight.effective_memory_available(meminfo=str(meminfo), cgroup=str(cgroup)) == 60_000_000 * 1024 + (cgroup / "memory.max").write_text("50000000000\n") + (cgroup / "memory.current").write_text("10000000000\n") + assert weight.effective_memory_available(meminfo=str(meminfo), cgroup=str(cgroup)) == 40_000_000_000 + (cgroup / "memory.current").write_text("60000000000\n") # over its limit: a known zero, not "unknown" + assert weight.effective_memory_available(meminfo=str(meminfo), cgroup=str(cgroup)) == 0 + meminfo.write_text("MemTotal: 1 kB\n") # no MemAvailable, cgroup only + (cgroup / "memory.current").write_text("10000000000\n") + assert weight.effective_memory_available(meminfo=str(meminfo), cgroup=str(cgroup)) == 40_000_000_000 + assert weight.effective_memory_available(meminfo=str(tmp_path / "absent"), cgroup=str(tmp_path / "absent")) is None + assert weight.effective_memory_available() == weight.effective_memory_available("/proc/meminfo", "/sys/fs/cgroup") + + +@pytest.mark.parametrize( + "available, admitted", + [ + (65_024 * 100 // 88, True), # exactly the 12% headroom: 65,024 <= 73,890 * 88 // 100 = 65,023 -> refused + (65_024 * 100 // 88 + 2, True), # one byte of headroom to spare + (65_024, False), # the table fits in RAM but not with headroom + (0, False), # a known zero is a hard refusal + (None, True), # unknown: best effort, admit + ], +) +def test_pinned_bank_is_admitted_by_effective_memory_before_it_is_allocated(tmp_path, monkeypatch, available, admitted): + _write_table(tmp_path) + monkeypatch.setattr(weight, "effective_memory_available", lambda: available) + if available == 65_024 * 100 // 88: + admitted = False # 73,890 * 88 // 100 = 65,023 < 65,024 + if admitted: + table = load_ple_table(str(tmp_path), _args(), pin=False) + assert table.bank.nbytes == 65_024 + return + ceiling = available * 88 // 100 + allocations = [] + real_bank = weight.HostBank + + class Bank(real_bank): + def __init__(self, *a, **k): + allocations.append(a) + super().__init__(*a, **k) + + monkeypatch.setattr(weight, "HostBank", Bank) + with pytest.raises(ValueError, match=rf"PLE table needs 65024 bytes of host memory for the pinned bank; effective memory available is {available} bytes, {ceiling} after the 12% headroom"): + load_ple_table(str(tmp_path), _args(), pin=False) + assert allocations == [] + + +def test_host_bank_discard_frees_the_mapping_and_forgets_it(): + from freetoken.moe.host_banks import HostBank + + before = len(_live_mmaps()) + bank = HostBank((1024, 16), torch.float8_e4m3fn) + assert len(_live_mmaps()) == before + 1 + view = bank.memoryview() + view.release() + bank.discard() + assert len(_live_mmaps()) == before + assert bank.tensor is None and bank.nbytes == 0 + bank.discard() # idempotent + + +@pytest.mark.parametrize( + "case, failure, message", + [ + ("a read fails", "read", r"disk went away"), + ("a shard changed under the read", "mutate", r"changed after preflight: st_ctime_ns"), + ("the pin fails", "pin", r"cudaHostRegister failed for 0\.0 GiB"), + ], +) +def test_failure_after_allocation_discards_the_bank_and_publishes_nothing(tmp_path, monkeypatch, case, failure, message): + from freetoken.moe.host_banks import HostBank, PinFailed + + _write_table(tmp_path) + layout = ple_table_layout(str(tmp_path)) + banks = [] + real_bank = weight.HostBank + real_read = weight.read_range_into + + class Bank(real_bank): + def __init__(self, *a, **k): + super().__init__(*a, **k) + banks.append(self) + + def read(buf, path, **kwargs): + if failure == "read": + raise OSError("disk went away") + if failure == "mutate" and not getattr(read, "done", False): + read.done = True + _flip_dtype_in_place(Path(layout.files[1].path)) + return real_read(buf, path, **kwargs) + + monkeypatch.setattr(weight, "HostBank", Bank) + monkeypatch.setattr(weight, "read_range_into", read) + if failure == "pin": + monkeypatch.setattr(weight.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(HostBank, "pin", lambda self: (_ for _ in ()).throw(PinFailed("cudaHostRegister failed for 0.0 GiB"))) + live = len(_live_mmaps()) + with pytest.raises((ValueError, OSError, PinFailed), match=message): + load_ple_table(str(tmp_path), _args(), pin=True) + assert len(banks) == 1 + assert banks[0].tensor is None and banks[0].nbytes == 0 # discarded ... + assert len(_live_mmaps()) == live # ... and its mapping is closed and forgotten