ExecuTorch for On-Device PyTorch Inference

#ai #machinelearning #llm #tutorial

ExecuTorch is PyTorch’s answer to “how do I run this model on a phone” that does not go through ONNX. The runtime is small and the API is short. Almost all of the difficulty is concentrated in one step, and it is the first one.

The shape of the toolchain

Three stages, and it is worth holding them separately because the error messages come from different layers:

Export. torch.export.export traces your nn.Module into a graph with no Python left in it. This is standard PyTorch, not ExecuTorch, and it is where models fail.
Lower. to_edge_transform_and_lower converts that graph to the Edge dialect and hands the parts a backend can accelerate to a partitioner.
Serialise. .to_executorch() produces a program you write to a .pte file — a single artefact containing the graph and the weights, which is what ships inside your app.

The runtime that loads a .pte is a C++ library you link into an Android or iOS application. There is no Python on the device, and the size of that runtime rather than of the framework is what you are adding to the binary — the model file itself is the part covered in bundling a quantized model inside a mobile app.

Export, which is where it breaks

torch.export requires a graph it can capture without running Python at inference time. Data-dependent control flow — an if on a tensor value, a loop whose trip count depends on the input, an early return on a threshold — cannot be captured, and the export raises rather than silently specialising. This is a feature: the alternative is a graph that is correct only for the example input you traced with.

Two practical consequences. First, shapes: by default the exported graph is specialised to the shapes of your sample inputs, and if you need a dimension to vary you must say so with a dynamic-shape specification rather than hoping. Second, the fixes are almost always in the model, not in the exporter — replace a Python if on tensor data with torch.where, hoist a configuration branch out of forward into construction, and replace a variable-length loop with a masked fixed-length one.

Lowering to a backend

Install the toolchain in a fresh environment: pip install executorch.
Put the model in eval() mode and prepare a tuple of sample inputs with the shapes you intend to run.
Export, lower with a partitioner for the target, and serialise to .pte.
Load the .pte back and check the outputs against eager PyTorch.

import torch
from executorch.exir import to_edge_transform_and_lower
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner

model = MyModel().eval()
sample_inputs = (torch.randn(1, 3, 224, 224),)

exported = torch.export.export(model, sample_inputs)

et_program = to_edge_transform_and_lower(
    exported,
    partitioner=[XnnpackPartitioner()],
).to_executorch()

with open("model.pte", "wb") as f:
    f.write(et_program.buffer)

The partitioner is the backend choice, and it is the only line that changes for a different target. XnnpackPartitioner targets Arm and x86 CPUs and is the portable default; there is a Core ML partitioner for Apple devices and a QualcommPartitioner for Qualcomm-powered Android phones. The current set is listed in PyTorch’s ExecuTorch getting-started guide.

Partial lowering is normal

A partitioner takes what it can accelerate and leaves the rest to the portable CPU kernels. That is not a failure, but it does mean the proportion of the graph a backend claimed is the number worth looking at — a model where the partitioner took two operators out of forty is going to behave like a CPU model with extra transfer overhead, for the same reason described in the Hexagon page.

Running the .pte

Before touching a device, load the file back in Python and compare against the eager model. Getting this wrong on a laptop is minutes; getting it wrong on a phone is an afternoon.

from executorch.runtime import Runtime

runtime = Runtime.get()
program = runtime.load_program("model.pte")
method = program.load_method("forward")

out_et = method.execute([sample_inputs[0]])[0]
out_eager = model(*sample_inputs)

print(torch.allclose(out_et, out_eager, atol=1e-4, rtol=1e-4))
print((out_et - out_eager).abs().max())

A small difference is expected — backends fuse operations and reassociate floating-point arithmetic, and reassociation is not exact. A large difference means the lowering changed the numerics, usually through a quantized kernel being applied where you did not intend one. Tighten the tolerance until it fails, then look at which output diverged.

Do this comparison on inputs from your real distribution rather than on torch.randn. Random tensors exercise a network’s middle range and almost nothing else, so a numerical check on noise will pass on a lowering that has broken the behaviour your product depends on. The sample inputs you exported with are a shape specification, not a test set, and the two should not be the same tensors.

Getting it onto a phone

The .pte is the portable artefact; the runtime is per-platform. On Android you link the ExecuTorch runtime through its AAR and call into it from Kotlin or through JNI; on iOS you add the runtime as a Swift package or framework and call it from Swift. The model file goes in as an asset, which puts it straight into the store size limits set out in the bundling page — a base module on Google Play is capped at 500 MB of compressed download, which a 1B model at 4 bits reaches on its own.

Two things to keep in view as you productionise. Backend selection is a build-time decision baked into the .pte, so a build lowered with QualcommPartitioner is not the build you ship to a MediaTek phone — either ship an Xnnpack build everywhere and accept CPU, or ship per-target artefacts and select at download time. And quantization is a separate concern that composes with this: you quantize before or during lowering, and the same accuracy verification discipline from the ONNX quantization page applies unchanged — compare against the unquantized model on real inputs, never on noise.

ExecuTorch reached a stable 1.x release and its partitioner module paths and runtime bindings have moved between versions. Pin the version you develop against and check the getting-started guide for that version rather than an older tutorial.