diff --git a/docs/source/quantization-overview.md b/docs/source/quantization-overview.md index c31c3ded837..921c0f5a022 100644 --- a/docs/source/quantization-overview.md +++ b/docs/source/quantization-overview.md @@ -73,3 +73,5 @@ Note that numerics on device can differ those in PyTorch even for unquantized mo ## 3. Lower the model The final step is to lower the quantized_model to the desired backend, as you would an unquantized one. See [backend-specific pages](backends-overview.md) for lowering information. + +For complex Quantization-Aware Training flows — where training loops, checkpointing, and job restarts need to live outside a single export call — see [Using Complex QAT Flows with Recipe-Based Lowering](using-complex-qat-flows-with-recipe-based-lowering.md). diff --git a/docs/source/usage.md b/docs/source/usage.md index 6ffc136093b..5ce6c9156dd 100644 --- a/docs/source/usage.md +++ b/docs/source/usage.md @@ -9,6 +9,7 @@ troubleshooting, and frequently asked questions. getting-started using-executorch-export +using-complex-qat-flows-with-recipe-based-lowering using-executorch-android using-executorch-ios using-executorch-cpp diff --git a/docs/source/using-complex-qat-flows-with-recipe-based-lowering.md b/docs/source/using-complex-qat-flows-with-recipe-based-lowering.md new file mode 100644 index 00000000000..17c7af3620a --- /dev/null +++ b/docs/source/using-complex-qat-flows-with-recipe-based-lowering.md @@ -0,0 +1,144 @@ +# Using Complex QAT Flows with Recipe-Based Lowering + +How to split the declarative export recipe around the quantization boundary so +that an arbitrary Quantization-Aware Training flow can live outside of — and +between the two halves of — the recipe pipeline. + +## Motivation + +`ExportRecipe` and `ExportSession` run the full export pipeline — source +transforms, quantize, torch.export, lower, to-executorch — as a single +declarative call. That works well for simple calibration flows. + +Complex QAT is different. A real training run involves iterative training +loops, periodic checkpointing, job restarts, accuracy evaluation on validation +datasets, rollback when a checkpoint regresses, and potentially separate +machines for training and for compilation. None of that can fit inside a single +`export()` call. + +## How It Works + +The recipe's `pipeline_stages` attribute lets you restrict an `ExportSession` +to a specific subset of stages. Use it to split the pipeline around the +`QUANTIZE` stage: run everything up to (but not including) quantization in one +recipe invocation, hand the result off to your QAT flow as a `.pt2` file, and +then continue with the post-quantize stages in a second recipe invocation. + +``` +ExportRecipe (pre-quantize slice) + | + v + .pt2 <-- cross-process / cross-machine handoff + | + v +Your QAT training loop (any framework, any duration) + | + v + .pt2 <-- cross-process / cross-machine handoff + | + v +ExportRecipe (post-quantize slice) + | + v + model.pte +``` + +The recipe still owns all the lowering and delegation work; you are only +carving out the `QUANTIZE` stage to replace it with your own flow. + +## The Method + +### Pre-quantize slice + +Run only the pre-quantize stages through `ExportSession`, then export the +resulting module to an ATEN-dialect `ExportedProgram` and save it as a `.pt2` +handoff file: + +```python +from executorch.export import export as et_export +from executorch.export.types import StageType + +recipe.pipeline_stages = [StageType.SOURCE_TRANSFORM] +sess = et_export(model, example_inputs=[inputs], export_recipe=recipe) + +transformed = sess.get_stage_artifacts()[StageType.SOURCE_TRANSFORM].data["forward"] + +ep = torch.export.export(transformed, inputs, strict=True) +torch.export.save(ep, "pre_qat.pt2") +``` + +> **Simplification:** if you do not need any pre-quantize recipe passes (such +> as source transforms provided by the recipe), you can skip this half entirely. +> Capture the eager model directly with `torch.export.export` and save the +> result as the handoff `.pt2`. No recipe is involved in that case. + +### Your QAT flow + +Load the handoff, prepare it for QAT, and train: + +```python +# Load the pre-quantize graph. +captured_gm = torch.export.load("pre_qat.pt2").module() + +# Prepare for QAT once. +from torchao.quantization.pt2e.quantize_pt2e import prepare_qat_pt2e +from torchao.quantization.pt2e import move_exported_model_to_train + +prepared = prepare_qat_pt2e(captured_gm, quantizer) +move_exported_model_to_train(prepared) + +# Train -- this is an open-ended process. +# It can be paused and resumed from checkpoints. +# It can run for hours, days, or weeks. +# It can be distributed across separate machines or jobs. +# Evaluate accuracy on a validation dataset periodically and +# roll back to an earlier checkpoint whenever accuracy regresses. +# Save checkpoints as you go: +# torch.save(prepared.state_dict(), "qat_ckpt.pt") +# Resume from a checkpoint: +# prepared.load_state_dict(torch.load("qat_ckpt.pt")) + +# Only once you are satisfied with the result: +from torchao.quantization.pt2e import move_exported_model_to_eval +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e + +move_exported_model_to_eval(prepared) +quantized = convert_pt2e(prepared) + +ep = torch.export.export(quantized, inputs, strict=True) +torch.export.save(ep, "post_qat.pt2") +``` + +### Post-quantize slice + +Continue the recipe pipeline from `TORCH_EXPORT` onward: + +```python +recipe.pipeline_stages = [ + StageType.TORCH_EXPORT, + StageType.TO_EDGE_TRANSFORM_AND_LOWER, + StageType.TO_EXECUTORCH, +] +quantized_gm = torch.export.load("post_qat.pt2").module() +sess = et_export(quantized_gm, example_inputs=[inputs], export_recipe=recipe) +sess.save_to_pte("model") # writes model.pte +``` + +> **Simplification:** if the recipe uses its default pipeline stages, you do +> not need to set `pipeline_stages` at all on the lowering side. Pass the +> `.pt2` file path directly to `export()` and `ExportSession` will +> automatically skip `SOURCE_TRANSFORM`, `QUANTIZE`, and `TORCH_EXPORT` when it +> detects an `ExportedProgram` input, picking up at +> `TO_EDGE_TRANSFORM_AND_LOWER`. This auto-skip only applies to the default +> pipeline; if the recipe defines a custom `pipeline_stages` list, set the +> post-quantize stages explicitly as shown above. + +## Example + +A runnable end-to-end example demonstrating both the full split and the +simplified path is available at +[`examples/export/qat_pipeline_split/`](https://github.com/pytorch/executorch/tree/main/examples/export/qat_pipeline_split). + +See the [example README](https://github.com/pytorch/executorch/blob/main/examples/export/qat_pipeline_split/README.md) +for full setup instructions, expected artifacts, and notes on adapting the +example to a different backend. diff --git a/examples/export/qat_pipeline_split/1_prepare.py b/examples/export/qat_pipeline_split/1_prepare.py new file mode 100644 index 00000000000..c3bbd55816e --- /dev/null +++ b/examples/export/qat_pipeline_split/1_prepare.py @@ -0,0 +1,148 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Stage 1: capture the model up to the quantize boundary. +# +# Two example paths (selected with --example): +# +# minimal (default) +# Export the eager nn.Module to an ATEN-dialect ExportedProgram with a +# single torch.export.export call. No recipe, no ExportSession. +# +# sliced +# Drive only the SOURCE_TRANSFORM stage through ExportSession, then +# export to ATEN and save. This shows how to restrict the pipeline to +# the stages that precede quantization. +# +# Both paths write a .pt2 file consumed by 2_qat.py. + +import argparse +import os + +import torch + +from model import get_example_inputs, get_model + + +def _build_recipe(): + from executorch.backends.xnnpack.recipes.xnnpack_recipe_types import ( + XNNPackRecipeType, + ) + from executorch.export.recipe import ExportRecipe + + return ExportRecipe.get_recipe(XNNPackRecipeType.PT2E_INT8_STATIC_PER_TENSOR) + + +def run_minimal(workdir: str) -> None: + print("[minimal] Capturing ATEN graph with torch.export.export (no recipe).") + + model = get_model() + ex = get_example_inputs() + + assert isinstance( + model, torch.nn.Module + ), f"Expected nn.Module from get_model(), got {type(model)}" + + ep = torch.export.export(model, ex, strict=True) + + assert ep is not None, "torch.export.export returned None" + print(f"[minimal] Graph nodes: {len(list(ep.graph.nodes))}") + + out = os.path.join(workdir, "stage1_minimal.pt2") + torch.export.save(ep, out) + print(f"[minimal] Saved ExportedProgram -> {out}") + + reloaded_gm = torch.export.load(out).module() + y = reloaded_gm(*ex) + assert y.shape == ( + 1, + 10, + ), f"[minimal] Unexpected output shape after reload: {y.shape}" + print(f"[minimal] Round-trip output shape: {y.shape} (assertion passed)") + + +def run_sliced(workdir: str) -> None: + print("[sliced] Running SOURCE_TRANSFORM stage only via ExportSession.") + + from executorch.export import export as et_export + from executorch.export.types import StageType + + recipe = _build_recipe() + + # Restrict the pipeline to the pre-quantize stages only. + recipe.pipeline_stages = [StageType.SOURCE_TRANSFORM] + + model = get_model() + ex = get_example_inputs() + + assert isinstance( + model, torch.nn.Module + ), f"Expected nn.Module from get_model(), got {type(model)}" + + sess = et_export(model, example_inputs=[ex], export_recipe=recipe) + + artifacts = sess.get_stage_artifacts() + assert ( + StageType.SOURCE_TRANSFORM in artifacts + ), "SOURCE_TRANSFORM artifact not found - did the stage run?" + + transformed = artifacts[StageType.SOURCE_TRANSFORM].data + assert isinstance( + transformed, dict + ), f"Expected method-keyed dict from SOURCE_TRANSFORM, got {type(transformed)}" + assert ( + "forward" in transformed + ), "'forward' method missing from SOURCE_TRANSFORM artifact" + + transformed_module = transformed["forward"] + assert isinstance( + transformed_module, torch.nn.Module + ), f"Expected nn.Module after SOURCE_TRANSFORM, got {type(transformed_module)}" + print( + f"[sliced] SOURCE_TRANSFORM produced: {type(transformed_module).__name__}" + " (pass-through for this PT2E recipe, as expected)" + ) + + ep = torch.export.export(transformed_module, ex, strict=True) + assert ep is not None, "torch.export.export returned None" + print(f"[sliced] Graph nodes: {len(list(ep.graph.nodes))}") + + out = os.path.join(workdir, "stage1_sliced.pt2") + torch.export.save(ep, out) + print(f"[sliced] Saved ExportedProgram -> {out}") + + reloaded_gm = torch.export.load(out).module() + y = reloaded_gm(*ex) + assert y.shape == ( + 1, + 10, + ), f"[sliced] Unexpected output shape after reload: {y.shape}" + print(f"[sliced] Round-trip output shape: {y.shape} (assertion passed)") + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Stage 1: capture the model up to the quantize boundary." + ) + parser.add_argument( + "--example", + choices=["minimal", "sliced"], + default="minimal", + ) + parser.add_argument("--workdir", required=True) + args = parser.parse_args() + + os.makedirs(args.workdir, exist_ok=True) + + if args.example == "minimal": + run_minimal(args.workdir) + else: + run_sliced(args.workdir) + + print(f"\nStage 1 done. Artifact written to: {args.workdir}") + + +if __name__ == "__main__": + main() diff --git a/examples/export/qat_pipeline_split/2_qat.py b/examples/export/qat_pipeline_split/2_qat.py new file mode 100644 index 00000000000..4b2f4009c79 --- /dev/null +++ b/examples/export/qat_pipeline_split/2_qat.py @@ -0,0 +1,177 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Stage 2: perform arbitrary Quantization-Aware Training (QAT). +# +# Loads the captured ATEN graph from stage 1, prepares it for QAT, runs dummy +# forward passes to exercise the observers, saves and restores a checkpoint, +# then converts to a quantized graph and saves it for stage 3. +# +# No real QAT is performed: the dummy forward passes only exist to show that +# observers collect statistics in training mode. + +import argparse +import copy +import os + +import torch + +from model import get_calibration_inputs, get_example_inputs + + +def _get_quantizer(): + from executorch.backends.xnnpack.recipes.xnnpack_recipe_types import ( + XNNPackRecipeType, + ) + from executorch.export.recipe import ExportRecipe + + recipe = ExportRecipe.get_recipe(XNNPackRecipeType.PT2E_INT8_STATIC_PER_TENSOR) + quantizers = recipe.quantization_recipe.quantizers + assert ( + quantizers and len(quantizers) > 0 + ), "Recipe carries no quantizers - cannot prepare for QAT" + return quantizers[0] + + +def _run_qat(captured_gm: torch.fx.GraphModule, workdir: str) -> torch.fx.GraphModule: + from torchao.quantization.pt2e import ( + move_exported_model_to_eval, + move_exported_model_to_train, + ) + from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_qat_pt2e + + calib = get_calibration_inputs(n=4) + quantizer = _get_quantizer() + + # prepare_qat_pt2e mutates the graph in-place; keep a pristine copy for the + # checkpoint restore demo below. + pristine_gm = copy.deepcopy(captured_gm) + prepared = prepare_qat_pt2e(captured_gm, quantizer) + + assert prepared is not None, "prepare_qat_pt2e returned None" + + # Verify that observer / fake-quant modules were inserted. + # Matched by call_module node target prefix because class identity differs + # between torchao and torch.ao namespaces. + obs_nodes = [ + n + for n in prepared.graph.nodes + if n.op == "call_module" + and isinstance(n.target, str) + and n.target.startswith("activation_post_process") + ] + assert len(obs_nodes) > 0, ( + "No activation_post_process nodes found after prepare_qat_pt2e - " + "the quantizer may not have annotated this graph" + ) + obs_type_names = ", ".join( + type(dict(prepared.named_modules())[n.target]).__name__ for n in obs_nodes[:3] + ) + print( + f" Inserted {len(obs_nodes)} observer/fake-quant nodes " + f"({obs_type_names}{'...' if len(obs_nodes) > 3 else ''})." + ) + + prepared = move_exported_model_to_train(prepared) + + print(" Running dummy training steps (observers collecting statistics)...") + for i, inputs in enumerate(calib): + _ = prepared(*inputs) + print(f" step {i + 1}/{len(calib)} input shape: {inputs[0].shape}") + + # Save observer state so training can be paused and resumed later. + ckpt_path = os.path.join(workdir, "qat_ckpt.pt") + torch.save(prepared.state_dict(), ckpt_path) + print(f" Checkpoint saved -> {ckpt_path}") + + # Restore into a fresh prepared module to demonstrate the pause/resume seam. + print(" Demonstrating checkpoint restore into a fresh prepared module...") + fresh_prepared = prepare_qat_pt2e(pristine_gm, quantizer) + fresh_prepared = move_exported_model_to_train(fresh_prepared) + + saved_state = torch.load(ckpt_path, weights_only=True) + fresh_prepared.load_state_dict(saved_state) + + restored_state = fresh_prepared.state_dict() + assert set(saved_state.keys()) == set( + restored_state.keys() + ), "State dict keys differ after checkpoint restore" + for key in saved_state: + assert torch.allclose( + saved_state[key], restored_state[key] + ), f"Tensor mismatch for key '{key}' after checkpoint restore" + print(" Checkpoint restore verified: all tensors match (assertion passed).") + + prepared = fresh_prepared + + prepared = move_exported_model_to_eval(prepared) + converted = convert_pt2e(prepared) + + assert converted is not None, "convert_pt2e returned None" + + # After convert_pt2e the quantize/dequantize ops are in the + # quantized_decomposed namespace; match by __name__ substring. + qdq_nodes = [ + n + for n in converted.graph.nodes + if n.op == "call_function" + and hasattr(n.target, "__name__") + and ( + "quantize_per_tensor" in n.target.__name__ + or "dequantize_per_tensor" in n.target.__name__ + or "quantize_per_channel" in n.target.__name__ + or "dequantize_per_channel" in n.target.__name__ + ) + ] + assert len(qdq_nodes) > 0, ( + "No quantize/dequantize nodes found after convert_pt2e - " + "conversion may have failed silently" + ) + print(f" convert_pt2e inserted {len(qdq_nodes)} quantize/dequantize ops.") + + return converted + + +def run(example: str, workdir: str) -> None: + pt2_in = os.path.join(workdir, f"stage1_{example}.pt2") + assert os.path.isfile(pt2_in), ( + f"Input file not found: {pt2_in} " + f"(run 1_prepare.py --example {example} first)" + ) + + print(f"[{example}] Loading captured graph from {pt2_in}") + captured_gm = torch.export.load(pt2_in).module() + + assert isinstance( + captured_gm, torch.fx.GraphModule + ), f"Expected GraphModule after load, got {type(captured_gm)}" + print( + f"[{example}] Loaded GraphModule with {len(list(captured_gm.graph.nodes))} nodes." + ) + + converted = _run_qat(captured_gm, workdir) + + ex = get_example_inputs() + ep = torch.export.export(converted, ex, strict=True) + out = os.path.join(workdir, f"stage2_{example}_quantized.pt2") + torch.export.save(ep, out) + print(f"[{example}] Saved quantized ExportedProgram -> {out}") + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Stage 2: perform arbitrary QAT on the captured graph from stage 1." + ) + parser.add_argument("--example", choices=["minimal", "sliced"], default="minimal") + parser.add_argument("--workdir", required=True) + args = parser.parse_args() + + run(args.example, args.workdir) + + print(f"\nStage 2 done. Artifacts in: {args.workdir}") + + +if __name__ == "__main__": + main() diff --git a/examples/export/qat_pipeline_split/3_lower.py b/examples/export/qat_pipeline_split/3_lower.py new file mode 100644 index 00000000000..828b06e25f7 --- /dev/null +++ b/examples/export/qat_pipeline_split/3_lower.py @@ -0,0 +1,184 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Stage 3: lower the quantized graph to an ExecuTorch .pte program. +# +# Two example paths (selected with --example): +# +# minimal (default) +# Pass the .pt2 file path directly to export(). ExportSession detects +# an ExportedProgram input and auto-skips SOURCE_TRANSFORM, QUANTIZE and +# TORCH_EXPORT; the effective pipeline starts at TO_EDGE_TRANSFORM_AND_LOWER. +# +# sliced +# Load the GraphModule and set pipeline_stages explicitly to +# [TORCH_EXPORT, TO_EDGE_TRANSFORM_AND_LOWER, TO_EXECUTORCH] - the +# mirror image of the [SOURCE_TRANSFORM] slice used in stage 1. +# +# Both paths assert delegation happened and write model.pte. +# The model is run in the separate stage 4 (4_run.py). + +import argparse +import os + +import torch + +from model import get_example_inputs + + +def _build_recipe(): + from executorch.backends.xnnpack.recipes.xnnpack_recipe_types import ( + XNNPackRecipeType, + ) + from executorch.export.recipe import ExportRecipe + + return ExportRecipe.get_recipe(XNNPackRecipeType.PT2E_INT8_STATIC_PER_TENSOR) + + +def _finish(sess, workdir: str) -> None: + """Assert delegation occurred and save model.pte.""" + from executorch.export.types import StageType + + pte_buffer = sess.get_pte_buffer() + assert ( + pte_buffer is not None and len(pte_buffer) > 0 + ), "get_pte_buffer() returned an empty buffer - lowering may have failed" + print(f" PTE buffer size: {len(pte_buffer)} bytes.") + + artifacts = sess.get_stage_artifacts() + lowering_stage = next( + ( + s + for s in ( + StageType.TO_EDGE_TRANSFORM_AND_LOWER, + StageType.TO_BACKEND, + ) + if s in artifacts + ), + None, + ) + assert ( + lowering_stage is not None + ), "No lowering stage artifact found - did TO_EDGE_TRANSFORM_AND_LOWER run?" + + delegation_info = artifacts[lowering_stage].get_context("delegation_info") + assert ( + delegation_info is not None + ), "delegation_info context is None - the lowering stage did not populate it" + + num_delegated = delegation_info.num_delegated_subgraphs + assert num_delegated > 0, ( + f"Expected at least one delegated subgraph, got {num_delegated}. " + "XNNPACK partitioner did not claim any nodes." + ) + print(f" Delegated subgraphs: {num_delegated} (assertion passed).") + + print("\n Delegation summary:") + sess.print_delegation_info() + + pte_path = os.path.join(workdir, "model.pte") + sess.save_to_pte(os.path.join(workdir, "model")) + assert os.path.isfile( + pte_path + ), f"Expected .pte at {pte_path} but file was not created" + print(f"\n Saved -> {pte_path} ({os.path.getsize(pte_path)} bytes)") + + +def run_minimal(workdir: str) -> None: + from executorch.export import export as et_export + from executorch.export.types import StageType + + pt2_in = os.path.join(workdir, "stage2_minimal_quantized.pt2") + assert os.path.isfile(pt2_in), ( + f"Input file not found: {pt2_in} " "(run 2_qat.py --example minimal first)" + ) + print(f"[minimal] Lowering quantized model from {pt2_in}") + + recipe = _build_recipe() + + # Pass the file path directly. ExportSession auto-skips SOURCE_TRANSFORM, + # QUANTIZE and TORCH_EXPORT for ExportedProgram input. + sess = et_export(pt2_in, export_recipe=recipe) + + artifacts = sess.get_stage_artifacts() + for skipped in ( + StageType.SOURCE_TRANSFORM, + StageType.QUANTIZE, + StageType.TORCH_EXPORT, + ): + assert skipped not in artifacts, ( + f"Stage {skipped} should have been skipped for ExportedProgram input " + "but an artifact was found" + ) + print( + "[minimal] Auto-skip verified: SOURCE_TRANSFORM / QUANTIZE / " + "TORCH_EXPORT absent from artifacts (assertion passed)." + ) + + _finish(sess, workdir) + + +def run_sliced(workdir: str) -> None: + from executorch.export import export as et_export + from executorch.export.types import StageType + + pt2_in = os.path.join(workdir, "stage2_sliced_quantized.pt2") + assert os.path.isfile(pt2_in), ( + f"Input file not found: {pt2_in} " "(run 2_qat.py --example sliced first)" + ) + print(f"[sliced] Lowering quantized model from {pt2_in}") + + quantized_gm = torch.export.load(pt2_in).module() + + assert isinstance( + quantized_gm, torch.fx.GraphModule + ), f"Expected GraphModule after load, got {type(quantized_gm)}" + print( + f"[sliced] Loaded GraphModule ({len(list(quantized_gm.graph.nodes))} nodes)." + ) + + recipe = _build_recipe() + + # Mirror image of the [SOURCE_TRANSFORM] slice in stage 1: together they + # cover the full pipeline with no overlap and no stage silently skipped. + recipe.pipeline_stages = [ + StageType.TORCH_EXPORT, + StageType.TO_EDGE_TRANSFORM_AND_LOWER, + StageType.TO_EXECUTORCH, + ] + + ex = get_example_inputs() + sess = et_export(quantized_gm, example_inputs=[ex], export_recipe=recipe) + + artifacts = sess.get_stage_artifacts() + for expected_stage in recipe.pipeline_stages: + assert ( + expected_stage in artifacts + ), f"Expected artifact for stage {expected_stage} but none was produced" + print( + "[sliced] All explicit pipeline stages produced artifacts (assertion passed)." + ) + + _finish(sess, workdir) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Stage 3: lower the quantized graph from stage 2 to an ExecuTorch .pte." + ) + parser.add_argument("--example", choices=["minimal", "sliced"], default="minimal") + parser.add_argument("--workdir", required=True) + args = parser.parse_args() + + if args.example == "minimal": + run_minimal(args.workdir) + else: + run_sliced(args.workdir) + + print(f"\nStage 3 done. model.pte written to: {args.workdir}") + + +if __name__ == "__main__": + main() diff --git a/examples/export/qat_pipeline_split/4_run.py b/examples/export/qat_pipeline_split/4_run.py new file mode 100644 index 00000000000..c639126cd56 --- /dev/null +++ b/examples/export/qat_pipeline_split/4_run.py @@ -0,0 +1,59 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Stage 4: run model.pte through the ExecuTorch runtime. +# +# Loads the .pte produced by stage 3 and executes it, asserting the output +# shape matches expectations. If the ExecuTorch pybindings are not available +# in this environment a warning is printed and the script exits cleanly; the +# .pte produced by stage 3 is still valid. + +import argparse +import os + +from model import get_example_inputs + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Stage 4: run model.pte through the ExecuTorch runtime." + ) + parser.add_argument("--workdir", required=True) + args = parser.parse_args() + + pte_path = os.path.join(args.workdir, "model.pte") + assert os.path.isfile( + pte_path + ), f"model.pte not found at {pte_path} (run 3_lower.py first)" + + try: + from executorch.runtime import Runtime + except ModuleNotFoundError: + print( + "WARNING: executorch.runtime is not available in this environment. " + "Build and install the ExecuTorch pybindings to run the .pte file. " + "Skipping runtime execution." + ) + return + + print(f"Loading {pte_path} ...") + runtime = Runtime.get() + program = runtime.load_program(pte_path) + method = program.load_method("forward") + + ex = get_example_inputs() + outputs = method.execute(ex) + + assert len(outputs) == 1, f"Expected 1 output tensor, got {len(outputs)}" + out_tensor = outputs[0] + assert out_tensor.shape == (1, 10), f"Unexpected output shape: {out_tensor.shape}" + print( + f"Runtime execution succeeded. Output shape: {out_tensor.shape} (assertion passed)" + ) + print("\nStage 4 done.") + + +if __name__ == "__main__": + main() diff --git a/examples/export/qat_pipeline_split/README.md b/examples/export/qat_pipeline_split/README.md new file mode 100644 index 00000000000..a03e6b069b3 --- /dev/null +++ b/examples/export/qat_pipeline_split/README.md @@ -0,0 +1,161 @@ +# QAT Pipeline Split Example + +Shows how to combine ExecuTorch's recipe-based (declarative) export with a +complex, multi-step Quantization-Aware Training flow that lives **outside** the +recipe. + +## The problem + +`ExportRecipe` / `ExportSession` run the full export pipeline — source +transforms, quantize, torch.export, lower, to-executorch — as a single +declarative call. That works well when PTQ calibration is sufficient. Real +QAT workflows are different: they involve training loops, checkpointing, job +restarts, and potentially separate machines for training and compilation. +Those workflows cannot fit inside a single `export()` call. + +## The solution: split the recipe around QUANTIZE + +The recipe's `pipeline_stages` attribute lets you restrict a session to a +specific subset of stages. This example uses that mechanism to insert an +arbitrary QAT step between the pre-quantize and post-quantize halves of the +recipe: + +``` +ExportRecipe (pre-quantize slice) + | + v + .pt2 <-- cross-process / cross-machine handoff + | + v +Your QAT training loop (any framework, any duration) + | + v + .pt2 <-- cross-process / cross-machine handoff + | + v +ExportRecipe (post-quantize slice) + | + v + model.pte +``` + +The recipe still does all the lowering and delegation work; you are only +carving out the QUANTIZE stage to replace it with your own flow. + +## Two example modes + +### `minimal` — let the recipe adapt to work you already did + +Stage 1 captures the eager model with a plain `torch.export.export` call, with +no recipe involved at all. After your QAT flow (stage 2), you hand the +quantized `.pt2` directly to `export()` with the full recipe. `ExportSession` +detects that the input is already an `ExportedProgram` and automatically skips +`SOURCE_TRANSFORM`, `QUANTIZE` and `TORCH_EXPORT` — the recipe picks up at +`TO_EDGE_TRANSFORM_AND_LOWER` and completes lowering as normal. + +Use this when you have no need for pre-quantize recipe passes and want the +simplest possible seam. + +### `sliced` — explicitly slice the recipe's pipeline_stages in two + +Stage 1 creates an `ExportRecipe`, restricts `pipeline_stages` to +`[SOURCE_TRANSFORM]`, and runs only that pre-quantize half of the pipeline. +After your QAT flow (stage 2), stage 3 picks up the same recipe, sets +`pipeline_stages` to `[TORCH_EXPORT, TO_EDGE_TRANSFORM_AND_LOWER, +TO_EXECUTORCH]`, and completes the pipeline. Together the two slices cover +every stage of the recipe with no overlap and no stage silently skipped. + +Use this when you need pre-quantize recipe passes to run before QAT (for +example, torchao source transforms that the recipe provides). + +### Which should I use? + +Use **`minimal`** when the recipe uses its default pipeline stages and you have +no need for pre-quantize recipe passes. The `minimal` mode works because +`ExportSession` auto-skips SOURCE_TRANSFORM, QUANTIZE and TORCH_EXPORT when it +receives an `ExportedProgram` input — this auto-skip only applies to the +default pipeline. If the recipe defines a custom `pipeline_stages` list, use +**`sliced`** instead. The +`sliced` mode is also required whenever pre-quantize recipe passes (such as +torchao source transforms) need to run before handing the graph to your +training loop. + +## Stage overview + +``` +Eager nn.Module + | + | Stage 1 -- 1_prepare.py + | minimal: torch.export.export (no recipe) + | sliced: ExportRecipe with pipeline_stages=[SOURCE_TRANSFORM] + v + stage1_*.pt2 + | + | Stage 2 -- 2_qat.py (runs in a separate process) + | prepare_qat_pt2e -> dummy train -> checkpoint -> restore -> convert_pt2e + | (stand-in for any real QAT training loop) + v + stage2_*_quantized.pt2 + | + | Stage 3 -- 3_lower.py + | minimal: ExportRecipe, full recipe (auto-skips pre-quantize stages) + | sliced: ExportRecipe with pipeline_stages=[TORCH_EXPORT, + | TO_EDGE_TRANSFORM_AND_LOWER, TO_EXECUTORCH] + v + model.pte + | + | Stage 4 -- 4_run.py + | executorch.runtime: load .pte, run forward, assert output shape + v + [1, 10] logits +``` + +## Files + +| File | Role | +|------|------| +| `model.py` | `SmallConvNet` definition (stage 1 only) and example-input helpers (all stages) | +| `1_prepare.py` | Stage 1: capture to ATEN, optionally via a recipe slice | +| `2_qat.py` | Stage 2: QAT + checkpoint save/restore demo | +| `3_lower.py` | Stage 3: lower to `.pte` with the export recipe | +| `4_run.py` | Stage 4: run `.pte` through the ExecuTorch runtime | +| `run.sh` | Bash orchestrator: drives all four stages in sequence | + +## Running + +```bash +# From the executorch root: +PYTHON=/path/to/python bash examples/export/qat_pipeline_split/run.sh \ + --example minimal \ + --workdir /tmp/qat_pipeline_split_minimal + +PYTHON=/path/to/python bash examples/export/qat_pipeline_split/run.sh \ + --example sliced \ + --workdir /tmp/qat_pipeline_split_sliced +``` + +Each stage can also be run independently: + +```bash +python 1_prepare.py --example minimal --workdir /tmp/out +python 2_qat.py --example minimal --workdir /tmp/out +python 3_lower.py --example minimal --workdir /tmp/out +python 4_run.py --workdir /tmp/out +``` + +## Expected artifacts in `--workdir` + +| File | Written by | +|------|-----------| +| `stage1_minimal.pt2` / `stage1_sliced.pt2` | stage 1 | +| `stage2_minimal_quantized.pt2` / `stage2_sliced_quantized.pt2` | stage 2 | +| `qat_ckpt.pt` | stage 2 (checkpoint demo) | +| `model.pte` | stage 3 | + +## Notes + +- The concrete backend is XNNPACK (`PT2E_INT8_STATIC_PER_TENSOR`). To adapt + to a different backend replace `_build_recipe()` in `1_prepare.py` and + `3_lower.py`, and `_get_quantizer()` in `2_qat.py`. +- Stage 4 requires the ExecuTorch pybindings (`executorch.runtime`). If they + are not installed a warning is printed and the stage exits cleanly. diff --git a/examples/export/qat_pipeline_split/model.py b/examples/export/qat_pipeline_split/model.py new file mode 100644 index 00000000000..56f429142b2 --- /dev/null +++ b/examples/export/qat_pipeline_split/model.py @@ -0,0 +1,47 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Model definition for the QAT pipeline split example. +# +# SmallConvNet is used only by stage 1 (1_prepare.py), which constructs the +# eager model before capture. Stages 2 and 3 reload the model from a saved +# .pt2 file and only import the example-input helpers below. + +from typing import List, Tuple + +import torch +import torch.nn as nn + + +class SmallConvNet(nn.Module): + """A small conv net for demonstrating PT2E quantization. + + Input: (N, 1, 28, 28) + Output: (N, 10) + """ + + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv2d(1, 8, kernel_size=3, padding=1) + self.pool = nn.AdaptiveAvgPool2d((4, 4)) + self.fc = nn.Linear(8 * 4 * 4, 10) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = torch.relu(self.conv(x)) + x = self.pool(x) + x = torch.flatten(x, 1) + return self.fc(x) + + +def get_model() -> SmallConvNet: + return SmallConvNet().eval() + + +def get_example_inputs() -> Tuple[torch.Tensor, ...]: + return (torch.randn(1, 1, 28, 28),) + + +def get_calibration_inputs(n: int = 4) -> List[Tuple[torch.Tensor, ...]]: + return [(torch.randn(1, 1, 28, 28),) for _ in range(n)] diff --git a/examples/export/qat_pipeline_split/run.sh b/examples/export/qat_pipeline_split/run.sh new file mode 100755 index 00000000000..b8e7a2031db --- /dev/null +++ b/examples/export/qat_pipeline_split/run.sh @@ -0,0 +1,123 @@ +#!/usr/bin/env bash +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Orchestrate the four-stage QAT pipeline split example. +# +# Each stage runs in its own Python process and communicates only through files +# in WORKDIR. This mirrors a real workflow where training and lowering happen +# on separate machines or at different times. +# +# Usage: +# ./run.sh [--example {minimal|sliced}] [--workdir PATH] +# +# Options: +# --example Which of the two example paths to run (default: minimal). +# minimal : capture with a plain torch.export one-liner, skip +# the recipe entirely before QAT, then lower with the +# full recipe afterwards. +# sliced : slice the recipe around the QUANTIZE stage so that +# only the pre-quantize stages run in stage 1 and only +# the post-quantize stages run in stage 3. +# --workdir Directory for intermediate .pt2 / .pte / checkpoint files +# (default: