Lowering a custom TorchAO QAT model to XNNPACK backend

[Lorenzo-Mazza]

Hi,

I trained a custom PyTorch model using the most recent TorchAO eager QAT workflow, cf. qat_workflow.:

What I did so far is i) I prepared the model using Int8DynamicActivationIntxWeight (using int4) config, ii) I trained the model with a standard pytorch training loop, iii) I saved the final checkpoint as a .ckpt file, before QAT convert.
I now want to deploy this model on an ARM device through ExecuTorch with the XNNPACK backend.
My understanding is that the next logical step should be:

quantize_(model, QATConfig(base_config, step="convert"))

After that, I would like to export/lower the converted model to a .pte file for XNNPACK. This is the step I cannot find a clear example for.

Here they describe pretty much what I have done, they reach the point where they have a trained model and then run the "convert" step from torchao. What I am missing is how to go from there to a lowered model to the correct XNNPACK backend.

Here it is said that "XNNPACK backend also supports quantizing models with the torchao quantize_ API", thus what I have in mind seems achievable. But there is no concrete example on how to do the actual lowering, and the link in the file is stale and sends to a 404 page. Likely, the missing link is referring to this page here, but in this torchao tutorial there is only a partial example showing the conversion happening through some off-the-shelf scripts for a specific standard model, namely

                                                                                                                                                                                                                                                              
# Convert checkpoint format for ExecuTorch                                                                                                                                                                                                                    
                                                                                                                                                                                                                                                              
python -m executorch.examples.models.phi_4_mini.convert_weights pytorch_model.bin pytorch_model_converted.bin
                                                                                                                              
# Export to PTE format with torchao optimizations preserved
                                                              
PARAMS="executorch/examples/models/phi_4_mini/config.json"
                                                              
python -m executorch.examples.models.llama.export_llama \
    --model "phi_4_mini" \
    --checkpoint "pytorch_model_converted.bin" \
    --params "$PARAMS" \
    -kv \
    --use_sdpa_with_kv_cache \
    -X \
    --metadata '{"get_bos_id":199999, "get_eos_ids":[200020,199999]}' \
    --max_seq_length 128 \
    --max_context_length 128 \
    --output_name="phi4-mini-8da4w.pte"
                                                              

This is not directly applicable to my custom model.
So my questions are:

• Is a custom non-LLM model trained with TorchAO eager QAT and Int8DynamicActivationIntxWeightConfig currently supported by ExecuTorch/XNNPACK?
• After running the TorchAO QAT convert step, what is the intended export/lowering API for a custom model?
• Is there a minimal reference example of: custom torch.nn.Module→ TorchAO QAT prepare→ train→ TorchAO QAT convert→ ExecuTorch/XNNPACK lowering→ .pte ?

Thanks,

Lorenzo Mazza

[rksharma-owg]

Yes, the standard export/lowering flow is usable for a custom model, but there is an important distinction in the current XNNPACK docs.

For a custom non-LLM model, the documented XNNPACK path is the PT2E flow:

1. Export the model with torch.export.export.
2. Prepare it with prepare_pt2e and XNNPACKQuantizer.
3. Calibrate if needed.
4. Run convert_pt2e.
5. Export the converted model and lower it with to_edge_transform_and_lower(..., partitioner=[XnnpackPartitioner()]).
6. Call .to_executorch() to produce the .pte program.

The current XNNPACK quantization guide says this path supports 8-bit symmetric weights with 8-bit asymmetric activations, including dynamic or static activations. It also explicitly says weight-only quantization is not currently supported by XNNPACK.

The Int8DynamicActivationIntxWeightConfig example in that same guide is under the LLM-oriented quantize_ section. I would not assume that an arbitrary custom model using int4 weights will lower successfully just because the TorchAO convert step completes. The backend still needs matching supported operators and quantization patterns.

So for a first working custom-model prototype, I would switch to XNNPACKQuantizer plus PT2E 8-bit quantization and use the generic lowering API. If int4 is a hard requirement, the next step is to check whether the resulting graph contains XNNPACK-supported int4 operators; otherwise it will need backend support or a different deployment path.