From 2d89a3b2812d2b22e29674688e2271d50297deed Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Mon, 26 Oct 2026 22:40:09 -0700 Subject: [PATCH 1/5] Initial export of the PTN samples stack [ghstack-poisoned] --- backends/native/README.md | 83 +++ .../test/runtime/deserialize/test_package.cpp | 26 +- exir/_serialize/BUCK | 15 + exir/_serialize/_ptn.py | 12 +- exir/_serialize/_ptn_samples.py | 404 ++++++++++++++ exir/_serialize/test/BUCK | 12 + exir/_serialize/test/test_ptn_samples.py | 520 ++++++++++++++++++ exir/native/BUCK | 1 + exir/native/__init__.py | 12 + exir/native/test/test_native.py | 50 +- 10 files changed, 1128 insertions(+), 7 deletions(-) create mode 100644 exir/_serialize/_ptn_samples.py create mode 100644 exir/_serialize/test/test_ptn_samples.py diff --git a/backends/native/README.md b/backends/native/README.md index a784537989f..086664de294 100644 --- a/backends/native/README.md +++ b/backends/native/README.md @@ -41,6 +41,89 @@ A `.ptn` is an uncompressed zip with fixed entry names: | `program.safetensors` | Constants. Byte-identical immutable tensors are stored once. | | `aliases.json` | Deduplicated key to owner key; present only when needed. | +## Samples (optional) + +A `.ptn` may carry recorded method calls: inputs, and optionally the outputs and +intermediate values they produced. A sample is an ordered sequence of calls, so a +stateful method such as an LLM's prefill followed by decode steps can be checked. They serve reference checks, layer-by-layer +numerical debugging, calibration and benchmarking. Program loaders ignore them. + +```python +from executorch.exir.native import Call, read_samples, Sample, SampleSet, to_native, write_samples + +to_native(ep).save("model.ptn") +golden = SampleSet([Sample([Call([x], [model(x)])])]) +write_samples("model.ptn", "forward", "golden", golden) +read_samples("model.ptn") # {("forward", "golden"): SampleSet(...)} +``` + +`write_samples` appends a set to any existing `.ptn`, not only one this process +saved, and replaces the file only once the set is complete. `read_samples` +validates every set against the rules below, except that `` exists. + +### Layout + +``` +samples///index.json +samples///data.safetensors +``` + +- `` names a program method. `` is an identifier with no meaning to + the format. Both match `[A-Za-z0-9_-]+`. +- Each set has exactly these two members, and nothing else lives under `samples/`. +- Members are stored (uncompressed). The `data.safetensors` data section starts + 64-byte aligned in the file, like `program.safetensors`. +- Adding a set only appends members; removing one means rewriting the zip. + +### `index.json` + +``` +{ "version": "1.0", required, "MAJOR.MINOR" + "description": str, optional + "source": str, optional, producer of the outputs + "tol": [atol, rtol], optional, default for every call + "samples": [Sample, ...] } required, non-empty + +Sample = { "calls": [Call, ...], required, non-empty, run in order + "tol": [atol, rtol] } optional, overrides the set's + +Call = { "in": [Value, ...], required, positional arguments + "out": [Value, ...], optional, positional outputs + "intermediates": {name: Key, ...}, optional + "tol": [atol, rtol] } optional, overrides the sample's + +Value = Key | number | true | false | [Value, ...] +Key = str naming a tensor in the set's data.safetensors +``` + +- Each sample starts from the program's initial state. Samples are independent and + may run in any order; the calls within a sample run in order with no reset + between them, so each sees the state the previous one left. +- A tensor's key is its JSON path in the index, e.g. `samples[0].calls[1].in[2]`. + Below a call that is `.in[k]`, `.in[k][l]` for an element of a list argument, + `.out[k]` or `.intermediates.`. Every key resolves, and every tensor is + referenced exactly once. +- A JSON integer is an int scalar, any other number a float. +- Intermediate names are top-level SSA value names (`[A-Za-z_][A-Za-z0-9_]*`). + Producers record what they can; consumers compare the names both sides have. +- Tolerances are finite and non-negative, and advisory: a consumer may override them. + `write_samples` warns when atol or rtol is above 1e-3. +- Readers reject another major `version` and accept any minor one: a minor bump + only adds fields, and readers ignore unknown fields. + +Not in v1: + +- Sequences that span methods, such as a separate `prefill` and `decode` sharing + state. +- Tensors shared across sets. +- Values inside nested graphs. +- Quantized method inputs and outputs (a `QuantSpec` on graph I/O). Integer and + float I/O of quantized models is supported; torch quantized dtypes and tensor + subclasses are not. + +TODO: specify quantized I/O as the storage form the method binds (packed for +sub-byte), with quant parameters read from the program rather than the sample. + ## Runtime The C++ runtime reads a `.ptn` into an in-memory graph and hands each method to diff --git a/backends/native/test/runtime/deserialize/test_package.cpp b/backends/native/test/runtime/deserialize/test_package.cpp index 3b29afc3d8b..c49791025f5 100644 --- a/backends/native/test/runtime/deserialize/test_package.cpp +++ b/backends/native/test/runtime/deserialize/test_package.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -56,7 +57,9 @@ class TempPackage { std::filesystem::path path_; public: - explicit TempPackage(std::vector tensors = {}) { + explicit TempPackage( + std::vector tensors = {}, + const std::vector>& extra = {}) { path_ = std::filesystem::temp_directory_path() / ("ptn_package_" + std::to_string(reinterpret_cast(this)) + ".ptn"); @@ -74,6 +77,9 @@ class TempPackage { add(archive.get(), kProgramEntry, as_bytes(program_)); add(archive.get(), kSafeTensorsEntry, ByteSpan(tensors)); add(archive.get(), kAliasesEntry, as_bytes(aliases_)); + for (const auto& [name, bytes] : extra) { + add(archive.get(), name.c_str(), as_bytes(bytes)); + } if (zip_close(archive.get()) != 0) { throw std::runtime_error("failed to close test package"); } @@ -135,6 +141,24 @@ TEST(PackageTest, LoadsConstantsOnDemand) { EXPECT_NO_THROW(package.verify()); } +TEST(PackageTest, IgnoresSampleMembers) { + const std::string index = + R"({"version":1,"samples":[{"calls":[{"in":["samples[0].calls[0].in[0]"]}]}]})"; + const std::vector data = make_safetensors( + R"({"samples[0].calls[0].in[0]":{"dtype":"U8","shape":[2],"data_offsets":[0,2]}})", + "in"); + const TempPackage file( + {}, + {{"samples/forward/golden/index.json", index}, + {"samples/forward/golden/data.safetensors", + std::string(data.begin(), data.end())}}); + const Package package = Package::load(file.path()); + + EXPECT_EQ(package.owner_keys(), (std::vector{"weight", "bias"})); + EXPECT_FALSE(package.constant_info("samples[0].calls[0].in[0]")); + EXPECT_NO_THROW(package.verify()); +} + TEST(PackageTest, ReportsMissingConstantsAndWrongDestinations) { const TempPackage file; const Package package = Package::load(file.path()); diff --git a/exir/_serialize/BUCK b/exir/_serialize/BUCK index 4d19c94ac2c..75797e1fd16 100644 --- a/exir/_serialize/BUCK +++ b/exir/_serialize/BUCK @@ -38,6 +38,21 @@ fbcode_target(_kind = runtime.python_library, ], ) +# Optional recorded method calls inside a .ptn; see backends/native/README.md. +fbcode_target(_kind = runtime.python_library, + name = "ptn_samples", + typing = True, + srcs = [ + "_ptn_samples.py", + ], + visibility = ["PUBLIC"], + deps = [ + ":ptn", + "fbsource//third-party/pypi/safetensors:safetensors", + "//caffe2:torch", + ], +) + # so we can have access to EXECUTORCH_CLIENTS list. fbcode_target(_kind = runtime.python_library, name = "lib", diff --git a/exir/_serialize/_ptn.py b/exir/_serialize/_ptn.py index 2d997957348..d06e0942c63 100644 --- a/exir/_serialize/_ptn.py +++ b/exir/_serialize/_ptn.py @@ -243,16 +243,20 @@ def _safetensors_header(owners: dict[str, torch.Tensor]) -> tuple[bytes, list[st return len(encoded).to_bytes(8, "little") + encoded, order -def _write_safetensors(pkg: zipfile.ZipFile, owners: dict[str, torch.Tensor]) -> None: - """Stream the safetensors member into `pkg` with its data section aligned.""" +def _write_safetensors( + pkg: zipfile.ZipFile, + owners: dict[str, torch.Tensor], + entry: str = SAFETENSORS_ENTRY, +) -> None: + """Stream a safetensors member into `pkg` with its data section aligned.""" header, order = _safetensors_header(owners) size = len(header) + sum(t.numel() * t.element_size() for t in owners.values()) zip64 = size > zipfile.ZIP64_LIMIT - zinfo = zipfile.ZipInfo(SAFETENSORS_ENTRY, time.localtime(time.time())[:6]) + zinfo = zipfile.ZipInfo(entry, time.localtime(time.time())[:6]) local_header_size = ( _LOCAL_HEADER_SIZE - + len(SAFETENSORS_ENTRY.encode("utf-8")) + + len(entry.encode("utf-8")) + (_ZIP64_LOCAL_EXTRA_SIZE if zip64 else 0) ) assert pkg.fp is not None diff --git a/exir/_serialize/_ptn_samples.py b/exir/_serialize/_ptn_samples.py new file mode 100644 index 00000000000..0325f0cc8ff --- /dev/null +++ b/exir/_serialize/_ptn_samples.py @@ -0,0 +1,404 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +""" +Optional recorded method calls stored inside a .ptn package. + +Each set lives under its own prefix, so adding one only appends members: + + samples///index.json which tensors and scalars each call used + samples///data.safetensors those tensors, aligned like constants + +Program loaders ignore these members. The format is specified in +backends/native/README.md; read_samples enforces it. +""" + +from __future__ import annotations + +import json +import math +import os +import re +import shutil +import struct +import tempfile +import warnings +import zipfile +from dataclasses import dataclass + +import torch +from executorch.exir._serialize._ptn import ( + _DATA_ALIGNMENT, + _DTYPE_CODES, + _LOCAL_HEADER_SIZE, + _write_safetensors, + PTG_ENTRY, +) + +_SAMPLES_DIR = "samples/" +_INDEX_ENTRY = "index.json" +_DATA_ENTRY = "data.safetensors" +_VERSION = "1.0" + +_NAME = re.compile(r"[A-Za-z0-9_-]+") +_VERSION_PATTERN = re.compile(r"(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)") +_SSA_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") + +Value = torch.Tensor | int | float | bool | list["Value"] +Tolerance = tuple[float, float] +_LOOSE_TOL = 1e-3 + + +@dataclass +class Call: + """One call of a method: positional inputs, and optionally what it produced.""" + + inputs: list[Value] + outputs: list[Value] | None = None + intermediates: dict[str, torch.Tensor] | None = None + tol: Tolerance | None = None + + +@dataclass +class Sample: + """Calls run in order from the program's initial state, with no reset between.""" + + calls: list[Call] + tol: Tolerance | None = None + + +@dataclass +class SampleSet: + samples: list[Sample] + description: str | None = None + source: str | None = None + tol: Tolerance | None = None + + +def _check_tol(tol: object, where: str) -> Tolerance: + if ( + not isinstance(tol, (list, tuple)) + or len(tol) != 2 + or not all( + isinstance(t, (int, float)) + and not isinstance(t, bool) + and math.isfinite(t) + and t >= 0 + for t in tol + ) + ): + raise ValueError(f"{where}: tol must be two finite numbers >= 0, got {tol!r}") + return (float(tol[0]), float(tol[1])) + + +def _encode_tol(tol: object, where: str, stacklevel: int) -> list[float]: + checked = _check_tol(tol, f"write_samples: {where}") + if max(checked) > _LOOSE_TOL: + warnings.warn( + f"write_samples: {where} tol (atol, rtol) = {checked} is above " + f"{_LOOSE_TOL} and may hide numerical bugs", + stacklevel=stacklevel, + ) + return list(checked) + + +def _encode(value: object, key: str, tensors: dict[str, torch.Tensor]) -> object: + """The index.json form of `value`; tensors are collected under their path key.""" + if isinstance(value, torch.Tensor): + if value.dtype not in _DTYPE_CODES: + raise KeyError( + f"write_samples: safetensors cannot represent dtype {value.dtype} " + f"at {key}" + ) + tensors[key] = value.detach().cpu().contiguous() + return key + if isinstance(value, (list, tuple)): + return [_encode(v, f"{key}[{i}]", tensors) for i, v in enumerate(value)] + if isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"write_samples: scalar at {key} is not finite") + if isinstance(value, (bool, int, float)): + return value + raise TypeError( + f"write_samples: unsupported value at {key}: {type(value).__name__}" + ) + + +def _encode_call( + call: Call, at: str, tensors: dict[str, torch.Tensor] +) -> dict[str, object]: + for field, values in (("inputs", call.inputs), ("outputs", call.outputs)): + if values is not None and not isinstance(values, (list, tuple)): + raise TypeError( + f"write_samples: {at} {field} must be a list, " + f"got {type(values).__name__}" + ) + entry: dict[str, object] = {"in": _encode(call.inputs, f"{at}.in", tensors)} + if call.outputs is not None: + entry["out"] = _encode(call.outputs, f"{at}.out", tensors) + if call.intermediates is not None: + intermediates = {} + for ssa, tensor in call.intermediates.items(): + if not _SSA_NAME.fullmatch(ssa): + raise ValueError( + f"write_samples: intermediate name {ssa!r} must match " + f"{_SSA_NAME.pattern}" + ) + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"write_samples: intermediate {ssa!r} must be a Tensor") + intermediates[ssa] = _encode(tensor, f"{at}.intermediates.{ssa}", tensors) + entry["intermediates"] = intermediates + if call.tol is not None: + entry["tol"] = _encode_tol(call.tol, at, stacklevel=5) + return entry + + +def _encode_set( + samples: SampleSet, tensors: dict[str, torch.Tensor] +) -> dict[str, object]: + index: dict[str, object] = {"version": _VERSION} + for field, text in ( + ("description", samples.description), + ("source", samples.source), + ): + if text is None: + continue + if not isinstance(text, str): + raise TypeError(f"write_samples: {field} must be a string") + index[field] = text + if samples.tol is not None: + index["tol"] = _encode_tol(samples.tol, "set", stacklevel=4) + entries = [] + for i, sample in enumerate(samples.samples): + if not sample.calls: + raise ValueError(f"write_samples: samples[{i}] needs at least one call") + calls = [] + for j, call in enumerate(sample.calls): + calls.append(_encode_call(call, f"samples[{i}].calls[{j}]", tensors)) + entry: dict[str, object] = {"calls": calls} + if sample.tol is not None: + entry["tol"] = _encode_tol(sample.tol, f"samples[{i}]", stacklevel=4) + entries.append(entry) + index["samples"] = entries + return index + + +def write_samples(path: str, method: str, name: str, samples: SampleSet) -> None: + """Append a sample set to an existing .ptn without touching its other members. + + The package is replaced only once the new set is complete, so a failure + leaves it as it was. This costs a copy of the package. + + Warns when a tolerance's atol or rtol is above 1e-3. + + Raises: + ValueError: If a name is invalid, a set or sample is empty, the set + already exists, the package has no program, or a scalar or tolerance + is out of range. + KeyError: If a tensor has a dtype safetensors cannot represent. + TypeError: If a value is not a tensor, scalar or list, or a call's inputs + or outputs are not a list. + """ + for label, value in (("method", method), ("set", name)): + if not _NAME.fullmatch(value): + raise ValueError( + f"write_samples: {label} name {value!r} must match {_NAME.pattern}" + ) + if not samples.samples: + raise ValueError("write_samples: a set needs at least one sample") + + tensors: dict[str, torch.Tensor] = {} + encoded = json.dumps(_encode_set(samples, tensors), indent=1, allow_nan=False) + + prefix = f"{_SAMPLES_DIR}{method}/{name}/" + with zipfile.ZipFile(path) as pkg: + names = pkg.namelist() + if PTG_ENTRY not in names: + raise ValueError(f"write_samples: {path} has no {PTG_ENTRY}") + if any(n.startswith(prefix) for n in names): + raise ValueError(f"write_samples: {prefix} already exists in {path}") + fd, staging = tempfile.mkstemp( + dir=os.path.dirname(os.path.abspath(path)), suffix=".ptn.tmp" + ) + os.close(fd) + try: + shutil.copyfile(path, staging) + shutil.copymode(path, staging) + with zipfile.ZipFile(staging, "a", zipfile.ZIP_STORED, allowZip64=True) as pkg: + _write_safetensors(pkg, tensors, prefix + _DATA_ENTRY) + pkg.writestr(prefix + _INDEX_ENTRY, encoded) + os.replace(staging, path) + except BaseException: + os.unlink(staging) + raise + + +def _data_offset(path: str, info: zipfile.ZipInfo) -> int: + """File offset of a stored member's data, read from its local header.""" + with open(path, "rb") as f: + f.seek(info.header_offset) + local = f.read(_LOCAL_HEADER_SIZE) + name_size, extra_size = struct.unpack(" None: + self.tensors = tensors + self.where = where + self.unreferenced: set[str] = set(tensors) + + def fail(self, message: str) -> ValueError: + return ValueError(f"{self.where}: {message}") + + def tol(self, entry: dict[str, object]) -> Tolerance | None: + return _check_tol(entry["tol"], self.where) if "tol" in entry else None + + def tensor(self, value: str, key: str) -> torch.Tensor: + if value != key: + raise self.fail(f"key {value!r} must be {key!r}") + if value not in self.tensors: + raise self.fail(f"no tensor named {value!r}") + self.unreferenced.discard(value) + return self.tensors[value] + + def value(self, value: object, key: str) -> Value: + if isinstance(value, str): + return self.tensor(value, key) + if isinstance(value, list): + return self.values(value, key) + if isinstance(value, (bool, int)) or ( + isinstance(value, float) and math.isfinite(value) + ): + return value + raise self.fail(f"unsupported value at {key}: {value!r}") + + def values(self, values: list[object], key: str) -> list[Value]: + return [self.value(v, f"{key}[{i}]") for i, v in enumerate(values)] + + def intermediates(self, raw: object, at: str) -> dict[str, torch.Tensor]: + if not isinstance(raw, dict): + raise self.fail(f"{at}.intermediates must be an object") + result = {} + for ssa, key in raw.items(): + if not _SSA_NAME.fullmatch(ssa) or not isinstance(key, str): + raise self.fail(f"bad intermediate {ssa!r} in {at}") + result[ssa] = self.tensor(key, f"{at}.intermediates.{ssa}") + return result + + def call(self, entry: object, at: str) -> Call: + if not isinstance(entry, dict) or not isinstance(entry.get("in"), list): + raise self.fail(f"{at} must be an object with an 'in' list") + if "out" in entry and not isinstance(entry["out"], list): + raise self.fail(f"{at}.out must be a list") + return Call( + inputs=self.values(entry["in"], f"{at}.in"), + outputs=self.values(entry["out"], f"{at}.out") if "out" in entry else None, + intermediates=( + self.intermediates(entry["intermediates"], at) + if "intermediates" in entry + else None + ), + tol=self.tol(entry), + ) + + def sample(self, entry: object, at: str) -> Sample: + calls = entry.get("calls") if isinstance(entry, dict) else None + if not isinstance(calls, list) or not calls: + raise self.fail(f"{at} must be an object with a non-empty 'calls' list") + assert isinstance(entry, dict) + return Sample( + calls=[self.call(c, f"{at}.calls[{j}]") for j, c in enumerate(calls)], + tol=self.tol(entry), + ) + + +def _decode_set( + index: object, tensors: dict[str, torch.Tensor], where: str +) -> SampleSet: + decoder = _SetDecoder(tensors, where) + if not isinstance(index, dict): + raise decoder.fail("index.json must contain a JSON object") + version = index.get("version") + match = _VERSION_PATTERN.fullmatch(version) if isinstance(version, str) else None + if match is None or match[1] != _VERSION.split(".")[0]: + raise decoder.fail(f"unsupported version {version!r}, expected {_VERSION}") + entries = index.get("samples") + if not isinstance(entries, list) or not entries: + raise decoder.fail("samples must be a non-empty list") + for field in ("description", "source"): + if field in index and not isinstance(index[field], str): + raise decoder.fail(f"{field} must be a string") + + samples = [decoder.sample(e, f"samples[{i}]") for i, e in enumerate(entries)] + if decoder.unreferenced: + raise decoder.fail(f"unreferenced tensors {sorted(decoder.unreferenced)}") + return SampleSet( + samples=samples, + description=index.get("description"), + source=index.get("source"), + tol=decoder.tol(index), + ) + + +def read_samples(path: str) -> dict[tuple[str, str], SampleSet]: + """Every sample set in a .ptn, keyed by (method, set), validated as read. + + A package without samples returns an empty dict. + + Raises: + ValueError: If any samples/ member, index or tensor violates the format. + """ + from safetensors import SafetensorError + from safetensors.torch import load as safetensors_load + + sets: dict[tuple[str, str], SampleSet] = {} + with zipfile.ZipFile(path) as pkg: + found: dict[tuple[str, str], dict[str, zipfile.ZipInfo]] = {} + for info in pkg.infolist(): + if not info.filename.startswith(_SAMPLES_DIR): + continue + parts = info.filename[len(_SAMPLES_DIR) :].split("/") + if ( + len(parts) != 3 + or not all(_NAME.fullmatch(p) for p in parts[:2]) + or parts[2] not in (_INDEX_ENTRY, _DATA_ENTRY) + ): + raise ValueError(f"read_samples: unexpected member {info.filename!r}") + if info.compress_type != zipfile.ZIP_STORED: + raise ValueError(f"read_samples: {info.filename!r} is compressed") + found.setdefault((parts[0], parts[1]), {})[parts[2]] = info + + for (method, name), members in sorted(found.items()): + where = f"read_samples: {_SAMPLES_DIR}{method}/{name}" + for entry in (_INDEX_ENTRY, _DATA_ENTRY): + if entry not in members: + raise ValueError(f"{where}: missing {entry}") + data = pkg.read(members[_DATA_ENTRY]) + header_size = int.from_bytes(data[:8], "little") + start = _data_offset(path, members[_DATA_ENTRY]) + 8 + header_size + if start % _DATA_ALIGNMENT: + raise ValueError( + f"{where}: tensor data is not {_DATA_ALIGNMENT}-byte aligned" + ) + try: + index = json.loads(pkg.read(members[_INDEX_ENTRY])) + tensors = safetensors_load(data) + except (ValueError, RecursionError, SafetensorError) as e: + raise ValueError(f"{where}: {e}") from e + sets[(method, name)] = _decode_set(index, tensors, where) + return sets + + +__all__ = [ + "Call", + "Sample", + "SampleSet", + "Value", + "read_samples", + "write_samples", +] diff --git a/exir/_serialize/test/BUCK b/exir/_serialize/test/BUCK index 673fbaa5414..5e1cd4c8150 100644 --- a/exir/_serialize/test/BUCK +++ b/exir/_serialize/test/BUCK @@ -81,3 +81,15 @@ fbcode_target(_kind = runtime.python_test, "//executorch/exir/_serialize:ptn", ], ) + +fbcode_target(_kind = runtime.python_test, + name = "test_ptn_samples", + srcs = [ + "test_ptn_samples.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir/_serialize:ptn", + "//executorch/exir/_serialize:ptn_samples", + ], +) diff --git a/exir/_serialize/test/test_ptn_samples.py b/exir/_serialize/test/test_ptn_samples.py new file mode 100644 index 00000000000..3e6e23f3ba4 --- /dev/null +++ b/exir/_serialize/test/test_ptn_samples.py @@ -0,0 +1,520 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import json +import os +import struct +import tempfile +import unittest +import warnings +import zipfile +from unittest import mock + +import torch +from executorch.exir._serialize._ptn import read_ptn, write_ptn +from executorch.exir._serialize._ptn_samples import ( + Call, + read_samples, + Sample, + SampleSet, + write_samples, +) + +_PTG = b"\x00\x00\x00\x00NPTG stand-in graph blob" +_INDEX = "samples/forward/golden/index.json" +_DATA = "samples/forward/golden/data.safetensors" + + +def _golden() -> SampleSet: + return SampleSet( + samples=[ + Sample([Call(inputs=[torch.randn(1, 3)], outputs=[torch.randn(1, 2)])]), + Sample([Call(inputs=[torch.randn(1, 3)], outputs=[torch.randn(1, 2)])]), + ], + description="two calls", + source="eager/cpu/f32", + tol=(1e-4, 1e-4), + ) + + +def _members(path: str) -> dict[str, bytes]: + with zipfile.ZipFile(path) as pkg: + return {name: pkg.read(name) for name in pkg.namelist()} + + +def _write_aligned(pkg: zipfile.ZipFile, name: str, data: bytes) -> None: + """Store `data` so it starts 64-byte aligned, padding the local header.""" + start = pkg.fp.tell() + 30 + len(name.encode()) + 6 + padding = (-start) % 64 + zinfo = zipfile.ZipInfo(name) + zinfo.extra = struct.pack(" None: + """Rebuild a package, replacing, adding or (for None) dropping entries.""" + entries = _members(path) + entries.update(overrides) + with zipfile.ZipFile(path, "w", zipfile.ZIP_STORED) as pkg: + for name, data in entries.items(): + if data is None: + continue + if name.endswith(".safetensors"): + _write_aligned(pkg, name, data) + else: + pkg.writestr(name, data) + + +def _data_start(path: str) -> int: + """File offset where the golden set's tensor data begins.""" + with zipfile.ZipFile(path) as pkg: + info = pkg.getinfo(_DATA) + data = pkg.read(_DATA) + with open(path, "rb") as f: + f.seek(info.header_offset) + name_size, extra_size = struct.unpack(" None: + index = json.loads(_members(path)[_INDEX]) + edit(index) + _rewrite(path, {_INDEX: json.dumps(index).encode()}) + + +class PtnSamplesTest(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + self.path = os.path.join(self._tmp.name, "m.ptn") + write_ptn(self.path, _PTG, {"w": torch.ones(4)}) + + def assertValuesEqual(self, actual, expected): + if isinstance(expected, torch.Tensor): + self.assertIsInstance(actual, torch.Tensor) + self.assertEqual(actual.dtype, expected.dtype) + self.assertTrue(torch.equal(actual, expected)) + elif isinstance(expected, list): + self.assertEqual(len(actual), len(expected)) + for a, e in zip(actual, expected): + self.assertValuesEqual(a, e) + else: + self.assertEqual(type(actual), type(expected)) + self.assertEqual(actual, expected) + + def assertSetsEqual(self, actual: SampleSet, expected: SampleSet): + self.assertEqual(actual.description, expected.description) + self.assertEqual(actual.source, expected.source) + self.assertEqual(actual.tol, expected.tol) + self.assertEqual(len(actual.samples), len(expected.samples)) + for a, e in zip(actual.samples, expected.samples): + self.assertEqual(a.tol, e.tol) + self.assertEqual(len(a.calls), len(e.calls)) + for ac, ec in zip(a.calls, e.calls): + self.assertCallsEqual(ac, ec) + + def assertCallsEqual(self, a: Call, e: Call): + self.assertValuesEqual(a.inputs, e.inputs) + self.assertEqual(a.outputs is None, e.outputs is None) + if e.outputs is not None: + self.assertValuesEqual(a.outputs, e.outputs) + self.assertEqual(a.intermediates is None, e.intermediates is None) + if e.intermediates is not None: + self.assertEqual(a.intermediates.keys(), e.intermediates.keys()) + for name in e.intermediates: + self.assertValuesEqual(a.intermediates[name], e.intermediates[name]) + self.assertEqual(a.tol, e.tol) + + def test_package_without_samples_reads_empty(self): + self.assertEqual(read_samples(self.path), {}) + + def test_golden_round_trips(self): + golden = _golden() + write_samples(self.path, "forward", "golden", golden) + + sets = read_samples(self.path) + self.assertEqual(list(sets), [("forward", "golden")]) + self.assertSetsEqual(sets[("forward", "golden")], golden) + + def test_intermediates_and_nested_tol_round_trip(self): + debug = SampleSet( + samples=[ + Sample( + [ + Call( + inputs=[torch.randn(2, 4)], + outputs=[torch.randn(2, 4)], + intermediates={ + "linear": torch.randn(2, 8), + "gelu_1": torch.randn(2, 8), + }, + tol=(1e-3, 5e-4), + ) + ], + tol=(5e-4, 5e-4), + ) + ], + tol=(1e-4, 1e-4), + ) + write_samples(self.path, "forward", "debug", debug) + + self.assertSetsEqual(read_samples(self.path)[("forward", "debug")], debug) + + def test_calls_round_trip_in_order(self): + prefill = Call( + [torch.arange(6).view(1, 6), torch.arange(6)], [torch.randn(1, 1, 5)] + ) + decode = [ + Call([torch.tensor([[t]]), torch.tensor([6 + i])], [torch.randn(1, 1, 5)]) + for i, t in enumerate((11, 12, 13)) + ] + golden = SampleSet([Sample([prefill, *decode]), Sample([prefill])]) + write_samples(self.path, "forward", "golden", golden) + + self.assertSetsEqual(read_samples(self.path)[("forward", "golden")], golden) + + def test_inputs_only_round_trips(self): + calib = SampleSet(samples=[Sample([Call([torch.randn(3)])]) for _ in range(3)]) + write_samples(self.path, "forward", "calibration", calib) + + self.assertSetsEqual(read_samples(self.path)[("forward", "calibration")], calib) + + def test_scalars_lists_and_dtypes_round_trip(self): + mixed = SampleSet( + samples=[ + Sample( + [ + Call( + inputs=[ + torch.arange(6, dtype=torch.int64).reshape(2, 3), + [ + torch.ones(2, dtype=torch.bfloat16), + torch.zeros(1, dtype=torch.bool), + ], + 3, + 2.0, + True, + ], + outputs=[torch.ones(2, dtype=torch.float16), 7], + ) + ] + ) + ] + ) + write_samples(self.path, "forward", "mixed", mixed) + + self.assertSetsEqual(read_samples(self.path)[("forward", "mixed")], mixed) + + def test_samples_may_differ_in_shape(self): + dynamic = SampleSet( + samples=[Sample([Call([torch.randn(1, n, 8)])]) for n in (3, 9)] + ) + write_samples(self.path, "forward", "dynamic", dynamic) + + samples = read_samples(self.path)[("forward", "dynamic")].samples + self.assertEqual([s.calls[0].inputs[0].shape[1] for s in samples], [3, 9]) + + def test_keys_are_json_paths(self): + write_samples( + self.path, + "forward", + "golden", + SampleSet( + samples=[ + Sample( + [ + Call( + inputs=[torch.ones(1), [torch.ones(1), 2]], + outputs=[torch.ones(1)], + intermediates={"linear": torch.ones(1)}, + ), + Call(inputs=[torch.ones(1)]), + ] + ) + ] + ), + ) + + calls = json.loads(_members(self.path)[_INDEX])["samples"][0]["calls"] + at = "samples[0].calls[0]" + self.assertEqual(calls[0]["in"], [f"{at}.in[0]", [f"{at}.in[1][0]", 2]]) + self.assertEqual(calls[0]["out"], [f"{at}.out[0]"]) + self.assertEqual( + calls[0]["intermediates"], {"linear": f"{at}.intermediates.linear"} + ) + self.assertEqual(calls[1], {"in": ["samples[0].calls[1].in[0]"]}) + + def test_sets_across_methods_coexist(self): + write_samples(self.path, "encode", "golden", _golden()) + write_samples(self.path, "decode", "golden", _golden()) + write_samples( + self.path, "decode", "bench", SampleSet([Sample([Call([torch.ones(1)])])]) + ) + + self.assertEqual( + sorted(read_samples(self.path)), + [("decode", "bench"), ("decode", "golden"), ("encode", "golden")], + ) + + def test_append_leaves_program_untouched(self): + before = _members(self.path) + write_samples(self.path, "forward", "golden", _golden()) + + after = _members(self.path) + for name, data in before.items(): + self.assertEqual(after[name], data, name) + ptg, constants = read_ptn(self.path) + self.assertEqual(ptg, _PTG) + self.assertTrue(torch.equal(constants["w"], torch.ones(4))) + + def test_members_are_stored_and_data_aligned(self): + # Vary the preceding index.json size so the data member starts at every + # offset modulo the alignment. + pad = SampleSet([Sample([Call([torch.ones(1)])])]) + for size in range(64): + with self.subTest(size=size): + write_ptn(self.path, _PTG, {"w": torch.ones(4)}) + pad.description = "x" * size + write_samples(self.path, "forward", "pad", pad) + write_samples(self.path, "forward", "golden", _golden()) + with zipfile.ZipFile(self.path) as pkg: + for name in (_INDEX, _DATA): + self.assertEqual( + pkg.getinfo(name).compress_type, zipfile.ZIP_STORED + ) + self.assertEqual(_data_start(self.path) % 64, 0) + read_samples(self.path) + + def test_reader_rejects_misaligned_data(self): + # Program sizes one byte apart cannot both leave the data aligned. + for size in (0, 1): + write_ptn(self.path, b"x" * size, {}) + write_samples(self.path, "forward", "golden", _golden()) + entries = _members(self.path) + with zipfile.ZipFile(self.path, "w", zipfile.ZIP_STORED) as pkg: + for name, data in entries.items(): + pkg.writestr(name, data) + if _data_start(self.path) % 64: + break + with self.assertRaisesRegex(ValueError, "not 64-byte aligned"): + read_samples(self.path) + + def test_failed_append_leaves_package_untouched(self): + write_samples(self.path, "forward", "golden", _golden()) + with open(self.path, "rb") as f: + before = f.read() + with mock.patch( + "executorch.exir._serialize._ptn_samples._write_safetensors", + side_effect=OSError("disk full"), + ): + with self.assertRaisesRegex(OSError, "disk full"): + write_samples(self.path, "forward", "other", _golden()) + + with open(self.path, "rb") as f: + self.assertEqual(f.read(), before) + self.assertEqual(os.listdir(self._tmp.name), ["m.ptn"]) + + def test_writer_rejects_existing_set(self): + write_samples(self.path, "forward", "golden", _golden()) + with self.assertRaisesRegex(ValueError, "already exists"): + write_samples(self.path, "forward", "golden", _golden()) + + def test_writer_rejects_bad_names(self): + for method, name in (("for/ward", "golden"), ("forward", "gol.den"), ("", "x")): + with self.subTest(method=method, name=name): + with self.assertRaisesRegex(ValueError, "must match"): + write_samples(self.path, method, name, _golden()) + bad = SampleSet( + [Sample([Call([torch.ones(1)], intermediates={"a.b": torch.ones(1)})])] + ) + with self.assertRaisesRegex(ValueError, "must match"): + write_samples(self.path, "forward", "debug", bad) + + def test_writer_rejects_bad_values(self): + cases = [ + (SampleSet([]), ValueError, "at least one sample"), + (SampleSet([Sample([])]), ValueError, "at least one call"), + (SampleSet([Sample([Call([float("nan")])])]), ValueError, "not finite"), + ( + SampleSet([Sample([Call([torch.ones(1)])])], tol=(-1.0, 0.0)), + ValueError, + "tol", + ), + ( + SampleSet([Sample([Call([torch.ones(1)], tol=(0.0, -1.0))])]), + ValueError, + r"samples\[0\]\.calls\[0\]: tol", + ), + (SampleSet([Sample([Call(["x"])])]), TypeError, "unsupported value"), + ( + SampleSet([Sample([Call([torch.ones(1)])])], description=1), + TypeError, + "description must be a string", + ), + ( + SampleSet([Sample([Call(torch.ones(1))])]), + TypeError, + "inputs must be a list", + ), + ( + SampleSet([Sample([Call([torch.ones(1)], torch.ones(1))])]), + TypeError, + "outputs must be a list", + ), + ( + SampleSet([Sample([Call([torch.ones(1, dtype=torch.complex128)])])]), + KeyError, + "cannot represent dtype", + ), + ] + for samples, error, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(error, message): + write_samples(self.path, "forward", "golden", samples) + self.assertEqual(read_samples(self.path), {}) + + def test_writer_warns_on_loose_tol(self): + one = [torch.ones(1)] + cases = [ + (SampleSet([Sample([Call(one)])], tol=(2.0, 0.05)), "set"), + (SampleSet([Sample([Call(one)], tol=(0.0, 2e-3))]), r"samples\[0\]"), + ( + SampleSet([Sample([Call(one), Call(one, tol=(2e-3, 0.0))])]), + r"samples\[0\]\.calls\[1\]", + ), + ] + for i, (loose, where) in enumerate(cases): + with self.assertWarnsRegex( + UserWarning, f"{where} tol .* above 0.001" + ) as caught: + write_samples(self.path, "forward", f"loose{i}", loose) + self.assertEqual(caught.filename, __file__) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + write_samples(self.path, "forward", "golden", _golden()) + tight = SampleSet([Sample([Call([torch.ones(1)], tol=(1e-3, 1e-3))])]) + write_samples(self.path, "forward", "tight", tight) + + def test_writer_requires_a_program(self): + path = os.path.join(self._tmp.name, "empty.ptn") + with zipfile.ZipFile(path, "w") as pkg: + pkg.writestr("other", b"") + with self.assertRaisesRegex(ValueError, "has no program.ptg"): + write_samples(path, "forward", "golden", _golden()) + + def test_reader_rejects_bad_layout(self): + cases = [ + ({_DATA: None}, "missing data.safetensors"), + ({_INDEX: None}, "missing index.json"), + ({"samples/forward/golden/notes.txt": b""}, "unexpected member"), + ({"samples/forward/stray.json": b""}, "unexpected member"), + ({"samples/for.ward/golden/index.json": b""}, "unexpected member"), + ({"samples/forward/golden/x/index.json": b""}, "unexpected member"), + ] + for overrides, message in cases: + with self.subTest(message=message): + write_ptn(self.path, _PTG, {}) + write_samples(self.path, "forward", "golden", _golden()) + _rewrite(self.path, overrides) + with self.assertRaisesRegex(ValueError, message): + read_samples(self.path) + + def test_reader_rejects_compressed_member(self): + write_samples(self.path, "forward", "golden", _golden()) + entries = _members(self.path) + with zipfile.ZipFile(self.path, "w") as pkg: + for name, data in entries.items(): + compression = ( + zipfile.ZIP_DEFLATED if name == _INDEX else zipfile.ZIP_STORED + ) + pkg.writestr(name, data, compress_type=compression) + with self.assertRaisesRegex(ValueError, "is compressed"): + read_samples(self.path) + + def test_reader_rejects_bad_index(self): + def set_field(field, value): + return lambda index: index.__setitem__(field, value) + + def set_sample_field(field, value): + return lambda index: index["samples"][0].__setitem__(field, value) + + def set_call_field(field, value): + return lambda index: index["samples"][0]["calls"][0].__setitem__( + field, value + ) + + def set_intermediates(raw): + return set_call_field("intermediates", raw) + + at = "samples[0].calls[0]" + cases = [ + (lambda index: index.clear(), "unsupported version"), + (set_field("version", "2.0"), "unsupported version"), + (set_field("version", 1), "unsupported version"), + (set_field("version", "1"), "unsupported version"), + (set_field("version", "1.01"), "unsupported version"), + (set_field("samples", []), "non-empty list"), + (set_field("tol", [1e-4]), "tol must be"), + (set_field("description", 1), "description must be a string"), + (set_field("source", None), "source must be a string"), + (set_sample_field("calls", []), "non-empty 'calls' list"), + (set_sample_field("tol", [1e-4]), "tol must be"), + (set_call_field("tol", [0.0, -1.0]), "tol must be"), + (set_call_field("in", None), "'in' list"), + (set_call_field("in", [f"{at}.in[0]", None]), "unsupported value"), + (set_call_field("in", [f"{at}.in[0]", float("nan")]), "unsupported value"), + (set_call_field("out", f"{at}.out[0]"), "out must be a list"), + (set_intermediates([]), "intermediates must be an object"), + (set_intermediates({"a.b": f"{at}.intermediates.a.b"}), "bad intermediate"), + (set_intermediates({"linear": 1}), "bad intermediate"), + ( + set_call_field("in", [f"{at}.in[7]"]), + "must be 'samples\\[0\\].calls\\[0\\].in\\[0\\]'", + ), + ( + set_call_field("out", [f"{at}.out[0]", f"{at}.out[1]"]), + "no tensor named", + ), + (set_call_field("out", [1]), "unreferenced tensors"), + ] + for edit, message in cases: + with self.subTest(message=message): + write_ptn(self.path, _PTG, {}) + write_samples(self.path, "forward", "golden", _golden()) + _edit_index(self.path, edit) + with self.assertRaisesRegex(ValueError, message): + read_samples(self.path) + + def test_reader_names_set_on_unparseable_members(self): + cases = [ + ({_INDEX: b'{"version": '}, "Expecting value"), + ({_INDEX: b"[" * 100_000}, ""), + ({_DATA: b"\x08" + bytes(7) + b"not json"}, ""), + ] + for overrides, message in cases: + with self.subTest(message=message): + write_ptn(self.path, _PTG, {}) + write_samples(self.path, "forward", "golden", _golden()) + _rewrite(self.path, overrides) + with self.assertRaisesRegex( + ValueError, f"samples/forward/golden.*{message}" + ): + read_samples(self.path) + + def test_reader_accepts_newer_minor_version_and_unknown_fields(self): + golden = _golden() + write_samples(self.path, "forward", "golden", golden) + + def add_fields(index): + index["version"] = "1.7" + index["future"] = {"anything": 1} + index["samples"][0]["future"] = True + index["samples"][0]["calls"][0]["future"] = True + + _edit_index(self.path, add_fields) + self.assertSetsEqual(read_samples(self.path)[("forward", "golden")], golden) diff --git a/exir/native/BUCK b/exir/native/BUCK index 0c83e5baaa8..c6dbeac4bff 100644 --- a/exir/native/BUCK +++ b/exir/native/BUCK @@ -20,5 +20,6 @@ fbcode_target( "//executorch/backends/native:lib", "//executorch/exir:lib", "//executorch/exir/_serialize:ptn", + "//executorch/exir/_serialize:ptn_samples", ], ) diff --git a/exir/native/__init__.py b/exir/native/__init__.py index 0498fba4e51..f7e3b0dbacd 100644 --- a/exir/native/__init__.py +++ b/exir/native/__init__.py @@ -6,9 +6,21 @@ # pyre-strict +from executorch.exir._serialize._ptn_samples import ( + Call, + read_samples, + Sample, + SampleSet, + write_samples, +) from executorch.exir.native.native import NativeProgramManager, to_native __all__ = [ + "Call", "NativeProgramManager", + "Sample", + "SampleSet", + "read_samples", "to_native", + "write_samples", ] diff --git a/exir/native/test/test_native.py b/exir/native/test/test_native.py index a45277bbc91..e122630497f 100644 --- a/exir/native/test/test_native.py +++ b/exir/native/test/test_native.py @@ -15,7 +15,6 @@ import torch import torch.nn as nn - from executorch.backends.native import get_default_compile_config from executorch.backends.native.partitioner import NativePartitioner from executorch.backends.native.passes import get_default_passes @@ -32,7 +31,14 @@ executorch_call_delegate, get_lowered_submodules, ) -from executorch.exir.native import to_native +from executorch.exir.native import ( + Call, + read_samples, + Sample, + SampleSet, + to_native, + write_samples, +) from executorch.exir.native.native import ( _check_delegate_contract, _validate_merged_program, @@ -433,5 +439,45 @@ def forward(self, x): ) +class SamplesTest(unittest.TestCase): + def test_samples_round_trip_on_saved_package(self): + model = nn.Linear(8, 4) + prefill_x, decode_x = torch.randn(2, 8), torch.randn(1, 8) + native = to_native( + { + "prefill": torch.export.export(model, (prefill_x,)), + "decode": torch.export.export(model, (decode_x,)), + } + ) + golden = { + ("decode", "golden"): SampleSet( + [Sample([Call([decode_x], [model(decode_x).detach()])])], + tol=(1e-5, 1e-5), + ), + ("prefill", "calib"): SampleSet([Sample([Call([prefill_x])])]), + } + with tempfile.TemporaryDirectory() as td: + path = os.path.join(td, "m.ptn") + native.save(path) + with zipfile.ZipFile(path) as pkg: + program = {n: pkg.read(n) for n in pkg.namelist()} + for (method, name), samples in golden.items(): + write_samples(path, method, name, samples) + with zipfile.ZipFile(path) as pkg: + self.assertEqual({n: pkg.read(n) for n in program}, program) + read = read_samples(path) + self.assertEqual(read.keys(), golden.keys()) + for key, expected in golden.items(): + self.assertEqual(read[key].tol, expected.tol) + for got, want in zip(read[key].samples, expected.samples, strict=True): + (got_call,), (want_call,) = got.calls, want.calls + torch.testing.assert_close( + got_call.inputs, want_call.inputs, rtol=0, atol=0 + ) + torch.testing.assert_close( + got_call.outputs, want_call.outputs, rtol=0, atol=0 + ) + + if __name__ == "__main__": unittest.main() From 109f726fe6d27e64337d839868b6e343864334d2 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Tue, 27 Oct 2026 14:56:15 -0700 Subject: [PATCH 2/5] Update [ghstack-poisoned] --- backends/native/README.md | 36 +-- .../test/runtime/deserialize/test_package.cpp | 6 +- exir/_serialize/_ptn_samples.py | 55 ++-- exir/_serialize/test/test_ptn_samples.py | 235 ++++++++++++------ exir/native/test/test_native.py | 32 ++- 5 files changed, 231 insertions(+), 133 deletions(-) diff --git a/backends/native/README.md b/backends/native/README.md index 086664de294..13583d2537b 100644 --- a/backends/native/README.md +++ b/backends/native/README.md @@ -44,32 +44,35 @@ A `.ptn` is an uncompressed zip with fixed entry names: ## Samples (optional) A `.ptn` may carry recorded method calls: inputs, and optionally the outputs and -intermediate values they produced. A sample is an ordered sequence of calls, so a -stateful method such as an LLM's prefill followed by decode steps can be checked. They serve reference checks, layer-by-layer -numerical debugging, calibration and benchmarking. Program loaders ignore them. +intermediate values they produced. A sample is an ordered sequence of calls, each +naming the method it runs, so a stateful sequence such as an LLM's prefill +followed by decode steps can be checked whether it uses one method or several. +Samples serve reference checks, layer-by-layer numerical debugging, calibration +and benchmarking. Program loaders ignore them. ```python from executorch.exir.native import Call, read_samples, Sample, SampleSet, to_native, write_samples to_native(ep).save("model.ptn") -golden = SampleSet([Sample([Call([x], [model(x)])])]) -write_samples("model.ptn", "forward", "golden", golden) -read_samples("model.ptn") # {("forward", "golden"): SampleSet(...)} +golden = SampleSet([Sample([Call("forward", [x], [model(x)])])]) +write_samples("model.ptn", "golden", golden) +read_samples("model.ptn") # {"golden": SampleSet(...)} ``` `write_samples` appends a set to any existing `.ptn`, not only one this process saved, and replaces the file only once the set is complete. `read_samples` -validates every set against the rules below, except that `` exists. +validates every set against the rules below, except that each call's method +exists in the program. ### Layout ``` -samples///index.json -samples///data.safetensors +samples//index.json +samples//data.safetensors ``` -- `` names a program method. `` is an identifier with no meaning to - the format. Both match `[A-Za-z0-9_-]+`. +- `` is an identifier with no meaning to the format, matching + `[A-Za-z0-9_-]+`. - Each set has exactly these two members, and nothing else lives under `samples/`. - Members are stored (uncompressed). The `data.safetensors` data section starts 64-byte aligned in the file, like `program.safetensors`. @@ -87,7 +90,8 @@ samples///data.safetensors Sample = { "calls": [Call, ...], required, non-empty, run in order "tol": [atol, rtol] } optional, overrides the set's -Call = { "in": [Value, ...], required, positional arguments +Call = { "method": str, required, a program method + "in": [Value, ...], required, positional arguments "out": [Value, ...], optional, positional outputs "intermediates": {name: Key, ...}, optional "tol": [atol, rtol] } optional, overrides the sample's @@ -98,13 +102,15 @@ Key = str naming a tensor in the set's data.safetensors - Each sample starts from the program's initial state. Samples are independent and may run in any order; the calls within a sample run in order with no reset - between them, so each sees the state the previous one left. + between them, so each sees the state the previous one left, whichever method + that call ran. - A tensor's key is its JSON path in the index, e.g. `samples[0].calls[1].in[2]`. Below a call that is `.in[k]`, `.in[k][l]` for an element of a list argument, `.out[k]` or `.intermediates.`. Every key resolves, and every tensor is referenced exactly once. - A JSON integer is an int scalar, any other number a float. -- Intermediate names are top-level SSA value names (`[A-Za-z_][A-Za-z0-9_]*`). +- Intermediate names are top-level SSA value names of the call's method + (`[A-Za-z_][A-Za-z0-9_]*`). Producers record what they can; consumers compare the names both sides have. - Tolerances are finite and non-negative, and advisory: a consumer may override them. `write_samples` warns when atol or rtol is above 1e-3. @@ -113,8 +119,6 @@ Key = str naming a tensor in the set's data.safetensors Not in v1: -- Sequences that span methods, such as a separate `prefill` and `decode` sharing - state. - Tensors shared across sets. - Values inside nested graphs. - Quantized method inputs and outputs (a `QuantSpec` on graph I/O). Integer and diff --git a/backends/native/test/runtime/deserialize/test_package.cpp b/backends/native/test/runtime/deserialize/test_package.cpp index c49791025f5..e81939fdd27 100644 --- a/backends/native/test/runtime/deserialize/test_package.cpp +++ b/backends/native/test/runtime/deserialize/test_package.cpp @@ -143,14 +143,14 @@ TEST(PackageTest, LoadsConstantsOnDemand) { TEST(PackageTest, IgnoresSampleMembers) { const std::string index = - R"({"version":1,"samples":[{"calls":[{"in":["samples[0].calls[0].in[0]"]}]}]})"; + R"({"version":1,"samples":[{"calls":[{"method":"forward","in":["samples[0].calls[0].in[0]"]}]}]})"; const std::vector data = make_safetensors( R"({"samples[0].calls[0].in[0]":{"dtype":"U8","shape":[2],"data_offsets":[0,2]}})", "in"); const TempPackage file( {}, - {{"samples/forward/golden/index.json", index}, - {"samples/forward/golden/data.safetensors", + {{"samples/golden/index.json", index}, + {"samples/golden/data.safetensors", std::string(data.begin(), data.end())}}); const Package package = Package::load(file.path()); diff --git a/exir/_serialize/_ptn_samples.py b/exir/_serialize/_ptn_samples.py index 0325f0cc8ff..afad11fb1a2 100644 --- a/exir/_serialize/_ptn_samples.py +++ b/exir/_serialize/_ptn_samples.py @@ -11,8 +11,8 @@ Each set lives under its own prefix, so adding one only appends members: - samples///index.json which tensors and scalars each call used - samples///data.safetensors those tensors, aligned like constants + samples//index.json each call's method, tensors and scalars + samples//data.safetensors those tensors, aligned like constants Program loaders ignore these members. The format is specified in backends/native/README.md; read_samples enforces it. @@ -58,6 +58,7 @@ class Call: """One call of a method: positional inputs, and optionally what it produced.""" + method: str inputs: list[Value] outputs: list[Value] | None = None intermediates: dict[str, torch.Tensor] | None = None @@ -137,7 +138,12 @@ def _encode_call( f"write_samples: {at} {field} must be a list, " f"got {type(values).__name__}" ) - entry: dict[str, object] = {"in": _encode(call.inputs, f"{at}.in", tensors)} + if not isinstance(call.method, str) or not call.method: + raise ValueError(f"write_samples: {at} method must be a non-empty string") + entry: dict[str, object] = { + "method": call.method, + "in": _encode(call.inputs, f"{at}.in", tensors), + } if call.outputs is not None: entry["out"] = _encode(call.outputs, f"{at}.out", tensors) if call.intermediates is not None: @@ -187,7 +193,7 @@ def _encode_set( return index -def write_samples(path: str, method: str, name: str, samples: SampleSet) -> None: +def write_samples(path: str, name: str, samples: SampleSet) -> None: """Append a sample set to an existing .ptn without touching its other members. The package is replaced only once the new set is complete, so a failure @@ -196,25 +202,22 @@ def write_samples(path: str, method: str, name: str, samples: SampleSet) -> None Warns when a tolerance's atol or rtol is above 1e-3. Raises: - ValueError: If a name is invalid, a set or sample is empty, the set - already exists, the package has no program, or a scalar or tolerance - is out of range. + ValueError: If the set name or a method name is invalid, a set or sample + is empty, the set already exists, the package has no program, or a + scalar or tolerance is out of range. KeyError: If a tensor has a dtype safetensors cannot represent. TypeError: If a value is not a tensor, scalar or list, or a call's inputs or outputs are not a list. """ - for label, value in (("method", method), ("set", name)): - if not _NAME.fullmatch(value): - raise ValueError( - f"write_samples: {label} name {value!r} must match {_NAME.pattern}" - ) + if not _NAME.fullmatch(name): + raise ValueError(f"write_samples: set name {name!r} must match {_NAME.pattern}") if not samples.samples: raise ValueError("write_samples: a set needs at least one sample") tensors: dict[str, torch.Tensor] = {} encoded = json.dumps(_encode_set(samples, tensors), indent=1, allow_nan=False) - prefix = f"{_SAMPLES_DIR}{method}/{name}/" + prefix = f"{_SAMPLES_DIR}{name}/" with zipfile.ZipFile(path) as pkg: names = pkg.namelist() if PTG_ENTRY not in names: @@ -293,9 +296,13 @@ def intermediates(self, raw: object, at: str) -> dict[str, torch.Tensor]: def call(self, entry: object, at: str) -> Call: if not isinstance(entry, dict) or not isinstance(entry.get("in"), list): raise self.fail(f"{at} must be an object with an 'in' list") + method = entry.get("method") + if not isinstance(method, str) or not method: + raise self.fail(f"{at}.method must be a non-empty string") if "out" in entry and not isinstance(entry["out"], list): raise self.fail(f"{at}.out must be a list") return Call( + method=method, inputs=self.values(entry["in"], f"{at}.in"), outputs=self.values(entry["out"], f"{at}.out") if "out" in entry else None, intermediates=( @@ -345,8 +352,8 @@ def _decode_set( ) -def read_samples(path: str) -> dict[tuple[str, str], SampleSet]: - """Every sample set in a .ptn, keyed by (method, set), validated as read. +def read_samples(path: str) -> dict[str, SampleSet]: + """Every sample set in a .ptn, keyed by set name, validated as read. A package without samples returns an empty dict. @@ -356,25 +363,25 @@ def read_samples(path: str) -> dict[tuple[str, str], SampleSet]: from safetensors import SafetensorError from safetensors.torch import load as safetensors_load - sets: dict[tuple[str, str], SampleSet] = {} + sets: dict[str, SampleSet] = {} with zipfile.ZipFile(path) as pkg: - found: dict[tuple[str, str], dict[str, zipfile.ZipInfo]] = {} + found: dict[str, dict[str, zipfile.ZipInfo]] = {} for info in pkg.infolist(): if not info.filename.startswith(_SAMPLES_DIR): continue parts = info.filename[len(_SAMPLES_DIR) :].split("/") if ( - len(parts) != 3 - or not all(_NAME.fullmatch(p) for p in parts[:2]) - or parts[2] not in (_INDEX_ENTRY, _DATA_ENTRY) + len(parts) != 2 + or not _NAME.fullmatch(parts[0]) + or parts[1] not in (_INDEX_ENTRY, _DATA_ENTRY) ): raise ValueError(f"read_samples: unexpected member {info.filename!r}") if info.compress_type != zipfile.ZIP_STORED: raise ValueError(f"read_samples: {info.filename!r} is compressed") - found.setdefault((parts[0], parts[1]), {})[parts[2]] = info + found.setdefault(parts[0], {})[parts[1]] = info - for (method, name), members in sorted(found.items()): - where = f"read_samples: {_SAMPLES_DIR}{method}/{name}" + for name, members in sorted(found.items()): + where = f"read_samples: {_SAMPLES_DIR}{name}" for entry in (_INDEX_ENTRY, _DATA_ENTRY): if entry not in members: raise ValueError(f"{where}: missing {entry}") @@ -390,7 +397,7 @@ def read_samples(path: str) -> dict[tuple[str, str], SampleSet]: tensors = safetensors_load(data) except (ValueError, RecursionError, SafetensorError) as e: raise ValueError(f"{where}: {e}") from e - sets[(method, name)] = _decode_set(index, tensors, where) + sets[name] = _decode_set(index, tensors, where) return sets diff --git a/exir/_serialize/test/test_ptn_samples.py b/exir/_serialize/test/test_ptn_samples.py index 3e6e23f3ba4..1d95645c4ab 100644 --- a/exir/_serialize/test/test_ptn_samples.py +++ b/exir/_serialize/test/test_ptn_samples.py @@ -24,15 +24,31 @@ ) _PTG = b"\x00\x00\x00\x00NPTG stand-in graph blob" -_INDEX = "samples/forward/golden/index.json" -_DATA = "samples/forward/golden/data.safetensors" +_INDEX = "samples/golden/index.json" +_DATA = "samples/golden/data.safetensors" def _golden() -> SampleSet: return SampleSet( samples=[ - Sample([Call(inputs=[torch.randn(1, 3)], outputs=[torch.randn(1, 2)])]), - Sample([Call(inputs=[torch.randn(1, 3)], outputs=[torch.randn(1, 2)])]), + Sample( + [ + Call( + "forward", + inputs=[torch.randn(1, 3)], + outputs=[torch.randn(1, 2)], + ) + ] + ), + Sample( + [ + Call( + "forward", + inputs=[torch.randn(1, 3)], + outputs=[torch.randn(1, 2)], + ) + ] + ), ], description="two calls", source="eager/cpu/f32", @@ -118,6 +134,7 @@ def assertSetsEqual(self, actual: SampleSet, expected: SampleSet): self.assertCallsEqual(ac, ec) def assertCallsEqual(self, a: Call, e: Call): + self.assertEqual(a.method, e.method) self.assertValuesEqual(a.inputs, e.inputs) self.assertEqual(a.outputs is None, e.outputs is None) if e.outputs is not None: @@ -134,11 +151,11 @@ def test_package_without_samples_reads_empty(self): def test_golden_round_trips(self): golden = _golden() - write_samples(self.path, "forward", "golden", golden) + write_samples(self.path, "golden", golden) sets = read_samples(self.path) - self.assertEqual(list(sets), [("forward", "golden")]) - self.assertSetsEqual(sets[("forward", "golden")], golden) + self.assertEqual(list(sets), ["golden"]) + self.assertSetsEqual(sets["golden"], golden) def test_intermediates_and_nested_tol_round_trip(self): debug = SampleSet( @@ -146,6 +163,7 @@ def test_intermediates_and_nested_tol_round_trip(self): Sample( [ Call( + "forward", inputs=[torch.randn(2, 4)], outputs=[torch.randn(2, 4)], intermediates={ @@ -160,28 +178,36 @@ def test_intermediates_and_nested_tol_round_trip(self): ], tol=(1e-4, 1e-4), ) - write_samples(self.path, "forward", "debug", debug) + write_samples(self.path, "debug", debug) - self.assertSetsEqual(read_samples(self.path)[("forward", "debug")], debug) + self.assertSetsEqual(read_samples(self.path)["debug"], debug) def test_calls_round_trip_in_order(self): prefill = Call( - [torch.arange(6).view(1, 6), torch.arange(6)], [torch.randn(1, 1, 5)] + "forward", + [torch.arange(6).view(1, 6), torch.arange(6)], + [torch.randn(1, 1, 5)], ) decode = [ - Call([torch.tensor([[t]]), torch.tensor([6 + i])], [torch.randn(1, 1, 5)]) + Call( + "forward", + [torch.tensor([[t]]), torch.tensor([6 + i])], + [torch.randn(1, 1, 5)], + ) for i, t in enumerate((11, 12, 13)) ] golden = SampleSet([Sample([prefill, *decode]), Sample([prefill])]) - write_samples(self.path, "forward", "golden", golden) + write_samples(self.path, "golden", golden) - self.assertSetsEqual(read_samples(self.path)[("forward", "golden")], golden) + self.assertSetsEqual(read_samples(self.path)["golden"], golden) def test_inputs_only_round_trips(self): - calib = SampleSet(samples=[Sample([Call([torch.randn(3)])]) for _ in range(3)]) - write_samples(self.path, "forward", "calibration", calib) + calib = SampleSet( + samples=[Sample([Call("forward", [torch.randn(3)])]) for _ in range(3)] + ) + write_samples(self.path, "calibration", calib) - self.assertSetsEqual(read_samples(self.path)[("forward", "calibration")], calib) + self.assertSetsEqual(read_samples(self.path)["calibration"], calib) def test_scalars_lists_and_dtypes_round_trip(self): mixed = SampleSet( @@ -189,6 +215,7 @@ def test_scalars_lists_and_dtypes_round_trip(self): Sample( [ Call( + "forward", inputs=[ torch.arange(6, dtype=torch.int64).reshape(2, 3), [ @@ -205,34 +232,34 @@ def test_scalars_lists_and_dtypes_round_trip(self): ) ] ) - write_samples(self.path, "forward", "mixed", mixed) + write_samples(self.path, "mixed", mixed) - self.assertSetsEqual(read_samples(self.path)[("forward", "mixed")], mixed) + self.assertSetsEqual(read_samples(self.path)["mixed"], mixed) def test_samples_may_differ_in_shape(self): dynamic = SampleSet( - samples=[Sample([Call([torch.randn(1, n, 8)])]) for n in (3, 9)] + samples=[Sample([Call("forward", [torch.randn(1, n, 8)])]) for n in (3, 9)] ) - write_samples(self.path, "forward", "dynamic", dynamic) + write_samples(self.path, "dynamic", dynamic) - samples = read_samples(self.path)[("forward", "dynamic")].samples + samples = read_samples(self.path)["dynamic"].samples self.assertEqual([s.calls[0].inputs[0].shape[1] for s in samples], [3, 9]) def test_keys_are_json_paths(self): write_samples( self.path, - "forward", "golden", SampleSet( samples=[ Sample( [ Call( + "forward", inputs=[torch.ones(1), [torch.ones(1), 2]], outputs=[torch.ones(1)], intermediates={"linear": torch.ones(1)}, ), - Call(inputs=[torch.ones(1)]), + Call("forward", inputs=[torch.ones(1)]), ] ) ] @@ -246,23 +273,30 @@ def test_keys_are_json_paths(self): self.assertEqual( calls[0]["intermediates"], {"linear": f"{at}.intermediates.linear"} ) - self.assertEqual(calls[1], {"in": ["samples[0].calls[1].in[0]"]}) - - def test_sets_across_methods_coexist(self): - write_samples(self.path, "encode", "golden", _golden()) - write_samples(self.path, "decode", "golden", _golden()) - write_samples( - self.path, "decode", "bench", SampleSet([Sample([Call([torch.ones(1)])])]) - ) - self.assertEqual( - sorted(read_samples(self.path)), - [("decode", "bench"), ("decode", "golden"), ("encode", "golden")], + calls[1], {"method": "forward", "in": ["samples[0].calls[1].in[0]"]} ) + def test_sample_spans_methods(self): + prefill = Call("prefill", [torch.arange(6).view(1, 6)], [torch.randn(1, 6, 5)]) + decode = Call("decode", [torch.tensor([[11]])], [torch.randn(1, 1, 5)]) + golden = SampleSet([Sample([prefill, decode])]) + write_samples(self.path, "golden", golden) + + calls = json.loads(_members(self.path)[_INDEX])["samples"][0]["calls"] + self.assertEqual([c["method"] for c in calls], ["prefill", "decode"]) + self.assertSetsEqual(read_samples(self.path)["golden"], golden) + + def test_sets_coexist(self): + write_samples(self.path, "golden", _golden()) + bench = SampleSet([Sample([Call("forward", [torch.ones(1)])])]) + write_samples(self.path, "bench", bench) + + self.assertEqual(sorted(read_samples(self.path)), ["bench", "golden"]) + def test_append_leaves_program_untouched(self): before = _members(self.path) - write_samples(self.path, "forward", "golden", _golden()) + write_samples(self.path, "golden", _golden()) after = _members(self.path) for name, data in before.items(): @@ -274,13 +308,13 @@ def test_append_leaves_program_untouched(self): def test_members_are_stored_and_data_aligned(self): # Vary the preceding index.json size so the data member starts at every # offset modulo the alignment. - pad = SampleSet([Sample([Call([torch.ones(1)])])]) + pad = SampleSet([Sample([Call("forward", [torch.ones(1)])])]) for size in range(64): with self.subTest(size=size): write_ptn(self.path, _PTG, {"w": torch.ones(4)}) pad.description = "x" * size - write_samples(self.path, "forward", "pad", pad) - write_samples(self.path, "forward", "golden", _golden()) + write_samples(self.path, "pad", pad) + write_samples(self.path, "golden", _golden()) with zipfile.ZipFile(self.path) as pkg: for name in (_INDEX, _DATA): self.assertEqual( @@ -293,7 +327,7 @@ def test_reader_rejects_misaligned_data(self): # Program sizes one byte apart cannot both leave the data aligned. for size in (0, 1): write_ptn(self.path, b"x" * size, {}) - write_samples(self.path, "forward", "golden", _golden()) + write_samples(self.path, "golden", _golden()) entries = _members(self.path) with zipfile.ZipFile(self.path, "w", zipfile.ZIP_STORED) as pkg: for name, data in entries.items(): @@ -304,7 +338,7 @@ def test_reader_rejects_misaligned_data(self): read_samples(self.path) def test_failed_append_leaves_package_untouched(self): - write_samples(self.path, "forward", "golden", _golden()) + write_samples(self.path, "golden", _golden()) with open(self.path, "rb") as f: before = f.read() with mock.patch( @@ -312,61 +346,88 @@ def test_failed_append_leaves_package_untouched(self): side_effect=OSError("disk full"), ): with self.assertRaisesRegex(OSError, "disk full"): - write_samples(self.path, "forward", "other", _golden()) + write_samples(self.path, "other", _golden()) with open(self.path, "rb") as f: self.assertEqual(f.read(), before) self.assertEqual(os.listdir(self._tmp.name), ["m.ptn"]) def test_writer_rejects_existing_set(self): - write_samples(self.path, "forward", "golden", _golden()) + write_samples(self.path, "golden", _golden()) with self.assertRaisesRegex(ValueError, "already exists"): - write_samples(self.path, "forward", "golden", _golden()) + write_samples(self.path, "golden", _golden()) def test_writer_rejects_bad_names(self): - for method, name in (("for/ward", "golden"), ("forward", "gol.den"), ("", "x")): - with self.subTest(method=method, name=name): + for name in ("gol/den", "gol.den", ""): + with self.subTest(name=name): with self.assertRaisesRegex(ValueError, "must match"): - write_samples(self.path, method, name, _golden()) + write_samples(self.path, name, _golden()) + no_method = SampleSet([Sample([Call("", [torch.ones(1)])])]) + with self.assertRaisesRegex(ValueError, "method must be a non-empty string"): + write_samples(self.path, "golden", no_method) bad = SampleSet( - [Sample([Call([torch.ones(1)], intermediates={"a.b": torch.ones(1)})])] + [ + Sample( + [ + Call( + "forward", + [torch.ones(1)], + intermediates={"a.b": torch.ones(1)}, + ) + ] + ) + ] ) with self.assertRaisesRegex(ValueError, "must match"): - write_samples(self.path, "forward", "debug", bad) + write_samples(self.path, "debug", bad) def test_writer_rejects_bad_values(self): cases = [ (SampleSet([]), ValueError, "at least one sample"), (SampleSet([Sample([])]), ValueError, "at least one call"), - (SampleSet([Sample([Call([float("nan")])])]), ValueError, "not finite"), ( - SampleSet([Sample([Call([torch.ones(1)])])], tol=(-1.0, 0.0)), + SampleSet([Sample([Call("forward", [float("nan")])])]), + ValueError, + "not finite", + ), + ( + SampleSet( + [Sample([Call("forward", [torch.ones(1)])])], tol=(-1.0, 0.0) + ), ValueError, "tol", ), ( - SampleSet([Sample([Call([torch.ones(1)], tol=(0.0, -1.0))])]), + SampleSet( + [Sample([Call("forward", [torch.ones(1)], tol=(0.0, -1.0))])] + ), ValueError, r"samples\[0\]\.calls\[0\]: tol", ), - (SampleSet([Sample([Call(["x"])])]), TypeError, "unsupported value"), ( - SampleSet([Sample([Call([torch.ones(1)])])], description=1), + SampleSet([Sample([Call("forward", ["x"])])]), + TypeError, + "unsupported value", + ), + ( + SampleSet([Sample([Call("forward", [torch.ones(1)])])], description=1), TypeError, "description must be a string", ), ( - SampleSet([Sample([Call(torch.ones(1))])]), + SampleSet([Sample([Call("forward", torch.ones(1))])]), TypeError, "inputs must be a list", ), ( - SampleSet([Sample([Call([torch.ones(1)], torch.ones(1))])]), + SampleSet([Sample([Call("forward", [torch.ones(1)], torch.ones(1))])]), TypeError, "outputs must be a list", ), ( - SampleSet([Sample([Call([torch.ones(1, dtype=torch.complex128)])])]), + SampleSet( + [Sample([Call("forward", [torch.ones(1, dtype=torch.complex128)])])] + ), KeyError, "cannot represent dtype", ), @@ -374,16 +435,28 @@ def test_writer_rejects_bad_values(self): for samples, error, message in cases: with self.subTest(message=message): with self.assertRaisesRegex(error, message): - write_samples(self.path, "forward", "golden", samples) + write_samples(self.path, "golden", samples) self.assertEqual(read_samples(self.path), {}) def test_writer_warns_on_loose_tol(self): one = [torch.ones(1)] cases = [ - (SampleSet([Sample([Call(one)])], tol=(2.0, 0.05)), "set"), - (SampleSet([Sample([Call(one)], tol=(0.0, 2e-3))]), r"samples\[0\]"), + (SampleSet([Sample([Call("forward", one)])], tol=(2.0, 0.05)), "set"), ( - SampleSet([Sample([Call(one), Call(one, tol=(2e-3, 0.0))])]), + SampleSet([Sample([Call("forward", one)], tol=(0.0, 2e-3))]), + r"samples\[0\]", + ), + ( + SampleSet( + [ + Sample( + [ + Call("forward", one), + Call("forward", one, tol=(2e-3, 0.0)), + ] + ) + ] + ), r"samples\[0\]\.calls\[1\]", ), ] @@ -391,41 +464,43 @@ def test_writer_warns_on_loose_tol(self): with self.assertWarnsRegex( UserWarning, f"{where} tol .* above 0.001" ) as caught: - write_samples(self.path, "forward", f"loose{i}", loose) + write_samples(self.path, f"loose{i}", loose) self.assertEqual(caught.filename, __file__) with warnings.catch_warnings(): warnings.simplefilter("error") - write_samples(self.path, "forward", "golden", _golden()) - tight = SampleSet([Sample([Call([torch.ones(1)], tol=(1e-3, 1e-3))])]) - write_samples(self.path, "forward", "tight", tight) + write_samples(self.path, "golden", _golden()) + tight = SampleSet( + [Sample([Call("forward", [torch.ones(1)], tol=(1e-3, 1e-3))])] + ) + write_samples(self.path, "tight", tight) def test_writer_requires_a_program(self): path = os.path.join(self._tmp.name, "empty.ptn") with zipfile.ZipFile(path, "w") as pkg: pkg.writestr("other", b"") with self.assertRaisesRegex(ValueError, "has no program.ptg"): - write_samples(path, "forward", "golden", _golden()) + write_samples(path, "golden", _golden()) def test_reader_rejects_bad_layout(self): cases = [ ({_DATA: None}, "missing data.safetensors"), ({_INDEX: None}, "missing index.json"), - ({"samples/forward/golden/notes.txt": b""}, "unexpected member"), - ({"samples/forward/stray.json": b""}, "unexpected member"), - ({"samples/for.ward/golden/index.json": b""}, "unexpected member"), - ({"samples/forward/golden/x/index.json": b""}, "unexpected member"), + ({"samples/golden/notes.txt": b""}, "unexpected member"), + ({"samples/stray.json": b""}, "unexpected member"), + ({"samples/gol.den/index.json": b""}, "unexpected member"), + ({"samples/golden/x/index.json": b""}, "unexpected member"), ] for overrides, message in cases: with self.subTest(message=message): write_ptn(self.path, _PTG, {}) - write_samples(self.path, "forward", "golden", _golden()) + write_samples(self.path, "golden", _golden()) _rewrite(self.path, overrides) with self.assertRaisesRegex(ValueError, message): read_samples(self.path) def test_reader_rejects_compressed_member(self): - write_samples(self.path, "forward", "golden", _golden()) + write_samples(self.path, "golden", _golden()) entries = _members(self.path) with zipfile.ZipFile(self.path, "w") as pkg: for name, data in entries.items(): @@ -466,6 +541,12 @@ def set_intermediates(raw): (set_sample_field("tol", [1e-4]), "tol must be"), (set_call_field("tol", [0.0, -1.0]), "tol must be"), (set_call_field("in", None), "'in' list"), + (set_call_field("method", ""), "method must be a non-empty string"), + (set_call_field("method", 1), "method must be a non-empty string"), + ( + lambda index: index["samples"][0]["calls"][0].pop("method"), + "method must be a non-empty string", + ), (set_call_field("in", [f"{at}.in[0]", None]), "unsupported value"), (set_call_field("in", [f"{at}.in[0]", float("nan")]), "unsupported value"), (set_call_field("out", f"{at}.out[0]"), "out must be a list"), @@ -485,7 +566,7 @@ def set_intermediates(raw): for edit, message in cases: with self.subTest(message=message): write_ptn(self.path, _PTG, {}) - write_samples(self.path, "forward", "golden", _golden()) + write_samples(self.path, "golden", _golden()) _edit_index(self.path, edit) with self.assertRaisesRegex(ValueError, message): read_samples(self.path) @@ -499,16 +580,14 @@ def test_reader_names_set_on_unparseable_members(self): for overrides, message in cases: with self.subTest(message=message): write_ptn(self.path, _PTG, {}) - write_samples(self.path, "forward", "golden", _golden()) + write_samples(self.path, "golden", _golden()) _rewrite(self.path, overrides) - with self.assertRaisesRegex( - ValueError, f"samples/forward/golden.*{message}" - ): + with self.assertRaisesRegex(ValueError, f"samples/golden.*{message}"): read_samples(self.path) def test_reader_accepts_newer_minor_version_and_unknown_fields(self): golden = _golden() - write_samples(self.path, "forward", "golden", golden) + write_samples(self.path, "golden", golden) def add_fields(index): index["version"] = "1.7" @@ -517,4 +596,4 @@ def add_fields(index): index["samples"][0]["calls"][0]["future"] = True _edit_index(self.path, add_fields) - self.assertSetsEqual(read_samples(self.path)[("forward", "golden")], golden) + self.assertSetsEqual(read_samples(self.path)["golden"], golden) diff --git a/exir/native/test/test_native.py b/exir/native/test/test_native.py index e122630497f..01a65e65b65 100644 --- a/exir/native/test/test_native.py +++ b/exir/native/test/test_native.py @@ -450,19 +450,26 @@ def test_samples_round_trip_on_saved_package(self): } ) golden = { - ("decode", "golden"): SampleSet( - [Sample([Call([decode_x], [model(decode_x).detach()])])], + "golden": SampleSet( + [ + Sample( + [ + Call("prefill", [prefill_x], [model(prefill_x).detach()]), + Call("decode", [decode_x], [model(decode_x).detach()]), + ] + ) + ], tol=(1e-5, 1e-5), ), - ("prefill", "calib"): SampleSet([Sample([Call([prefill_x])])]), + "calib": SampleSet([Sample([Call("prefill", [prefill_x])])]), } with tempfile.TemporaryDirectory() as td: path = os.path.join(td, "m.ptn") native.save(path) with zipfile.ZipFile(path) as pkg: program = {n: pkg.read(n) for n in pkg.namelist()} - for (method, name), samples in golden.items(): - write_samples(path, method, name, samples) + for name, samples in golden.items(): + write_samples(path, name, samples) with zipfile.ZipFile(path) as pkg: self.assertEqual({n: pkg.read(n) for n in program}, program) read = read_samples(path) @@ -470,13 +477,14 @@ def test_samples_round_trip_on_saved_package(self): for key, expected in golden.items(): self.assertEqual(read[key].tol, expected.tol) for got, want in zip(read[key].samples, expected.samples, strict=True): - (got_call,), (want_call,) = got.calls, want.calls - torch.testing.assert_close( - got_call.inputs, want_call.inputs, rtol=0, atol=0 - ) - torch.testing.assert_close( - got_call.outputs, want_call.outputs, rtol=0, atol=0 - ) + for got_call, want_call in zip(got.calls, want.calls, strict=True): + self.assertEqual(got_call.method, want_call.method) + torch.testing.assert_close( + got_call.inputs, want_call.inputs, rtol=0, atol=0 + ) + torch.testing.assert_close( + got_call.outputs, want_call.outputs, rtol=0, atol=0 + ) if __name__ == "__main__": From 12fe7d86405fc8dc379a8e08a8311671d8df639c Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Tue, 27 Oct 2026 18:38:15 -0700 Subject: [PATCH 3/5] Update [ghstack-poisoned] --- backends/native/README.md | 11 ++--- .../test/runtime/deserialize/test_package.cpp | 2 +- exir/_serialize/_ptn_samples.py | 34 ++++++++++----- exir/_serialize/test/test_ptn_samples.py | 43 +++++++++++++++++++ 4 files changed, 73 insertions(+), 17 deletions(-) diff --git a/backends/native/README.md b/backends/native/README.md index 13583d2537b..f973b6005a3 100644 --- a/backends/native/README.md +++ b/backends/native/README.md @@ -62,7 +62,7 @@ read_samples("model.ptn") # {"golden": SampleSet(...)} `write_samples` appends a set to any existing `.ptn`, not only one this process saved, and replaces the file only once the set is complete. `read_samples` validates every set against the rules below, except that each call's method -exists in the program. +exists in the program and its intermediate names are values of that method. ### Layout @@ -105,9 +105,9 @@ Key = str naming a tensor in the set's data.safetensors between them, so each sees the state the previous one left, whichever method that call ran. - A tensor's key is its JSON path in the index, e.g. `samples[0].calls[1].in[2]`. - Below a call that is `.in[k]`, `.in[k][l]` for an element of a list argument, - `.out[k]` or `.intermediates.`. Every key resolves, and every tensor is - referenced exactly once. + Below a call that is `.in[k]`, `.out[k]` or `.intermediates.`, with + `[l]` appended for each level of list nesting. Every key resolves, and every + tensor is referenced exactly once. - A JSON integer is an int scalar, any other number a float. - Intermediate names are top-level SSA value names of the call's method (`[A-Za-z_][A-Za-z0-9_]*`). @@ -115,7 +115,8 @@ Key = str naming a tensor in the set's data.safetensors - Tolerances are finite and non-negative, and advisory: a consumer may override them. `write_samples` warns when atol or rtol is above 1e-3. - Readers reject another major `version` and accept any minor one: a minor bump - only adds fields, and readers ignore unknown fields. + only adds fields, and readers ignore unknown fields and the tensors only those + fields reference. Not in v1: diff --git a/backends/native/test/runtime/deserialize/test_package.cpp b/backends/native/test/runtime/deserialize/test_package.cpp index e81939fdd27..0afa787051a 100644 --- a/backends/native/test/runtime/deserialize/test_package.cpp +++ b/backends/native/test/runtime/deserialize/test_package.cpp @@ -143,7 +143,7 @@ TEST(PackageTest, LoadsConstantsOnDemand) { TEST(PackageTest, IgnoresSampleMembers) { const std::string index = - R"({"version":1,"samples":[{"calls":[{"method":"forward","in":["samples[0].calls[0].in[0]"]}]}]})"; + R"({"version":"1.0","samples":[{"calls":[{"method":"forward","in":["samples[0].calls[0].in[0]"]}]}]})"; const std::vector data = make_safetensors( R"({"samples[0].calls[0].in[0]":{"dtype":"U8","shape":[2],"data_offsets":[0,2]}})", "in"); diff --git a/exir/_serialize/_ptn_samples.py b/exir/_serialize/_ptn_samples.py index afad11fb1a2..54e486ca951 100644 --- a/exir/_serialize/_ptn_samples.py +++ b/exir/_serialize/_ptn_samples.py @@ -116,7 +116,7 @@ def _encode(value: object, key: str, tensors: dict[str, torch.Tensor]) -> object f"write_samples: safetensors cannot represent dtype {value.dtype} " f"at {key}" ) - tensors[key] = value.detach().cpu().contiguous() + tensors[key] = value.detach().resolve_conj().resolve_neg().cpu().contiguous() return key if isinstance(value, (list, tuple)): return [_encode(v, f"{key}[{i}]", tensors) for i, v in enumerate(value)] @@ -213,6 +213,7 @@ def write_samples(path: str, name: str, samples: SampleSet) -> None: raise ValueError(f"write_samples: set name {name!r} must match {_NAME.pattern}") if not samples.samples: raise ValueError("write_samples: a set needs at least one sample") + path = os.path.realpath(path) tensors: dict[str, torch.Tensor] = {} encoded = json.dumps(_encode_set(samples, tensors), indent=1, allow_nan=False) @@ -224,9 +225,7 @@ def write_samples(path: str, name: str, samples: SampleSet) -> None: raise ValueError(f"write_samples: {path} has no {PTG_ENTRY}") if any(n.startswith(prefix) for n in names): raise ValueError(f"write_samples: {prefix} already exists in {path}") - fd, staging = tempfile.mkstemp( - dir=os.path.dirname(os.path.abspath(path)), suffix=".ptn.tmp" - ) + fd, staging = tempfile.mkstemp(dir=os.path.dirname(path), suffix=".ptn.tmp") os.close(fd) try: shutil.copyfile(path, staging) @@ -342,7 +341,8 @@ def _decode_set( raise decoder.fail(f"{field} must be a string") samples = [decoder.sample(e, f"samples[{i}]") for i, e in enumerate(entries)] - if decoder.unreferenced: + newer_minor = int(match[2]) > int(_VERSION.split(".")[1]) + if decoder.unreferenced and not newer_minor: raise decoder.fail(f"unreferenced tensors {sorted(decoder.unreferenced)}") return SampleSet( samples=samples, @@ -378,14 +378,28 @@ def read_samples(path: str) -> dict[str, SampleSet]: raise ValueError(f"read_samples: unexpected member {info.filename!r}") if info.compress_type != zipfile.ZIP_STORED: raise ValueError(f"read_samples: {info.filename!r} is compressed") - found.setdefault(parts[0], {})[parts[1]] = info + members = found.setdefault(parts[0], {}) + if parts[1] in members: + raise ValueError(f"read_samples: duplicate member {info.filename!r}") + members[parts[1]] = info for name, members in sorted(found.items()): where = f"read_samples: {_SAMPLES_DIR}{name}" for entry in (_INDEX_ENTRY, _DATA_ENTRY): if entry not in members: raise ValueError(f"{where}: missing {entry}") - data = pkg.read(members[_DATA_ENTRY]) + try: + data = pkg.read(members[_DATA_ENTRY]) + index = json.loads(pkg.read(members[_INDEX_ENTRY])) + tensors = safetensors_load(data) + except ( + ValueError, + KeyError, + RecursionError, + SafetensorError, + zipfile.BadZipFile, + ) as e: + raise ValueError(f"{where}: {e}") from e header_size = int.from_bytes(data[:8], "little") start = _data_offset(path, members[_DATA_ENTRY]) + 8 + header_size if start % _DATA_ALIGNMENT: @@ -393,11 +407,9 @@ def read_samples(path: str) -> dict[str, SampleSet]: f"{where}: tensor data is not {_DATA_ALIGNMENT}-byte aligned" ) try: - index = json.loads(pkg.read(members[_INDEX_ENTRY])) - tensors = safetensors_load(data) - except (ValueError, RecursionError, SafetensorError) as e: + sets[name] = _decode_set(index, tensors, where) + except (RecursionError, OverflowError) as e: raise ValueError(f"{where}: {e}") from e - sets[name] = _decode_set(index, tensors, where) return sets diff --git a/exir/_serialize/test/test_ptn_samples.py b/exir/_serialize/test/test_ptn_samples.py index 1d95645c4ab..3861be4f23d 100644 --- a/exir/_serialize/test/test_ptn_samples.py +++ b/exir/_serialize/test/test_ptn_samples.py @@ -305,6 +305,14 @@ def test_append_leaves_program_untouched(self): self.assertEqual(ptg, _PTG) self.assertTrue(torch.equal(constants["w"], torch.ones(4))) + def test_append_through_symlink_updates_target(self): + link = os.path.join(self._tmp.name, "link.ptn") + os.symlink(self.path, link) + write_samples(link, "golden", _golden()) + + self.assertTrue(os.path.islink(link)) + self.assertEqual(list(read_samples(self.path)), ["golden"]) + def test_members_are_stored_and_data_aligned(self): # Vary the preceding index.json size so the data member starts at every # offset modulo the alignment. @@ -511,6 +519,16 @@ def test_reader_rejects_compressed_member(self): with self.assertRaisesRegex(ValueError, "is compressed"): read_samples(self.path) + def test_reader_rejects_duplicate_member(self): + write_samples(self.path, "golden", _golden()) + index = _members(self.path)[_INDEX] + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with zipfile.ZipFile(self.path, "a") as pkg: + pkg.writestr(_INDEX, index) + with self.assertRaisesRegex(ValueError, "duplicate member"): + read_samples(self.path) + def test_reader_rejects_bad_index(self): def set_field(field, value): return lambda index: index.__setitem__(field, value) @@ -530,6 +548,7 @@ def set_intermediates(raw): cases = [ (lambda index: index.clear(), "unsupported version"), (set_field("version", "2.0"), "unsupported version"), + (set_field("tol", [10**400, 0.0]), "too large"), (set_field("version", 1), "unsupported version"), (set_field("version", "1"), "unsupported version"), (set_field("version", "1.01"), "unsupported version"), @@ -575,6 +594,16 @@ def test_reader_names_set_on_unparseable_members(self): cases = [ ({_INDEX: b'{"version": '}, "Expecting value"), ({_INDEX: b"[" * 100_000}, ""), + ( + { + _INDEX: b'{"version": "1.0", "samples": [{"calls": [{"method": ' + + b'"forward", "in": ' + + b"[" * 600 + + b"]" * 600 + + b"}]}]}" + }, + "recursion", + ), ({_DATA: b"\x08" + bytes(7) + b"not json"}, ""), ] for overrides, message in cases: @@ -597,3 +626,17 @@ def add_fields(index): _edit_index(self.path, add_fields) self.assertSetsEqual(read_samples(self.path)["golden"], golden) + + def test_newer_minor_may_reference_tensors_from_unknown_fields(self): + write_samples(self.path, "golden", _golden()) + + def move_outputs_to_unknown_field(index): + index["version"] = "1.1" + call = index["samples"][0]["calls"][0] + call["future"] = call.pop("out") + + _edit_index(self.path, move_outputs_to_unknown_field) + self.assertIsNone(read_samples(self.path)["golden"].samples[0].calls[0].outputs) + _edit_index(self.path, lambda index: index.__setitem__("version", "1.0")) + with self.assertRaisesRegex(ValueError, "unreferenced tensors"): + read_samples(self.path) From deafada9b6b5ab48444bb822f8666b2ebae6d5ce Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Tue, 27 Oct 2026 19:59:27 -0700 Subject: [PATCH 4/5] Update [ghstack-poisoned] --- exir/_serialize/_ptn_samples.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/exir/_serialize/_ptn_samples.py b/exir/_serialize/_ptn_samples.py index 54e486ca951..bc4ec875c46 100644 --- a/exir/_serialize/_ptn_samples.py +++ b/exir/_serialize/_ptn_samples.py @@ -352,7 +352,7 @@ def _decode_set( ) -def read_samples(path: str) -> dict[str, SampleSet]: +def read_samples(path: str) -> dict[str, SampleSet]: # noqa: C901 """Every sample set in a .ptn, keyed by set name, validated as read. A package without samples returns an empty dict. From a1b24480108e2d8b6fbf72fb2222402558940b3e Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Tue, 27 Oct 2026 21:18:48 -0700 Subject: [PATCH 5/5] Update [ghstack-poisoned] --- exir/_serialize/test/test_ptn_samples.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/exir/_serialize/test/test_ptn_samples.py b/exir/_serialize/test/test_ptn_samples.py index 3861be4f23d..faeca978583 100644 --- a/exir/_serialize/test/test_ptn_samples.py +++ b/exir/_serialize/test/test_ptn_samples.py @@ -7,6 +7,7 @@ import json import os import struct +import sys import tempfile import unittest import warnings @@ -611,8 +612,15 @@ def test_reader_names_set_on_unparseable_members(self): write_ptn(self.path, _PTG, {}) write_samples(self.path, "golden", _golden()) _rewrite(self.path, overrides) - with self.assertRaisesRegex(ValueError, f"samples/golden.*{message}"): - read_samples(self.path) + limit = sys.getrecursionlimit() + sys.setrecursionlimit(1000) + try: + with self.assertRaisesRegex( + ValueError, f"samples/golden.*{message}" + ): + read_samples(self.path) + finally: + sys.setrecursionlimit(limit) def test_reader_accepts_newer_minor_version_and_unknown_fields(self): golden = _golden()