Add documentation and example of arbitrary QAT flow used with recipe-based export.

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).

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

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.

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()