From 3016b89300c097841ad6d9fd24f821f7e22d5ae8 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Tue, 3 Nov 2026 06:24:34 -0800 Subject: [PATCH] [executorch][native] Record dynamic-shape linears as chained bundled samples Pull Request resolved: https://github.com/pytorch/executorch/pull/24405 Engines resize dynamic inputs between inferences, so the op samples now cover a shape changing mid-run. A case built with `chain` is exported with its `Dim`s and recorded as one sample of four calls. Each `Dim` starts mid-range, where export traces it, visits both bounds and returns, every other `Dim` in the opposite order. The harness already runs a sample's calls in order on one executable from one context, resizing inputs between them, so only its header comment changes. The new linear cases are the XNNPACK delegate's dynamic linear tests, each also followed by a ReLU: one and two batch dims, one `Dim` on two dims or two inputs, and sequential and parallel linears, with batch dims down to 1, which tracing mid-range keeps export from specializing. The cases every quantization scheme fits take 256 input features over up to 64 tokens: `dynamic_tokens` in 2-D and, with a static batch of one as in an LLM, in 3-D, two batch dims sharing one `Dim` of up to 16, and the sequential and parallel linears. Authored with Claude Code. Differential Revision: [D128770754](https://our.internmc.facebook.com/intern/diff/D128770754/) **NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D128770754/)! ghstack-source-id: 463657750 --- backends/native/test/samples/cases.py | 32 ++++++- backends/native/test/samples/linear.py | 92 ++++++++++++++++++- backends/native/test/samples/record.py | 33 +++++-- .../native/test/samples/sample_test_main.cpp | 5 +- backends/native/test/samples/test_record.py | 44 ++++++++- 5 files changed, 191 insertions(+), 15 deletions(-) diff --git a/backends/native/test/samples/cases.py b/backends/native/test/samples/cases.py index 82b97316b24..9a613a14db9 100644 --- a/backends/native/test/samples/cases.py +++ b/backends/native/test/samples/cases.py @@ -4,7 +4,8 @@ # 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.""" +"""Op cases: a model and the inputs of each sample it is recorded with, or of +the calls of one sample when its shapes are dynamic.""" import inspect from collections.abc import Callable @@ -13,6 +14,7 @@ from typing import TYPE_CHECKING import torch +from torch.export import Dim if TYPE_CHECKING: from facto.inputgen.specs.model import Spec @@ -30,6 +32,9 @@ class Case: model: torch.nn.Module # Positional float32 inputs of each sample; the count may vary by case. inputs: list[Inputs] + # `{dim: Dim}` of each positional input. A dynamic case's inputs are the + # calls of one sample, which run in order on one executable. + dynamic_shapes: tuple[dict[int, Dim], ...] | None = None def outputs(result) -> list[torch.Tensor]: @@ -46,6 +51,31 @@ def draws(count: int, *shapes: tuple[int, ...]) -> list[Inputs]: ] +def chain(name: str, model: torch.nn.Module, *shapes: tuple[int | Dim, ...]) -> Case: + """A dynamic case with standard-normal inputs of `shapes`, from a fixed seed. + Each Dim starts mid-range, which export traces, then visits both bounds and + returns, every other Dim in the opposite order.""" + dims = list(dict.fromkeys(d for s in shapes for d in s if not isinstance(d, int))) + walks = {} + for k, d in enumerate(dims): + mid = (d.min + d.max + 1) // 2 + walks[d] = (mid, d.min, d.max, mid) if k % 2 == 0 else (mid, d.max, d.min, mid) + generator = torch.Generator().manual_seed(0) + inputs = [ + tuple( + torch.randn( + [walks[d][step] if d in walks else d for d in s], generator=generator + ) + for s in shapes + ) + for step in range(4) + ] + dynamic_shapes = tuple( + {i: d for i, d in enumerate(s) if not isinstance(d, int)} for s in shapes + ) + return Case(name, model, inputs, dynamic_shapes) + + class _SuiteRunner: """Stands in for the backend test suite's runner, keeping what a test lowers.""" diff --git a/backends/native/test/samples/linear.py b/backends/native/test/samples/linear.py index 9beea6408fa..6b6a22b1438 100644 --- a/backends/native/test/samples/linear.py +++ b/backends/native/test/samples/linear.py @@ -4,12 +4,13 @@ # 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, FACTO's, and shapes and inputs that -stress accumulation and quantization.""" +"""Linear cases: the backend test suite's, FACTO's, shapes and inputs that +stress accumulation and quantization, and dynamic shapes.""" import torch from executorch.backends.native.test.samples.cases import ( Case, + chain, draws, facto_cases, suite_cases, @@ -19,6 +20,7 @@ from executorch.extension.llm.export.gguf import QK_K from facto.inputgen.argument.type import ArgType from facto.inputgen.specs.model import ConstraintProducer as cp, InPosArg, OutArg, Spec +from torch.export import Dim def _unit(*constraints) -> list: @@ -90,6 +92,91 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, ...]: return tuple(m(x) for m in self.linears) +class _Pair(torch.nn.Module): + """Two modules, each on its own input.""" + + def __init__(self, first: torch.nn.Module, second: torch.nn.Module) -> None: + super().__init__() + self.first, self.second = first, second + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, ...]: + return self.first(x), self.second(y) + + +_XNNPACK_SHAPES = ((3, 4, 4), (4, 37, 17), (4, 17, 37)) + + +def _batch_dims(count: int, batch: int) -> tuple[Dim, ...]: + return tuple(Dim(f"batch{i}", min=1, max=batch) for i in range(count)) + + +def _structural( + suffix: str, k: int, hidden: int, n: int, batch: Dim, batch2: Dim +) -> list[Case]: + """Parallel, sequential and both, each linear `k` to `hidden` but the last of + a sequence, `hidden` to `n`.""" + + def sequential() -> torch.nn.Module: + return torch.nn.Sequential( + torch.nn.Linear(k, hidden), torch.nn.Linear(hidden, n) + ) + + return [ + chain( + f"dynamic_parallel{suffix}", + _Pair(torch.nn.Linear(k, hidden), torch.nn.Linear(k, hidden)), + (batch, k), + (batch, k), + ), + chain(f"dynamic_sequential{suffix}", sequential(), (batch, k)), + chain( + f"dynamic_parallel_and_sequential{suffix}", + _Pair(torch.nn.Linear(k, hidden), sequential()), + (batch, k), + (batch2, k), + ), + ] + + +def _dynamic_cases() -> list[Case]: + """The XNNPACK delegate's dynamic linear tests, with batch dims down to 1, + and LLM token shapes every quantization scheme fits.""" + cases = [] + for batch, k, n in _XNNPACK_SHAPES: + for batch_dims in (1, 2): + for bias in (True, False): + name = f"dynamic_{batch_dims}batch_{k}_{n}" + ( + "" if bias else "_no_bias" + ) + shape = (*_batch_dims(batch_dims, batch), k) + cases.append(chain(name, torch.nn.Linear(k, n, bias=bias), shape)) + batch, batch2 = Dim("batch", min=1, max=100), Dim("batch2", min=1, max=100) + cases += [ + chain("dynamic_two_batch", torch.nn.Linear(14, 15), (batch, batch, 14)), + *_structural("", 4, 5, 3, batch, batch2), + ] + tokens = Dim("tokens", min=1, max=64) + cases.append(chain("dynamic_tokens", torch.nn.Linear(256, 64), (tokens, 256))) + for batch, k, n in _XNNPACK_SHAPES: + for batch_dims in (1, 2): + for bias in (True, False): + name = f"dynamic_{batch_dims}batch_{k}_{n}_relu" + ( + "" if bias else "_no_bias" + ) + relu = torch.nn.Sequential( + torch.nn.Linear(k, n, bias=bias), torch.nn.ReLU() + ) + cases.append(chain(name, relu, (*_batch_dims(batch_dims, batch), k))) + rows = Dim("rows", min=1, max=16) + return [ + *cases, + # A static batch of one, as an LLM's prefill and decode. + chain("dynamic_tokens_3d", torch.nn.Linear(256, 64), (1, tokens, 256)), + chain("dynamic_two_batch_256", torch.nn.Linear(256, 64), (rows, rows, 256)), + *_structural("_256", 256, 256, 64, tokens, Dim("tokens2", min=1, max=64)), + ] + + def cases() -> list[Case]: suite = suite_cases(test_linear, samples=4) torch.manual_seed(0) @@ -127,4 +214,5 @@ def cases() -> list[Case]: _facto_case, lambda i, args: i % 5 == 0 or args[1].size(1) % QK_K == 0, ), + *_dynamic_cases(), ] diff --git a/backends/native/test/samples/record.py b/backends/native/test/samples/record.py index 90a57daa399..9465509795e 100644 --- a/backends/native/test/samples/record.py +++ b/backends/native/test/samples/record.py @@ -11,9 +11,10 @@ `precise`, the same math in float64 with quantized weights decoded in float32, whatever the method's dtype, and activations unquantized, floating outputs stored as float32. Weight-only schemes add `golden-int8_act`, the eager outputs -with the same weights and int8 per-token dynamic activations. Nothing in it -depends on an engine or a device: consumers judge an engine by its distance -from `precise` relative to eager's. +with the same weights and int8 per-token dynamic activations. A dynamic case +is exported with its dynamic shapes and recorded as one sample whose calls +change them. 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 @@ -45,6 +46,13 @@ 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 _samples(calls: list[Call], chained: bool) -> list[Sample]: + """One sample per call, or one running all the calls when they are chained.""" + if chained: + return [Sample(calls)] if calls else [] + return [Sample([call]) for call in calls] + + def record( cases: list[Case], scheme: Scheme, @@ -69,20 +77,27 @@ def record( 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]) + chained = case.dynamic_shapes is not None + programs[name] = torch.export.export( + method, + inputs[0], + dynamic_shapes=({"args": case.dynamic_shapes} if chained else None), + ) + golden_calls, precise_calls, act_calls = [], [], [] with torch.no_grad(): for x in inputs: - golden.append(Sample([Call(name, list(x), outputs(method(*x)))])) + golden_calls.append(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)])) + precise_calls.append(Call(name, list(x), ref)) if scheme.weight_only: with int8_activations(method): - golden_act.append( - Sample([Call(name, list(x), outputs(method(*x)))]) - ) + act_calls.append(Call(name, list(x), outputs(method(*x)))) + golden += _samples(golden_calls, chained) + precise += _samples(precise_calls, chained) + golden_act += _samples(act_calls, chained) if not programs: logger.info(f"{path.name}: no case fits the scheme, not written") path.unlink(missing_ok=True) diff --git a/backends/native/test/samples/sample_test_main.cpp b/backends/native/test/samples/sample_test_main.cpp index e44d987343c..18e76fde8ec 100644 --- a/backends/native/test/samples/sample_test_main.cpp +++ b/backends/native/test/samples/sample_test_main.cpp @@ -6,8 +6,9 @@ // A gtest main that runs the bundled samples of .ptn packages on the engine the // binary links (see PtnRunner.h): one test per package and method, grouping -// samples by their first call. Each sample runs from a fresh context, and each -// method output is judged over all its samples by OutputCheck. +// samples by their first call. Each sample runs from a fresh context, its calls +// in order on one executable per method, resizing dynamic inputs between them. +// Each method output is judged over all its calls by OutputCheck. // // Usage: // --packages=a.ptn,dir [--rmse_ratio --max_ratio --rtol --sigmas] diff --git a/backends/native/test/samples/test_record.py b/backends/native/test/samples/test_record.py index fd5d28d3b2f..710e89b436c 100644 --- a/backends/native/test/samples/test_record.py +++ b/backends/native/test/samples/test_record.py @@ -12,7 +12,12 @@ 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.cases import ( + Case, + chain, + draws, + suite_cases, +) from executorch.backends.native.test.samples.record import record from executorch.backends.native.test.samples.schemes import ( Gguf, @@ -22,6 +27,7 @@ Torchao, ) from executorch.exir.native import read_samples +from torch.export import Dim from torchao.quantization import IntxWeightOnlyConfig, quantize_ from torchao.quantization.granularity import PerGroup @@ -105,6 +111,42 @@ def test_record_adds_the_8da4w_golden_to_weight_only_packages(self) -> None: got.calls[0].outputs, want.calls[0].outputs, rtol=0, atol=0 ) + def test_chain_visits_each_dims_bounds_starting_mid_range(self) -> None: + a, b = Dim("a", min=1, max=8), Dim("b", min=2, max=4) + case = chain("pair", torch.nn.Identity(), (a, 3), (b, a)) + self.assertEqual( + [[x.shape for x in xs] for xs in case.inputs], + [ + [(5, 3), (3, 5)], + [(1, 3), (4, 1)], + [(8, 3), (2, 8)], + [(5, 3), (3, 5)], + ], + ) + self.assertEqual(case.dynamic_shapes, ({0: a}, {0: b, 1: a})) + + def test_record_chains_a_dynamic_case_into_one_sample(self) -> None: + case = chain("linear", torch.nn.Linear(32, 4), (Dim("batch", min=1, max=6), 32)) + scheme = schemes()["int4wo_g32"] + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "linear.ptn" + record([case], scheme, path, dtypes=["f32"]) + sets = read_samples(str(path)) + self.assertEqual(set(sets), {"golden", "precise", "golden-int8_act"}) + for samples in sets.values(): + (sample,) = samples.samples + self.assertEqual( + [call.inputs[0].shape for call in sample.calls], + [(4, 32), (1, 32), (6, 32), (4, 32)], + ) + model = copy.deepcopy(case.model).eval() + scheme.quantize(model) + with torch.no_grad(): + for call, x in zip( + sets["golden"].samples[0].calls, case.inputs, strict=True + ): + torch.testing.assert_close(call.outputs[0], model(*x), 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: