diff --git a/backends/native/README.md b/backends/native/README.md index a784537989f..f973b6005a3 100644 --- a/backends/native/README.md +++ b/backends/native/README.md @@ -41,6 +41,94 @@ 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, 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("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 each call's method +exists in the program and its intermediate names are values of that method. + +### Layout + +``` +samples//index.json +samples//data.safetensors +``` + +- `` 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`. +- 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 = { "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 + +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, 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]`, `.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_]*`). + 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 and the tensors only those + fields reference. + +Not in v1: + +- 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..0afa787051a 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.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"); + const TempPackage file( + {}, + {{"samples/golden/index.json", index}, + {"samples/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..bc4ec875c46 --- /dev/null +++ b/exir/_serialize/_ptn_samples.py @@ -0,0 +1,423 @@ +# 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 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. +""" + +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.""" + + method: str + 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().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)] + 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__}" + ) + 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: + 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, 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 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. + """ + 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") + path = os.path.realpath(path) + + tensors: dict[str, torch.Tensor] = {} + encoded = json.dumps(_encode_set(samples, tensors), indent=1, allow_nan=False) + + prefix = f"{_SAMPLES_DIR}{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(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") + 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=( + 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)] + 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, + description=index.get("description"), + source=index.get("source"), + tol=decoder.tol(index), + ) + + +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. + + 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[str, SampleSet] = {} + with zipfile.ZipFile(path) as pkg: + 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) != 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") + 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}") + 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: + raise ValueError( + f"{where}: tensor data is not {_DATA_ALIGNMENT}-byte aligned" + ) + try: + sets[name] = _decode_set(index, tensors, where) + except (RecursionError, OverflowError) as e: + raise ValueError(f"{where}: {e}") from e + 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..faeca978583 --- /dev/null +++ b/exir/_serialize/test/test_ptn_samples.py @@ -0,0 +1,650 @@ +# 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 sys +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/golden/index.json" +_DATA = "samples/golden/data.safetensors" + + +def _golden() -> SampleSet: + return SampleSet( + samples=[ + 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", + 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.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: + 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, "golden", golden) + + sets = read_samples(self.path) + self.assertEqual(list(sets), ["golden"]) + self.assertSetsEqual(sets["golden"], golden) + + def test_intermediates_and_nested_tol_round_trip(self): + debug = SampleSet( + samples=[ + Sample( + [ + Call( + "forward", + 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, "debug", debug) + + self.assertSetsEqual(read_samples(self.path)["debug"], debug) + + def test_calls_round_trip_in_order(self): + prefill = Call( + "forward", + [torch.arange(6).view(1, 6), torch.arange(6)], + [torch.randn(1, 1, 5)], + ) + decode = [ + 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, "golden", golden) + + self.assertSetsEqual(read_samples(self.path)["golden"], golden) + + def test_inputs_only_round_trips(self): + 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)["calibration"], calib) + + def test_scalars_lists_and_dtypes_round_trip(self): + mixed = SampleSet( + samples=[ + Sample( + [ + Call( + "forward", + 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, "mixed", mixed) + + self.assertSetsEqual(read_samples(self.path)["mixed"], mixed) + + def test_samples_may_differ_in_shape(self): + dynamic = SampleSet( + samples=[Sample([Call("forward", [torch.randn(1, n, 8)])]) for n in (3, 9)] + ) + write_samples(self.path, "dynamic", dynamic) + + 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, + "golden", + SampleSet( + samples=[ + Sample( + [ + Call( + "forward", + inputs=[torch.ones(1), [torch.ones(1), 2]], + outputs=[torch.ones(1)], + intermediates={"linear": torch.ones(1)}, + ), + Call("forward", 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], {"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, "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_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. + 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, "pad", pad) + write_samples(self.path, "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, "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, "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, "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, "golden", _golden()) + with self.assertRaisesRegex(ValueError, "already exists"): + write_samples(self.path, "golden", _golden()) + + def test_writer_rejects_bad_names(self): + for name in ("gol/den", "gol.den", ""): + with self.subTest(name=name): + with self.assertRaisesRegex(ValueError, "must match"): + 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( + "forward", + [torch.ones(1)], + intermediates={"a.b": torch.ones(1)}, + ) + ] + ) + ] + ) + with self.assertRaisesRegex(ValueError, "must match"): + 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("forward", [float("nan")])])]), + ValueError, + "not finite", + ), + ( + SampleSet( + [Sample([Call("forward", [torch.ones(1)])])], tol=(-1.0, 0.0) + ), + ValueError, + "tol", + ), + ( + SampleSet( + [Sample([Call("forward", [torch.ones(1)], tol=(0.0, -1.0))])] + ), + ValueError, + r"samples\[0\]\.calls\[0\]: tol", + ), + ( + 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("forward", torch.ones(1))])]), + TypeError, + "inputs must be a list", + ), + ( + SampleSet([Sample([Call("forward", [torch.ones(1)], torch.ones(1))])]), + TypeError, + "outputs must be a list", + ), + ( + SampleSet( + [Sample([Call("forward", [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, "golden", samples) + self.assertEqual(read_samples(self.path), {}) + + def test_writer_warns_on_loose_tol(self): + one = [torch.ones(1)] + cases = [ + (SampleSet([Sample([Call("forward", one)])], tol=(2.0, 0.05)), "set"), + ( + 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\]", + ), + ] + for i, (loose, where) in enumerate(cases): + with self.assertWarnsRegex( + UserWarning, f"{where} tol .* above 0.001" + ) as caught: + write_samples(self.path, f"loose{i}", loose) + self.assertEqual(caught.filename, __file__) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + 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, "golden", _golden()) + + def test_reader_rejects_bad_layout(self): + cases = [ + ({_DATA: None}, "missing data.safetensors"), + ({_INDEX: None}, "missing index.json"), + ({"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, "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, "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_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) + + 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("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"), + (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("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"), + (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, "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}, ""), + ( + { + _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: + with self.subTest(message=message): + write_ptn(self.path, _PTG, {}) + write_samples(self.path, "golden", _golden()) + _rewrite(self.path, overrides) + 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() + write_samples(self.path, "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)["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) 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..01a65e65b65 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,53 @@ 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 = { + "golden": SampleSet( + [ + Sample( + [ + Call("prefill", [prefill_x], [model(prefill_x).detach()]), + Call("decode", [decode_x], [model(decode_x).detach()]), + ] + ) + ], + tol=(1e-5, 1e-5), + ), + "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 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) + 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): + 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__": unittest.main()