diff --git a/backends/native/test/CMakeLists.txt b/backends/native/test/CMakeLists.txt index 73ea13f321c..eb404c4cba9 100644 --- a/backends/native/test/CMakeLists.txt +++ b/backends/native/test/CMakeLists.txt @@ -29,6 +29,24 @@ target_compile_options( $<$>:-fexceptions> ) +# Op packages with bundled samples, one per op and quantization scheme. +set(_native_op_samples ${CMAKE_CURRENT_BINARY_DIR}/op_samples) +file(GLOB _native_op_samples_srcs CONFIGURE_DEPENDS + ${CMAKE_CURRENT_SOURCE_DIR}/samples/*.py +) +add_custom_command( + OUTPUT ${_native_op_samples}.stamp + COMMAND + ${CMAKE_COMMAND} -E env "PYTHONPATH=${EXECUTORCH_ROOT}/src" + ${PYTHON_EXECUTABLE} -m executorch.backends.native.test.samples.generate + --out ${_native_op_samples} + COMMAND ${CMAKE_COMMAND} -E touch ${_native_op_samples}.stamp + DEPENDS ${_native_op_samples_srcs} + ${EXECUTORCH_ROOT}/backends/test/suite/operators/test_linear.py + VERBATIM +) +add_custom_target(native_op_samples DEPENDS ${_native_op_samples}.stamp) + if(EXECUTORCH_BUILD_NATIVE_VULKAN) add_executable( native_vulkan_passes_test diff --git a/backends/native/test/samples/BUCK b/backends/native/test/samples/BUCK new file mode 100644 index 00000000000..97b4b66b7db --- /dev/null +++ b/backends/native/test/samples/BUCK @@ -0,0 +1,50 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +fbcode_target( + _kind = runtime.python_library, + name = "record", + srcs = [ + "cases.py", + "linear.py", + "ops.py", + "record.py", + "schemes.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/test/suite/operators:test_operators", + "//executorch/exir/native:native", + "//pytorch/ao:torchao", + ], +) + +fbcode_target( + _kind = runtime.python_binary, + name = "generate", + srcs = ["generate.py"], + main_function = "executorch.backends.native.test.samples.generate.main", + deps = [":record"], +) + +fbcode_target( + _kind = runtime.python_test, + name = "test_record", + srcs = ["test_record.py"], + deps = [ + ":record", + "//caffe2:torch", + "//executorch/exir/native:native", + ], +) + +# Op packages with bundled samples, one per op and quantization scheme. +fbcode_target( + _kind = runtime.genrule, + name = "op_samples", + out = "op_samples", + cmd = "$(exe :generate) --out $OUT", + visibility = ["//executorch/backends/native/test/..."], +) diff --git a/backends/native/test/samples/cases.py b/backends/native/test/samples/cases.py new file mode 100644 index 00000000000..b08baba963a --- /dev/null +++ b/backends/native/test/samples/cases.py @@ -0,0 +1,87 @@ +# 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. + +"""Op cases: a model and the inputs of each sample it is recorded with.""" + +import inspect +from dataclasses import dataclass +from types import ModuleType + +import torch + +Inputs = tuple[torch.Tensor, ...] + +# Caps a suite case's input and output elements over its samples, keeping at +# least one sample: a large output alone has enough elements for statistics. +_ELEMENT_BUDGET = 1 << 16 + + +@dataclass +class Case: + name: str + model: torch.nn.Module + # Positional float32 inputs of each sample; the count may vary by case. + inputs: list[Inputs] + + +def outputs(result) -> list[torch.Tensor]: + return list(result) if isinstance(result, (tuple, list)) else [result] + + +def draws(count: int, *shapes: tuple[int, ...]) -> list[Inputs]: + """`count` samples of standard-normal inputs of the given shapes, from a fixed + seed, so equal calls return equal samples.""" + generator = torch.Generator().manual_seed(0) + return [ + tuple(torch.randn(shape, generator=generator) for shape in shapes) + for _ in range(count) + ] + + +class _SuiteRunner: + """Stands in for the backend test suite's runner, keeping what a test lowers.""" + + _flow = None + + def __init__(self) -> None: + self.calls: list[tuple[torch.nn.Module, Inputs]] = [] + + def lower_and_run_model(self, model, inputs, **_) -> None: + self.calls.append((model, tuple(inputs))) + + +def suite_cases(module: ModuleType, samples: int) -> list[Case]: + """One case per model a backend test suite operator module lowers, with + inputs from up to `samples` runs under different seeds.""" + cases = [] + for test_name, test in sorted(vars(module).items()): + if not test_name.startswith("test_"): + continue + runs = [] + for seed in range(samples): + torch.manual_seed(seed) + runner = _SuiteRunner() + if "dtype" in inspect.signature(test).parameters: + test(runner, torch.float32) + else: + test(runner) + runs.append(runner.calls) + lowered = [[(type(m), [x.shape for x in xs]) for m, xs in run] for run in runs] + if any(seed_lowered != lowered[0] for seed_lowered in lowered): + raise ValueError(f"{test_name} lowers different models across seeds") + for i, (model, inputs) in enumerate(runs[0]): + name = f"{test_name.removeprefix('test_')}_{i}" + try: + results = outputs(model(*inputs)) + except RuntimeError as e: + e.add_note(f"while recording suite case {name}") + raise + numel = sum(t.numel() for t in (*inputs, *results)) + count = max(1, min(samples, _ELEMENT_BUDGET // max(1, numel))) + cases.append(Case(name, model, [run[i][1] for run in runs[:count]])) + if not cases: + raise ValueError(f"no test in {module.__name__} lowers a model") + return cases diff --git a/backends/native/test/samples/generate.py b/backends/native/test/samples/generate.py new file mode 100644 index 00000000000..665a17aa22b --- /dev/null +++ b/backends/native/test/samples/generate.py @@ -0,0 +1,50 @@ +# 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. + +"""Write _.ptn for each op and quantization scheme that fits a case +of the op.""" + +import argparse +import logging +import multiprocessing +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path + +from executorch.backends.native.test.samples.ops import ops +from executorch.backends.native.test.samples.record import record +from executorch.backends.native.test.samples.schemes import schemes + + +def _record(name: str, scheme: str, out: Path) -> None: + logging.basicConfig(level=logging.INFO, force=True) + op = ops()[name] + record(op.cases(), schemes()[scheme], out / f"{name}_{scheme}.ptn", op.dtypes) + + +def main() -> None: + all_ops, all_schemes = ops(), schemes() + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--ops", nargs="+", choices=all_ops, default=list(all_ops)) + parser.add_argument( + "--schemes", nargs="+", choices=all_schemes, default=list(all_schemes) + ) + args = parser.parse_args() + args.out.mkdir(parents=True, exist_ok=True) + # Forked workers can't use the CUDA state torch may have initialized here. + spawn = multiprocessing.get_context("spawn") + with ProcessPoolExecutor(mp_context=spawn) as pool: + jobs = [ + pool.submit(_record, name, scheme, args.out) + for name in args.ops + for scheme in args.schemes + ] + for job in jobs: + job.result() + + +if __name__ == "__main__": + main() diff --git a/backends/native/test/samples/linear.py b/backends/native/test/samples/linear.py new file mode 100644 index 00000000000..24c39fe2ccc --- /dev/null +++ b/backends/native/test/samples/linear.py @@ -0,0 +1,36 @@ +# 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. + +"""Linear cases: the backend test suite's, and shapes and inputs that stress +accumulation and quantization.""" + +import torch +from executorch.backends.native.test.samples.cases import Case, draws, suite_cases +from executorch.backends.test.suite.operators import test_linear + + +def cases() -> list[Case]: + suite = suite_cases(test_linear, samples=4) + torch.manual_seed(0) + # One channel far above the rest, as in LLM activations. + outlier = draws(4, (16, 256)) + for (x,) in outlier: + x[:, 3] *= 32 + zero_row = draws(2, (8, 256)) + for (x,) in zero_row: + x[2] = 0 + return suite + [ + Case("decode", torch.nn.Linear(2048, 512), draws(8, (1, 2048))), + Case("prefill", torch.nn.Linear(1024, 256), draws(2, (64, 1024))), + Case("large_k", torch.nn.Linear(8192, 32), draws(4, (4, 8192))), + Case( + "leading_batch_no_bias", + torch.nn.Linear(256, 96, bias=False), + draws(4, (2, 7, 256)), + ), + Case("outlier_channel", torch.nn.Linear(256, 128), outlier), + Case("zero_row", torch.nn.Linear(256, 64), zero_row), + ] diff --git a/backends/native/test/samples/ops.py b/backends/native/test/samples/ops.py new file mode 100644 index 00000000000..16332b665ee --- /dev/null +++ b/backends/native/test/samples/ops.py @@ -0,0 +1,24 @@ +# 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. + +"""Recorded ops, keyed by the name their packages are written under.""" + +from collections.abc import Callable +from dataclasses import dataclass + +from executorch.backends.native.test.samples import linear +from executorch.backends.native.test.samples.cases import Case +from executorch.backends.native.test.samples.record import DTYPES + + +@dataclass(frozen=True) +class Op: + cases: Callable[[], list[Case]] + dtypes: tuple[str, ...] = tuple(DTYPES) + + +def ops() -> dict[str, Op]: + return {"linear": Op(linear.cases)} diff --git a/backends/native/test/samples/record.py b/backends/native/test/samples/record.py new file mode 100644 index 00000000000..6846890f081 --- /dev/null +++ b/backends/native/test/samples/record.py @@ -0,0 +1,92 @@ +# 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. + +"""Record op cases into .ptn packages with bundled samples. + +A package holds one method per case and dtype under one quantization scheme, +and two sample sets over the same inputs: `golden`, the eager outputs, and +`precise`, the same math in float64 with weights dequantized and activations +unquantized, floating outputs stored as float32. Nothing in it depends on an +engine or a device: consumers judge an engine by its distance from `precise` +relative to eager's. +""" + +import copy +import logging +from collections.abc import Iterable +from pathlib import Path + +import torch +from executorch.backends.native.test.samples.cases import Case, Inputs, outputs +from executorch.backends.native.test.samples.schemes import Scheme +from executorch.exir.native import Call, Sample, SampleSet, to_native, write_samples +from torch.func import functional_call + +logger: logging.Logger = logging.getLogger(__name__) + +DTYPES = {"f32": torch.float32, "f16": torch.float16, "bf16": torch.bfloat16} + + +class _Method(torch.nn.ModuleDict): + """One model keyed by its method name, keeping constant FQNs unique across a + package's methods.""" + + def forward(self, *args): + (model,) = self.values() + return model(*args) + + +def _cast(inputs: Inputs, dtype: torch.dtype) -> Inputs: + return tuple(t.to(dtype) if t.is_floating_point() else t for t in inputs) + + +def record( + cases: list[Case], + scheme: Scheme, + path: Path, + dtypes: Iterable[str] = tuple(DTYPES), +) -> None: + """Writes no package, removing any at `path`, when the scheme fits no case.""" + programs, golden, precise = {}, [], [] + for case in cases: + if not scheme.fits(case.model): + logger.info(f"{path.name}: {case.name} does not fit the scheme") + continue + for suffix in dtypes: + name = f"{case.name}_{suffix}" + if name in programs: + raise ValueError(f"{path.name}: duplicate method {name}") + dtype = DTYPES[suffix] + model = copy.deepcopy(case.model).to(dtype).eval() + scheme.quantize(model) + method = _Method({name: model}) + state = { + k: scheme.decode(v).double() for k, v in method.state_dict().items() + } + inputs = [_cast(x, dtype) for x in case.inputs] + programs[name] = torch.export.export(method, inputs[0]) + with torch.no_grad(): + for x in inputs: + golden.append(Sample([Call(name, list(x), outputs(method(*x)))])) + ref = functional_call(method, state, _cast(x, torch.float64)) + ref = [ + t.float() if t.is_floating_point() else t for t in outputs(ref) + ] + precise.append(Sample([Call(name, list(x), ref)])) + if not programs: + logger.info(f"{path.name}: no case fits the scheme, not written") + path.unlink(missing_ok=True) + return + to_native(programs).save(str(path)) + source = f"torch {torch.__version__}" + write_samples(str(path), "golden", SampleSet(golden, "eager", source)) + write_samples( + str(path), + "precise", + SampleSet( + precise, "float64, weights dequantized, activations unquantized", source + ), + ) diff --git a/backends/native/test/samples/schemes.py b/backends/native/test/samples/schemes.py new file mode 100644 index 00000000000..8966e7facad --- /dev/null +++ b/backends/native/test/samples/schemes.py @@ -0,0 +1,91 @@ +# 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. + +"""Quantization schemes. Every op is recorded under every scheme, keeping the +cases the scheme fits.""" + +from dataclasses import dataclass + +import torch +from torchao.core.config import AOBaseConfig +from torchao.quantization import ( + Int8DynamicActivationIntxWeightConfig, + IntxWeightOnlyConfig, + quantize_, +) +from torchao.quantization.granularity import PerGroup +from torchao.utils import TorchAOBaseTensor + + +class Scheme: + """No quantization.""" + + def fits(self, model: torch.nn.Module) -> bool: + return True + + def quantize(self, model: torch.nn.Module) -> None: + pass + + def decode(self, value: torch.Tensor) -> torch.Tensor: + """A state tensor of a quantized model as a plain tensor.""" + return value + + +def _fits(model: torch.nn.Module, group: int | None) -> bool: + """Whether `model` has an nn.Linear and every one's in_features split into + groups of `group`, or form one group when it is None.""" + linears = [m for m in model.modules() if isinstance(m, torch.nn.Linear)] + for m in linears: + size = m.in_features if group is None else group + # torchao reduces a one-element-group weight to a single scale. + if size < 2 or m.in_features % size: + return False + return bool(linears) + + +@dataclass(frozen=True) +class Torchao(Scheme): + """A torchao config applied to every nn.Linear.""" + + config: AOBaseConfig + + def fits(self, model: torch.nn.Module) -> bool: + granularity = ( + self.config.granularity + if isinstance(self.config, IntxWeightOnlyConfig) + else self.config.weight_granularity + ) + return _fits( + model, granularity.group_size if isinstance(granularity, PerGroup) else None + ) + + def quantize(self, model: torch.nn.Module) -> None: + quantize_(model, self.config) + + def decode(self, value: torch.Tensor) -> torch.Tensor: + return value.dequantize() if isinstance(value, TorchAOBaseTensor) else value + + +def schemes() -> dict[str, Scheme]: + return { + "float": Scheme(), + "int4wo_g32": Torchao( + IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32)) + ), + "int4wo_g128": Torchao( + IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(128)) + ), + "8da4w_g32": Torchao( + Int8DynamicActivationIntxWeightConfig( + weight_dtype=torch.int4, weight_granularity=PerGroup(32) + ) + ), + "8da4w_g128": Torchao( + Int8DynamicActivationIntxWeightConfig( + weight_dtype=torch.int4, weight_granularity=PerGroup(128) + ) + ), + } diff --git a/backends/native/test/samples/test_record.py b/backends/native/test/samples/test_record.py new file mode 100644 index 00000000000..391ecd9b52a --- /dev/null +++ b/backends/native/test/samples/test_record.py @@ -0,0 +1,90 @@ +# 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 tempfile +import unittest +from pathlib import Path +from types import ModuleType + +import torch +import torch.nn.functional as F +from executorch.backends.native.test.samples.cases import Case, draws, suite_cases +from executorch.backends.native.test.samples.record import record +from executorch.backends.native.test.samples.schemes import Scheme, schemes +from executorch.exir.native import read_samples + + +class RecordTest(unittest.TestCase): + def test_suite_cases_raises_on_a_model_eager_rejects(self) -> None: + module = ModuleType("suite") + module.test_mismatch = lambda runner: runner.lower_and_run_model( + torch.nn.Linear(4, 2), (torch.randn(1, 3),) + ) + with self.assertRaises(RuntimeError) as raised: + suite_cases(module, samples=1) + self.assertIn( + "while recording suite case mismatch_0", raised.exception.__notes__ + ) + + def test_suite_cases_raises_when_seeds_lower_different_shapes(self) -> None: + module = ModuleType("suite") + sizes = iter([2, 3]) + module.test_varying = lambda runner: runner.lower_and_run_model( + torch.nn.Identity(), (torch.randn(next(sizes)),) + ) + with self.assertRaises(ValueError): + suite_cases(module, samples=2) + + def test_suite_cases_raises_when_no_test_lowers_a_model(self) -> None: + with self.assertRaises(ValueError): + suite_cases(ModuleType("suite"), samples=1) + + def test_record_raises_on_duplicate_case_names(self) -> None: + case = Case("same", torch.nn.Linear(2, 2), draws(1, (1, 2))) + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(ValueError): + record([case, case], Scheme(), Path(tmp) / "dup.ptn") + + def test_record_writes_no_package_when_the_scheme_fits_no_case(self) -> None: + case = Case("relu", torch.nn.ReLU(), draws(1, (1, 32))) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "relu.ptn" + path.write_bytes(b"stale") + record([case], schemes()["int4wo_g32"], path) + self.assertFalse(path.exists()) + + def test_record_golden_is_eager_and_precise_is_float64(self) -> None: + model = torch.nn.Linear(16, 4) + (x,) = draws(1, (3, 16))[0] + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "linear.ptn" + record([Case("linear", model, [(x,)])], Scheme(), path, dtypes=["bf16"]) + sets = read_samples(str(path)) + x, w, b = (t.detach().bfloat16() for t in (x, model.weight, model.bias)) + torch.testing.assert_close( + sets["golden"].samples[0].calls[0].outputs[0], + F.linear(x, w, b), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + sets["precise"].samples[0].calls[0].outputs[0], + F.linear(x.double(), w.double(), b.double()).float(), + rtol=0, + atol=0, + ) + + def test_record_keeps_integer_outputs_precise(self) -> None: + class Argmax(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.argmax(-1) + + case = Case("argmax", Argmax(), draws(1, (2, 8))) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "argmax.ptn" + record([case], Scheme(), path, dtypes=["f32"]) + precise = read_samples(str(path))["precise"].samples + self.assertEqual(precise[0].calls[0].outputs[0].dtype, torch.int64)